Fixed collision movie maker to work with the new NetCDF files

examples/symba_energy_momentum/collision_movie.py

@@ -10,15 +10,21 @@
 xmax = 20.0
 ymin = -20.0
 ymax = 20.0
+framejump = 1
+ncutoff = 1e20
+radscale = 2000
 
 cases = ['supercat_head', 'supercat_off', 'disruption_head', 'disruption_off']
+#cases = ['supercat_head']
 
-def scale_sim(ds):
-
+def scale_sim(ds, Rcb):
+    dst0 = ds.isel(time=0)
+    dst0 = dst0.where(dst0['radius'] < Rcb, drop=True)
+    dst0 = dst0.where(dst0.id > 0, drop=True)
+    GMtot = dst0['Gmass'].sum(skipna=True, dim="id")
     dsscale = ds.where(ds.id > 0, drop=True) # Remove the central body
-
-    GMtot = dsscale['Gmass'].sum(skipna=True, dim="id").isel(time=0)
-    rscale = ds['radius'].sel(id=1, time=0)
+
+    rscale = ds['radius'].sel(id=1, time=0).values
     dsscale['radius'] /= rscale
 
     dsscale['radmarker'] = dsscale['radius'].fillna(0)
@@ -64,12 +70,12 @@ class AnimatedScatter(object):
     """An animated scatter plot using matplotlib.animations.FuncAnimation."""
 
     def __init__(self, ds, param):
-
         frame = 0
-        nframes = ds['time'].size
-        self.ds = scale_sim(ds)
+        nframes = int(ds['time'].size / framejump)
         self.param = param
+        self.Rcb = ds['radius'].sel(id=0, time=0).values
         self.rot_angle = {}
+        self.ds = scale_sim(ds, self.Rcb)
 
         self.clist = {'Initial conditions' : 'xkcd:windows blue',
                       'Disruption' : 'xkcd:baby poop',
@@ -94,8 +100,9 @@ def __init__(self, ds, param):
         # Then setup FuncAnimation.
         self.ani = animation.FuncAnimation(fig, self.update, interval=1, frames=nframes,
                                           init_func=self.setup_plot, blit=False)
-        self.ani.save(animfile, fps=60, dpi=300,
-                      extra_args=['-vcodec', 'libx264'])
+        self.ani.save(animfile, fps=60, dpi=300, extra_args=['-vcodec', 'mpeg4'])
+        #self.ani.save(animfile, fps=60, dpi=300, extra_args=['-vcodec', 'libx264'])
+        print(f"Finished writing {animfile}")
 
     def plot_pl_circles(self, pl, radmarker):
         patches = []
@@ -239,14 +246,15 @@ def update(self,frame):
     def data_stream(self, frame=0):
         while True:
             d = self.ds.isel(time=frame)
-            radius = d['radmarker'].values
+            #d = d.where(d['radius'] > 0.0, drop=True)
+            radius = d['radius'].values
             Gmass = d['Gmass'].values
             x = d['xhxb'].values
             y = d['pyb'].values
             vhx = d['vhxb'].values
             vhy = d['vhyb'].values
             name = d['id'].values
-            npl = d.id.count().values
+            npl = d['npl'].values[0]
             id = d['id'].values
             rotx = d['rotx'].values
             roty = d['roty'].values
@@ -256,7 +264,7 @@ def data_stream(self, frame=0):
             origin = d['origin_type'].values
 
             t = self.ds.coords['time'].values[frame]
-            self.mask = np.logical_not(np.isnan(x))
+            self.mask = np.logical_not(radius > self.Rcb)
 
             x = np.nan_to_num(x, copy=False)
             y = np.nan_to_num(y, copy=False)