Change to brackets for timesteps on cartoons, fix pop saving for complicated sampling params

cogsworth/plot.py

@@ -243,14 +243,18 @@ def plot_cartoon_evolution(bpp, bin_num, label_type="long", plot_title="Cartoon
     # annotate the time on the left side of the binary
     for time, inds in zip(times, row_inds):
         ax.annotate(f'{time:1.2e} Myr' if time > 1e4 else f'{time:1.2f} Myr',
-                    xy=(-offset - 0.3, total - np.mean(inds) * y_sep_mult), ha="right", va="center",
+                    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,
                     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 double arrowed line connecting them
+        # if there's more than one ind then plot a bracketed line connecting them
         if len(inds) > 1:
-            ax.annotate("", xy=(-offset - 0.4, total - inds[0] * y_sep_mult + 0.3),
-                        xytext=(-offset - 0.4, total - inds[-1] * y_sep_mult - 0.3),
-                        arrowprops=dict(arrowstyle="<|-|>", lw=1, color="black"), zorder=-2)
+            ax.annotate('', xy=(-offset - 0.35, total - np.mean(inds) * y_sep_mult),
+                        xytext=(-offset - 0.4, total - np.mean(inds) * y_sep_mult),
+                        ha='center', va='center',
+                        # 2.5 is a magic number here to make the bracket the right size
+                        arrowprops=dict(arrowstyle=f'-[, widthB={y_sep_mult * len(inds) * 2.5}, lengthB=1',
+                                        lw=1.5, color='k'))
 
     period_offset = 0.2
 

cogsworth/pop.py

@@ -1429,8 +1429,7 @@ def save(self, file_name, overwrite=False):
 
             # save sampling params
             d = file.create_dataset("sampling_params", data=[])
-            for key in self.sampling_params:
-                d.attrs[key] = self.sampling_params[key]
+            d.attrs["dict"] = yaml.dump(self.sampling_params, default_flow_style=None)
 
 
 def load(file_name, parts=["initial_binaries", "initial_galaxy", "stellar_evolution"]):
@@ -1470,9 +1469,7 @@ def load(file_name, parts=["initial_binaries", "initial_galaxy", "stellar_evolut
         for key in file["BSE_settings"].attrs:
             BSE_settings[key] = file["BSE_settings"].attrs[key]
 
-        # load in sampling params
-        for key in file["sampling_params"].attrs:
-            sampling_params[key] = file["sampling_params"].attrs[key]
+        sampling_params = yaml.load(file["sampling_params"].attrs["dict"], Loader=yaml.Loader)
 
     with h5.File(file_name, 'r') as f:
         galactic_potential = potential_from_dict(yaml.load(f.attrs["potential_dict"], Loader=yaml.Loader))

cogsworth/tests/test_pop.py

@@ -73,6 +73,22 @@ def test_io(self):
 
         os.remove("testing-pop-io.h5")
 
+    def test_save_complicated_sampling(self):
+        """Check that you can save a population with complicated sampling params"""
+        p = pop.Population(2, processes=1, 
+                           sampling_params={"qmin": 0.5, "porb_model": {
+            "min": 0.15,
+            "max": 5,
+            "slope": 0.0
+        }})
+        p.create_population()
+
+        p.save("testing-pop-io-sampling", overwrite=True)
+
+        p_loaded = pop.load("testing-pop-io-sampling", parts=["initial_binaries"])
+
+        self.assertTrue(np.all(p.initC == p_loaded.initC))
+
     def test_lazy_io(self):
         """Check that a population can be saved and re-loaded lazily"""
         p = pop.Population(2, processes=1, bcm_timestep_conditions=[['dtp=100000.0']],

docs/case_studies/binaries_and_potentials.ipynb

@@ -140,34 +140,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 65,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "<Population - 60481 evolved systems - galactic_potential=NullPotential, sfh_model=Wagg2022>"
-      ]
-     },
-     "execution_count": 65,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "p_no_potential"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 20,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "71504it [00:54, 1304.08it/s]                            \n"
+      "71504it [00:41, 1743.30it/s]                            \n"
      ]
     }
    ],
@@ -198,7 +178,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -220,7 +200,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 21,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -460,7 +440,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.11"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

docs/modules/changelog.rst

@@ -11,6 +11,7 @@ Major release to go with the release paper submission! 🎉
 - New feature: ``plot_cartoon_binary`` will now adjust the width of the binary based on the orbital period and label simultaneous timesteps more clearly
 - Major change: calls like ``p.bpp`` will now raise an error if sampling is not yet done to avoid confusion
 - Bug fix: Can now save and load unevolved populations
+- Bug fix: Saving ``sampling_params`` now works correctly when you have a ``sampling_params`` object that includes a dictionary (nested dictionaries were causing crashes before)
 
 1.2.0
 =====