diff --git a/examples/Fragmentation/Fragmentation_Movie.py b/examples/Fragmentation/Fragmentation_Movie.py index ba98c6900..f5f9cd792 100644 --- a/examples/Fragmentation/Fragmentation_Movie.py +++ b/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) diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index e315cd244..78d71c820 100644 --- a/python/swiftest/swiftest/io.py +++ b/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 diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 7afc72b76..3c9105554 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/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.