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)