Added new ability to process multiple encounter files on the Python side

examples/Fragmentation/Fragmentation_Movie.py

@@ -180,6 +180,6 @@ def data_stream(self, frame=0):
        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, 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)
+       sim.run(dt=1e-4, 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

@@ -800,6 +800,35 @@ def swifter2xr(param, verbose=True):
         if verbose: print(f"Successfully converted {ds.sizes['time']} output frames.")
     return ds
 
+def process_netcdf_intput(ds, param):
+    """
+    Performs several tasks to convert raw NetCDF files output by the Fortran program into a form that
+    is used by the Python side. These include:
+    - Ensuring all types are correct
+    - Removing any bad id values (empty id slots)
+    - Swapping the id and name dimension if the names are unique
+
+    Parameters
+    ----------
+    ds : Xarray dataset
+
+    Returns
+    -------
+    ds : xarray dataset
+    """
+
+    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"})
+        ds = ds.reset_coords("id")
+
+    return ds
+
 def swiftest2xr(param, verbose=True):
     """
     Converts a Swiftest binary data file into an xarray DataSet.
@@ -814,19 +843,11 @@ def swiftest2xr(param, verbose=True):
     xarray dataset
     """
 
+
     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"})
-           ds = ds.reset_coords("id")
+        ds = process_netcdf_intput(ds, param)
     else:
         print(f"Error encountered. OUT_TYPE {param['OUT_TYPE']} not recognized.")
         return None

python/swiftest/swiftest/simulation_class.py

@@ -21,6 +21,7 @@
 import datetime
 import xarray as xr
 import numpy as np
+from functools import partial
 import numpy.typing as npt
 import shutil
 import subprocess
@@ -355,7 +356,6 @@ def __init__(self,read_param: bool = False, read_old_output_file: bool = False,
                 # overriding everything with defaults when there are no arguments passed to Simulation()
                 kwargs['param_file'] = self.param_file
                 param_file_found = True
-
             else:
                 param_file_found = False
 
@@ -2729,8 +2729,10 @@ def read_output_file(self,read_init_cond : bool = True):
         # This is done to handle cases where the method is called from a different working directory than the simulation
         # results
 
-        # EXPERIMENTAL
-        read_encounter = True
+        if "ENCOUNTER_SAVE" in self.param or "FRAGMENTATION_SAVE" in self.param:
+            read_encounter = self.param["ENCOUNTER_SAVE"] != "NONE" or self.param["FRAGMENTATION_SAVE"] != "NONE"
+        else:
+            read_encounter = False
         param_tmp = self.param.copy()
         param_tmp['BIN_OUT'] = os.path.join(self.simdir, self.param['BIN_OUT'])
         if self.codename == "Swiftest":
@@ -2746,10 +2748,7 @@ def read_output_file(self,read_init_cond : bool = True):
                 else:
                     self.ic = self.data.isel(time=0)
             if read_encounter:
-                param_tmp['BIN_OUT'] = self.simdir / "encounter.nc"
-                if self.verbose:
-                    print("Reading encounter history file as .enc")
-                self.enc = io.swiftest2xr(param_tmp, verbose=self.verbose)
+                self.read_encounter()
 
         elif self.codename == "Swifter":
             self.data = io.swifter2xr(param_tmp, verbose=self.verbose)
@@ -2760,6 +2759,22 @@ def read_output_file(self,read_init_cond : bool = True):
             warnings.warn('Cannot process unknown code type. Call the read_param method with a valid code name. Valid options are "Swiftest", "Swifter", or "Swift".',stacklevel=2)
         return
 
+    def read_encounter(self):
+        if self.verbose:
+            print("Reading encounter history file as .enc")
+        enc_files = self.simdir.glob("**/enc_*.nc")
+
+        # This is needed in order to pass the param argument down to the io.process_netcdf_input function
+        def _preprocess(ds, param):
+            return io.process_netcdf_input(ds,param)
+        partial_func = partial(_precprocess, param=self.param)
+
+        self.enc = xr.open_mfdataset(enc_files,parallel=True,combine="nested", concat_dim="encounter",preprocess=partial_func)
+
+        return
+
+
+
     def follow(self, codestyle="Swifter"):
         """
         An implementation of the Swift tool_follow algorithm. Under development. Currently only for Swift simulations.