Skip to content
This repository was archived by the owner on Aug 28, 2024. It is now read-only.

Commit

Permalink
Merge branch 'debug'
Browse files Browse the repository at this point in the history
  • Loading branch information
@daminton
daminton committed Nov 30, 2022
2 parents 0b42bed + 2006549 commit 555cb8f
Show file tree
Hide file tree
Showing 2 changed files with 134 additions and 82 deletions.
210 changes: 130 additions & 80 deletions examples/Fragmentation/Fragmentation_Movie.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand All @@ -48,16 +44,16 @@
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]),
np.array([-2.562596e-04, -6.280005, 0.0])],
"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]),
Expand All @@ -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'])

# ----------------------------------------------------------------------------------------------------------------------
# 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)
6 changes: 4 additions & 2 deletions python/swiftest/swiftest/simulation_class.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down

0 comments on commit 555cb8f

Please sign in to comment.