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

Commit

Permalink
Merge branch 'fraggle_simplification' into debug
Browse files Browse the repository at this point in the history
  • Loading branch information
@daminton
daminton committed Dec 31, 2022
2 parents d395773 + 28e5cde commit 1bbc416
Show file tree
Hide file tree
Showing 115 changed files with 18,599 additions and 18,814 deletions.
74 changes: 52 additions & 22 deletions examples/Fragmentation/Fragmentation_Movie.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,39 +36,61 @@
# ----------------------------------------------------------------------------------------------------------------------
# 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 Disruption", "Off-axis Supercatastrophic", "Hit and Run"]
available_movie_styles = ["disruption_headon", "disruption_off_axis", "supercatastrophic_headon", "supercatastrophic_off_axis","hitandrun_disrupt", "hitandrun_pure"]
movie_title_list = ["Head-on Disruption", "Off-axis Disruption", "Head-on Supercatastrophic", "Off-axis Supercatastrophic", "Hit and Run w/ Runner Disruption", "Pure Hit and Run"]
movie_titles = dict(zip(available_movie_styles, movie_title_list))
num_movie_frames = 1200

# These initial conditions were generated by trial and error
names = ["Target","Projectile"]
pos_vectors = {"disruption_headon" : [np.array([1.0, -5.0e-05, 0.0]),
np.array([1.0, 5.0e-05 ,0.0])],
"disruption_off_axis" : [np.array([1.0, -5.0e-05, 0.0]),
np.array([1.0, 5.0e-05 ,0.0])],
"supercatastrophic_headon": [np.array([1.0, -5.0e-05, 0.0]),
np.array([1.0, 5.0e-05, 0.0])],
"supercatastrophic_off_axis": [np.array([1.0, -5.0e-05, 0.0]),
np.array([1.0, 5.0e-05, 0.0])],
"hitandrun" : [np.array([1.0, -4.2e-05, 0.0]),
np.array([1.0, 4.2e-05, 0.0])]
"hitandrun_disrupt" : [np.array([1.0, -4.2e-05, 0.0]),
np.array([1.0, 4.2e-05, 0.0])],
"hitandrun_pure" : [np.array([1.0, -4.2e-05, 0.0]),
np.array([1.0, 4.2e-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([0.5, -6.28, 0.0])],
"hitandrun" : [np.array([0.0, 6.28, 0.0]),
np.array([-1.45, -6.28, 0.00])]
vel_vectors = {"disruption_headon" : [np.array([ 0.00, 6.280005, 0.0]),
np.array([ 0.00, -6.280005, 0.0])],
"disruption_off_axis" : [np.array([ 0.00, 6.280005, 0.0]),
np.array([ 0.50, -6.280005, 0.0])],
"supercatastrophic_headon": [np.array([ 0.00, 6.28, 0.0]),
np.array([ 0.00, -6.28, 0.0])],
"supercatastrophic_off_axis": [np.array([ 0.00, 6.28, 0.0]),
np.array([ 0.50, -6.28, 0.0])],
"hitandrun_disrupt" : [np.array([ 0.00, 6.28, 0.0]),
np.array([-1.45, -6.28, 0.0])],
"hitandrun_pure" : [np.array([ 0.00, 6.28, 0.0]),
np.array([-1.51, -6.28, 0.0])]
}

rot_vectors = {"disruption_headon" : [np.array([0.0, 0.0, 0.0]),
np.array([0.0, 0.0, 0.0])],
"disruption_off_axis": [np.array([0.0, 0.0, -6.0e4]),
np.array([0.0, 0.0, 1.0e5])],
"supercatastrophic_headon": [np.array([0.0, 0.0, 0.0]),
np.array([0.0, 0.0, 0.0])],
"supercatastrophic_off_axis": [np.array([0.0, 0.0, -6.0e4]),
np.array([0.0, 0.0, 1.0e5])],
"hitandrun" : [np.array([0.0, 0.0, 6.0e4]),
"hitandrun_disrupt" : [np.array([0.0, 0.0, 6.0e4]),
np.array([0.0, 0.0, 1.0e5])],
"hitandrun_pure" : [np.array([0.0, 0.0, 6.0e4]),
np.array([0.0, 0.0, 1.0e5])]
}

body_Gmass = {"disruption_headon" : [1e-7, 1e-10],
"disruption_off_axis" : [1e-7, 1e-10],
"supercatastrophic_headon" : [1e-7, 1e-8],
"supercatastrophic_off_axis": [1e-7, 1e-8],
"hitandrun" : [1e-7, 7e-10]
"hitandrun_disrupt" : [1e-7, 7e-10],
"hitandrun_pure" : [1e-7, 7e-10]
}

density = 3000 * swiftest.AU2M**3 / swiftest.MSun
Expand All @@ -77,7 +99,8 @@
for k,v in body_Gmass.items():
body_radius[k] = [((Gmass/GU)/(4./3.*np.pi*density))**(1./3.) for Gmass in v]

body_radius["hitandrun"] = [7e-6, 3.25e-6]
body_radius["hitandrun_disrupt"] = [7e-6, 3.25e-6]
body_radius["hitandrun_pure"] = [7e-6, 3.25e-6]

# ----------------------------------------------------------------------------------------------------------------------
# Define the animation class that will generate the movies of the fragmentation outcomes
Expand All @@ -103,6 +126,10 @@ def encounter_combiner(sim):

# The following will combine the two datasets along the time dimension, sort the time dimension, and then fill in any time gaps with interpolation
ds = xr.combine_nested([data,enc],concat_dim='time').sortby("time").interpolate_na(dim="time")

# Interpolate in time to make a smooth, constant time step dataset
smooth_time = np.linspace(start=tgood.isel(time=0), stop=ds.time[-1], num=num_movie_frames)
ds = ds.interp(time=smooth_time)

return ds

Expand All @@ -126,8 +153,7 @@ def __init__(self, sim, animfile, title, style, nskip=1):
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 = 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}")

Expand Down Expand Up @@ -185,18 +211,22 @@ def data_stream(self, frame=0):

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")
print("2. Off-axis disruption")
print("3. Head-on supercatastrophic")
print("4. Off-axis supercatastrophic")
print("5. Hit and run with disruption of the runner")
print("6. Pure hit and run")
print("7. All of the above")
user_selection = int(input("? "))

if user_selection > 0 and user_selection < 4:
if user_selection > 0 and user_selection < 7:
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:
print(f"Generating {movie_titles[style]}")
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(simdir=style, rotation=True, init_cond_format = "XV", compute_conservation_values=True)
Expand All @@ -206,8 +236,8 @@ def data_stream(self, frame=0):
# 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, encounter_save="both", gmtiny=gmtiny, minimum_fragment_gmass=minimum_fragment_gmass, verbose=False)
sim.run(dt=1e-3, tstop=1.0e-3, istep_out=1, dump_cadence=0)
sim.set_parameter(collision_model="fraggle", encounter_save="both", gmtiny=gmtiny, minimum_fragment_gmass=minimum_fragment_gmass, verbose=False)
sim.run(dt=5e-4, tstop=5.0e-4, istep_out=1, dump_cadence=0)

print("Generating animation")
anim = AnimatedScatter(sim,movie_filename,movie_titles[style],style,nskip=1)
anim = AnimatedScatter(sim,movie_filename,movie_titles[style],style,nskip=1)
49 changes: 31 additions & 18 deletions python/swiftest/swiftest/simulation_class.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -224,16 +224,18 @@ def __init__(self,read_param: bool = False, read_old_output: bool = False, simdi
general_relativity : bool, default True
Include the post-Newtonian correction in acceleration calculations.
Parameter input file equivalent: `GR`
fragmentation : bool, default True
If set to True, this turns on the Fraggle fragment generation code and `rotation` must also be True.
collision_model: {"MERGE","BOUNCE","FRAGGLE"}, default "MERGE"
This is used to set the collision/fragmentation model. [TODO: DESCRIBE THESE]
This argument only applies to Swiftest-SyMBA simulations. It will be ignored otherwise.
Parameter input file equivalent: `FRAGMENTATION`
Parameter input file equivalent: `COLLISION_MODEL`
minimum_fragment_gmass : float, optional
If fragmentation is turned on, this sets the mimimum G*mass of a collisional fragment that can be generated.
If fragmentation is turned on, this sets the mimimum G*mass of a collisional fragment that can be generated if a
fragmentation model is enabled. Ignored otherwise.
*Note.* Only set one of minimum_fragment_gmass or minimum_fragment_mass
Parameter input file equivalent: None
minimum_fragment_mass : float, optional
If fragmentation is turned on, this sets the mimimum mass of a collisional fragment that can be generated.
If fragmentation is turned on, this sets the mimimum mass of a collisional fragment that can be generated. if a
fragmentation model is enabled. Ignored otherwise
*Note.* Only set one of minimum_fragment_gmass or minimum_fragment_mass
Parameter input file equivalent: `MIN_GMFRAG`
rotation : bool, default False
Expand Down Expand Up @@ -777,7 +779,7 @@ def set_parameter(self, verbose: bool = True, **kwargs):
"mtiny": None,
"close_encounter_check": True,
"general_relativity": True,
"fragmentation": False,
"collision_model": "MERGE",
"minimum_fragment_mass": None,
"minimum_fragment_gmass": 0.0,
"rotation": False,
Expand Down Expand Up @@ -1013,7 +1015,7 @@ def get_integrator(self,arg_list: str | List[str] | None = None, verbose: bool |
def set_feature(self,
close_encounter_check: bool | None = None,
general_relativity: bool | None = None,
fragmentation: bool | None = None,
collision_model: Literal["MERGE","BOUNCE","FRAGGLE"] | None = None,
minimum_fragment_gmass: float | None = None,
minimum_fragment_mass: float | None = None,
rotation: bool | None = None,
Expand Down Expand Up @@ -1046,15 +1048,20 @@ def set_feature(self,
*WARNING*: Enabling this feature could lead to very large files.
general_relativity : bool, optional
Include the post-Newtonian correction in acceleration calculations.
fragmentation : bool, optional
If set to True, this turns on the Fraggle fragment generation code and `rotation` must also be True.
collision_model: {"MERGE","BOUNCE","FRAGGLE"}, default "MERGE"
This is used to set the collision/fragmentation model. [TODO: DESCRIBE THESE]
This argument only applies to Swiftest-SyMBA simulations. It will be ignored otherwise.
Parameter input file equivalent: `COLLISION_MODEL`
minimum_fragment_gmass : float, optional
If fragmentation is turned on, this sets the mimimum G*mass of a collisional fragment that can be generated.
If fragmentation is turned on, this sets the mimimum G*mass of a collisional fragment that can be generated if a
fragmentation model is enabled. Ignored otherwise.
*Note.* Only set one of minimum_fragment_gmass or minimum_fragment_mass
Parameter input file equivalent: None
minimum_fragment_mass : float, optional
If fragmentation is turned on, this sets the mimimum mass of a collisional fragment that can be generated.
If fragmentation is turned on, this sets the mimimum mass of a collisional fragment that can be generated. if a
fragmentation model is enabled. Ignored otherwise
*Note.* Only set one of minimum_fragment_gmass or minimum_fragment_mass
Parameter input file equivalent: `MIN_GMFRAG`
rotation : bool, optional
If set to True, this turns on rotation tracking and radius, rotation vector, and moments of inertia values
must be included in the initial conditions.
Expand Down Expand Up @@ -1122,13 +1129,16 @@ def set_feature(self,
self.param["GR"] = general_relativity
update_list.append("general_relativity")

if fragmentation is not None:
fragmentation_models = ["FRAGGLE"]
if collision_model is not None:
collision_model = collision_model.upper()
fragmentation = collision_model in fragmentation_models
if self.codename != "Swiftest" and self.integrator != "symba" and fragmentation:
warnings.warn("Fragmentation is only available on Swiftest SyMBA.",stacklevel=2)
self.param['FRAGMENTATION'] = False
self.param['COLLISION_MODEL'] = "MERGE"
else:
self.param['FRAGMENTATION'] = fragmentation
update_list.append("fragmentation")
self.param['COLLISION_MODEL'] = collision_model
update_list.append("collision_model")
if fragmentation:
if "MIN_GMFRAG" not in self.param and minimum_fragment_mass is None and minimum_fragment_gmass is None:
warnings.warn("Minimum fragment mass is not set. Set it using minimum_fragment_gmass or minimum_fragment_mass",stacklevel=2)
Expand All @@ -1151,7 +1161,7 @@ def set_feature(self,
self.param['ROTATION'] = rotation
update_list.append("rotation")

if self.param['FRAGMENTATION'] and not self.param['ROTATION']:
if self.param['COLLISION_MODEL'] == "FRAGGLE" and not self.param['ROTATION']:
self.param['ROTATION'] = True
update_list.append("rotation")

Expand Down Expand Up @@ -1230,7 +1240,7 @@ def get_feature(self, arg_list: str | List[str] | None = None, verbose: bool | N
----------
arg_list: str | List[str], optional
A single string or list of strings containing the names of the features to extract. Default is all of:
["close_encounter_check", "general_relativity", "fragmentation", "rotation", "compute_conservation_values"]
["close_encounter_check", "general_relativity", "collision_model", "rotation", "compute_conservation_values"]
verbose: bool, optional
If passed, it will override the Simulation object's verbose flag
**kwargs
Expand All @@ -1245,7 +1255,7 @@ def get_feature(self, arg_list: str | List[str] | None = None, verbose: bool | N
"""

valid_var = {"close_encounter_check": "CHK_CLOSE",
"fragmentation": "FRAGMENTATION",
"collision_model": "COLLISION_MODEL",
"encounter_save": "ENCOUNTER_SAVE",
"minimum_fragment_gmass": "MIN_GMFRAG",
"rotation": "ROTATION",
Expand Down Expand Up @@ -2765,6 +2775,9 @@ def _preprocess(ds, param):
# Remove any overlapping time values
tgood,tid = np.unique(self.encounters.time,return_index=True)
self.encounters = self.encounters.isel(time=tid)
# Remove any NaN values
tgood=self.encounters.time.where(~np.isnan(self.encounters.time),drop=True)
self.encounters = self.encounters.sel(time=tgood)

return

Expand Down
Loading

0 comments on commit 1bbc416

Please sign in to comment.