diff --git a/examples/Fragmentation/Fragmentation_Movie.py b/examples/Fragmentation/Fragmentation_Movie.py index ab666fb00..c260c8116 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)) @@ -48,8 +44,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([1.0, 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]), @@ -57,7 +53,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]), @@ -79,73 +75,127 @@ 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 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 - -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. - - # 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.scatter(xhx, xhy, s=(5000000000 * Gmass)) - - plt.tight_layout() - -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=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) - - # 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) - ani.save(movie_filename, fps=60, dpi=300, extra_args=['-vcodec', 'libx264']) \ No newline at end of file + +# ---------------------------------------------------------------------------------------------------------------------- +# 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.""" + + 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("x") + ax.set_ylabel("y") + 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(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 + + 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, + + 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 + Gmass = ds['Gmass'].values + x = ds['xhx'].values + y = ds['xhy'].values + point_rad = 2 * radius * self.ax_pt_size + 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) 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: