Improved the movie maker for Chambers do to a vs e and a vs. i (and p…

…otentially others)

examples/Chambers2013/.gitignore

@@ -1,4 +1,4 @@
 *
 !.gitignore
 !init_cond.py
-!aescattermovie.py
+!scattermovie.py

examples/Chambers2013/aescattermovie.py → examples/Chambers2013/scattermovie.py

@@ -6,21 +6,26 @@
 import matplotlib.colors as mcolors
 
 titletext = "Chambers (2013)"
-animation_file = "Chambers2013-aescatter.mp4"
-radscale = 2000
-AU = 1.0
-xmin = 0.0
-xmax = 2.25
-ymin = 0.0
-ymax = 1.0
+valid_plot_styles = ["aescatter", "aiscatter"]
+xlim={"aescatter" : (0.0, 2.25),
+      "aiscatter" : (0.0, 2.25)}
+ylim={"aescatter" : (0.0, 1.0),
+      "aiscatter" : (0.0, 40.0)}
+xlabel={"aescatter": "Semimajor axis (AU)",
+        "aiscatter": "Semimajor axis (AU)"}
+ylabel={"aescatter": "Eccentricity",
+        "aiscatter": "Inclination (deg)"}
+
+
+plot_style = valid_plot_styles[1]
 framejump = 1
-ncutoff = 1e20
+animation_file = f"Chambers2013-{plot_style}.mp4"
 
 class AnimatedScatter(object):
     """An animated scatter plot using matplotlib.animations.FuncAnimation."""
     def __init__(self, ds, param):
 
-        frame = 0
+        self.radscale = 2000
         nframes = int(ds['time'].size / framejump)
         self.ds = ds
         self.param = param
@@ -37,38 +42,38 @@ def __init__(self, ds, param):
         plt.tight_layout(pad=0)
         # set up the figure
         self.ax = plt.Axes(fig, [0.1, 0.15, 0.8, 0.75])
-        self.ax.set_xlim(xmin, xmax)
-        self.ax.set_ylim(ymin, ymax)
+        self.ax.set_xlim(xlim[plot_style])
+        self.ax.set_ylim(ylim[plot_style])
         fig.add_axes(self.ax)
         self.ani = animation.FuncAnimation(fig, self.update, interval=1, frames=nframes, init_func=self.setup_plot, blit=True)
-        self.ani.save(animation_file, fps=30, dpi=300, extra_args=['-vcodec', 'libx264'])
+        self.ani.save(animation_file, fps=60, dpi=300, extra_args=['-vcodec', 'libx264'])
         print(f'Finished writing {animation_file}')
 
     def scatters(self, pl, radmarker, origin):
         scat = []
         for key, value in self.clist.items():
             idx = origin == key
-            s = self.ax.scatter(pl[idx, 0], pl[idx, 1], marker='o', s=radmarker[idx], c=value, alpha=0.75, label=key)
+            s = self.ax.scatter(pl[idx, 0], pl[idx, 1], marker='o', s=radmarker[idx], c=value, alpha=0.75, label=key, animated=True)
             scat.append(s)
         return scat
 
     def setup_plot(self):
         # First frame
         """Initial drawing of the scatter plot."""
-        t, name, Gmass, radius, npl, pl, radmarker, origin = next(self.data_stream(0))
+        t, npl, pl, radmarker, origin = next(self.data_stream(0))
 
         # set up the figure
         self.ax.margins(x=10, y=1)
-        self.ax.set_xlabel("Semimajor Axis (AU)", fontsize='16', labelpad=1)
-        self.ax.set_ylabel("Eccentricity", fontsize='16', labelpad=1)
+        self.ax.set_xlabel(xlabel[plot_style], fontsize='16', labelpad=1)
+        self.ax.set_ylabel(ylabel[plot_style], fontsize='16', labelpad=1)
 
         title = self.ax.text(0.50, 1.05, "", bbox={'facecolor': 'w', 'alpha': 0.5, 'pad': 5}, transform=self.ax.transAxes,
-                        ha="center")
+                        ha="center", animated=True)
 
         title.set_text(f"{titletext} -  Time = ${t*1e-6:6.2f}$ My with ${npl:4.0f}$ particles")
         self.slist = self.scatters(pl, radmarker, origin)
         self.slist.append(title)
-        leg = plt.legend(loc="upper right", scatterpoints=1, fontsize=10)
+        leg = plt.legend(loc="upper left", scatterpoints=1, fontsize=10)
         for i,l in enumerate(leg.legendHandles):
            leg.legendHandles[i]._sizes = [20]
         return self.slist
@@ -78,25 +83,23 @@ def data_stream(self, frame=0):
             d = self.ds.isel(time = frame)
             name_good = d['name'].where(d['name'] != "Sun", drop=True)
             d = d.sel(name=name_good)
-            d['radmarker'] = (d['radius'] / self.Rcb) * radscale
-            radius = d['radmarker'].values
-
-            radius = d['radmarker'].values
-            Gmass = d['Gmass'].values
-            a = d['a'].values 
-            e = d['e'].values
-            name = d['name'].values
-            npl = d['npl'].values
-            radmarker = d['radmarker']
-            origin = d['origin_type']
+            d['radmarker'] = (d['radius'] / self.Rcb) * self.radscale
 
-            t = self.ds.coords['time'].values[frame]
+            t = d['time'].values
+            npl = d['npl'].values
+            radmarker = d['radmarker'].values
+            origin = d['origin_type'].values
+
+            if plot_style == "aescatter":
+               pl = np.c_[d['a'].values,d['e'].values]
+            elif plot_style == "aiscatter":
+               pl = np.c_[d['a'].values,d['inc'].values]
 
-            yield t, name, Gmass, radius, npl, np.c_[a, e], radmarker, origin
+            yield t, npl, pl, radmarker, origin
 
     def update(self,frame):
         """Update the scatter plot."""
-        t, name, Gmass, radius, npl, pl, radmarker, origin = next(self.data_stream(framejump * frame))
+        t,  npl, pl, radmarker, origin = next(self.data_stream(framejump * frame))
 
         self.slist[-1].set_text(f"{titletext} - Time = ${t*1e-6:6.3f}$ My with ${npl:4.0f}$ particles")