Merge pull request #151 from TomWagg/plots

Plot improvements

.github/workflows/test_cogsworth.yml

@@ -39,8 +39,10 @@ jobs:
       run: |
         # test with as many cores as possible for speed
         pytest cogsworth -n=auto --cov=./ --cov-report=xml
-    - name: Get Coverage
-      uses: orgoro/coverage@v3
+    - name: Upload coverage reports to Codecov
+      uses: codecov/codecov-action@v5
       with:
-        coverageFile: ./coverage.xml
-        token: ${{ secrets.GITHUB_TOKEN }}
+        fail_ci_if_error: false
+        files: ./cov.xml
+        token: ${{ secrets.CODECOV_TOKEN }}
+        verbose: true

cogsworth/_version.py

@@ -1 +1 @@
-__version__ = "2.0.1"
+__version__ = "2.0.2"

cogsworth/hydro/potential.py

@@ -48,7 +48,7 @@ def get_snapshot_potential(snap, components=[{"label": "star", "attr": "s", "r_s
                                                     nmax=nmax, lmax=lmax, r_s=comp["r_s"], skip_m=True)
 
         # add the SCF potential to the composite potential
-        pot[comp["label"]] = gp.scf.SCFPotential(m=subsnap["mass"].sum(), r_s=comp["r_s"],
+        pot[comp["label"]] = gp.scf.SCFPotential(m=subsnap["mass"].sum().tolist(), r_s=comp["r_s"],
                                                  Snlm=Snlm, Tnlm=Tnlm, units=galactic)
 
     # save the potential to a file if requested

cogsworth/hydro/rewind.py

@@ -43,10 +43,10 @@ def rewind_to_formation(subsnap, pot, dt=-1 * u.Myr, processes=1):
     final_time = subsnap.properties["time"].in_units("Myr") * u.Myr
 
     # convert the final particles to a phase-space position
-    wf = gd.PhaseSpacePosition(pos=np.transpose(subsnap["pos"].in_units("kpc")) * u.kpc,
-                               vel=np.transpose(subsnap["vel"].in_units("km s**-1")) * u.km / u.s)
+    wf = gd.PhaseSpacePosition(pos=np.transpose(subsnap["pos"].in_units("kpc").tolist()) * u.kpc,
+                               vel=np.transpose(subsnap["vel"].in_units("km s**-1").tolist()) * u.km / u.s)
     # store their formation times in Myr
-    tforms = subsnap["tform"].in_units("Myr") * u.Myr
+    tforms = subsnap["tform"].in_units("Myr").tolist() * u.Myr
 
     def args(wf, tforms):       # pragma: no cover
         for w, t in zip(wf, tforms):
@@ -64,16 +64,16 @@ def args(wf, tforms):       # pragma: no cover
 
     # check if the formation masses are available, warn if not
     if "massform" not in subsnap.all_keys():
-        mass = subsnap["mass"].in_units("Msol")
+        mass = subsnap["mass"].in_units("Msol").tolist()
         logging.getLogger("cogsworth").info("Formation masses (`massform`) not found, using present day masses instead")
     else:
-        mass = subsnap["massform"].in_units("Msol")
+        mass = subsnap["massform"].in_units("Msol").tolist()
 
     # create a df of the initial particles and return
     init_particles = np.zeros((len(subsnap), 10))
     init_particles[:, 0] = mass
     init_particles[:, 1] = subsnap["metals"] if np.ndim(subsnap["metals"]) == 1 else subsnap["metals"][:, 0]
-    init_particles[:, 2] = subsnap["tform"].in_units("Gyr")
+    init_particles[:, 2] = subsnap["tform"].in_units("Gyr").tolist()
     init_particles[:, 3] = subsnap["iord"]
 
     for i in range(len(w0)):

cogsworth/hydro/utils.py

@@ -79,8 +79,8 @@ def prepare_snapshot(path, halo_params={"mode": "ssc"}):
                     raise e
 
     # orient the snapshot face-on and side-on
-    side_on = pynbody.analysis.angmom.sideon(snap, cen=(0, 0, 0))
-    face_on = pynbody.analysis.angmom.faceon(snap, cen=(0, 0, 0))
+    side_on = pynbody.analysis.angmom.sideon(snap, already_centered=True)
+    face_on = pynbody.analysis.angmom.faceon(snap, already_centered=True)
 
     return snap
 

cogsworth/plot.py

@@ -166,7 +166,10 @@ def _rlof_path(centre, width, height, m=1.5, flip=False):
 
 
 def plot_cartoon_evolution(bpp, bin_num, label_type="long", plot_title="Cartoon Binary Evolution",
-                           y_sep_mult=1.5, offset=0.2, s_base=1000, fig=None, ax=None, show=True):    # pragma: no cover
+                           y_sep_mult=1.5, offset=0.2, s_base=1000,
+                           time_fs_mult=1.0, mass_fs_mult=1.0, kstar_fs_mult=1.0,
+                           porb_fs_mult=1.0, label_fs_mult=1.0,
+                           fig=None, ax=None, show=True):    # pragma: no cover
     """Plot COSMIC bpp output as a cartoon evolution
 
     Parameters
@@ -185,6 +188,16 @@ def plot_cartoon_evolution(bpp, bin_num, label_type="long", plot_title="Cartoon
         Offset from the centre for each of the stars (larger=wider binaries)
     s_base : `float`, optional
         Base scatter point size for the stars
+    time_fs_mult : `float`, optional
+        Multiplier for the time annotation fontsize
+    mass_fs_mult : `float`, optional
+        Multiplier for the mass annotation fontsize
+    kstar_fs_mult : `float`, optional
+        Multiplier for the kstar annotation fontsize
+    porb_fs_mult : `float`, optional
+        Multiplier for the porb annotation fontsize
+    label_fs_mult : `float`, optional
+        Multiplier for the evolution label fontsize
     fig : :class:`~matplotlib.pyplot.figure`, optional
         Figure on which to plot, by default will create a new one
     ax : :class:`~matplotlib.pyplot.axis`, optional
@@ -245,7 +258,7 @@ def plot_cartoon_evolution(bpp, bin_num, label_type="long", plot_title="Cartoon
         ax.annotate(f'{time:1.2e} Myr' if time > 1e4 else f'{time:1.2f} Myr',
                     xy=(-offset - 0.3 - (0.12 if len(inds) > 1 else 0),
                         total - np.mean(inds) * y_sep_mult), ha="right", va="center",
-                    fontsize=0.4*fs, fontweight="bold", zorder=-1,
+                    fontsize=0.4*fs*time_fs_mult, fontweight="bold", zorder=-1,
                     bbox=dict(boxstyle="round,pad=0.2", fc="white", ec="white") if len(inds) > 1 else None)
         # if there's more than one ind then plot a bracketed line connecting them
         if len(inds) > 1:
@@ -297,7 +310,7 @@ def plot_cartoon_evolution(bpp, bin_num, label_type="long", plot_title="Cartoon
         # check if either star is now a massless remnant
         mr_1 = k1["short"] == "MR"
         mr_2 = k2["short"] == "MR"
-        ks_fontsize = 0.3 * fs
+        ks_fontsize = 0.3 * fs * kstar_fs_mult
 
         # start an evolution label variable
         evol_label = et[label_type]
@@ -309,7 +322,7 @@ def plot_cartoon_evolution(bpp, bin_num, label_type="long", plot_title="Cartoon
             evol_label = f'{which_star} evolved to\n{to_type}'
 
         # annotate the evolution label on the right side of the binary
-        ax.annotate(evol_label, xy=(0.5, total - i), va="center")
+        ax.annotate(evol_label, xy=(0.5, total - i), va="center", fontsize=0.4*fs*label_fs_mult)
 
         # if we've got a common envelope then draw an ellipse behind the binary
         if common_envelope:
@@ -336,7 +349,7 @@ def plot_cartoon_evolution(bpp, bin_num, label_type="long", plot_title="Cartoon
 
             # annotate the correct mass
             ax.annotate(f'{row["mass_1"] if mr_2 else row["mass_2"]:1.2f} ' + r'$\rm M_{\odot}$',
-                        xy=(0, total - i - 0.45), ha="center", va="top", fontsize=0.3*fs,
+                        xy=(0, total - i - 0.45), ha="center", va="top", fontsize=0.3*fs*mass_fs_mult,
                         bbox=dict(boxstyle="round,pad=0.2", fc="white", ec="none", alpha=0.7)
                         if et_ind in [15, 16] else None)
 
@@ -358,13 +371,14 @@ def plot_cartoon_evolution(bpp, bin_num, label_type="long", plot_title="Cartoon
             mass_y_offset = 0.35 if not (rlof and not common_envelope) else 0.5
             ax.annotate(f'{row["mass_1"]:1.2f} ' + r'$\rm M_{\odot}$',
                         xy=(0 - offset * contact_adjust - off_s, total - i - mass_y_offset),
-                        ha="left" if common_envelope else "center", va="top", fontsize=0.3*fs,
+                        ha="left" if common_envelope else "center", va="top", fontsize=0.3*fs*mass_fs_mult,
                         rotation=45 if contact else 0,
                         bbox=dict(boxstyle="round,pad=0.2", fc="white", ec="none", alpha=0.7)
                         if et_ind in [15, 16] else None)
             ax.annotate(f'{row["mass_2"]:1.2f} ' + r'$\rm M_{\odot}$',
                         xy=(0 + offset * contact_adjust + off_s, total - i - mass_y_offset),
-                        ha="right" if common_envelope else "center", va="top", fontsize=0.3*fs, zorder=1000,
+                        ha="right" if common_envelope else "center", va="top", fontsize=0.3*fs*mass_fs_mult,
+                        zorder=1000,
                         rotation=45 if contact else 0,
                         bbox=dict(boxstyle="round,pad=0.2", fc="white", ec="none", alpha=0.7)
                         if et_ind in [15, 16] else None)
@@ -398,7 +412,7 @@ def plot_cartoon_evolution(bpp, bin_num, label_type="long", plot_title="Cartoon
                 p_lab = f'{row["porb"]:1.2e} days' if row["porb"] > 10000 or row["porb"] < 1\
                     else f'{row["porb"]:1.0f} days'
                 ax.annotate(p_lab, xy=(x, total - i + 0.05), ha="center", va="bottom",
-                            fontsize=0.2*fs if row["porb"] > 10000 or row["porb"] < 1 else 0.3*fs)
+                            fontsize=0.2*fs*porb_fs_mult if row["porb"] > 10000 or row["porb"] < 1 else 0.3*fs*porb_fs_mult)
 
             # for non-common-envelope RLOF, plot a RLOF teardrop in the background
             if rlof and not common_envelope:

cogsworth/utils.py

@@ -13,8 +13,8 @@
     {'long': 'Hertzsprung Gap', 'short': 'HG', 'colour': (0.939608, 0.471373, 0.094902, 1.0)},
     {'long': 'First Giant Branch', 'short': 'FGB', 'colour': (0.716186, 0.833203, 0.916155, 1.0)},
     {'long': 'Core Helium Burning', 'short': 'CHeB', 'colour': (0.29098, 0.59451, 0.78902, 1.0)},
-    {'long': 'Early Asymptotic Giant Branch', 'short': 'EAGB', 'colour': (0.294902, 0.690196, 0.384314, 1.0)},
-    {'long': 'Thermally Pulsing Asymptotic Giant Branch', 'short': 'TPAGB',
+    {'long': 'Early AGB', 'short': 'EAGB', 'colour': (0.294902, 0.690196, 0.384314, 1.0)},
+    {'long': 'Thermally Pulsing AGB', 'short': 'TPAGB',
      'colour': (0.723122, 0.889612, 0.697178, 1.0)},
     {'long': 'Helium Main Sequence', 'short': 'HeMS', 'colour': (0.254627, 0.013882, 0.615419, 1.0)},
     {'long': 'Helium Hertsprung Gap', 'short': 'HeHG', 'colour': (0.562738, 0.051545, 0.641509, 1.0)},

docs/modules/changelog.rst

@@ -4,6 +4,11 @@ Full changelog
 
 This page tracks all of the changes that have been made to ``cogsworth``. We follow the standard versioning convention of A.B.C, where C is a patch/bugfix, B is a large bugfix or new feature and A is a major new breaking change. B/C are backwards compatible but A changes may be breaking.
 
+2.0.2
+=====
+- Add some more flexbility to the ``plot_cartoon_binary`` function in terms of fontsize
+- Allow use of pynbody 2.0 with cogsworth
+
 2.0.1
 =====
 

setup.cfg

@@ -53,7 +53,7 @@ actions =
 lisa =
     legwork >= 0.4.6
 hydro = 
-    pynbody
+    pynbody >= 2.0.0
 test = 
     %(observables)s
     %(lisa)s