Add translation and cartoon helpers

cogsworth/__init__.py

@@ -1,8 +1,8 @@
-from . import galaxy, kicks, pop, events, classify, observables
+from . import galaxy, kicks, pop, events, classify, observables, utils
 from ._version import __version__
 from .citations import CITATIONS
 
 __bibtex__ = __citation__ = CITATIONS["general"]["cogsworth"]["bibtex"]
 __uri__ = "https://cogsworth.readthedocs.io/"
 __author__ = "Tom Wagg"
-__email__ = "tomjwagg@gmail.com"
+__email__ = "tomjwagg@gmail.com"

cogsworth/pop.py

@@ -21,6 +21,7 @@
 from cogsworth.classify import determine_final_classes
 from cogsworth.observables import get_photometry
 from cogsworth.tests.optional_deps import check_dependencies
+from cogsworth.utils import translate_COSMIC_tables, plot_cartoon_evolution
 
 from cogsworth.citations import CITATIONS
 
@@ -950,6 +951,38 @@ def plot_map(self, nside=128, coord="C",
         if show:
             plt.show()
 
+    def translate_tables(self, **kwargs):
+        """Translate the COSMIC BSE tables to human readable format
+
+        Parameters
+        ----------
+
+        **kwargs : `various`
+            Any arguments to pass to :func:`~cogsworth.utils.translate_COSMIC_tables`
+        """
+        self._bpp = translate_COSMIC_tables(self._bpp, **kwargs)
+        self._final_bpp = translate_COSMIC_tables(self._final_bpp, **kwargs)
+
+        kwargs.update({"evol_type": False})
+        self._bcm = translate_COSMIC_tables(self._bcm, **kwargs)
+
+    def plot_cartoon_binary(self, bin_num, **kwargs):
+        """Plot a cartoon of the evolution of a single binary
+
+        Parameters
+        ----------
+        bin_num : `int`
+            Which binary to plot
+        **kwargs : `various`
+            Keyword arguments to pass, see :func:`~cogsworth.utils.plot_cartoon_evolution` for options
+
+        Returns
+        -------
+        fig, ax : :class:`~matplotlib.pyplot.figure`, :class:`~matplotlib.pyplot.axis`
+            Figure and axis of the plot
+        """
+        return plot_cartoon_evolution(self.bpp, bin_num, **kwargs)
+
     def save(self, file_name, overwrite=False):
         """Save a Population to disk
 

cogsworth/tests/test_pop.py

@@ -321,3 +321,23 @@ def test_sampling(self):
             p.sample_initial_binaries()
             q = p._initial_binaries["mass_2"] / p._initial_binaries["mass_1"]
             self.assertTrue(min(q) >= qmin)
+
+    def test_translation(self):
+        """Ensure that COSMIC tables are being translated properly"""
+        p = pop.Population(10)
+        p.perform_stellar_evolution()
+        p.translate_tables(replace_columns=False, label_type="short")
+
+        self.assertTrue(p.bpp["kstar_1"].dtype == np.float64)
+        self.assertTrue((p.bpp["kstar_1_str"][p.bpp["kstar_1"] == 1] == "MS").all())
+
+        p.translate_tables(replace_columns=True)
+        self.assertFalse(p.bpp["kstar_1"].dtype == np.float64)
+
+    def test_cartoon(self):
+        """Ensure that the cartoon plot works"""
+        p = pop.Population(10, final_kstar1=[14])
+        p.perform_stellar_evolution()
+
+        for bin_num in p.bin_nums:
+            p.plot_cartoon_binary(bin_num, show=False)