From c25ac25c14781fef11575ce6be5e6aafe8e00b9d Mon Sep 17 00:00:00 2001 From: David A Minton Date: Thu, 5 Jan 2023 14:56:38 -0500 Subject: [PATCH] Fixed issues getting the before/after collision snapshots properly recorded --- python/swiftest/swiftest/io.py | 2 - python/swiftest/swiftest/simulation_class.py | 53 +++++++++----------- src/collision/collision_io.f90 | 3 +- src/collision/collision_util.f90 | 17 +++++++ 4 files changed, 42 insertions(+), 33 deletions(-) diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index ade7cac00..d03204904 100644 --- a/python/swiftest/swiftest/io.py +++ b/python/swiftest/swiftest/io.py @@ -804,8 +804,6 @@ def process_netcdf_input(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 ---------- diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index cfe88f26f..8a27de3d0 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -211,12 +211,12 @@ def __init__(self,read_param: bool = False, read_old_output: bool = False, simdi but not each other (SyMBA only). *Note.* Only mtiny or gmtiny is accepted, not both. Parameter input file equivalent: `GMTINY` close_encounter_check : bool, default True - Check for close encounters between bodies. If set to True, then the radii of massive bodies must be included + Check for close encounter 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", "BOTH"}, default "NONE" Indicate if and how encounter data should be saved. If set to "TRAJECTORY", the position and velocity vectors - of all bodies undergoing close encounters are saved at each intermediate step to the encounter files. + of all bodies undergoing close encounter are saved at each intermediate step to the encounter files. If set to "CLOSEST", the position and velocities at the point of closest approach between pairs of bodies are computed and stored to the encounter files. If set to "BOTH", then this stores the values that would be computed in "TRAJECTORY" and "CLOSEST". If set to "NONE" no trajectory information is saved. @@ -316,8 +316,8 @@ def __init__(self,read_param: bool = False, read_old_output: bool = False, simdi self.param = {} self.data = xr.Dataset() self.init_cond = xr.Dataset() - self.encounters = xr.Dataset() - self.collisions = xr.Dataset() + self.encounter = xr.Dataset() + self.collision = xr.Dataset() self.simdir = Path(simdir) if self.simdir.exists(): @@ -1035,11 +1035,11 @@ def set_feature(self, Parameters ---------- close_encounter_check : bool, optional - Check for close encounters between bodies. If set to True, then the radii of massive bodies must be included + Check for close encounter between bodies. If set to True, then the radii of massive bodies must be included in initial conditions. encounter_save : {"NONE","TRAJECTORY","CLOSEST","BOTH"}, default "NONE" Indicate if and how encounter data should be saved. If set to "TRAJECTORY", the position and velocity vectors - of all bodies undergoing close encounters are saved at each intermediate step to the encounter files. + of all bodies undergoing close encounter are saved at each intermediate step to the encounter files. If set to "CLOSEST", the position and velocities at the point of closest approach between pairs of bodies are computed and stored to the encounter files. If set to "BOTH", then this stores the values that would be computed in "TRAJECTORY" and "CLOSEST". If set to "NONE" no trajectory information is saved. @@ -2738,8 +2738,8 @@ def read_output_file(self,read_init_cond : bool = True): else: self.init_cond = self.data.isel(time=0) - self.read_encounters() - self.read_collisions() + self.read_encounter() + self.read_collision() if self.verbose: print("Finished reading Swiftest dataset files.") @@ -2752,13 +2752,13 @@ 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_encounters(self): + def read_encounter(self): enc_files = glob(f"{self.simdir}{os.path.sep}encounter_*.nc") if len(enc_files) == 0: return if self.verbose: - print("Reading encounter history file as .encounters") + print("Reading encounter history file as .encounter") enc_files.sort() @@ -2767,33 +2767,28 @@ def _preprocess(ds, param): return io.process_netcdf_input(ds,param) partial_func = partial(_preprocess, param=self.param) - self.encounters = xr.open_mfdataset(enc_files,parallel=True,combine="nested",concat_dim="time",join="left",preprocess=partial_func,mask_and_scale=True) - self.encounters = io.process_netcdf_input(self.encounters, self.param) + self.encounter = xr.open_mfdataset(enc_files,parallel=True,combine="nested",concat_dim="time",join="left",preprocess=partial_func,mask_and_scale=True) + self.encounter = io.process_netcdf_input(self.encounter, self.param) # Remove any overlapping time values - tgood,tid = np.unique(self.encounters.time,return_index=True) - self.encounters = self.encounters.isel(time=tid) + tgood,tid = np.unique(self.encounter.time,return_index=True) + self.encounter = self.encounter.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) + tgood=self.encounter.time.where(~np.isnan(self.encounter.time),drop=True) + self.encounter = self.encounter.sel(time=tgood) return - def read_collisions(self): - col_files = glob(f"{self.simdir}{os.path.sep}collision_*.nc") - if len(col_files) == 0: - return + def read_collision(self): + + col_file = self.simdir / "collision.nc" + if not os.path.exists(col_file): + return - col_files.sort() if self.verbose: - print("Reading collision history file as .collisions") - - # 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(_preprocess, param=self.param) + print("Reading collision history file as .collision") - self.collisions = xr.open_mfdataset(col_files,parallel=True, combine="nested", concat_dim="collision", preprocess=partial_func,mask_and_scale=True) - self.collisions = io.process_netcdf_input(self.collisions, self.param) + self.collision = xr.open_dataset(col_file) + self.collision = io.process_netcdf_input(self.collision, self.param) return diff --git a/src/collision/collision_io.f90 b/src/collision/collision_io.f90 index a35b7229c..feda0d029 100644 --- a/src/collision/collision_io.f90 +++ b/src/collision/collision_io.f90 @@ -105,11 +105,10 @@ module subroutine collision_io_netcdf_dump(self, param) call nc%open(param) - do i = 1, self%nframes + do i = 1, self%iframe if (allocated(self%frame(i)%item)) then select type(snapshot => self%frame(i)%item) class is (collision_snapshot) - param%ioutput = i call snapshot%write_frame(self,param) end select else diff --git a/src/collision/collision_util.f90 b/src/collision/collision_util.f90 index 882bb8f29..fda192a5b 100644 --- a/src/collision/collision_util.f90 +++ b/src/collision/collision_util.f90 @@ -695,6 +695,7 @@ module subroutine collision_util_snapshot(self, param, nbody_system, t, arg) snapshot%collider%pe(1) = nbody_system%pe snapshot%collider%be(1) = nbody_system%be snapshot%collider%te(1) = nbody_system%te + case("after") ! Get record the energy of the sytem after the collision call nbody_system%get_energy_and_momentum(param) @@ -707,6 +708,22 @@ module subroutine collision_util_snapshot(self, param, nbody_system, t, arg) snapshot%collider%be(2) = nbody_system%be snapshot%collider%te(2) = nbody_system%te + select type(before_snap => snapshot%collider%before ) + class is (swiftest_nbody_system) + select type(before_orig => nbody_system%collider%before) + class is (swiftest_nbody_system) + call move_alloc(before_orig%pl, before_snap%pl) + end select + end select + + select type(after_snap => snapshot%collider%after ) + class is (swiftest_nbody_system) + select type(after_orig => nbody_system%collider%after) + class is (swiftest_nbody_system) + call move_alloc(after_orig%pl, after_snap%pl) + end select + end select + ! Save the snapshot for posterity call collision_util_save_snapshot(nbody_system%collision_history,snapshot) deallocate(snapshot)