Squashed commit of the following:

commit d584541 Author: João Faria <[email protected]> Date: Thu Jul 18 11:50:31 2024 +0200 artifact naming per platform; retention 2 days commit 79cb25d Author: João Faria <[email protected]> Date: Thu Jul 18 11:45:07 2024 +0200 update and bug fixes in phase_plot commit 6b21891 Author: João Faria <[email protected]> Date: Thu Jul 18 11:38:42 2024 +0200 always commit a2ac1ab Author: João Faria <[email protected]> Date: Thu Jul 18 11:36:25 2024 +0200 no warnings commit 38dd769 Author: João Faria <[email protected]> Date: Thu Jul 18 11:35:28 2024 +0200 try adding artifacts commit 839f7e3 Author: João Faria <[email protected]> Date: Thu Jul 18 11:25:24 2024 +0200 separate tests commit 0041d83 Author: João Faria <[email protected]> Date: Thu Jul 18 11:17:30 2024 +0200 try a more explicit owner commit 9ee1ca1 Author: João Faria <[email protected]> Date: Thu Jul 18 10:41:52 2024 +0200 none owner to force copy? commit eff2640 Author: João Faria <[email protected]> Date: Thu Jul 18 09:33:32 2024 +0200 test phase plots against baseline commit eaca351 Author: João Faria <[email protected]> Date: Thu Jul 18 09:30:35 2024 +0200 update pip and several python versions for windows commit 5f97c4b Author: João Faria <[email protected]> Date: Thu Jul 18 09:30:01 2024 +0200 update test keplerian commit 5bc2f90 Author: João Faria <[email protected]> Date: Wed Jul 17 19:24:41 2024 +0200 explicit rv_policy copy commit ad52117 Author: João Faria <[email protected]> Date: Wed Jul 17 19:10:49 2024 +0200 debug
kima-org · Jul 18, 2024 · 22dde1a · 22dde1a
1 parent 7abc523
commit 22dde1a
Show file tree

Hide file tree

Showing 9 changed files with 138 additions and 63 deletions.
diff --git a/.github/workflows/pip-extensive.yml b/.github/workflows/pip-extensive.yml
@@ -16,6 +16,7 @@ jobs:
         experimental: [false]
         include:
           - platform: windows-latest
+            python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
             experimental: true
 
     steps:
@@ -32,7 +33,7 @@ jobs:
 
     - name: Dependencies
       run: |
-        pip install pip -U
+        python -m pip install pip -U
         pip install pytest
 
     - name: Check versions

diff --git a/.github/workflows/pip.yml b/.github/workflows/pip.yml
@@ -53,4 +53,12 @@ jobs:
 
     - name: Test
       run: |
-        python -m pytest -s tests/
+        python -m pytest -s tests/ 
+        python -m pytest -s tests/ --mpl  --mpl-results-path=results --mpl-generate-summary=html
+
+    - uses: actions/upload-artifact@v4
+      if: always()
+      with:
+        name: results-${{ matrix.platform }}
+        path: results
+        retention-days: 2
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -8,13 +8,13 @@ if(WIN32)
   add_compile_definitions(_USE_MATH_DEFINES)
 endif()
 
-if (MSVC)
-    # warning level 3
-    add_compile_options(/W3)
-else()
-    # additional warnings
-    add_compile_options(-Wall -Wextra -Wpedantic)
-endif()
+# if (MSVC)
+#     # warning level 3
+#     add_compile_options(/W3)
+# else()
+#     # additional warnings
+#     add_compile_options(-Wall -Wextra -Wpedantic)
+# endif()
 
 # add_compile_options(-Wall -Wextra -Wpedantic)
 

diff --git a/src/kima/kepler.cpp b/src/kima/kepler.cpp
@@ -681,8 +681,8 @@ namespace brandt
         else
         {
             // If cos(E) = -1, E = pi and tan(0.5*E) -> inf and f = E = pi
-            *sinf = 0;
-            *cosf = -1;
+            *sinf = 0.0;
+            *cosf = -1.0;
         }
     }
 
@@ -1190,7 +1190,12 @@ NB_MODULE(kepler, m) {
                            const double &P, const double &K, const double &ecc, 
                            const double &w, const double &M0, const double &M0_epoch) {
         size_t size = t.size();
-        auto v = brandt::keplerian(t, P, K, ecc, w, M0, M0_epoch);
-        return nb::ndarray<nb::numpy, double>(v.data(), {size}, nb::handle());
-    }, "t"_a, "P"_a, "K"_a, "ecc"_a, "w"_a, "M0"_a, "M0_epoch"_a, KEPLERIAN_DOC);
+        struct Temp { std::vector<double> v; };
+        Temp *temp = new Temp();
+        temp->v = brandt::keplerian(t, P, K, ecc, w, M0, M0_epoch);
+        nb::capsule owner(temp, [](void *p) noexcept { delete (Temp *) p; });
+        return nb::ndarray<nb::numpy, double>(temp->v.data(), {size}, owner);
+    }, "t"_a, "P"_a, "K"_a, "ecc"_a, "w"_a, "M0"_a, "M0_epoch"_a, KEPLERIAN_DOC
+    // nb::rv_policy::copy
+    );
 }
diff --git a/src/kima/pykima/display.py b/src/kima/pykima/display.py
@@ -2244,9 +2244,11 @@ def corner_known_object(res, star_mass=1.0, adda=False, **kwargs):
 
 def phase_plot_logic(res, sample, sort_by_decreasing_K=False, sort_by_increasing_P=False):
     from string import ascii_lowercase
+    letters = ascii_lowercase[1:]
+
     nplanets = int(sample[res.indices['np']])
 
-    params = {ascii_lowercase[1+i]: {} for i in range(nplanets)}
+    params = {letters[i]: {} for i in range(nplanets)}
     for i, k in enumerate(params.keys()):
         params[k]['P'] = P = sample[res.indices['planets.P']][i]
         params[k]['K'] = K = sample[res.indices['planets.K']][i]
@@ -2258,8 +2260,8 @@ def phase_plot_logic(res, sample, sort_by_decreasing_K=False, sort_by_increasing
 
     pj = 0
     if res.KO:
+        ko = {letters[i]: {} for i in range(nplanets, nplanets + res.nKO)}
         nplanets += res.nKO
-        ko = {ascii_lowercase[i]: {} for i in range(nplanets, nplanets + res.nKO)}
         for i, k in enumerate(ko.keys()):
             ko[k]['P'] = P = sample[res.indices['KOpars']][i]
             ko[k]['K'] = K = sample[res.indices['KOpars']][i + res.nKO]
@@ -2272,16 +2274,17 @@ def phase_plot_logic(res, sample, sort_by_decreasing_K=False, sort_by_increasing
         params.update(ko)
 
     if res.TR:
+        tr = {letters[i]: {} for i in range(nplanets, nplanets + res.nTR)}
         nplanets += res.nTR
-        tr = {ascii_lowercase[i]: {} for i in range(nplanets, nplanets + res.nTR)}
         for i, k in enumerate(tr.keys()):
             tr[k]['P'] = P = sample[res.indices['TRpars']][i]
             tr[k]['K'] = K = sample[res.indices['TRpars']][i + res.nTR]
-            tr[k]['φ'] = φ = sample[res.indices['TRpars']][i + 2 * res.nTR]
+            tr[k]['Tc'] = Tc = sample[res.indices['TRpars']][i + 2 * res.nTR]
             tr[k]['e'] = e = sample[res.indices['TRpars']][i + 3 * res.nTR]
             tr[k]['w'] = w = sample[res.indices['TRpars']][i + 4 * res.nTR]
-            tr[k]['Tp'] = res.M0_epoch - (P * φ) / (2*np.pi)
+            # tr[k]['Tp'] = res.M0_epoch - (P * φ) / (2*np.pi)
             tr[k]['index'] = -pj - 1
+            pj += 1
         params.update(tr)
 
     keys = list(params.keys())
@@ -2295,20 +2298,44 @@ def phase_plot_logic(res, sample, sort_by_decreasing_K=False, sort_by_increasing
     return nplanets, params, keys
 
 
-def phase_plot(res,
-               sample,
-               highlight=None,
-               only=None,
-               phase_axs=None,
-               add_titles=True,
-               sharey=False,
-               highlight_points=None,
-               sort_by_increasing_P=False,
-               sort_by_decreasing_K=True,
-               show_gls_residuals=False,
+def phase_plot(res, sample, phase_axs=None,
+               sort_by_increasing_P=False, sort_by_decreasing_K=True,
+               highlight=None, highlight_points=None, only=None,
+               add_titles=True, sharey=False, show_gls_residuals=False,
                **kwargs):
-    """ Plot the phase curves given the solution in `sample` """
-    # this is probably the most complicated function in the whole package!!
+    """
+    Plot the planet phase curves, GP, and residuals, for a given `sample`.
+    
+    Args:
+        res (kima.KimaResults):
+            The `KimaResults` instance
+        sample (array):
+            Array with one posterior sample
+        phase_axs (list[matplotlib.axes.Axes]):
+            One or more axes for the phase plot(s)
+        sort_by_increasing_P (bool):
+            Sort the planets by increasing period
+        sort_by_decreasing_K (bool):
+            Sort the planets by decreasing semi-amplitude
+        highlight (list):
+            Highlight all data points from a specific instrument
+        highlight_points (list):
+            Highlight specific data points by index
+        only (list):
+            Only show data from specific instrument(s)
+        add_titles (bool):
+            Add titles to each phase plot
+        sharey (bool):
+            Share the y-axis of the phase plots
+        show_gls_residuals (bool):
+            Add a panel with the Lomb-Scargle periodogram of the residuals
+        **kwargs (dict):
+            Keyword arguments passed to `plt.errorbar`
+    
+    Warning:
+        This is probably the most complicated function in the whole package! For
+        one, the layout of the axes in the figure may not always be optimal.
+    """
 
     if res.max_components == 0 and not res.KO and not res.TR:
         print('Model has no planets! phase_plot() doing nothing...')
@@ -2373,24 +2400,29 @@ def phase_plot(res,
     if show_gls_residuals:
         fs[0] += 2
 
+    # axes layout:
+    # up to 3 planets: 1 row
+    # up to 6 planets: 2 rows
+
+
     fig = plt.figure(constrained_layout=True, figsize=fs)
     nrows = {
-        1: 2,
-        2: 2,
-        3: 2,
-        4: 4,
-        5: 4,
-        6: 4,
-        7: 5,
-        8: 5,
-        9: 5,
-        10: 5
+        1: 1, 2: 1, 3: 1,
+        4: 2, 5: 2, 6: 2,
+        7: 4, 8: 4, 9: 4,
     }[nplanets]
 
+    # at least the residuals plot
+    nrows += 1
+
     if res.has_gp:
         nrows += 1
 
-    ncols = nplanets if nplanets <= 3 else 3
+    ncols = {
+        1: 1, 2: 2, 3: 3,
+        4: 2,
+        5: 3, 6: 3, 7: 3, 8: 3, 9: 3
+    }[nplanets]
     hr = [2] * (nrows - 1) + [1]
     wr = None
 
@@ -2399,32 +2431,58 @@ def phase_plot(res,
         ncols += 1
         # fig.set_size_inches(fs[0], fs[1])
 
-    gs = gridspec.GridSpec(nrows, ncols, figure=fig, height_ratios=hr,
-                           width_ratios=wr)
-    gs_indices = {i: (i // 3, i % 3) for i in range(50)}
-    axs = []
+    gs = gridspec.GridSpec(nrows, ncols, figure=fig,
+                           height_ratios=hr, width_ratios=wr)
+    # gs_indices = {i: (i // 3, i % 3) for i in range(50)}
+
+    if phase_axs is None:
+        if nplanets == 1:
+            axs = [fig.add_subplot(gs[0, 0])]
+        elif nplanets == 2:
+            axs = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1])]
+        elif nplanets == 3:
+            axs = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1]), fig.add_subplot(gs[0, 2])]
+        elif nplanets == 4:
+            axs = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1]),
+                   fig.add_subplot(gs[1, 0]), fig.add_subplot(gs[1, 1])]
+        elif nplanets == 5:
+            axs = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1]), fig.add_subplot(gs[0, 2]), 
+                   fig.add_subplot(gs[1, 0]), fig.add_subplot(gs[1, 1])]
+        elif nplanets == 6:
+            axs = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1]), fig.add_subplot(gs[0, 2]),
+                   fig.add_subplot(gs[1, 0]), fig.add_subplot(gs[1, 1]), fig.add_subplot(gs[1, 2])]
+        elif nplanets == 7:
+            axs = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1]), fig.add_subplot(gs[0, 2]),
+                   fig.add_subplot(gs[1, 0]), fig.add_subplot(gs[1, 1]), fig.add_subplot(gs[1, 2]),
+                   fig.add_subplot(gs[2, 0])]
+        elif nplanets == 8:
+            axs = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1]), fig.add_subplot(gs[0, 2]),
+                   fig.add_subplot(gs[1, 0]), fig.add_subplot(gs[1, 1]), fig.add_subplot(gs[1, 2]),
+                   fig.add_subplot(gs[2, 0]), fig.add_subplot(gs[2, 1])]
+        elif nplanets == 9:
+            axs = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1]), fig.add_subplot(gs[0, 2]),
+                   fig.add_subplot(gs[1, 0]), fig.add_subplot(gs[1, 1]), fig.add_subplot(gs[1, 2]),
+                   fig.add_subplot(gs[2, 0]), fig.add_subplot(gs[2, 1]), fig.add_subplot(gs[2, 2])]
+        else:
+            raise NotImplementedError
+    else:
+        axs = phase_axs
+
+    colors = kwargs.get('colors', None)
 
     # for each planet in this sample
     for i, letter in enumerate(keys):
-        if phase_axs is None:
-            ax = fig.add_subplot(gs[gs_indices[i]])
-        else:
-            try:
-                ax = phase_axs[i]
-            except IndexError:
-                continue
-
-        axs.append(ax)
+        ax = axs[i]
 
         ax.axvline(0.5, ls='--', color='k', alpha=0.2, zorder=-5)
         ax.axhline(0.0, ls='--', color='k', alpha=0.2, zorder=-5)
 
         P = params[letter]['P']
-        Tp = params[letter]['Tp']
+        # Tp = params[letter]['Tp']
 
         # plot the keplerian curve in phase (3 times)
         phase = np.linspace(0, 1, 200)
-        tt = phase * P + Tp
+        tt = phase * P + res.M0_epoch
 
         # keplerian for this planet
         planet_index = params[letter]['index']
@@ -2454,13 +2512,16 @@ def phase_plot(res,
         if res.multi:
             for k in range(1, res.n_instruments + 1):
                 m = obs == k
-                phase = ((t[m] - Tp) / P) % 1.0
+                phase = ((t[m] - res.M0_epoch) / P) % 1.0
 
                 yy = (y - vv)[m]
                 ee = e[m].copy()
 
-                # one color for each instrument
-                color = f'C{k-1}'
+                if colors is None:
+                    # one color for each instrument
+                    color = f'C{k-1}'
+                else:
+                    color = colors[k - 1]
 
                 for j in (-1, 0, 1):
                     label = res.instruments[k - 1] if j == 0 else ''
@@ -2492,7 +2553,7 @@ def phase_plot(res,
                                     alpha=alpha, **hlkw)
 
         else:
-            phase = ((t - Tp) / P) % 1.0
+            phase = ((t - res.M0_epoch) / P) % 1.0
             yy = y - vv
 
             for j in (-1, 0, 1):

diff --git a/tests/baseline/test_phase_plot_include0-True-1-True-1-0.png b/tests/baseline/test_phase_plot_include0-True-1-True-1-0.png
diff --git a/tests/baseline/test_phase_plot_include1-True-2-True-3-1.png b/tests/baseline/test_phase_plot_include1-True-2-True-3-1.png
diff --git a/tests/baseline/test_phase_plot_include2-True-1-True-2-1.png b/tests/baseline/test_phase_plot_include2-True-1-True-2-1.png
diff --git a/tests/test_keplerian.py b/tests/test_keplerian.py
@@ -56,4 +56,4 @@ def test_speed():
         equality_check=None,
         n_range=[2**k for k in range(2, 10)],
     )
-    # plt.show()
+    plt.show()