From 6e728a37e71007df4b9a75ea87cc60b3a4e18bcc Mon Sep 17 00:00:00 2001 From: David Minton Date: Fri, 3 Sep 2021 16:40:57 -0400 Subject: [PATCH] Fixed collision movie maker to work with the new NetCDF files --- .../symba_energy_momentum/collision_movie.py | 34 ++++++++++++------- 1 file changed, 21 insertions(+), 13 deletions(-) diff --git a/examples/symba_energy_momentum/collision_movie.py b/examples/symba_energy_momentum/collision_movie.py index 56ed66b9b..1a17034c5 100755 --- a/examples/symba_energy_momentum/collision_movie.py +++ b/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)