From 0d41a93a3ed191b57a136ee050ae69ac66ebb216 Mon Sep 17 00:00:00 2001 From: David A Minton Date: Tue, 29 Nov 2022 17:16:39 -0500 Subject: [PATCH 1/6] Tweaked fragmentation movie script to make the movies a bit more fast-paced --- examples/Fragmentation/Fragmentation_Movie.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/examples/Fragmentation/Fragmentation_Movie.py b/examples/Fragmentation/Fragmentation_Movie.py index ab666fb00..08f872d0c 100644 --- a/examples/Fragmentation/Fragmentation_Movie.py +++ b/examples/Fragmentation/Fragmentation_Movie.py @@ -135,7 +135,7 @@ def animate(i,ds,movie_title): minimum_fragment_gmass = 0.2 * body_Gmass[style][1] # Make the minimum fragment mass a fraction of the smallest body gmtiny = 0.99 * body_Gmass[style][1] # Make GMTINY just smaller than the smallest original body. This will prevent runaway collisional cascades sim.set_parameter(fragmentation = True, gmtiny=gmtiny, minimum_fragment_gmass=minimum_fragment_gmass) - sim.run(dt=1e-8, tstop=1.e-5) + sim.run(dt=1e-8, tstop=2.e-5) # Calculate the number of frames in the dataset. nframes = int(sim.data['time'].size) @@ -147,5 +147,6 @@ def animate(i,ds,movie_title): # Set up the figure and the animation. fig, ax = plt.subplots(figsize=(4,4)) # Generate the movie. - ani = animation.FuncAnimation(fig, animate, fargs=(sim.data, movie_titles[style]), interval=1, frames=nframes, repeat=False) + nskip = 10 + ani = animation.FuncAnimation(fig, animate, fargs=(sim.data, movie_titles[style]), interval=1, frames=range(0,nframes,nskip), repeat=False) ani.save(movie_filename, fps=60, dpi=300, extra_args=['-vcodec', 'libx264']) \ No newline at end of file From 626faf1f3944655ea8f22a0224e85779f67586f9 Mon Sep 17 00:00:00 2001 From: David A Minton Date: Tue, 29 Nov 2022 18:09:50 -0500 Subject: [PATCH 2/6] Fixed movie maker so that circles are the actual radius of the bodies --- examples/Fragmentation/Fragmentation_Movie.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/examples/Fragmentation/Fragmentation_Movie.py b/examples/Fragmentation/Fragmentation_Movie.py index 08f872d0c..d290267b7 100644 --- a/examples/Fragmentation/Fragmentation_Movie.py +++ b/examples/Fragmentation/Fragmentation_Movie.py @@ -49,7 +49,7 @@ "supercatastrophic_off_axis": [np.array([1.0, -4.2e-05, 0.0]), np.array([1.0, 4.2e-05, 0.0])], "hitandrun" : [np.array([1.0, -4.2e-05, 0.0]), - np.array([1.0, 4.2e-05, 0.0])] + np.array([0.9999999, 4.2e-05, 0.0])] } vel_vectors = {"disruption_headon" : [np.array([-2.562596e-04, 6.280005, 0.0]), @@ -57,7 +57,7 @@ "supercatastrophic_off_axis": [np.array([0.0, 6.28, 0.0]), np.array([1.0, -6.28, 0.0])], "hitandrun" : [np.array([0.0, 6.28, 0.0]), - np.array([-1.5, -6.28, 0.0])] + np.array([-0.1, -6.28, 0.0])] } rot_vectors = {"disruption_headon" : [np.array([0.0, 0.0, 0.0]), @@ -92,6 +92,8 @@ def center(xhx, xhy, xhz, Gmass): z_com = np.sum(Gmass * xhz) / np.sum(Gmass) return x_com, y_com, z_com +figsize = (4,4) + def animate(i,ds,movie_title): # Calculate the position and mass of all bodies in the system at time i and store as a numpy array. @@ -99,6 +101,7 @@ def animate(i,ds,movie_title): xhy = ds['xhy'].isel(time=i).dropna(dim='name').values xhz = ds['xhx'].isel(time=i).dropna(dim='name').values Gmass = ds['Gmass'].isel(time=i).dropna(dim='name').values[1:] # Drop the Sun from the numpy array. + radius = ds['radius'].isel(time=i).dropna(dim='name').values[1:] # Drop the Sun from the numpy array. # Calculate the center of mass of the system at time i. While the center of mass relative to the # colliding bodies does not change, the center of mass of the collision will move as the bodies @@ -117,7 +120,9 @@ def animate(i,ds,movie_title): ax.set_xticks([]) ax.set_yticks([]) - ax.scatter(xhx, xhy, s=(5000000000 * Gmass)) + ax_pt_size = figsize[0] * 72 / (2 * scale_frame) + point_rad = 2 * radius * ax_pt_size + ax.scatter(xhx, xhy, s=point_rad**2) plt.tight_layout() @@ -135,7 +140,7 @@ def animate(i,ds,movie_title): minimum_fragment_gmass = 0.2 * body_Gmass[style][1] # Make the minimum fragment mass a fraction of the smallest body gmtiny = 0.99 * body_Gmass[style][1] # Make GMTINY just smaller than the smallest original body. This will prevent runaway collisional cascades sim.set_parameter(fragmentation = True, gmtiny=gmtiny, minimum_fragment_gmass=minimum_fragment_gmass) - sim.run(dt=1e-8, tstop=2.e-5) + sim.run(dt=1e-8, tstop=1.e-5) # Calculate the number of frames in the dataset. nframes = int(sim.data['time'].size) @@ -145,8 +150,8 @@ def animate(i,ds,movie_title): scale_frame = abs(sim.data['xhy'].isel(time=0, name=1).values) + abs(sim.data['xhy'].isel(time=0, name=2).values) # Set up the figure and the animation. - fig, ax = plt.subplots(figsize=(4,4)) + fig, ax = plt.subplots(figsize=figsize) # Generate the movie. - nskip = 10 + nskip = 1 ani = animation.FuncAnimation(fig, animate, fargs=(sim.data, movie_titles[style]), interval=1, frames=range(0,nframes,nskip), repeat=False) ani.save(movie_filename, fps=60, dpi=300, extra_args=['-vcodec', 'libx264']) \ No newline at end of file From 9ab21056ea30a6f47e6e33060b299efa3ff06ffd Mon Sep 17 00:00:00 2001 From: David A Minton Date: Tue, 29 Nov 2022 20:37:50 -0500 Subject: [PATCH 3/6] Re-wrote the animation in proper OOP style. Figure is no longer re-drawn each frame --- examples/Fragmentation/Fragmentation_Movie.py | 130 ++++++++++-------- 1 file changed, 76 insertions(+), 54 deletions(-) diff --git a/examples/Fragmentation/Fragmentation_Movie.py b/examples/Fragmentation/Fragmentation_Movie.py index d290267b7..c1f1eebd3 100644 --- a/examples/Fragmentation/Fragmentation_Movie.py +++ b/examples/Fragmentation/Fragmentation_Movie.py @@ -48,8 +48,8 @@ np.array([1.0, 1.807993e-05 ,0.0])], "supercatastrophic_off_axis": [np.array([1.0, -4.2e-05, 0.0]), np.array([1.0, 4.2e-05, 0.0])], - "hitandrun" : [np.array([1.0, -4.2e-05, 0.0]), - np.array([0.9999999, 4.2e-05, 0.0])] + "hitandrun" : [np.array([1.0, -2.0e-05, 0.0]), + np.array([0.999999, 2.0e-05, 0.0])] } vel_vectors = {"disruption_headon" : [np.array([-2.562596e-04, 6.280005, 0.0]), @@ -85,46 +85,80 @@ print("Generating all movie styles") movie_styles = available_movie_styles.copy() -# Define a function to calculate the center of mass of the system. -def center(xhx, xhy, xhz, Gmass): - x_com = np.sum(Gmass * xhx) / np.sum(Gmass) - y_com = np.sum(Gmass * xhy) / np.sum(Gmass) - z_com = np.sum(Gmass * xhz) / np.sum(Gmass) - return x_com, y_com, z_com - figsize = (4,4) +class AnimatedScatter(object): + """An animated scatter plot using matplotlib.animations.FuncAnimation.""" + + def __init__(self, sim, animfile, title, nskip=1): + nframes = int(sim.data['time'].size) + self.sim = sim + self.title = title + self.body_color_list = {'Initial conditions': 'xkcd:windows blue', + 'Disruption': 'xkcd:baby poop', + 'Supercatastrophic': 'xkcd:shocking pink', + 'Hit and run fragment': 'xkcd:blue with a hint of purple', + 'Central body': 'xkcd:almost black'} + + # Set up the figure and axes... + self.fig, self.ax = self.setup_plot() + + # Then setup FuncAnimation. + self.ani = animation.FuncAnimation(self.fig, self.update_plot, interval=1, frames=range(0,nframes,nskip), + blit=True) + self.ani.save(animfile, fps=60, dpi=300, extra_args=['-vcodec', 'libx264']) + print(f"Finished writing {animfile}") + + def setup_plot(self): + fig = plt.figure(figsize=figsize, dpi=300) + plt.tight_layout(pad=0) + + + # Calculate the distance along the y-axis between the colliding bodies at the start of the simulation. + # This will be used to scale the axis limits on the movie. + scale_frame = abs(sim.data['xhy'].isel(time=0, name=1).values) \ + + abs( sim.data['xhy'].isel(time=0, name=2).values) + ax = plt.Axes(fig, [0.1, 0.1, 0.8, 0.8]) + self.ax_pt_size = figsize[0] * 0.8 * 72 / (2 * scale_frame) + ax.set_xlim(-scale_frame, scale_frame) + ax.set_ylim(-scale_frame, scale_frame) + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_xlabel("xhx") + ax.set_ylabel("xhy") + ax.set_title(self.title) + fig.add_axes(ax) + + self.scatter_artist = ax.scatter([], [], animated=True) + return fig, ax + + def update_plot(self, frame): + # Define a function to calculate the center of mass of the system. + def center(frame): + ds = self.sim.data.isel(time=frame).where(self.sim.data['name'] != "Sun", drop=True) + Gmass = ds['Gmass'].values + xhx = ds['xhx'].values + xhy = ds['xhy'].values + x_com = np.sum(Gmass * xhx) / np.sum(Gmass) + y_com = np.sum(Gmass * xhy) / np.sum(Gmass) + return x_com, y_com + + x, y, point_rad = next(self.data_stream(frame)) + x_com, y_com = center(frame) + self.scatter_artist.set_offsets(np.c_[x - x_com, y - y_com]) + self.scatter_artist.set_sizes(point_rad) + return self.scatter_artist, + + def data_stream(self, frame=0): + while True: + ds = self.sim.data.isel(time=frame) + ds = ds.where(ds['name'] != "Sun", drop=True) + radius = ds['radius'].values + x = ds['xhx'].values + y = ds['xhy'].values + point_rad = 2 * radius * self.ax_pt_size + yield x, y, point_rad + -def animate(i,ds,movie_title): - - # Calculate the position and mass of all bodies in the system at time i and store as a numpy array. - xhx = ds['xhx'].isel(time=i).dropna(dim='name').values - xhy = ds['xhy'].isel(time=i).dropna(dim='name').values - xhz = ds['xhx'].isel(time=i).dropna(dim='name').values - Gmass = ds['Gmass'].isel(time=i).dropna(dim='name').values[1:] # Drop the Sun from the numpy array. - radius = ds['radius'].isel(time=i).dropna(dim='name').values[1:] # Drop the Sun from the numpy array. - - # Calculate the center of mass of the system at time i. While the center of mass relative to the - # colliding bodies does not change, the center of mass of the collision will move as the bodies - # orbit the system center of mass. - x_com, y_com, z_com = center(xhx, xhy, xhz, Gmass) - - # Create the figure and plot the bodies as points. - fig.clear() - ax = fig.add_subplot(111) - ax.set_title(movie_title) - ax.set_xlabel("xhx") - ax.set_ylabel("xhy") - ax.set_xlim(x_com - scale_frame, x_com + scale_frame) - ax.set_ylim(y_com - scale_frame, y_com + scale_frame) - ax.grid(False) - ax.set_xticks([]) - ax.set_yticks([]) - - ax_pt_size = figsize[0] * 72 / (2 * scale_frame) - point_rad = 2 * radius * ax_pt_size - ax.scatter(xhx, xhy, s=point_rad**2) - - plt.tight_layout() for style in movie_styles: param_file = Path(style) / "param.in" @@ -140,18 +174,6 @@ def animate(i,ds,movie_title): minimum_fragment_gmass = 0.2 * body_Gmass[style][1] # Make the minimum fragment mass a fraction of the smallest body gmtiny = 0.99 * body_Gmass[style][1] # Make GMTINY just smaller than the smallest original body. This will prevent runaway collisional cascades sim.set_parameter(fragmentation = True, gmtiny=gmtiny, minimum_fragment_gmass=minimum_fragment_gmass) - sim.run(dt=1e-8, tstop=1.e-5) - - # Calculate the number of frames in the dataset. - nframes = int(sim.data['time'].size) - - # Calculate the distance along the y-axis between the colliding bodies at the start of the simulation. - # This will be used to scale the axis limits on the movie. - scale_frame = abs(sim.data['xhy'].isel(time=0, name=1).values) + abs(sim.data['xhy'].isel(time=0, name=2).values) + sim.run(dt=1e-8, tstop=2.e-5) - # Set up the figure and the animation. - fig, ax = plt.subplots(figsize=figsize) - # Generate the movie. - nskip = 1 - ani = animation.FuncAnimation(fig, animate, fargs=(sim.data, movie_titles[style]), interval=1, frames=range(0,nframes,nskip), repeat=False) - ani.save(movie_filename, fps=60, dpi=300, extra_args=['-vcodec', 'libx264']) \ No newline at end of file + anim = AnimatedScatter(sim,movie_filename,movie_titles[style],nskip=10) From c0d25249df32bee28bcd5d80cbef416c09e3ae80 Mon Sep 17 00:00:00 2001 From: David A Minton Date: Wed, 30 Nov 2022 09:29:49 -0500 Subject: [PATCH 4/6] Fixed issue where units were being recomputed even if they hadn't changed --- python/swiftest/swiftest/simulation_class.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index aeedc8ac6..d5c6e21b8 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -1772,7 +1772,10 @@ def set_unit_system(self, if all(key in self.param for key in ["MU2KG","DU2M","TU2S"]): self.GU = constants.GC * self.param["TU2S"] ** 2 * self.param["MU2KG"] / self.param["DU2M"] ** 3 - if recompute_unit_values: + if recompute_unit_values and \ + MU2KG_old != self.param['MU2KG'] or \ + DU2M_old != self.param['DU2M'] or \ + TU2S_old != self.param['TU2S']: self.update_param_units(MU2KG_old, DU2M_old, TU2S_old) unit_dict = self.get_unit_system(update_list, verbose) @@ -1868,7 +1871,6 @@ def update_param_units(self, MU2KG_old, DU2M_old, TU2S_old): if MU2KG_old is not None: for k in mass_keys: if k in self.param: - print(f"param['{k}']: {self.param[k]}") self.param[k] *= MU2KG_old / self.param['MU2KG'] if DU2M_old is not None: From 2730cfcc29f7cfb8a1f8056b91c8465dfd75098f Mon Sep 17 00:00:00 2001 From: David A Minton Date: Wed, 30 Nov 2022 09:30:32 -0500 Subject: [PATCH 5/6] Restructured movie code a bit and added some more clarifying comments. Fixed bug that was causing bodies not to be drawn if there was a fragmentation event --- examples/Fragmentation/Fragmentation_Movie.py | 106 +++++++++++------- 1 file changed, 64 insertions(+), 42 deletions(-) diff --git a/examples/Fragmentation/Fragmentation_Movie.py b/examples/Fragmentation/Fragmentation_Movie.py index c1f1eebd3..8c2d2ca9d 100644 --- a/examples/Fragmentation/Fragmentation_Movie.py +++ b/examples/Fragmentation/Fragmentation_Movie.py @@ -32,13 +32,9 @@ import matplotlib.animation as animation from pathlib import Path -print("Select a fragmentation movie to generate.") -print("1. Head-on disruption") -print("2. Off-axis supercatastrophic") -print("3. Hit and run") -print("4. All of the above") -user_selection = int(input("? ")) - +# ---------------------------------------------------------------------------------------------------------------------- +# Define the names and initial conditions of the various fragmentation simulation types +# ---------------------------------------------------------------------------------------------------------------------- available_movie_styles = ["disruption_headon", "supercatastrophic_off_axis", "hitandrun"] movie_title_list = ["Head-on Disrutption", "Off-axis Supercatastrophic", "Hit and Run"] movie_titles = dict(zip(available_movie_styles, movie_title_list)) @@ -79,12 +75,10 @@ for k,v in body_Gmass.items(): body_radius[k] = [((Gmass/GU)/(4./3.*np.pi*density))**(1./3.) for Gmass in v] -if user_selection > 0 and user_selection < 4: - movie_styles = [available_movie_styles[user_selection-1]] -else: - print("Generating all movie styles") - movie_styles = available_movie_styles.copy() +# ---------------------------------------------------------------------------------------------------------------------- +# Define the animation class that will generate the movies of the fragmentation outcomes +# ---------------------------------------------------------------------------------------------------------------------- figsize = (4,4) class AnimatedScatter(object): """An animated scatter plot using matplotlib.animations.FuncAnimation.""" @@ -133,17 +127,16 @@ def setup_plot(self): def update_plot(self, frame): # Define a function to calculate the center of mass of the system. - def center(frame): - ds = self.sim.data.isel(time=frame).where(self.sim.data['name'] != "Sun", drop=True) - Gmass = ds['Gmass'].values - xhx = ds['xhx'].values - xhy = ds['xhy'].values - x_com = np.sum(Gmass * xhx) / np.sum(Gmass) - y_com = np.sum(Gmass * xhy) / np.sum(Gmass) + def center(Gmass, x, y): + x = x[~np.isnan(x)] + y = y[~np.isnan(y)] + Gmass = Gmass[~np.isnan(Gmass)] + x_com = np.sum(Gmass * x) / np.sum(Gmass) + y_com = np.sum(Gmass * y) / np.sum(Gmass) return x_com, y_com - x, y, point_rad = next(self.data_stream(frame)) - x_com, y_com = center(frame) + Gmass, x, y, point_rad = next(self.data_stream(frame)) + x_com, y_com = center(Gmass, x, y) self.scatter_artist.set_offsets(np.c_[x - x_com, y - y_com]) self.scatter_artist.set_sizes(point_rad) return self.scatter_artist, @@ -153,27 +146,56 @@ def data_stream(self, frame=0): ds = self.sim.data.isel(time=frame) ds = ds.where(ds['name'] != "Sun", drop=True) radius = ds['radius'].values + Gmass = ds['Gmass'].values x = ds['xhx'].values y = ds['xhy'].values point_rad = 2 * radius * self.ax_pt_size - yield x, y, point_rad - - - -for style in movie_styles: - param_file = Path(style) / "param.in" - - movie_filename = f"{style}.mp4" - - # Pull in the Swiftest output data from the parameter file and store it as a Xarray dataset. - sim = swiftest.Simulation(param_file=param_file, rotation=True, init_cond_format = "XV", compute_conservation_values=True) - sim.add_solar_system_body("Sun") - sim.add_body(Gmass=body_Gmass[style], radius=body_radius[style], xh=pos_vectors[style], vh=vel_vectors[style], rot=rot_vectors[style]) - - # Set fragmentation parameters - minimum_fragment_gmass = 0.2 * body_Gmass[style][1] # Make the minimum fragment mass a fraction of the smallest body - gmtiny = 0.99 * body_Gmass[style][1] # Make GMTINY just smaller than the smallest original body. This will prevent runaway collisional cascades - sim.set_parameter(fragmentation = True, gmtiny=gmtiny, minimum_fragment_gmass=minimum_fragment_gmass) - sim.run(dt=1e-8, tstop=2.e-5) - - anim = AnimatedScatter(sim,movie_filename,movie_titles[style],nskip=10) + yield Gmass, x, y, point_rad + +if __name__ == "__main__": + + print("Select a fragmentation movie to generate.") + print("1. Head-on disruption") + print("2. Off-axis supercatastrophic") + print("3. Hit and run") + print("4. All of the above") + user_selection = int(input("? ")) + + if user_selection > 0 and user_selection < 4: + movie_styles = [available_movie_styles[user_selection-1]] + else: + print("Generating all movie styles") + movie_styles = available_movie_styles.copy() + + for style in movie_styles: + param_file = Path(style) / "param.in" + bin_file = Path(style) / "bin.nc" + if bin_file.exists(): + user_selection = input(f"An older simulation of {movie_titles[style]} exists. Do you want to re-run it? [y/N] ") + if user_selection == "": + run_new = False + else: + try: + run_new = swiftest.io.str2bool(user_selection) + except: + run_new = False + else: + run_new = True + + movie_filename = f"{style}.mp4" + + # Pull in the Swiftest output data from the parameter file and store it as a Xarray dataset. + if run_new: + sim = swiftest.Simulation(param_file=param_file, rotation=True, init_cond_format = "XV", compute_conservation_values=True) + sim.add_solar_system_body("Sun") + sim.add_body(Gmass=body_Gmass[style], radius=body_radius[style], xh=pos_vectors[style], vh=vel_vectors[style], rot=rot_vectors[style]) + + # Set fragmentation parameters + minimum_fragment_gmass = 0.2 * body_Gmass[style][1] # Make the minimum fragment mass a fraction of the smallest body + gmtiny = 0.99 * body_Gmass[style][1] # Make GMTINY just smaller than the smallest original body. This will prevent runaway collisional cascades + sim.set_parameter(fragmentation = True, gmtiny=gmtiny, minimum_fragment_gmass=minimum_fragment_gmass, verbose=False) + sim.run(dt=1e-8, tstop=2.e-5) + else: + sim = swiftest.Simulation(param_file=param_file, read_old_output_file=True) + + anim = AnimatedScatter(sim,movie_filename,movie_titles[style],nskip=10) From 20065498b827e101f88e49941a3f818ec2ca2410 Mon Sep 17 00:00:00 2001 From: David A Minton Date: Wed, 30 Nov 2022 10:26:33 -0500 Subject: [PATCH 6/6] Updated axes labels to x and y --- examples/Fragmentation/Fragmentation_Movie.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/Fragmentation/Fragmentation_Movie.py b/examples/Fragmentation/Fragmentation_Movie.py index 8c2d2ca9d..c260c8116 100644 --- a/examples/Fragmentation/Fragmentation_Movie.py +++ b/examples/Fragmentation/Fragmentation_Movie.py @@ -117,8 +117,8 @@ def setup_plot(self): ax.set_ylim(-scale_frame, scale_frame) ax.set_xticks([]) ax.set_yticks([]) - ax.set_xlabel("xhx") - ax.set_ylabel("xhy") + ax.set_xlabel("x") + ax.set_ylabel("y") ax.set_title(self.title) fig.add_axes(ax)