Merge pull request #133 from TomWagg/docs

New tutorials and examples

cogsworth/plot.py

@@ -438,6 +438,7 @@ def plot_cartoon_evolution(bpp, bin_num, label_type="long", plot_title="Cartoon
 def plot_galactic_orbit(primary_orbit, secondary_orbit=None,
                         t_min=0 * u.Myr, t_max=np.inf * u.Myr, show_start=True,
                         primary_kwargs={}, secondary_kwargs={}, start_kwargs={},
+                        show_legend=True,
                         fig=None, axes=None, show=True):
     """Plot the galactic orbit of a binary system. This provides a wrapper around the gala
     :class:`~gala.dynamics.Orbit` method :meth:`~gala.dynamics.Orbit.plot`.
@@ -507,6 +508,11 @@ def plot_galactic_orbit(primary_orbit, secondary_orbit=None,
 
         primary_orbit[0].plot(axes=fig.axes, **full_start_kwargs)
 
+    if show_legend:     # pragma: no cover
+        handles, labels = axes[-1].get_legend_handles_labels()
+        fig.legend(handles, labels, loc='upper center', ncol=5, fontsize=16)
+        fig.subplots_adjust(top=0.9)
+
     # show the plot if desired
     if show:
         plt.show()

cogsworth/pop.py

@@ -334,6 +334,10 @@ def __getitem__(self, ind):
                 new_pop._final_vel = self._final_vel[all_inds]
         return new_pop
 
+    def copy(self):
+        """Create a copy of the population"""
+        return self[:]
+
     def get_citations(self, filename=None):
         """Print the citations for the packages/papers used in the population"""
         # ask users for a filename to save the bibtex to
@@ -1217,7 +1221,8 @@ def plot_cartoon_binary(self, bin_num, **kwargs):
         """
         return plot_cartoon_evolution(self.bpp, bin_num, **kwargs)
 
-    def plot_orbit(self, bin_num, show_sn=True, sn_kwargs={}, show=True, **kwargs):
+    def plot_orbit(self, bin_num, show_sn=True, sn_kwargs={}, show=True,
+                   show_legend=True, **kwargs):
         """Plot the Galactic orbit of a binary system
 
         Parameters
@@ -1228,6 +1233,10 @@ def plot_orbit(self, bin_num, show_sn=True, sn_kwargs={}, show=True, **kwargs):
             Whether to show the position of the SN, by default True
         sn_kwargs : `dict`, optional
             Any additional arguments to pass to the SN scatter plot call, by default {}
+        show : `bool`, optional
+            Whether to immediately show the plot
+        show_legend : `bool`, optional
+            Whether to show the legend
         **kwargs : `various`
             Any additional arguments to pass to :func:`cogsworth.plot.plot_galactic_orbit`
 
@@ -1252,8 +1261,14 @@ def plot_orbit(self, bin_num, show_sn=True, sn_kwargs={}, show=True, **kwargs):
         primary_orbit = self.primary_orbits[ind]
         secondary_orbit = self.secondary_orbits[ind] if disrupted else None
 
+        fig, axes = kwargs.pop("fig", None), kwargs.pop("axes", None)
+        if fig is None or axes is None:
+            fig, axes = plt.subplots(1, 3, figsize=(25, 7))
+            fig.subplots_adjust(wspace=0.35)
+
         # create the orbit plot
-        fig, axes = plot_galactic_orbit(primary_orbit, secondary_orbit, show=False, **kwargs)
+        fig, axes = plot_galactic_orbit(primary_orbit, secondary_orbit, show=False,
+                                        show_legend=False, fig=fig, axes=axes, **kwargs)
 
         # extra bit if you want to show the SN locations
         if show_sn:
@@ -1266,10 +1281,11 @@ def plot_orbit(self, bin_num, show_sn=True, sn_kwargs={}, show=True, **kwargs):
 
             # loop over the primary and secondary
             rows = self.bpp.loc[bin_num]
-            for mask, orbit, colour in zip([((rows["evol_type"] == 15)
-                                             | ((rows["evol_type"] == 16) & (rows["sep"] == 0.0))),
-                                            rows["evol_type"] == 16],
-                                           [primary_orbit, secondary_orbit], colours):
+            for mask, orbit, colour, l in zip([((rows["evol_type"] == 15)
+                                               | ((rows["evol_type"] == 16) & (rows["sep"] == 0.0))),
+                                              rows["evol_type"] == 16],
+                                              [primary_orbit, secondary_orbit], colours,
+                                              ["Primary", "Secondary"]):
                 # if there is no SN or no orbit then skip
                 if not np.any(mask) or orbit is None:           # pragma: no cover
                     continue
@@ -1283,7 +1299,7 @@ def plot_orbit(self, bin_num, show_sn=True, sn_kwargs={}, show=True, **kwargs):
                     "marker": (10, 2, 0),
                     "color": colour,
                     "s": 100,
-                    "label": "SN position"
+                    "label": f"{l} supernova"
                 }
                 full_sn_kwargs.update(sn_kwargs)
 
@@ -1292,6 +1308,11 @@ def plot_orbit(self, bin_num, show_sn=True, sn_kwargs={}, show=True, **kwargs):
                 axes[1].scatter(sn_pos.x, sn_pos.z, **full_sn_kwargs)
                 axes[2].scatter(sn_pos.y, sn_pos.z, **full_sn_kwargs)
 
+        if show_legend:
+            handles, labels = axes[-1].get_legend_handles_labels()
+            fig.legend(handles, labels, loc='upper center', ncol=5, fontsize=16)
+            fig.subplots_adjust(top=0.9)
+
         if show:
             plt.show()
 

cogsworth/tests/test_pop.py

@@ -368,6 +368,8 @@ def test_indexing(self):
         p._final_pos = np.zeros(len(p._orbits))
         p._final_vel = np.zeros(len(p._orbits))
 
+        p.copy()
+
         for ind in inds:
             p_ind = p[ind]
             if isinstance(ind, slice):