Added new parameters for controlling encounter and fragmentation output

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)

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"})

python/swiftest/swiftest/simulation_class.py

@@ -213,13 +213,24 @@ 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`
         fragmentation : bool, default True
             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