Added rotation to scatter movie, and also now loop over each plot sty…

…le to make all 3 plots in one go (a vs e, a vs. i, and a vs. rot period)

examples/Chambers2013/scattermovie.py

@@ -36,36 +36,41 @@
 plt.switch_backend('agg')
 
 titletext = "Chambers (2013)"
-valid_plot_styles = ["aescatter", "aiscatter"]
-xlim={"aescatter" : (0.0, 2.5),
-      "aiscatter" : (0.0, 2.5)}
+valid_plot_styles = ["aescatter", "aiscatter", "arotscatter"]
+xlim={"aescatter" : (0.0, 3.0),
+      "aiscatter" : (0.0, 3.0), 
+      "arotscatter" : (0.0, 2.5)}
 ylim={"aescatter" : (0.0, 1.0),
-      "aiscatter" : (0.0, 40.0)}
+      "aiscatter" : (0.0, 40.0),
+      "arotscatter" : (1.0, 10000.0)}
 xlabel={"aescatter": "Semimajor axis (AU)",
-        "aiscatter": "Semimajor axis (AU)"}
+        "aiscatter": "Semimajor axis (AU)",
+        "arotscatter": "Semimajor axis (AU)"}
 ylabel={"aescatter": "Eccentricity",
-        "aiscatter": "Inclination (deg)"}
+        "aiscatter": "Inclination (deg)",
+        "arotscatter": "Rotation period (h)"}
 
+YR2HR = 365.25 * 24
+ROT2PERIOD = YR2HR * 360.0
 
-plot_style = valid_plot_styles[0]
 framejump = 1
-animation_file = f"Chambers2013-{plot_style}.mp4"
-
 origin_types = ["Initial conditions", "Merger", "Disruption", "Supercatastrophic", "Hit and run fragmentation"]
 
+
 class AnimatedScatter(object):
     """An animated scatter plot using matplotlib.animations.FuncAnimation."""
     def __init__(self, ds, param):
 
         self.radscale = 2000
         nframes = int(ds['time'].size / framejump)
+        ds['rot_mag'] = swiftest.tool.magnitude(ds,"rot")
+        ds['rot_mag'] = ROT2PERIOD / ds['rot_mag']
+        self.Rcb = ds['radius'].sel(name="Sun").isel(time=0).values[()]
         self.ds = ds
         self.param = param
-        self.Rcb = self.ds['radius'].sel(name="Sun").isel(time=0).values[()]
         colors = ["k", "xkcd:faded blue", "xkcd:marigold", "xkcd:shocking pink", "xkcd:baby poop green"]
         self.clist = dict(zip(origin_types,colors))
 
-
         # Setup the figure and axes...
         fig = plt.figure(figsize=(8,4.5), dpi=300)
         plt.tight_layout(pad=0)
@@ -104,7 +109,10 @@ def init_plot(self):
 
         leg = plt.legend(loc="upper left", scatterpoints=1, fontsize=10)
         for i,l in enumerate(leg.legendHandles):
-           leg.legendHandles[i]._sizes = [20]      
+            leg.legendHandles[i]._sizes = [20]
+
+        if plot_style == "arotscatter":
+            self.ax.set_yscale('log')
 
         return self.artists
 
@@ -123,6 +131,8 @@ def get_data(self, frame=0):
             pl = np.c_[d['a'].values,d['e'].values]
         elif plot_style == "aiscatter":
             pl = np.c_[d['a'].values,d['inc'].values]
+        elif plot_style == "arotscatter":
+            pl = np.c_[d['a'].values,d['rot_mag'].values]
 
         return t, npl, pl, radmarker, origin
 
@@ -143,6 +153,8 @@ def update(self,frame):
         return self.artists
 
 sim = swiftest.Simulation(read_data=True)
-print('Making animation')
-anim = AnimatedScatter(sim.data,sim.param)
-print('Animation finished')
+for plot_style in valid_plot_styles:
+    animation_file = f"Chambers2013-{plot_style}.mp4"
+    print('Making animation')
+    anim = AnimatedScatter(sim.data,sim.param)
+    print('Animation finished')