From 3425a4261c41cf41fd125958a98c2df3d3341bae Mon Sep 17 00:00:00 2001 From: David A Minton Date: Tue, 6 Dec 2022 10:03:08 -0500 Subject: [PATCH] Added new parameters for controlling encounter and fragmentation output --- examples/Fragmentation/Fragmentation_Movie.py | 4 +- python/swiftest/swiftest/io.py | 53 ++++--------- python/swiftest/swiftest/simulation_class.py | 78 +++++++++++++++++++ 3 files changed, 96 insertions(+), 39 deletions(-) diff --git a/examples/Fragmentation/Fragmentation_Movie.py b/examples/Fragmentation/Fragmentation_Movie.py index a43555f08..ba98c6900 100644 --- a/examples/Fragmentation/Fragmentation_Movie.py +++ b/examples/Fragmentation/Fragmentation_Movie.py @@ -179,7 +179,7 @@ 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, gmtiny=gmtiny, minimum_fragment_gmass=minimum_fragment_gmass, verbose=False) - sim.run(dt=1e-4, tstop=2.0e-3, istep_out=1, dump_cadence=0) + sim.set_parameter(fragmentation=True, fragmentation_save="TRAJECTORY", gmtiny=gmtiny, minimum_fragment_gmass=minimum_fragment_gmass, verbose=False) + sim.run(dt=1e-5, tstop=2.0e-3, istep_out=1, dump_cadence=0) anim = AnimatedScatter(sim,movie_filename,movie_titles[style],nskip=1) diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index 40ce7ae33..e315cd244 100644 --- a/python/swiftest/swiftest/io.py +++ b/python/swiftest/swiftest/io.py @@ -31,7 +31,11 @@ "INTERACTION_LOOPS", "ENCOUNTER_CHECK", "TSTART", - "DUMP_CADENCE") + "DUMP_CADENCE", + "ENCOUNTER_SAVE", + "FRAGMENTATION_SAVE") + + # This list defines features that are booleans, so must be converted to/from string when writing/reading from file bool_param = ["RESTART", @@ -51,7 +55,10 @@ float_param = ["T0", "TSTART", "TSTOP", "DT", "CHK_RMIN", "CHK_RMAX", "CHK_EJECT", "CHK_QMIN", "DU2M", "MU2KG", "TU2S", "MIN_GMFRAG", "GMTINY"] -upper_str_param = ["OUT_TYPE","OUT_FORM","OUT_STAT","IN_TYPE","IN_FORM"] +upper_str_param = ["OUT_TYPE","OUT_FORM","OUT_STAT","IN_TYPE","IN_FORM","ENCOUNTER_SAVE","FRAGMENTATION_SAVE", "CHK_QMIN_COORD"] +lower_str_param = ["NC_IN", "PL_IN", "TP_IN", "CB_IN", "CHK_QMIN_RANGE"] + +param_keys = ['! VERSION'] + int_param + float_param + upper_str_param + lower_str_param+ bool_param # This defines Xarray Dataset variables that are strings, which must be processed due to quirks in how NetCDF-Fortran # handles strings differently than Python's Xarray. @@ -417,52 +424,22 @@ def write_labeled_param(param, param_file_name): Prints a text file containing the parameter information. """ outfile = open(param_file_name, 'w') - keylist = ['! VERSION', - 'T0', - 'TSTART', - 'TSTOP', - 'DT', - 'ISTEP_OUT', - 'DUMP_CADENCE', - 'NC_IN', - 'PL_IN', - 'TP_IN', - 'CB_IN', - 'IN_TYPE', - 'IN_FORM', - 'BIN_OUT', - 'OUT_FORM', - 'OUT_TYPE', - 'OUT_STAT', - 'CHK_QMIN', - 'CHK_RMIN', - 'CHK_RMAX', - 'CHK_EJECT', - 'CHK_QMIN_COORD', - 'CHK_QMIN_RANGE', - 'MU2KG', - 'TU2S', - 'DU2M', - 'GMTINY', - 'FRAGMENTATION', - 'MIN_GMFRAG', - 'RESTART'] ptmp = param.copy() # Print the list of key/value pairs in the preferred order - for key in keylist: + for key in param_keys: val = ptmp.pop(key, None) if val is not None: if type(val) is bool: - print(f"{key:<16} {bool2yesno(val)}", file=outfile) + print(f"{key:<32} {bool2yesno(val)}", file=outfile) else: - print(f"{key:<16} {val}", file=outfile) + print(f"{key:<32} {val}", file=outfile) # Print the remaining key/value pairs in whatever order for key, val in ptmp.items(): if val != "": if type(val) is bool: - print(f"{key:<16} {bool2yesno(val)}", file=outfile) + print(f"{key:<32} {bool2yesno(val)}", file=outfile) else: - print(f"{key:<16} {val}", file=outfile) + print(f"{key:<32} {val}", file=outfile) outfile.close() return @@ -840,10 +817,12 @@ def swiftest2xr(param, verbose=True): if ((param['OUT_TYPE'] == 'NETCDF_DOUBLE') or (param['OUT_TYPE'] == 'NETCDF_FLOAT')): if verbose: print('\nCreating Dataset from NetCDF file') ds = xr.open_dataset(param['BIN_OUT'], mask_and_scale=False) + if param['OUT_TYPE'] == "NETCDF_DOUBLE": ds = fix_types(ds,ftype=np.float64) elif param['OUT_TYPE'] == "NETCDF_FLOAT": ds = fix_types(ds,ftype=np.float32) + ds = ds.where(ds.id >=0 ,drop=True) # Check if the name variable contains unique values. If so, make name the dimension instead of id if len(np.unique(ds['name'])) == len(ds['name']): ds = ds.swap_dims({"id" : "name"}) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 727e47d07..a3723553d 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -213,6 +213,11 @@ def __init__(self,read_param: bool = False, read_old_output_file: bool = False, Check for close encounters between bodies. If set to True, then the radii of massive bodies must be included in initial conditions. Parameter input file equivalent: `CHK_CLOSE` + encounter_save : {"NONE","TRAJECTORY","CLOSEST"}, default "NONE" + Indicate if and how encounter data should be saved. If set to "TRAJECTORY" the full close encounter + trajectories are saved to file. If set to "CLOSEST" only the trajectories at the time of closest approach + are saved. If set to "NONE" no trajectory information is saved. + *WARNING*: Enabling this feature could lead to very large files. general_relativity : bool, default True Include the post-Newtonian correction in acceleration calculations. Parameter input file equivalent: `GR` @@ -220,6 +225,12 @@ def __init__(self,read_param: bool = False, read_old_output_file: bool = False, If set to True, this turns on the Fraggle fragment generation code and `rotation` must also be True. This argument only applies to Swiftest-SyMBA simulations. It will be ignored otherwise. Parameter input file equivalent: `FRAGMENTATION` + fragmentation_save : {"NONE","TRAJECTORY","CLOSEST"}, default "NONE" + Indicate if and how fragmentation data should be saved. If set to "TRAJECTORY" the full close encounter + trajectories associated with each collision are saved to file. If set to "CLOSEST" only the trajectories + at a the time the collision occurs are saved. If set to "NONE" no trajectory information is saved (collision + details are still logged fraggle.log). + *WARNING*: Enabling this feature could lead to very large files. minimum_fragment_gmass : float, optional If fragmentation is turned on, this sets the mimimum G*mass of a collisional fragment that can be generated. *Note.* Only set one of minimum_fragment_gmass or minimum_fragment_mass @@ -480,6 +491,9 @@ def run(self,**kwargs): warnings.warn(msg,stacklevel=2) return + if not self.restart: + self.clean() + print(f"Running a {self.codename} {self.integrator} run from tstart={self.param['TSTART']} {self.TU_name} to tstop={self.param['TSTOP']} {self.TU_name}") self._run_swiftest_driver() @@ -1010,6 +1024,8 @@ def set_feature(self, tides: bool | None = None, interaction_loops: Literal["TRIANGULAR", "FLAT", "ADAPTIVE"] | None = None, encounter_check_loops: Literal["TRIANGULAR", "SORTSWEEP", "ADAPTIVE"] | None = None, + encounter_save: Literal["NONE", "TRAJECTORY", "CLOSEST"] | None = None, + fragmentation_save: Literal["NONE", "TRAJECTORY", "CLOSEST"] | None = None, verbose: bool | None = None, **kwargs: Any ): @@ -1021,11 +1037,22 @@ def set_feature(self, close_encounter_check : bool, optional Check for close encounters between bodies. If set to True, then the radii of massive bodies must be included in initial conditions. + encounter_save : {"NONE","TRAJECTORY","CLOSEST"}, default "NONE" + Indicate if and how encounter data should be saved. If set to "TRAJECTORY" the full close encounter + trajectories are saved to file. If set to "CLOSEST" only the trajectories at the time of closest approach + are saved. If set to "NONE" no trajectory information is saved. + *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. This argument only applies to Swiftest-SyMBA simulations. It will be ignored otherwise. + fragmentation_save : {"NONE","TRAJECTORY","CLOSEST"}, default "NONE" + Indicate if and how fragmentation data should be saved. If set to "TRAJECTORY" the full close encounter + trajectories associated with each collision are saved to file. If set to "CLOSEST" only the trajectories + at a the time the collision occurs are saved. If set to "NONE" no trajectory information is saved (collision + details are still logged fraggle.log). + *WARNING*: Enabling this feature could lead to very large files. minimum_fragment_gmass : float, optional If fragmentation is turned on, this sets the mimimum G*mass of a collisional fragment that can be generated. *Note.* Only set one of minimum_fragment_gmass or minimum_fragment_mass @@ -1179,6 +1206,33 @@ def set_feature(self, self.param["ENCOUNTER_CHECK"] = encounter_check_loops update_list.append("encounter_check_loops") + if encounter_save is not None: + valid_vals = ["NONE", "TRAJECTORY", "CLOSEST"] + encounter_save = encounter_save.upper() + if encounter_save not in valid_vals: + msg = f"{encounter_save} is not a valid option for encounter_save." + msg += f"\nMust be one of {valid_vals}" + warnings.warn(msg,stacklevel=2) + if "ENCOUNTER_SAVE" not in self.param: + self.param["ENCOUNTER_SAVE"] = valid_vals[0] + else: + self.param["ENCOUNTER_SAVE"] = encounter_save + update_list.append("encounter_save") + + + if fragmentation_save is not None: + fragmentation_save = fragmentation_save.upper() + valid_vals = ["NONE", "TRAJECTORY", "CLOSEST"] + if fragmentation_save not in valid_vals: + msg = f"{fragmentation_save} is not a valid option for fragmentation_save." + msg += f"\nMust be one of {valid_vals}" + warnings.warn(msg,stacklevel=2) + if "FRAGMENTATION_SAVE" not in self.param: + self.param["FRAGMENTATION_SAVE"] = valid_vals[0] + else: + self.param["FRAGMENTATION_SAVE"] = fragmentation_save + update_list.append("fragmentation_save") + self.param["TIDES"] = False feature_dict = self.get_feature(update_list, verbose) @@ -1210,6 +1264,8 @@ def get_feature(self, arg_list: str | List[str] | None = None, verbose: bool | N valid_var = {"close_encounter_check": "CHK_CLOSE", "fragmentation": "FRAGMENTATION", + "encounter_save": "ENCOUNTER_SAVE", + "fragmentation_save": "FRAGMENTATION_SAVE", "minimum_fragment_gmass": "MIN_GMFRAG", "rotation": "ROTATION", "general_relativity": "GR", @@ -2860,3 +2916,25 @@ def initial_conditions_from_bin(self, framenum=-1, new_param=None, new_param_fil self.write_param(new_param_file, param=new_param) return frame + + def clean(self): + """ + Cleans up simulation directory by deleting old files (dump, logs, and output files). + """ + old_files = [self.simdir / self.param['BIN_OUT'], + self.simdir / "fraggle.log", + self.simdir / "swiftest.log", + ] + glob_files = [self.simdir.glob("**/dump_param?.in")] \ + + [self.simdir.glob("**/dump_bin?.nc")] \ + + [self.simdir.glob("**/enc*.nc")] \ + + [self.simdir.glob("**/frag*.nc")] + + for f in old_files: + if f.exists(): + os.remove(f) + for g in glob_files: + for f in g: + if f.exists(): + os.remove(f) + return