From 615730a5e460ef83932ed95e006f59082f01dd7a Mon Sep 17 00:00:00 2001 From: David Minton Date: Mon, 7 Nov 2022 19:22:22 -0500 Subject: [PATCH 01/17] Made sure that MIN_GMFRAG was cast as a float --- python/swiftest/swiftest/io.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index 4678b4784..044f1eff5 100644 --- a/python/swiftest/swiftest/io.py +++ b/python/swiftest/swiftest/io.py @@ -105,6 +105,8 @@ def read_swiftest_param(param_file_name, param, verbose=True): param['ENCOUNTER_CHECK'] = param['ENCOUNTER_CHECK'].upper() if 'GMTINY' in param: param['GMTINY'] = real2float(param['GMTINY']) + if 'MIN_GMFRAG' in param: + param['MIN_GMFRAG'] = real2float(param['MIN_GMFRAG']) except IOError: print(f"{param_file_name} not found.") return param From 810f48129e0e6ede6e5c511430ba4305fde8e937 Mon Sep 17 00:00:00 2001 From: David Minton Date: Mon, 7 Nov 2022 19:22:49 -0500 Subject: [PATCH 02/17] Added new method that converts the units inside the param dictionary if the unit system changes --- python/swiftest/swiftest/simulation_class.py | 115 ++++++++++++++++--- 1 file changed, 98 insertions(+), 17 deletions(-) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 5ad945d69..b38c7b34d 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -28,27 +28,20 @@ def __init__(self, codename="Swiftest", param_file="param.in", readbin=False, ve self.ds = xr.Dataset() self.param = { '! VERSION': f"Swiftest parameter input", - 'T0': "0.0", - 'TSTOP': "0.0", - 'DT': "0.0", + 'T0': 0.0, + 'TSTART' : 0.0, + 'TSTOP': 0.0, + 'DT': 0.0, 'IN_FORM': "XV", 'IN_TYPE': "NETCDF_DOUBLE", 'NC_IN' : "init_cond.nc", - 'CB_IN' : "cb.in", - 'PL_IN' : "pl.in", - 'TP_IN' : "tp.in", - 'ISTEP_OUT': "1", - 'ISTEP_DUMP': "1", + 'ISTEP_OUT': 1, + 'ISTEP_DUMP': 1, 'BIN_OUT': "bin.nc", 'OUT_TYPE': 'NETCDF_DOUBLE', 'OUT_FORM': "XVEL", 'OUT_STAT': "REPLACE", - 'CHK_RMAX': "-1.0", - 'CHK_EJECT': "-1.0", - 'CHK_RMIN': "-1.0", - 'CHK_QMIN': "-1.0", 'CHK_QMIN_COORD': "HELIO", - 'CHK_QMIN_RANGE': "-1.0 -1.0", 'ENC_OUT': "", 'EXTRA_FORCE': "NO", 'DISCARD_OUT': "", @@ -66,6 +59,7 @@ def __init__(self, codename="Swiftest", param_file="param.in", readbin=False, ve } self.codename = codename self.verbose = verbose + self.set_simulation_range(rmin=constants.RSun, rmax=1000.0) self.set_unit_system() # If the parameter file is in a different location than the current working directory, we will need @@ -132,6 +126,11 @@ def set_unit_system(self,MU="Msun",DU="AU",TU="YR",MU2KG=None,DU2M=None,TU2S=Non Sets the values of MU2KG, DU2M, and TU2S in the param dictionary to the appropriate units. Also computes the gravitational constant, GU """ + # Save the previously set values of the unit conversion factors so we can update parameters as needed + MU2KG_old = self.param.pop('MU2KG',None) + DU2M_old = self.param.pop('DU2M',None) + TU2S_old = self.param.pop('TU2S',None) + if MU2KG is not None: self.param['MU2KG'] = MU2KG self.MU_name = None @@ -178,13 +177,13 @@ def set_unit_system(self,MU="Msun",DU="AU",TU="YR",MU2KG=None,DU2M=None,TU2S=Non self.param['TU2S'] = TU2S self.TU_name = None else: - if TU.upper() == "YR": + if TU.upper() == "YR" or TU.upper() == "Y" or TU.upper() == "YEAR" or TU.upper() == "YEARS": self.param['TU2S'] = constants.YR2S self.TU_name = "y" - elif TU.upper() == "DAY" or TU.upper() == "D" or TU.upper() == "JD": + elif TU.upper() == "DAY" or TU.upper() == "D" or TU.upper() == "JD" or TU.upper() == "DAYS": self.param['TU2S'] = constants.JD2S self.TU_name = "Day" - elif TU.upper() == "S": + elif TU.upper() == "S" or TU.upper() == "SECONDS" or TU.upper() == "SEC": self.param['TU2S'] = 1.0 self.TU_name = "s" else: @@ -195,6 +194,8 @@ def set_unit_system(self,MU="Msun",DU="AU",TU="YR",MU2KG=None,DU2M=None,TU2S=Non self.VU_name = f"{self.DU_name}/{self.TU_name}" self.GC = constants.GC * self.param['TU2S']**2 * self.param['MU2KG'] / self.param['DU2M']**3 + self.update_param_units(MU2KG_old, DU2M_old, TU2S_old) + if self.verbose: if self.MU_name is None or self.DU_name is None or self.TU_name is None: print(f"Custom units set.") @@ -205,7 +206,87 @@ def set_unit_system(self,MU="Msun",DU="AU",TU="YR",MU2KG=None,DU2M=None,TU2S=Non print(f"Units set to {self.MU_name}-{self.DU_name}-{self.TU_name}") return - + + def update_param_units(self,MU2KG_old,DU2M_old,TU2S_old): + """ + Updates the values of parameters that have units when the units have changed. + + Parameters + ---------- + MU2KG_old : Old value of the mass unit conversion factor + DU2M_old : Old value of the distance unit conversion factor + TU2S_old : Old value of the time unit conversion factor + + Returns + ------- + Updates the set of param dictionary values for the new unit system + + """ + + mass_keys = ['GMTINY', 'MIN_GMFRAG'] + distance_keys = ['CHK_QMIN','CHK_RMIN','CHK_RMAX', 'CHK_EJECT'] + time_keys = ['T0','TSTOP','DT'] + + if MU2KG_old is not None: + for k in mass_keys: + if k in self.param: + print(f"param['{k}']: {self.param[k]}") + self.param[k] *= self.param['MU2KG'] / MU2KG_old + + if DU2M_old is not None: + for k in distance_keys: + if k in self.param: + self.param[k] *= self.param['DU2M'] / DU2M_old + + CHK_QMIN_RANGE = self.param.pop('CHK_QMIN_RANGE', None) + if CHK_QMIN_RANGE is not None: + CHK_QMIN_RANGE = CHK_QMIN_RANGE.split(" ") + for i, v in enumerate(CHK_QMIN_RANGE): + CHK_QMIN_RANGE[i] = float(CHK_QMIN_RANGE[i]) * self.param['DU2M'] / DU2M_old + self.param['CHK_QMIN_RANGE'] = f"{CHK_QMIN_RANGE[0]} {CHK_QMIN_RANGE[1]}" + + if TU2S_old is not None: + for k in time_keys: + if k in self.param: + self.param[k] *= self.param['TU2S'] / TU2S_old + + return + + + def set_simulation_range(self,rmin=None,rmax=None): + """ + Sets the minimum and maximum distances of the simulation. + + Parameters + ---------- + rmin : float + Minimum distance of the simulation (CHK_QMIN, CHK_RMIN, CHK_QMIN_RANGE[0]) + rmax : float + Maximum distance of the simulation (CHK_RMAX, CHK_QMIN_RANGE[1]) + + Returns + ------- + Sets the appropriate param dictionary values. + + """ + CHK_QMIN_RANGE = self.param.pop('CHK_QMIN_RANGE',None) + if CHK_QMIN_RANGE is None: + CHK_QMIN_RANGE = [-1, -1] + else: + CHK_QMIN_RANGE = CHK_QMIN_RANGE.split(" ") + if rmin is not None: + self.param['CHK_QMIN'] = rmin + self.param['CHK_RMIN'] = rmin + CHK_QMIN_RANGE[0] = rmin + if rmax is not None: + self.param['CHK_RMAX'] = rmax + self.param['CHK_EJECT'] = rmax + CHK_QMIN_RANGE[1] = rmax + + self.param['CHK_QMIN_RANGE'] =f"{CHK_QMIN_RANGE[0]} {CHK_QMIN_RANGE[1]}" + + return + def add(self, plname, date=date.today().isoformat(), idval=None): """ Adds a solar system body to an existing simulation DataSet. From 585629db5e493d3505fc20aa473ec64614cbd65a Mon Sep 17 00:00:00 2001 From: David Minton Date: Tue, 8 Nov 2022 12:42:25 -0500 Subject: [PATCH 03/17] Restructured the Simulation class to accept argument-based parameters. In the new structure, the param dictionary values are set internally by what arguments are passed. This allows for arguments to be sanitized and checked for consistency and compatibility in a more robust way. --- .../Basic_Simulation/initial_conditions.py | 39 +- python/swiftest/swiftest/io.py | 3 - python/swiftest/swiftest/simulation_class.py | 529 ++++++++++++++---- 3 files changed, 433 insertions(+), 138 deletions(-) diff --git a/examples/Basic_Simulation/initial_conditions.py b/examples/Basic_Simulation/initial_conditions.py index 640330f1f..dceb33cc7 100644 --- a/examples/Basic_Simulation/initial_conditions.py +++ b/examples/Basic_Simulation/initial_conditions.py @@ -30,7 +30,7 @@ from numpy.random import default_rng # Initialize the simulation object as a variable -sim = swiftest.Simulation() +sim = swiftest.Simulation(init_cond_file_type="ASCII") # Add parameter attributes to the simulation object sim.param['T0'] = 0.0 @@ -38,40 +38,9 @@ sim.param['DT'] = 0.005 sim.param['ISTEP_OUT'] = 200 sim.param['ISTEP_DUMP'] = 200 -sim.param['OUT_FORM'] = 'XVEL' -sim.param['OUT_TYPE'] = 'NETCDF_DOUBLE' -sim.param['OUT_STAT'] = 'REPLACE' -sim.param['IN_FORM'] = 'EL' -sim.param['IN_TYPE'] = 'ASCII' -sim.param['PL_IN'] = 'pl.in' -sim.param['TP_IN'] = 'tp.in' -sim.param['CB_IN'] = 'cb.in' -sim.param['BIN_OUT'] = 'out.nc' -sim.param['CHK_QMIN'] = swiftest.RSun / swiftest.AU2M -sim.param['CHK_RMIN'] = swiftest.RSun / swiftest.AU2M -sim.param['CHK_RMAX'] = 1000.0 -sim.param['CHK_EJECT'] = 1000.0 -sim.param['CHK_QMIN_COORD'] = 'HELIO' -sim.param['CHK_QMIN_RANGE'] = f'{swiftest.RSun / swiftest.AU2M} 1000.0' -sim.param['MU2KG'] = swiftest.MSun -sim.param['TU2S'] = swiftest.YR2S -sim.param['DU2M'] = swiftest.AU2M -sim.param['EXTRA_FORCE'] = 'NO' -sim.param['BIG_DISCARD'] = 'NO' -sim.param['CHK_CLOSE'] = 'YES' -sim.param['GR'] = 'YES' -sim.param['INTERACTION_LOOPS'] = 'ADAPTIVE' -sim.param['ENCOUNTER_CHECK'] = 'ADAPTIVE' -sim.param['RHILL_PRESENT'] = 'YES' -sim.param['FRAGMENTATION'] = 'YES' -sim.param['ROTATION'] = 'YES' -sim.param['ENERGY'] = 'YES' sim.param['GMTINY'] = 1e-6 sim.param['MIN_GMFRAG'] = 1e-9 -# Set gravitational units of the system -GU = swiftest.GC / (sim.param['DU2M'] ** 3 / (sim.param['MU2KG'] * sim.param['TU2S'] ** 2)) - # Add the modern planets and the Sun using the JPL Horizons Database sim.add("Sun", idval=0, date="2022-08-08") sim.add("Mercury", idval=1, date="2022-08-08") @@ -95,9 +64,9 @@ capom_pl = default_rng().uniform(0.0, 360.0, npl) omega_pl = default_rng().uniform(0.0, 360.0, npl) capm_pl = default_rng().uniform(0.0, 360.0, npl) -GM_pl = (np.array([6e23, 8e23, 1e24, 3e24, 5e24]) / sim.param['MU2KG']) * GU -R_pl = np.full(npl, (3 * (GM_pl / GU) / (4 * np.pi * density_pl)) ** (1.0 / 3.0)) -Rh_pl = a_pl * ((GM_pl) / (3 * GU)) ** (1.0 / 3.0) +GM_pl = (np.array([6e23, 8e23, 1e24, 3e24, 5e24]) / sim.param['MU2KG']) * sim.GU +R_pl = np.full(npl, (3 * (GM_pl / sim.GU) / (4 * np.pi * density_pl)) ** (1.0 / 3.0)) +Rh_pl = a_pl * ((GM_pl) / (3 * sim.GU)) ** (1.0 / 3.0) Ip1_pl = np.array([0.4, 0.4, 0.4, 0.4, 0.4]) Ip2_pl = np.array([0.4, 0.4, 0.4, 0.4, 0.4]) Ip3_pl = np.array([0.4, 0.4, 0.4, 0.4, 0.4]) diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index 044f1eff5..af74e34a4 100644 --- a/python/swiftest/swiftest/io.py +++ b/python/swiftest/swiftest/io.py @@ -94,9 +94,6 @@ def read_swiftest_param(param_file_name, param, verbose=True): param['CHK_CLOSE'] = param['CHK_CLOSE'].upper() param['RHILL_PRESENT'] = param['RHILL_PRESENT'].upper() param['FRAGMENTATION'] = param['FRAGMENTATION'].upper() - if param['FRAGMENTATION'] == 'YES' and param['PARTICLE_OUT'] == '': - if param['OUT_TYPE'] == 'REAL8' or param['OUT_TYPE'] == 'REAL4': - param['PARTICLE_OUT'] = 'particle.dat' param['ROTATION'] = param['ROTATION'].upper() param['TIDES'] = param['TIDES'].upper() param['ENERGY'] = param['ENERGY'].upper() diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index b38c7b34d..38fa44d65 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -9,6 +9,7 @@ You should have received a copy of the GNU General Public License along with Swiftest. If not, see: https://www.gnu.org/licenses. """ +from __future__ import annotations from swiftest import io from swiftest import init_cond @@ -19,12 +20,166 @@ import numpy as np import os import shutil +from typing import ( + Literal, + Dict, +) + class Simulation: """ This is a class that defines the basic Swift/Swifter/Swiftest simulation object """ - def __init__(self, codename="Swiftest", param_file="param.in", readbin=False, verbose=True): + def __init__(self, + codename: Literal["Swiftest", "Swifter", "Swift"] = "Swiftest", + param_file: os.PathLike | str ="param.in", + read_param: bool = False, + init_cond_file_type: Literal["NETCDF_DOUBLE", "NETCDF_FLOAT", "ASCII"] = "NETCDF_DOUBLE", + init_cond_file_name: str | os.PathLike | Dict[str, str] | Dict[str, os.PathLike] | None = None, + init_cond_format: Literal["EL", "XV"] = "EL", + output_file_type: Literal["NETCDF_DOUBLE","NETCDF_FLOAT","REAL4","REAL8","XDR4","XDR8"] = "NETCDF_DOUBLE", + output_file_name: os.PathLike | str | None = None, + output_format: Literal["XV","XVEL"] = "XVEL", + read_old_output_file: bool = False, + MU: str = "MSUN", + DU: str = "AU", + TU: str = "Y", + MU2KG: float | None = None, + DU2M: float | None = None, + TU2S: float | None = None, + rmin: float = constants.RSun / constants.AU2M, + rmax: float = 10000.0, + close_encounter_check: bool =True, + general_relativity: bool =True, + fragmentation: bool =True, + rotation: bool = True, + compute_conservation_values: bool = False, + interaction_loops: Literal["TRIANGULAR","FLAT","ADAPTIVE"] = "TRIANGULAR", + encounter_check_loops: Literal["TRIANGULAR","SORTSWEEP","ADAPTIVE"] = "TRIANGULAR", + verbose: bool = True + ): + """ + + Parameters + ---------- + codename : {"Swiftest", "Swifter", "Swift"}, default "Swiftest" + Name of the n-body integrator that will be used. + param_file : str, path-like, or file-lke, default "param.in" + Name of the parameter input file that will be passed to the integrator. + read_param : bool, default False + If true, read in a pre-existing parameter input file given by the argument `param_file`. Otherwise, create + a new parameter file using the arguments passed to Simulation. + > *Note:* If set to true, the parameters defined in the input file will override any passed into the + > arguments of Simulation. + init_cond_file_type : {"NETCDF_DOUBLE", "NETCDF_FLOAT", "ASCII"}, default "NETCDF_DOUBLE" + The file type containing initial conditions for the simulation: + * NETCDF_DOUBLE: A single initial conditions input file in NetCDF file format of type NETCDF_DOUBLE + * NETCDF_FLOAT: A single initial conditions input file in NetCDF file format of type NETCDF_FLOAT + * ASCII : Three initial conditions files in ASCII format. The individual files define the central body, + massive body, and test particle initial conditions. + init_cond_file_name : str, path-like, or dict, optional + Name of the input initial condition file or files. Whether to pass a single file name or a dictionary + depends on the argument passed to `init_cond_file_type`: If `init_cond_file_type={"NETCDF_DOUBLE","NETCDF_FLOAT"}`, + then this will be a single file name. If `init_cond_file_type="ASCII"` then this must be a dictionary where: + ```init_cond_file_name = { + "CB" : *path to central body initial conditions file* (Swiftest only), + "PL" : *path to massive body initial conditions file*, + "TP" : *path to test particle initial conditions file* + }``` + If no file name is provided then the following default file names will be used. + * NETCDF_DOUBLE, NETCDF_FLOAT: `init_cond_file_name = "init_cond.nc"` + * ASCII: `init_cond_file_name = {"CB" : "cb.in", "PL" : "pl.in", "TP" : "tp.in"}` + init_cond_format : {"EL", "XV"}, default "EL" + Indicates whether the input initial conditions are given as orbital elements or cartesian position and + velocity vectors. + > *Note:* If `codename` is "Swift" or "Swifter", EL initial conditions are converted to XV. + output_file_type : {"NETCDF_DOUBLE", "NETCDF_FLOAT","REAL4","REAL8","XDR4","XDR8"}, default "NETCDF_DOUBLE" + The file type for the outputs of the simulation. Compatible file types depend on the `codename` argument. + * Swiftest: Only "NETCDF_DOUBLE" or "NETCDF_FLOAT" supported. + * Swifter: Only "REAL4","REAL8","XDR4" or "XDR8" supported. + * Swift: Only "REAL4" supported. + output_file_name : str or path-like, optional + Name of output file to generate. If not supplied, then one of the default file names are used, depending on + the value passed to `output_file_type`. If one of the NetCDF types are used, the default is "bin.nc". + Otherwise, the default is "bin.dat". + output_format : {"XV","XVEL"}, default "XVEL" + Specifies the format for the data saved to the output file. If "XV" then cartesian position and velocity + vectors for all bodies are stored. If "XVEL" then the orbital elements are also stored. + read_old_output_file : bool, default False + If true, read in a pre-existing binary input file given by the argument `output_file_name`. + MU : str, default "MSUN" + The mass unit system to use. Case-insensitive valid options are: + * "Msun" : Solar mass + * "Mearth" : Earth mass + * "kg" : kilograms + * "g" : grams + DU : str, optional + The distance unit system to use. Case-insensitive valid options are: + * "AU" : Astronomical Unit + * "Rearth" : Earth radius + * "m" : meter + * "cm" : centimeter + TU : str, optional + The time unit system to use. Case-insensitive valid options are: + * "YR" : Year + * "DAY" : Julian day + * "d" : Julian day + * "JD" : Julian day + * "s" : second + MU2KG: float, optional + The conversion factor to multiply by the mass unit that would convert it to kilogram. + Setting this overrides MU + DU2M : float, optional + The conversion factor to multiply by the distance unit that would convert it to meter. + Setting this overrides DU + TU2S : float, optional + The conversion factor to multiply by the time unit that would convert it to seconds. + Setting this overrides TU + rmin : float, default value is the radius of the Sun in the unit system defined by the unit input arguments. + Minimum distance of the simulation (CHK_QMIN, CHK_RMIN, CHK_QMIN_RANGE[0]) + rmax : float, default value is 10000 AU in the unit system defined by the unit input arguments. + Maximum distance of the simulation (CHK_RMAX, CHK_QMIN_RANGE[1]) + 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 + in initial conditions. + general_relativity : bool, default True + Include the post-Newtonian correction in acceleration calculations. + 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. + rotation : bool, default True + If set to True, this turns on rotation tracking and radius, rotation vector, and moments of inertia values + must be included in the initial conditions. + This argument only applies to Swiftest-SyMBA simulations. It will be ignored otherwise. + compute_conservation_values : bool, default False + Turns on the computation of energy, angular momentum, and mass conservation and reports the values + every output step of a running simulation. + interaction_loops : {"TRIANGULAR","FLAT","ADAPTIVE"}, default "TRIANGULAR" + *Swiftest Experimental feature* + Specifies which algorithm to use for the computation of body-body gravitational forces. + * "TRIANGULAR" : Upper-triangular double-loops . + * "FLAT" : Body-body interation pairs are flattened into a 1-D array. + * "ADAPTIVE" : Periodically times the TRIANGULAR and FLAT methods and determines which one to use based on + the wall time to complete the loop. *Note:* Using ADAPTIVE means that bit-identical repeatability cannot + be assured, as the choice of algorithm depends on possible external factors that can influence the wall + time calculation. The exact floating-point results of the interaction will be different between the two + algorithm types. + encounter_check_loops : {"TRIANGULAR","SORTSWEEP","ADAPTIVE"}, default "TRIANGULAR" + *Swiftest Experimental feature* + Specifies which algorithm to use for checking whether bodies are in a close encounter state or not. + * "TRIANGULAR" : Upper-triangular double-loops. + * "SORTSWEEP" : A Sort-Sweep algorithm is used to reduce the population of potential close encounter bodies. + This algorithm is still in development, and does not necessarily speed up the encounter checking. + Use with caution. + * "ADAPTIVE" : Periodically times the TRIANGULAR and SORTSWEEP methods and determines which one to use based + on the wall time to complete the loop. *Note:* Using ADAPTIVE means that bit-identical repeatability cannot + be assured, as the choice of algorithm depends on possible external factors that can influence the wall + time calculation. The exact floating-point results of the interaction will be different between the two + algorithm types. + verbose : bool, default True + If set to True, then more information is printed by Simulation methods as they are executed. Setting to + False suppresses most messages other than errors. + """ self.ds = xr.Dataset() self.param = { '! VERSION': f"Swiftest parameter input", @@ -32,42 +187,48 @@ def __init__(self, codename="Swiftest", param_file="param.in", readbin=False, ve 'TSTART' : 0.0, 'TSTOP': 0.0, 'DT': 0.0, - 'IN_FORM': "XV", - 'IN_TYPE': "NETCDF_DOUBLE", - 'NC_IN' : "init_cond.nc", 'ISTEP_OUT': 1, 'ISTEP_DUMP': 1, - 'BIN_OUT': "bin.nc", - 'OUT_TYPE': 'NETCDF_DOUBLE', - 'OUT_FORM': "XVEL", - 'OUT_STAT': "REPLACE", 'CHK_QMIN_COORD': "HELIO", - 'ENC_OUT': "", 'EXTRA_FORCE': "NO", - 'DISCARD_OUT': "", - 'PARTICLE_OUT' : "", 'BIG_DISCARD': "NO", 'CHK_CLOSE': "YES", 'RHILL_PRESENT': "YES", - 'FRAGMENTATION': "NO", - 'ROTATION': "NO", + 'FRAGMENTATION': "YES", + 'ROTATION': "YES", 'TIDES': "NO", - 'ENERGY': "NO", + 'ENERGY': "YES", 'GR': "YES", - 'INTERACTION_LOOPS': "TRIANGULAR", - 'ENCOUNTER_CHECK': "TRIANGULAR" + 'INTERACTION_LOOPS': interaction_loops, + 'ENCOUNTER_CHECK': encounter_check_loops } self.codename = codename self.verbose = verbose - self.set_simulation_range(rmin=constants.RSun, rmax=1000.0) - self.set_unit_system() + + self.set_distance_range(rmin=rmin, rmax=rmax) + + self.set_unit_system(MU=MU, DU=DU, TU=TU, + MU2KG=MU2KG, DU2M=DU2M, TU2S=TU2S, + recompute_values=True) + + self.set_init_cond_files(init_cond_file_type=init_cond_file_type, + init_cond_file_name=init_cond_file_name, + init_cond_format=init_cond_format) + + self.set_output_files(output_file_type=output_file_type, + output_file_name=output_file_name, + output_format=output_format ) # If the parameter file is in a different location than the current working directory, we will need # to use it to properly open bin files self.sim_dir = os.path.dirname(os.path.realpath(param_file)) - if os.path.exists(param_file): - self.read_param(param_file, codename=codename, verbose=self.verbose) - if readbin: + if read_param: + if os.path.exists(param_file): + self.read_param(param_file, codename=codename, verbose=self.verbose) + else: + print(f"{param_file} not found.") + + if read_old_output_file: binpath = os.path.join(self.sim_dir,self.param['BIN_OUT']) if os.path.exists(binpath): self.bin2xr() @@ -76,7 +237,163 @@ def __init__(self, codename="Swiftest", param_file="param.in", readbin=False, ve return - def set_unit_system(self,MU="Msun",DU="AU",TU="YR",MU2KG=None,DU2M=None,TU2S=None): + def set_init_cond_files(self, + init_cond_file_type: Literal["NETCDF_DOUBLE", "NETCDF_FLOAT", "ASCII"], + init_cond_file_name: str | os.PathLike | Dict[str, str] | Dict[str, os.PathLike] | None, + init_cond_format: Literal["EL", "XV"]): + """ + Sets the initial condition file parameters in the parameters dictionary. + + Parameters + ---------- + init_cond_file_type : {"NETCDF_DOUBLE", "NETCDF_FLOAT", "ASCII"} + The file type containing initial conditions for the simulation: + * NETCDF_DOUBLE: A single initial conditions input file in NetCDF file format of type NETCDF_DOUBLE + * NETCDF_FLOAT: A single initial conditions input file in NetCDF file format of type NETCDF_FLOAT + * ASCII : Three initial conditions files in ASCII format. The individual files define the central body, + massive body, and test particle initial conditions. + init_cond_file_name : str, path-like, or dict, optional + Name of the input initial condition file or files. Whether to pass a single file name or a dictionary + depends on the argument passed to `init_cond_file_type`: If `init_cond_file_type={"NETCDF_DOUBLE","NETCDF_FLOAT"}`, + then this will be a single file name. If `init_cond_file_type="ASCII"` then this must be a dictionary where: + ```init_cond_file_name = { + "CB" : *path to central body initial conditions file* (Swiftest only), + "PL" : *path to massive body initial conditions file*, + "TP" : *path to test particle initial conditions file* + }``` + If no file name is provided then the following default file names will be used. + * NETCDF_DOUBLE, NETCDF_FLOAT: `init_cond_file_name = "init_cond.nc"` + * ASCII: `init_cond_file_name = {"CB" : "cb.in", "PL" : "pl.in", "TP" : "tp.in"}` + init_cond_format : {"EL", "XV"} + Indicates whether the input initial conditions are given as orbital elements or cartesian position and + velocity vectors. + > *Note:* If `codename` is "Swift" or "Swifter", EL initial conditions are converted to XV. + + Returns + ------- + Sets the appropriate initial conditions file parameters inside the self.param dictionary. + + """ + + def ascii_file_input_error_msg(codename): + print(f"Error in set_init_cond_files: init_cond_file_name must be a dictionary of the form: ") + print('{') + if codename == "Swiftest": + print('"CB" : *path to central body initial conditions file*,') + print('"PL" : *path to massive body initial conditions file*,') + print('"TP" : *path to test particle initial conditions file*') + print('}') + return + + if self.codename == "Swiftest": + init_cond_keys = ["CB", "PL", "TP"] + else: + init_cond_keys = ["PL", "TP"] + if init_cond_file_type != "ASCII": + print(f"{init_cond_file_type} is not supported by {self.codename}. Using ASCII instead") + init_cond_file_type="ASCII" + if init_cond_format != "XV": + print(f"{init_cond_format} is not supported by {self.codename}. Using XV instead") + init_cond_format = "XV" + + self.param["IN_TYPE"] = init_cond_file_type + self.param["IN_FORM"] = init_cond_format + + if init_cond_file_type == "ASCII": + if init_cond_file_name is None: + # No file names passed, so we will just use the defaults. + for key in init_cond_keys: + self.param[f"{key}_IN"] = f"{key.lower()}.in" + elif type(init_cond_file_name) is not dict: + # Oops, accidentally passed a single string or path-like instead of the expected dictionary for ASCII + # input type. + ascii_file_input_error_msg(self.codename) + elif not all(key in init_cond_file_name for key in init_cond_keys): + # This is the case where the dictionary doesn't have all the keys we expect. Print an error message. + ascii_file_input_error_msg(self.codename) + else: + # A valid initial conditions file dictionary was passed. + for key in init_cond_keys: + self.param[f"{key}_IN"] = init_cond_file_name[key] + else: + if init_cond_file_name is None: + # No file names passed, so we will just use the default. + self.param["NC_IN"] = "init_cond.nc" + elif type(init_cond_file_name) is dict: + # Oops, accidentally passed a dictionary instead of the expected single string or path-like for NetCDF + # input type. + print(f"Only a single input file is used for NetCDF files") + else: + self.param["NC_IN"] = init_cond_file_name + + return + + def set_output_files(self, + output_file_type: Literal[ "NETCDF_DOUBLE", "NETCDF_FLOAT", "REAL4", "REAL8", "XDR4", "XDR8"], + output_file_name: os.PathLike | str | None, + output_format: Literal["XV", "XVEL"] + ): + """ + Sets the output file parameters in the parameters dictionary. + + Parameters + ---------- + output_file_type : {"NETCDF_DOUBLE", "NETCDF_FLOAT","REAL4","REAL8","XDR4","XDR8"}, default "NETCDF_DOUBLE" + The file type for the outputs of the simulation. Compatible file types depend on the `codename` argument. + * Swiftest: Only "NETCDF_DOUBLE" or "NETCDF_FLOAT" supported. + * Swifter: Only "REAL4","REAL8","XDR4" or "XDR8" supported. + * Swift: Only "REAL4" supported. + output_file_name : str or path-like, optional + Name of output file to generate. If not supplied, then one of the default file names are used, depending on + the value passed to `output_file_type`. If one of the NetCDF types are used, the default is "bin.nc". + Otherwise, the default is "bin.dat". + output_format : {"XV","XVEL"}, default "XVEL" + Specifies the format for the data saved to the output file. If "XV" then cartesian position and velocity + vectors for all bodies are stored. If "XVEL" then the orbital elements are also stored. + + Returns + ------- + Sets the appropriate initial conditions file parameters inside the self.param dictionary. + + """ + + if self.codename == "Swiftest": + if output_file_type not in ["NETCDF_DOUBLE", "NETCDF_FLOAT"]: + print(f"{output_file_type} is not compatible with Swiftest. Setting to NETCDF_DOUBLE") + output_file_type = "NETCDF_DOUBLE" + elif self.codename == "Swifter": + if output_file_type not in ["REAL4","REAL8","XDR4","XDR8"]: + print(f"{output_file_type} is not compatible with Swifter. Setting to REAL8") + output_file_type = "REAL8" + elif self.codename == "Swift": + if output_file_type not in ["REAL4"]: + print(f"{output_file_type} is not compatible with Swift. Setting to REAL4") + output_file_type = "REAL4" + + self.param['OUT_TYPE'] = output_file_type + if output_file_name is None: + if output_file_type in ["NETCDF_DOUBLE", "NETCDF_FLOAT"]: + self.param['BIN_OUT'] = "bin.nc" + else: + self.param['BIN_OUT'] = "bin.dat" + else: + self.param['BIN_OUT'] = output_file_name + + if output_format != "XV" and self.codename != "Swiftest": + print(f"{output_format} is not compatible with {self.codename}. Setting to XV") + output_format = "XV" + self.param['OUT_FORM'] = output_format + + return + + def set_unit_system(self, + MU: str | None = None, + DU: str | None = None, + TU: str | None = None, + MU2KG: float | None = None, + DU2M: float | None = None, + TU2S: float | None = None, + recompute_values: bool = False): """ Setter for setting the unit conversion between one of the standard sets. @@ -90,111 +407,121 @@ def set_unit_system(self,MU="Msun",DU="AU",TU="YR",MU2KG=None,DU2M=None,TU2S=Non Parameters ---------- - MU : str, default "Msun" + MU : str, optional The mass unit system to use. Case-insensitive valid options are: "Msun" : Solar mass "Mearth" : Earth mass "kg" : kilograms "g" : grams - - DU : str, default "AU" + DU : str, optional The distance unit system to use. Case-insensitive valid options are: "AU" : Astronomical Unit "Rearth" : Earth radius "m" : meter "cm" : centimeter - - TU : str, default "YR" + TU : str, optional The time unit system to use. Case-insensitive valid options are: "YR" : Year "DAY" : Julian day "d" : Julian day "JD" : Julian day "s" : second - - MU2KG : float, default `None` - The conversion factor to multiply by the mass unit that would convert it to kilogram. Setting this overrides MU - - DU2M : float, default `None` - The conversion factor to multiply by the distance unit that would convert it to meter. Setting this overrides DU - - TU2S : float, default `None` - The conversion factor to multiply by the time unit that would convert it to seconds. Setting this overrides TU + MU2KG : float, optional + The conversion factor to multiply by the mass unit that would convert it to kilogram. + Setting this overrides MU + DU2M : float, optional + The conversion factor to multiply by the distance unit that would convert it to meter. + Setting this overrides DU + TU2S : float, optional + The conversion factor to multiply by the time unit that would convert it to seconds. + Setting this overrides TU + recompute_values : bool, default False + Recompute all values into the new unit system. + >*Note:* This is a destructive operation, however if not executed then the values contained in the parameter + > file and input/output data files computed previously may not be consistent with the new unit conversion + > factors. Returns ---------- - Sets the values of MU2KG, DU2M, and TU2S in the param dictionary to the appropriate units. Also computes the gravitational constant, GU + Sets the values of MU2KG, DU2M, and TU2S in the param dictionary to the appropriate units. Also computes the + gravitational constant, self.GU and recomput """ - # Save the previously set values of the unit conversion factors so we can update parameters as needed - MU2KG_old = self.param.pop('MU2KG',None) - DU2M_old = self.param.pop('DU2M',None) - TU2S_old = self.param.pop('TU2S',None) + MU2KG_old = None + DU2M_old = None + TU2S_old = None - if MU2KG is not None: - self.param['MU2KG'] = MU2KG - self.MU_name = None - else: - if MU.upper() == "MSUN": - self.param['MU2KG'] = constants.MSun - self.MU_name = "MSun" - elif MU.upper() == "MEARTH": - self.param['MU2KG'] = constants.MEarth - self.MU_name = "MEarth" - elif MU.upper() == "KG": - self.param['MU2KG'] = 1.0 - self.MU_name = "kg" - elif MU.upper() == "G": - self.param['MU2KG'] = 1000.0 - self.MU_name = "g" + if MU2KG is not None or MU is not None: + MU2KG_old = self.param.pop('MU2KG',None) + if MU2KG is not None: + self.param['MU2KG'] = MU2KG + self.MU_name = None else: - print(f"{MU} not a recognized unit system. Using MSun as a default.") - self.param['MU2KG'] = constants.MSun - self.MU_name = "MSun" - - if DU2M is not None: - self.param['DU2M'] = DU2M - self.DU_name = None - else: - if DU.upper() == "AU": - self.param['DU2M'] = constants.AU2M - self.DU_name = "AU" - elif DU.upper() == "REARTH": - self.param['DU2M'] = constants.REarth - self.DU_name = "REarth" - elif DU.upper() == "M": - self.param['DU2M'] = 1.0 - self.DU_name = "m" - elif DU.upper() == "CM": - self.param['DU2M'] = 100.0 - self.DU_name = "cm" + if MU.upper() == "MSUN": + self.param['MU2KG'] = constants.MSun + self.MU_name = "MSun" + elif MU.upper() == "MEARTH": + self.param['MU2KG'] = constants.MEarth + self.MU_name = "MEarth" + elif MU.upper() == "KG": + self.param['MU2KG'] = 1.0 + self.MU_name = "kg" + elif MU.upper() == "G": + self.param['MU2KG'] = 1000.0 + self.MU_name = "g" + else: + print(f"{MU} not a recognized unit system. Using MSun as a default.") + self.param['MU2KG'] = constants.MSun + self.MU_name = "MSun" + + if DU2M is not None or DU is not None: + DU2M_old = self.param.pop('DU2M',None) + if DU2M is not None: + self.param['DU2M'] = DU2M + self.DU_name = None else: - print(f"{DU} not a recognized unit system. Using AU as a default.") - self.param['DU2M'] = constants.AU2M - self.DU_name = "AU" - - if TU2S is not None: - self.param['TU2S'] = TU2S - self.TU_name = None - else: - if TU.upper() == "YR" or TU.upper() == "Y" or TU.upper() == "YEAR" or TU.upper() == "YEARS": - self.param['TU2S'] = constants.YR2S - self.TU_name = "y" - elif TU.upper() == "DAY" or TU.upper() == "D" or TU.upper() == "JD" or TU.upper() == "DAYS": - self.param['TU2S'] = constants.JD2S - self.TU_name = "Day" - elif TU.upper() == "S" or TU.upper() == "SECONDS" or TU.upper() == "SEC": - self.param['TU2S'] = 1.0 - self.TU_name = "s" + if DU.upper() == "AU": + self.param['DU2M'] = constants.AU2M + self.DU_name = "AU" + elif DU.upper() == "REARTH": + self.param['DU2M'] = constants.REarth + self.DU_name = "REarth" + elif DU.upper() == "M": + self.param['DU2M'] = 1.0 + self.DU_name = "m" + elif DU.upper() == "CM": + self.param['DU2M'] = 100.0 + self.DU_name = "cm" + else: + print(f"{DU} not a recognized unit system. Using AU as a default.") + self.param['DU2M'] = constants.AU2M + self.DU_name = "AU" + + if TU2S is not None or TU is not None: + TU2S_old = self.param.pop('TU2S',None) + if TU2S is not None: + self.param['TU2S'] = TU2S + self.TU_name = None else: - print(f"{TU} not a recognized unit system. Using YR as a default.") - self.param['TU2S'] = constants.YR2S - self.TU_name = "y" + if TU.upper() == "YR" or TU.upper() == "Y" or TU.upper() == "YEAR" or TU.upper() == "YEARS": + self.param['TU2S'] = constants.YR2S + self.TU_name = "y" + elif TU.upper() == "DAY" or TU.upper() == "D" or TU.upper() == "JD" or TU.upper() == "DAYS": + self.param['TU2S'] = constants.JD2S + self.TU_name = "Day" + elif TU.upper() == "S" or TU.upper() == "SECONDS" or TU.upper() == "SEC": + self.param['TU2S'] = 1.0 + self.TU_name = "s" + else: + print(f"{TU} not a recognized unit system. Using YR as a default.") + self.param['TU2S'] = constants.YR2S + self.TU_name = "y" self.VU_name = f"{self.DU_name}/{self.TU_name}" - self.GC = constants.GC * self.param['TU2S']**2 * self.param['MU2KG'] / self.param['DU2M']**3 + self.GU = constants.GC * self.param['TU2S']**2 * self.param['MU2KG'] / self.param['DU2M']**3 - self.update_param_units(MU2KG_old, DU2M_old, TU2S_old) + if recompute_values: + self.update_param_units(MU2KG_old, DU2M_old, TU2S_old) if self.verbose: if self.MU_name is None or self.DU_name is None or self.TU_name is None: @@ -253,7 +580,7 @@ def update_param_units(self,MU2KG_old,DU2M_old,TU2S_old): return - def set_simulation_range(self,rmin=None,rmax=None): + def set_distance_range(self,rmin=None,rmax=None): """ Sets the minimum and maximum distances of the simulation. @@ -287,6 +614,8 @@ def set_simulation_range(self,rmin=None,rmax=None): return + + def add(self, plname, date=date.today().isoformat(), idval=None): """ Adds a solar system body to an existing simulation DataSet. From ccb405feb287cb0fe1dbd09668003d4c7486c7bc Mon Sep 17 00:00:00 2001 From: David Minton Date: Tue, 8 Nov 2022 14:58:59 -0500 Subject: [PATCH 04/17] Updated parameter "YES" and "NO" values to boolean on the Python side. Wrote a feature setter and getter that sets the boolean values of the simulation parameters --- python/swiftest/swiftest/io.py | 179 ++++++++++++------ python/swiftest/swiftest/simulation_class.py | 180 +++++++++++++++++-- 2 files changed, 289 insertions(+), 70 deletions(-) diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index af74e34a4..b12d32cd3 100644 --- a/python/swiftest/swiftest/io.py +++ b/python/swiftest/swiftest/io.py @@ -17,8 +17,72 @@ import tempfile import re -newfeaturelist = ("FRAGMENTATION", "ROTATION", "TIDES", "ENERGY", "GR", "YARKOVSKY", "YORP", "IN_FORM") +newfeaturelist = ("FRAGMENTATION", "ROTATION", "TIDES", "ENERGY", "GR", "YARKOVSKY", "YORP", "IN_FORM", "SEED", "INTERACTION_LOOPS", "ENCOUNTER_CHECK") string_varnames = ["name", "particle_type", "status", "origin_type"] +bool_param = ["CHK_CLOSE", "EXTRA_FORCE", "RHILL_PRESENT", "BIG_DISCARD", "FRAGMENTATION", "ROTATION", "TIDES", "ENERGY", "GR", "YARKOVSKY", "YORP"] + +def bool2yesno(boolval): + """ + Converts a boolean into a string of either "YES" or "NO". + + Parameters + ---------- + boolval : bool + Input value + + Returns + ------- + {"YES","NO"} + + """ + if boolval: + return "YES" + else: + return "NO" + +def bool2tf(boolval): + """ + Converts a boolean into a string of either "T" or "F". + + Parameters + ---------- + boolval : bool + Input value + + Returns + ------- + {"T","F"} + + """ + if boolval: + return "T" + else: + return "F" + +def str2bool(input_str): + """ + Converts a string into an equivalent boolean. + + Parameters + ---------- + input_str : {"YES", "Y", "T", "TRUE", ".TRUE.", "NO", "N", "F", "FALSE", ".FALSE."} + Input string. Input is case-insensitive. + + Returns + ------- + {True, False} + + """ + valid_true = ["YES", "Y", "T", "TRUE", ".TRUE."] + valid_false = ["NO", "N", "F", "FALSE", ".FALSE."] + if input_str.upper() in valid_true: + return True + elif input_str.lower() in valid_false: + return False + else: + raise ValueError(f"{input_str} cannot is not recognized as boolean") + + def real2float(realstr): """ @@ -27,7 +91,7 @@ def real2float(realstr): Parameters ---------- - realstr : string + realstr : str Fortran-generated ASCII string of a real value. Returns @@ -89,21 +153,15 @@ def read_swiftest_param(param_file_name, param, verbose=True): param['DU2M'] = real2float(param['DU2M']) param['MU2KG'] = real2float(param['MU2KG']) param['TU2S'] = real2float(param['TU2S']) - param['EXTRA_FORCE'] = param['EXTRA_FORCE'].upper() - param['BIG_DISCARD'] = param['BIG_DISCARD'].upper() - param['CHK_CLOSE'] = param['CHK_CLOSE'].upper() - param['RHILL_PRESENT'] = param['RHILL_PRESENT'].upper() - param['FRAGMENTATION'] = param['FRAGMENTATION'].upper() - param['ROTATION'] = param['ROTATION'].upper() - param['TIDES'] = param['TIDES'].upper() - param['ENERGY'] = param['ENERGY'].upper() - param['GR'] = param['GR'].upper() param['INTERACTION_LOOPS'] = param['INTERACTION_LOOPS'].upper() param['ENCOUNTER_CHECK'] = param['ENCOUNTER_CHECK'].upper() if 'GMTINY' in param: param['GMTINY'] = real2float(param['GMTINY']) if 'MIN_GMFRAG' in param: param['MIN_GMFRAG'] = real2float(param['MIN_GMFRAG']) + for b in bool_param: + if b in param: + param[b] = str2bool(param[b]) except IOError: print(f"{param_file_name} not found.") return param @@ -139,7 +197,7 @@ def read_swifter_param(param_file_name, verbose=True): 'OUT_STAT': "NEW", 'J2': "0.0", 'J4': "0.0", - 'CHK_CLOSE': 'NO', + 'CHK_CLOSE': False, 'CHK_RMIN': "-1.0", 'CHK_RMAX': "-1.0", 'CHK_EJECT': "-1.0", @@ -147,9 +205,9 @@ def read_swifter_param(param_file_name, verbose=True): 'CHK_QMIN_COORD': "HELIO", 'CHK_QMIN_RANGE': "", 'ENC_OUT': "", - 'EXTRA_FORCE': 'NO', - 'BIG_DISCARD': 'NO', - 'RHILL_PRESENT': 'NO', + 'EXTRA_FORCE': False, + 'BIG_DISCARD': False, + 'RHILL_PRESENT': False, 'C': "-1.0", } @@ -179,10 +237,9 @@ def read_swifter_param(param_file_name, verbose=True): param['CHK_RMAX'] = real2float(param['CHK_RMAX']) param['CHK_EJECT'] = real2float(param['CHK_EJECT']) param['CHK_QMIN'] = real2float(param['CHK_QMIN']) - param['EXTRA_FORCE'] = param['EXTRA_FORCE'].upper() - param['BIG_DISCARD'] = param['BIG_DISCARD'].upper() - param['CHK_CLOSE'] = param['CHK_CLOSE'].upper() - param['RHILL_PRESENT'] = param['RHILL_PRESENT'].upper() + for b in bool_param: + if b in param: + param[b] = str2bool(param[b]) if param['C'] != '-1.0': param['C'] = real2float(param['C']) else: @@ -349,10 +406,18 @@ def write_labeled_param(param, param_file_name): # Print the list of key/value pairs in the preferred order for key in keylist: val = ptmp.pop(key, None) - if val is not None: print(f"{key:<16} {val}", file=outfile) + if val is not None: + if type(val) is bool: + print(f"{key:<16} {bool2yesno(val)}", file=outfile) + else: + print(f"{key:<16} {val}", file=outfile) # Print the remaining key/value pairs in whatever order for key, val in ptmp.items(): - if val != "": print(f"{key:<16} {val}", file=outfile) + if val != "": + if type(val) is bool: + print(f"{key:<16} {bool2yesno(val)}", file=outfile) + else: + print(f"{key:<16} {val}", file=outfile) outfile.close() return @@ -483,12 +548,12 @@ def make_swiftest_labels(param): tlab.append('capm') plab = tlab.copy() plab.append('Gmass') - if param['CHK_CLOSE'] == 'YES': + if param['CHK_CLOSE']: plab.append('radius') - if param['RHILL_PRESENT'] == 'YES': + if param['RHILL_PRESENT']: plab.append('rhill') clab = ['Gmass', 'radius', 'j2rp2', 'j4rp4'] - if param['ROTATION'] == 'YES': + if param['ROTATION']: clab.append('Ip1') clab.append('Ip2') clab.append('Ip3') @@ -501,7 +566,7 @@ def make_swiftest_labels(param): plab.append('rotx') plab.append('roty') plab.append('rotz') - if param['TIDES'] == 'YES': + if param['TIDES']: clab.append('k2') clab.append('Q') plab.append('k2') @@ -588,14 +653,14 @@ def swiftest_stream(f, param): Rcb = f.read_reals(np.float64) J2cb = f.read_reals(np.float64) J4cb = f.read_reals(np.float64) - if param['ROTATION'] == 'YES': + if param['ROTATION']: Ipcbx = f.read_reals(np.float64) Ipcby = f.read_reals(np.float64) Ipcbz = f.read_reals(np.float64) rotcbx = f.read_reals(np.float64) rotcby = f.read_reals(np.float64) rotcbz = f.read_reals(np.float64) - if param['TIDES'] == 'YES': + if param['TIDES']: k2cb = f.read_reals(np.float64) Qcb = f.read_reals(np.float64) if npl[0] > 0: @@ -619,17 +684,17 @@ def swiftest_stream(f, param): p11 = f.read_reals(np.float64) p12 = f.read_reals(np.float64) GMpl = f.read_reals(np.float64) - if param['RHILL_PRESENT'] == 'YES': + if param['RHILL_PRESENT']: rhill = f.read_reals(np.float64) Rpl = f.read_reals(np.float64) - if param['ROTATION'] == 'YES': + if param['ROTATION']: Ipplx = f.read_reals(np.float64) Ipply = f.read_reals(np.float64) Ipplz = f.read_reals(np.float64) rotplx = f.read_reals(np.float64) rotply = f.read_reals(np.float64) rotplz = f.read_reals(np.float64) - if param['TIDES'] == 'YES': + if param['TIDES']: k2pl = f.read_reals(np.float64) Qpl = f.read_reals(np.float64) if ntp[0] > 0: @@ -679,14 +744,14 @@ def swiftest_stream(f, param): tpid = np.empty(0) tpnames = np.empty(0) cvec = np.array([Mcb, Rcb, J2cb, J4cb]) - if param['RHILL_PRESENT'] == 'YES': + if param['RHILL_PRESENT']: if npl > 0: pvec = np.vstack([pvec, rhill]) - if param['ROTATION'] == 'YES': + if param['ROTATION']: cvec = np.vstack([cvec, Ipcbx, Ipcby, Ipcbz, rotcbx, rotcby, rotcbz]) if npl > 0: pvec = np.vstack([pvec, Ipplx, Ipply, Ipplz, rotplx, rotply, rotplz]) - if param['TIDES'] == 'YES': + if param['TIDES']: cvec = np.vstack([cvec, k2cb, Qcb]) if npl > 0: pvec = np.vstack([pvec, k2pl, Qpl]) @@ -1018,14 +1083,14 @@ def swiftest_xr2infile(ds, param, in_type="NETCDF_DOUBLE", infile_name=None,fram tp = frame.where(np.isnan(frame['Gmass']), drop=True).drop_vars(['Gmass', 'radius', 'j2rp2', 'j4rp4'],errors="ignore") GMSun = np.double(cb['Gmass']) - if param['CHK_CLOSE'] == 'YES': + if param['CHK_CLOSE']: RSun = np.double(cb['radius']) else: RSun = param['CHK_RMIN'] J2 = np.double(cb['j2rp2']) J4 = np.double(cb['j4rp4']) cbname = cb['name'].values[0] - if param['ROTATION'] == 'YES': + if param['ROTATION']: Ip1cb = np.double(cb['Ip1']) Ip2cb = np.double(cb['Ip2']) Ip3cb = np.double(cb['Ip3']) @@ -1042,7 +1107,7 @@ def swiftest_xr2infile(ds, param, in_type="NETCDF_DOUBLE", infile_name=None,fram print(RSun, file=cbfile) print(J2, file=cbfile) print(J4, file=cbfile) - if param['ROTATION'] == 'YES': + if param['ROTATION']: print(Ip1cb, Ip2cb, Ip3cb, file=cbfile) print(rotxcb, rotycb, rotzcb, file=cbfile) cbfile.close() @@ -1051,11 +1116,11 @@ def swiftest_xr2infile(ds, param, in_type="NETCDF_DOUBLE", infile_name=None,fram print(pl.id.count().values, file=plfile) for i in pl.id: pli = pl.sel(id=i) - if param['RHILL_PRESENT'] == 'YES': + if param['RHILL_PRESENT']: print(pli['name'].values[0], pli['Gmass'].values[0], pli['rhill'].values[0], file=plfile) else: print(pli['name'].values[0], pli['Gmass'].values[0], file=plfile) - if param['CHK_CLOSE'] == 'YES': + if param['CHK_CLOSE']: print(pli['radius'].values[0], file=plfile) if param['IN_FORM'] == 'XV': print(pli['xhx'].values[0], pli['xhy'].values[0], pli['xhz'].values[0], file=plfile) @@ -1065,7 +1130,7 @@ def swiftest_xr2infile(ds, param, in_type="NETCDF_DOUBLE", infile_name=None,fram print(pli['capom'].values[0], pli['omega'].values[0], pli['capm'].values[0], file=plfile) else: print(f"{param['IN_FORM']} is not a valid input format type.") - if param['ROTATION'] == 'YES': + if param['ROTATION']: print(pli['Ip1'].values[0], pli['Ip2'].values[0], pli['Ip3'].values[0], file=plfile) print(pli['rotx'].values[0], pli['roty'].values[0], pli['rotz'].values[0], file=plfile) plfile.close() @@ -1114,7 +1179,7 @@ def swifter_xr2infile(ds, param, framenum=-1): tp = frame.where(np.isnan(frame['Gmass']), drop=True).drop_vars(['Gmass', 'radius', 'j2rp2', 'j4rp4']) GMSun = np.double(cb['Gmass']) - if param['CHK_CLOSE'] == 'YES': + if param['CHK_CLOSE']: RSun = np.double(cb['radius']) else: RSun = param['CHK_RMIN'] @@ -1130,11 +1195,11 @@ def swifter_xr2infile(ds, param, framenum=-1): print('0.0 0.0 0.0', file=plfile) for i in pl.id: pli = pl.sel(id=i) - if param['RHILL_PRESENT'] == "YES": + if param['RHILL_PRESENT']: print(i.values, pli['Gmass'].values, pli['rhill'].values, file=plfile) else: print(i.values, pli['Gmass'].values, file=plfile) - if param['CHK_CLOSE'] == "YES": + if param['CHK_CLOSE']: print(pli['radius'].values, file=plfile) print(pli['xhx'].values, pli['xhy'].values, pli['xhz'].values, file=plfile) print(pli['vhx'].values, pli['vhy'].values, pli['vhz'].values, file=plfile) @@ -1181,10 +1246,10 @@ def swift2swifter(swift_param, plname="", tpname="", conversion_questions={}): intxt = input("Is this a SyMBA input file with RHILL values in pl.in? (y/N)> ") if intxt.upper() == 'Y': isSyMBA = True - swifter_param['RHILL_PRESENT'] = 'YES' + swifter_param['RHILL_PRESENT'] = True else: isSyMBA = False - swifter_param['RHILL_PRESENT'] = 'NO' + swifter_param['RHILL_PRESENT'] = False isDouble = conversion_questions.get('DOUBLE', None) if not isDouble: @@ -1212,9 +1277,9 @@ def swift2swifter(swift_param, plname="", tpname="", conversion_questions={}): swifter_param['OUT_FORM'] = 'XV' if swift_param['LCLOSE'] == "T": - swifter_param['CHK_CLOSE'] = "YES" + swifter_param['CHK_CLOSE'] = True else: - swifter_param['CHK_CLOSE'] = "NO" + swifter_param['CHK_CLOSE'] = False swifter_param['CHK_RMIN'] = swift_param['RMIN'] swifter_param['CHK_RMAX'] = swift_param['RMAX'] @@ -1253,17 +1318,17 @@ def swift2swifter(swift_param, plname="", tpname="", conversion_questions={}): if not intxt: intxt = input("EXTRA_FORCE: Use additional user-specified force routines? (y/N)> ") if intxt.upper() == 'Y': - swifter_param['EXTRA_FORCE'] = 'YES' + swifter_param['EXTRA_FORCE'] = True else: - swifter_param['EXTRA_FORCE'] = 'NO' + swifter_param['EXTRA_FORCE'] = False intxt = conversion_questions.get('BIG_DISCARD', None) if not intxt: intxt = input("BIG_DISCARD: include data for all bodies > GMTINY for each discard record? (y/N)> ") if intxt.upper() == 'Y': - swifter_param['BIG_DISCARD'] = 'YES' + swifter_param['BIG_DISCARD'] = True else: - swifter_param['BIG_DISCARD'] = 'NO' + swifter_param['BIG_DISCARD'] = False # Convert the PL file if plname == '': @@ -1309,11 +1374,11 @@ def swift2swifter(swift_param, plname="", tpname="", conversion_questions={}): else: if swift_param['LCLOSE'] == "T": plrad = real2float(i_list[1]) - if swifter_param['RHILL_PRESENT'] == 'YES': + if swifter_param['RHILL_PRESENT']: print(n + 1, GMpl, rhill, file=plnew) else: print(n + 1, GMpl, file=plnew) - if swifter_param['CHK_CLOSE'] == 'YES': + if swifter_param['CHK_CLOSE']: print(plrad, file=plnew) line = plold.readline() i_list = [i for i in re.split(' +|\t',line) if i.strip()] @@ -1433,12 +1498,12 @@ def swifter2swiftest(swifter_param, plname="", tpname="", cbname="", conversion_ i_list = [i for i in re.split(' +|\t',line) if i.strip()] idnum = int(i_list[0]) GMpl = real2float(i_list[1]) - if swifter_param['RHILL_PRESENT'] == 'YES': + if swifter_param['RHILL_PRESENT']: rhill = real2float(i_list[2]) print(idnum, GMpl, rhill, file=plnew) else: print(idnum, GMpl, file=plnew) - if swifter_param['CHK_CLOSE'] == 'YES': + if swifter_param['CHK_CLOSE']: line = plold.readline() i_list = [i for i in re.split(' +|\t',line) if i.strip()] plrad = real2float(i_list[0]) @@ -1608,7 +1673,7 @@ def swifter2swiftest(swifter_param, plname="", tpname="", cbname="", conversion_ # Remove the unneeded parameters if 'C' in swiftest_param: - swiftest_param['GR'] = 'YES' + swiftest_param['GR'] = True swiftest_param.pop('C', None) swiftest_param.pop('J2', None) swiftest_param.pop('J4', None) @@ -1697,7 +1762,7 @@ def swiftest2swifter_param(swiftest_param, J2=0.0, J4=0.0): TU2S = swifter_param.pop("TU2S", 1.0) GR = swifter_param.pop("GR", None) if GR is not None: - if GR == 'YES': + if GR: swifter_param['C'] = swiftest.einsteinC * np.longdouble(TU2S) / np.longdouble(DU2M) for key in newfeaturelist: tmp = swifter_param.pop(key, None) @@ -1769,7 +1834,7 @@ def swifter2swift_param(swifter_param, J2=0.0, J4=0.0): swift_param['BINARY_OUTPUTFILE'] = swifter_param['BIN_OUT'] swift_param['STATUS_FLAG_FOR_OPEN_STATEMENTS'] = swifter_param['OUT_STAT'] - if swifter_param['CHK_CLOSE'] == "YES": + if swifter_param['CHK_CLOSE']: swift_param['LCLOSE'] = "T" else: swift_param['LCLOSE'] = "F" diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 38fa44d65..f95e8b426 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -23,6 +23,7 @@ from typing import ( Literal, Dict, + List ) @@ -49,11 +50,14 @@ def __init__(self, TU2S: float | None = None, rmin: float = constants.RSun / constants.AU2M, rmax: float = 10000.0, - close_encounter_check: bool =True, - general_relativity: bool =True, - fragmentation: bool =True, + close_encounter_check: bool = True, + general_relativity: bool = True, + fragmentation: bool = True, rotation: bool = True, compute_conservation_values: bool = False, + extra_force: bool = False, + big_discard: bool = False, + rhill_present: bool = False, interaction_loops: Literal["TRIANGULAR","FLAT","ADAPTIVE"] = "TRIANGULAR", encounter_check_loops: Literal["TRIANGULAR","SORTSWEEP","ADAPTIVE"] = "TRIANGULAR", verbose: bool = True @@ -154,6 +158,12 @@ def __init__(self, compute_conservation_values : bool, default False Turns on the computation of energy, angular momentum, and mass conservation and reports the values every output step of a running simulation. + extra_force: bool, default False + Turns on user-defined force function. + big_discard: bool, default False + Includes big bodies when performing a discard (Swifter only) + rhill_present: bool, default False + Include the Hill's radius with the input files. interaction_loops : {"TRIANGULAR","FLAT","ADAPTIVE"}, default "TRIANGULAR" *Swiftest Experimental feature* Specifies which algorithm to use for the computation of body-body gravitational forces. @@ -190,15 +200,6 @@ def __init__(self, 'ISTEP_OUT': 1, 'ISTEP_DUMP': 1, 'CHK_QMIN_COORD': "HELIO", - 'EXTRA_FORCE': "NO", - 'BIG_DISCARD': "NO", - 'CHK_CLOSE': "YES", - 'RHILL_PRESENT': "YES", - 'FRAGMENTATION': "YES", - 'ROTATION': "YES", - 'TIDES': "NO", - 'ENERGY': "YES", - 'GR': "YES", 'INTERACTION_LOOPS': interaction_loops, 'ENCOUNTER_CHECK': encounter_check_loops } @@ -217,7 +218,17 @@ def __init__(self, self.set_output_files(output_file_type=output_file_type, output_file_name=output_file_name, - output_format=output_format ) + output_format=output_format) + + self.set_features(close_encounter_check=close_encounter_check, + general_relativity=general_relativity, + fragmentation=fragmentation, + rotation=rotation, + compute_conservation_values=compute_conservation_values, + extra_force=extra_force, + big_discard=big_discard, + rhill_present=rhill_present, + verbose = False) # If the parameter file is in a different location than the current working directory, we will need # to use it to properly open bin files @@ -236,6 +247,149 @@ def __init__(self, print(f"BIN_OUT file {binpath} not found.") return + def set_features(self, + close_encounter_check: bool | None = None, + general_relativity: bool | None = None, + fragmentation: bool | None = None, + rotation: bool | None = None, + compute_conservation_values: bool | None=None, + extra_force: bool | None = None, + big_discard: bool | None = None, + rhill_present: bool | None = None, + tides: bool | None = None, + verbose: bool | None = None, + ): + """ + Turns on or off various features of a simulation. + + 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 + in initial conditions. + 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. + rotation : bool, optional + If set to True, this turns on rotation tracking and radius, rotation vector, and moments of inertia values + must be included in the initial conditions. + This argument only applies to Swiftest-SyMBA simulations. It will be ignored otherwise. + compute_conservation_values : bool, optional + Turns on the computation of energy, angular momentum, and mass conservation and reports the values + every output step of a running simulation. + extra_force: bool, optional + Turns on user-defined force function. + big_discard: bool, optional + Includes big bodies when performing a discard (Swifter only) + rhill_present: bool, optional + Include the Hill's radius with the input files. + tides: bool, optional + Turns on tidal model (IN DEVELOPMENT - IGNORED) + Yarkovsky: bool, optional + Turns on Yarkovsky model (IN DEVELOPMENT - IGNORED) + YORP: bool, optional + Turns on YORP model (IN DEVELOPMENT - IGNORED) + verbose: bool, optional + If passed, it will override the Simulation object's verbose flag + + Returns + ------- + Sets the appropriate parameters in the self.param dictionary + + """ + + update_list = [] + if close_encounter_check is not None: + self.param["CHK_CLOSE"] = close_encounter_check + update_list.append("close_encounter_check") + + if general_relativity is not None: + self.param["GR"] = general_relativity + update_list.append("general_relativity") + + if fragmentation is not None: + self.param['FRAGMENTATION'] = fragmentation + update_list.append("fragmentation") + + if rotation is not None: + self.param['ROTATION'] = rotation + update_list.append("rotation") + + if compute_conservation_values is not None: + self.param["ENERGY"] = compute_conservation_values + update_list.append("compute_conservation_values") + + if extra_force is not None: + self.param["EXTRA_FORCE"] = extra_force + update_list.append("extra_force") + + if big_discard is not None: + if self.codename != "Swifter": + self.param["BIG_DISCARD"] = False + else: + self.param["BIG_DISCARD"] = big_discard + update_list.append("big_discard") + + if rhill_present is not None: + self.param["RHILL_PRESENT"] = rhill_present + update_list.append("rhill_present") + + if verbose is None: + verbose = self.verbose + if verbose: + if len(update_list) > 0: + feature_dict = self.get_features(update_list) + return + + + def get_features(self,feature: str | List[str] | None = None): + """ + + Returns a subset of the parameter dictionary containing the current value of the feature boolean values. + If the verbose option is set in the Simulation object, then it will also print the values. + + Parameters + ---------- + feature : str | List[str], optional + A single string or list of strings containing the names of the features to extract. Default is all of: + ["close_encounter_check", "general_relativity", "fragmentation", "rotation", "compute_conservation_values"] + + Returns + ------- + feature_dict : dict + A dictionary containing the requested features. + + """ + + valid_features = {"close_encounter_check": "CHK_CLOSE", + "general_relativity": "GR", + "fragmentation": "FRAGMENTATION", + "rotation": "ROTATION", + "compute_conservation_values": "ENERGY", + "extra_force": "EXTRA_FORCE", + "big_discard": "BIG_DISCARD", + "rhill_present": "RHILL_PRESENT" + } + + if feature is None: + feature_list = valid_features.keys() + elif type(feature) is str: + feature_list = [feature] + else: + # Only allow features to be checked if they are valid. Otherwise ignore. + feature_list = [feat for feat in feature if feat in set(valid_features.keys())] + + # Extract the feature dictionary + feature_dict = {valid_features[feat]:self.param[valid_features[feat]] for feat in feature_list} + + if self.verbose: + print("Simulation feature parameters:") + for key,val in feature_dict.items(): + print(f"{key:<16} {val}") + + return feature_dict def set_init_cond_files(self, init_cond_file_type: Literal["NETCDF_DOUBLE", "NETCDF_FLOAT", "ASCII"], From 2ac1e39d5a4fd012f381c50fecd27387f4369c2f Mon Sep 17 00:00:00 2001 From: David Minton Date: Tue, 8 Nov 2022 15:00:25 -0500 Subject: [PATCH 05/17] Changed from "_features" to "_feature" before it gets too late to change it --- python/swiftest/swiftest/simulation_class.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index f95e8b426..622e36fe2 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -220,7 +220,7 @@ def __init__(self, output_file_name=output_file_name, output_format=output_format) - self.set_features(close_encounter_check=close_encounter_check, + self.set_feature(close_encounter_check=close_encounter_check, general_relativity=general_relativity, fragmentation=fragmentation, rotation=rotation, @@ -247,7 +247,7 @@ def __init__(self, print(f"BIN_OUT file {binpath} not found.") return - def set_features(self, + def set_feature(self, close_encounter_check: bool | None = None, general_relativity: bool | None = None, fragmentation: bool | None = None, @@ -340,11 +340,11 @@ def set_features(self, verbose = self.verbose if verbose: if len(update_list) > 0: - feature_dict = self.get_features(update_list) + feature_dict = self.get_feature(update_list) return - def get_features(self,feature: str | List[str] | None = None): + def get_feature(self,feature: str | List[str] | None = None): """ Returns a subset of the parameter dictionary containing the current value of the feature boolean values. From a784afc04e82a7fc32c53cdd42fde3424f84dfcc Mon Sep 17 00:00:00 2001 From: David Minton Date: Tue, 8 Nov 2022 17:56:20 -0500 Subject: [PATCH 06/17] Added simulation time setter and getter methods --- .../Basic_Simulation/initial_conditions.py | 7 +- python/swiftest/swiftest/io.py | 37 ++- python/swiftest/swiftest/simulation_class.py | 236 ++++++++++++++++-- 3 files changed, 252 insertions(+), 28 deletions(-) diff --git a/examples/Basic_Simulation/initial_conditions.py b/examples/Basic_Simulation/initial_conditions.py index dceb33cc7..2b3693b0c 100644 --- a/examples/Basic_Simulation/initial_conditions.py +++ b/examples/Basic_Simulation/initial_conditions.py @@ -32,12 +32,9 @@ # Initialize the simulation object as a variable sim = swiftest.Simulation(init_cond_file_type="ASCII") +sim.set_simulation_time(tstart=0.0, tstop=10.0, dt=0.005, tstep_out=1.0,verbose=True) + # Add parameter attributes to the simulation object -sim.param['T0'] = 0.0 -sim.param['TSTOP'] = 10 -sim.param['DT'] = 0.005 -sim.param['ISTEP_OUT'] = 200 -sim.param['ISTEP_DUMP'] = 200 sim.param['GMTINY'] = 1e-6 sim.param['MIN_GMFRAG'] = 1e-9 diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index b12d32cd3..28303f8f8 100644 --- a/python/swiftest/swiftest/io.py +++ b/python/swiftest/swiftest/io.py @@ -17,9 +17,38 @@ import tempfile import re -newfeaturelist = ("FRAGMENTATION", "ROTATION", "TIDES", "ENERGY", "GR", "YARKOVSKY", "YORP", "IN_FORM", "SEED", "INTERACTION_LOOPS", "ENCOUNTER_CHECK") +# This defines features that are new in Swiftest and not in Swifter (for conversion between param.in files) +newfeaturelist = ("RESTART", + "FRAGMENTATION", + "ROTATION", + "TIDES", + "ENERGY", + "GR", + "YARKOVSKY", + "YORP", + "IN_FORM", + "SEED", + "INTERACTION_LOOPS", + "ENCOUNTER_CHECK", + "TSTART") + +# This list defines features that are booleans, so must be converted to/from string when writing/reading from file +bool_param = ["RESTART", + "CHK_CLOSE", + "EXTRA_FORCE", + "RHILL_PRESENT", + "BIG_DISCARD", + "FRAGMENTATION", + "ROTATION", + "TIDES", + "ENERGY", + "GR", + "YARKOVSKY", + "YORP"] + +# 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. string_varnames = ["name", "particle_type", "status", "origin_type"] -bool_param = ["CHK_CLOSE", "EXTRA_FORCE", "RHILL_PRESENT", "BIG_DISCARD", "FRAGMENTATION", "ROTATION", "TIDES", "ENERGY", "GR", "YARKOVSKY", "YORP"] def bool2yesno(boolval): """ @@ -144,6 +173,7 @@ def read_swiftest_param(param_file_name, param, verbose=True): param['IN_TYPE'] = param['IN_TYPE'].upper() param['IN_FORM'] = param['IN_FORM'].upper() param['T0'] = real2float(param['T0']) + param['TSTART'] = real2float(param['TSTART']) param['TSTOP'] = real2float(param['TSTOP']) param['DT'] = real2float(param['DT']) param['CHK_RMIN'] = real2float(param['CHK_RMIN']) @@ -401,7 +431,8 @@ def write_labeled_param(param, param_file_name): 'CHK_QMIN_RANGE', 'MU2KG', 'TU2S', - 'DU2M' ] + 'DU2M', + 'RESTART'] ptmp = param.copy() # Print the list of key/value pairs in the preferred order for key in keylist: diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 622e36fe2..bda2d521d 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -58,6 +58,7 @@ def __init__(self, extra_force: bool = False, big_discard: bool = False, rhill_present: bool = False, + restart: bool = False, interaction_loops: Literal["TRIANGULAR","FLAT","ADAPTIVE"] = "TRIANGULAR", encounter_check_loops: Literal["TRIANGULAR","SORTSWEEP","ADAPTIVE"] = "TRIANGULAR", verbose: bool = True @@ -164,6 +165,10 @@ def __init__(self, Includes big bodies when performing a discard (Swifter only) rhill_present: bool, default False Include the Hill's radius with the input files. + restart : bool, default False + If true, will restart an old run. The file given by `output_file_name` must exist before the run can + execute. If false, will start a new run. If the file given by `output_file_name` exists, it will be replaced + when the run is executed. interaction_loops : {"TRIANGULAR","FLAT","ADAPTIVE"}, default "TRIANGULAR" *Swiftest Experimental feature* Specifies which algorithm to use for the computation of body-body gravitational forces. @@ -205,6 +210,7 @@ def __init__(self, } self.codename = codename self.verbose = verbose + self.restart = restart self.set_distance_range(rmin=rmin, rmax=rmax) @@ -221,14 +227,15 @@ def __init__(self, output_format=output_format) self.set_feature(close_encounter_check=close_encounter_check, - general_relativity=general_relativity, - fragmentation=fragmentation, - rotation=rotation, - compute_conservation_values=compute_conservation_values, - extra_force=extra_force, - big_discard=big_discard, - rhill_present=rhill_present, - verbose = False) + general_relativity=general_relativity, + fragmentation=fragmentation, + rotation=rotation, + compute_conservation_values=compute_conservation_values, + extra_force=extra_force, + big_discard=big_discard, + rhill_present=rhill_present, + restart=restart, + verbose = False) # If the parameter file is in a different location than the current working directory, we will need # to use it to properly open bin files @@ -247,18 +254,198 @@ def __init__(self, print(f"BIN_OUT file {binpath} not found.") return + def set_simulation_time(self, + t0: float | None = None, + tstart: float | None = None, + tstop: float | None = None, + dt: float | None = None, + istep_out : int | None = None, + tstep_out : float | None = None, + istep_dump : int | None = None, + verbose : bool | None = None, + ): + """ + + Parameters + ---------- + t0 : float, optional + The reference time for the start of the simulation. Defaults is 0.0 + tstart : float, optional + The start time for a restarted simulation. For a new simulation, tstart will be set to t0 automatically. + tstop : float, optional + The stopping time for a simulation. `tstop` must be greater than `tstart`. + dt : float, optional + The step size of the simulation. `dt` must be less than or equal to `tstop-dstart`. + istep_out : int, optional + The number of time steps between outputs to file. *Note*: only `istep_out` or `toutput` can be set. + tstep_out : float, optional + The approximate time between when outputs are written to file. Passing this computes + `istep_out = floor(tstep_out/dt)`. *Note*: only `istep_out` or `toutput` can be set. + istep_dump : int, optional + The anumber of time steps between outputs to dump file. If not set, this will be set to the value of + `istep_out` (or the equivalent value determined by `tstep_out`) + verbose: bool, optional + If passed, it will override the Simulation object's verbose flag + + Returns + ------- + Sets the appropriate parameter variables in the parameter dictionary. + + """ + + update_list = [] + + + if t0 is None: + t0 = self.param.pop("T0", None) + if t0 is None: + t0 = 0.0 + else: + update_list.append("t0") + + if tstart is None: + tstart = self.param.pop("TSTART", None) + if tstart is None: + tstart = t0 + else: + update_list.append("tstart") + + self.param['T0'] = t0 + self.param['TSTART'] = tstart + + if tstop is None: + tstop = self.param.pop("TSTOP", None) + if tstop is None: + print("Error! tstop is not set!") + return + else: + update_list.append("tstop") + + if tstop <= tstart: + print("Error! tstop must be greater than tstart.") + return + + self.param['TSTOP'] = tstop + + if dt is None: + dt = self.param.pop("DT", None) + if dt is None: + print("Error! dt is not set!") + return + else: + update_list.append("dt") + + if dt > (tstop - tstart): + print("Error! dt must be smaller than tstop-tstart") + print(f"Setting dt = {tstop-tstart} instead of {dt}") + dt = tstop - tstart + + self.param['DT'] = dt + + if istep_out is None and tstep_out is None: + istep_out = self.param.pop("ISTEP_OUT", None) + if istep_out is None: + print("Error! istep_out is not set") + return + + if istep_out is not None and tstep_out is not None: + print("Error! istep_out and tstep_out cannot both be set") + return + + if tstep_out is not None: + istep_out = int(np.floor(tstep_out / dt)) + + self.param['ISTEP_OUT'] = istep_out + + if istep_dump is None: + istep_dump = istep_out + + self.param['ISTEP_DUMP'] = istep_dump + + update_list.extend(["istep_out", "tstep_out", "istep_dump"]) + + if verbose is None: + verbose = self.verbose + if verbose: + if len(update_list) > 0: + time_dict = self.get_simulation_time(update_list) + + return + + def get_simulation_time(self, param_key: str | List[str] | None = None): + """ + + Returns a subset of the parameter dictionary containing the current simulation time parameters. + If the verbose option is set in the Simulation object, then it will also print the values. + + Parameters + ---------- + param_key : str | List[str], optional + A single string or list of strings containing the names of the simulation time parameters to extract. + Default is all of: + ["t0", "tstart", "tstop", "dt", "istep_out", "tstep_out", "istep_dump"] + + Returns + ------- + time_dict : dict + A dictionary containing the requested parameters + + """ + + valid_var = {"t0": "T0", + "tstart": "TSTART", + "tstop": "TSTOP", + "dt": "DT", + "istep_out": "ISTEP_OUT", + "istep_dump": "ISTEP_DUMP", + } + + units = {"T0" : self.TU_name, + "TSTART" : self.TU_name, + "TSTOP" : self.TU_name, + "DT" : self.TU_name, + "TSTEP_OUT" : self.TU_name, + "ISTEP_OUT" : "", + "ISTEP_DUMP" : ""} + + if "tstep_out" in param_key: + istep_out = self.param['ISTEP_OUT'] + dt = self.param['DT'] + tstep_out = istep_out * dt + else: + tstep_out = None + + if param_key is None: + param_list = valid_var.keys() + elif type(param_key) is str: + param_list = [param_key] + else: + param_list = [k for k in param_key if k in set(valid_var.keys())] + + time_dict = {valid_var[k]: self.param[valid_var[k]] for k in param_list} + + if self.verbose: + print("Simulation feature parameters:") + for key, val in time_dict.items(): + print(f"{key:<16} {val} {units[key]}") + if tstep_out is not None: + print(f"{'TSTEP_OUT':<16} {tstep_out} {units['TSTEP_OUT']}") + + return time_dict + def set_feature(self, - close_encounter_check: bool | None = None, - general_relativity: bool | None = None, - fragmentation: bool | None = None, - rotation: bool | None = None, - compute_conservation_values: bool | None=None, - extra_force: bool | None = None, - big_discard: bool | None = None, - rhill_present: bool | None = None, - tides: bool | None = None, - verbose: bool | None = None, - ): + close_encounter_check: bool | None = None, + general_relativity: bool | None = None, + fragmentation: bool | None = None, + rotation: bool | None = None, + compute_conservation_values: bool | None=None, + extra_force: bool | None = None, + big_discard: bool | None = None, + rhill_present: bool | None = None, + restart: bool | None = None, + tides: bool | None = None, + verbose: bool | None = None, + ): """ Turns on or off various features of a simulation. @@ -291,6 +478,10 @@ def set_feature(self, Turns on Yarkovsky model (IN DEVELOPMENT - IGNORED) YORP: bool, optional Turns on YORP model (IN DEVELOPMENT - IGNORED) + restart : bool, default False + If true, will restart an old run. The file given by `output_file_name` must exist before the run can + execute. If false, will start a new run. If the file given by `output_file_name` exists, it will be replaced + when the run is executed. verbose: bool, optional If passed, it will override the Simulation object's verbose flag @@ -336,6 +527,10 @@ def set_feature(self, self.param["RHILL_PRESENT"] = rhill_present update_list.append("rhill_present") + if restart is not None: + self.param["RESTART"] = restart + update_list.append("restart") + if verbose is None: verbose = self.verbose if verbose: @@ -370,7 +565,8 @@ def get_feature(self,feature: str | List[str] | None = None): "compute_conservation_values": "ENERGY", "extra_force": "EXTRA_FORCE", "big_discard": "BIG_DISCARD", - "rhill_present": "RHILL_PRESENT" + "rhill_present": "RHILL_PRESENT", + "restart" : "RESTART" } if feature is None: From b6d3018bbe03236bfbac8a3751a37e0e3e79f974 Mon Sep 17 00:00:00 2001 From: David Minton Date: Tue, 8 Nov 2022 17:57:49 -0500 Subject: [PATCH 07/17] Removed unnecessary verbose flag --- examples/Basic_Simulation/initial_conditions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/Basic_Simulation/initial_conditions.py b/examples/Basic_Simulation/initial_conditions.py index 2b3693b0c..c9d0823af 100644 --- a/examples/Basic_Simulation/initial_conditions.py +++ b/examples/Basic_Simulation/initial_conditions.py @@ -32,7 +32,7 @@ # Initialize the simulation object as a variable sim = swiftest.Simulation(init_cond_file_type="ASCII") -sim.set_simulation_time(tstart=0.0, tstop=10.0, dt=0.005, tstep_out=1.0,verbose=True) +sim.set_simulation_time(tstart=0.0, tstop=10.0, dt=0.005, tstep_out=1.0) # Add parameter attributes to the simulation object sim.param['GMTINY'] = 1e-6 From 6aa38058d294cb64352696778d60dbe8cf1a3934 Mon Sep 17 00:00:00 2001 From: David Minton Date: Tue, 8 Nov 2022 17:59:24 -0500 Subject: [PATCH 08/17] Set tides parameter to be False so at least it's in the param list --- python/swiftest/swiftest/simulation_class.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index bda2d521d..d01da849a 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -531,6 +531,8 @@ def set_feature(self, self.param["RESTART"] = restart update_list.append("restart") + self.param["TIDES"] = False + if verbose is None: verbose = self.verbose if verbose: From 163f321560a28724c206d5e8bd95050652b06d55 Mon Sep 17 00:00:00 2001 From: David Minton Date: Tue, 8 Nov 2022 19:41:20 -0500 Subject: [PATCH 09/17] Switched to using boolean True instead of YES in the initial conditions generator --- python/swiftest/swiftest/init_cond.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/python/swiftest/swiftest/init_cond.py b/python/swiftest/swiftest/init_cond.py index 287e5ecae..0166e78a8 100644 --- a/python/swiftest/swiftest/init_cond.py +++ b/python/swiftest/swiftest/init_cond.py @@ -146,7 +146,7 @@ def solar_system_horizons(plname, idval, param, ephemerides_start_date, ds): Rpl = Rcb J2 = J2RP2 J4 = J4RP4 - if param['ROTATION'] == 'YES': + if param['ROTATION']: Ip1 = [Ipsun[0]] Ip2 = [Ipsun[1]] Ip3 = [Ipsun[2]] @@ -202,19 +202,19 @@ def solar_system_horizons(plname, idval, param, ephemerides_start_date, ds): if ispl: GMpl = [] GMpl.append(GMcb[0] / MSun_over_Mpl[plname]) - if param['CHK_CLOSE'] == 'YES': + if param['CHK_CLOSE']: Rpl = [] Rpl.append(planetradius[plname] * DCONV) else: Rpl = None # Generate planet value vectors - if (param['RHILL_PRESENT'] == 'YES'): + if (param['RHILL_PRESENT']): rhill = [] rhill.append(pldata[plname].elements()['a'][0] * DCONV * (3 * MSun_over_Mpl[plname]) ** (-THIRDLONG)) else: rhill = None - if (param['ROTATION'] == 'YES'): + if (param['ROTATION']): Ip1 = [] Ip2 = [] Ip3 = [] @@ -304,7 +304,7 @@ def vec2xr(param, idvals, namevals, v1, v2, v3, v4, v5, v6, GMpl=None, Rpl=None, else: iscb = False - if param['ROTATION'] == 'YES': + if param['ROTATION']: if Ip1 is None: Ip1 = np.full_like(v1, 0.4) if Ip2 is None: @@ -325,10 +325,10 @@ def vec2xr(param, idvals, namevals, v1, v2, v3, v4, v5, v6, GMpl=None, Rpl=None, else: ispl = False - if ispl and param['CHK_CLOSE'] == 'YES' and Rpl is None: + if ispl and param['CHK_CLOSE'] and Rpl is None: print("Massive bodies need a radius value.") return None - if ispl and rhill is None and param['RHILL_PRESENT'] == 'YES': + if ispl and rhill is None and param['RHILL_PRESENT']: print("rhill is required.") return None @@ -342,7 +342,7 @@ def vec2xr(param, idvals, namevals, v1, v2, v3, v4, v5, v6, GMpl=None, Rpl=None, if iscb: label_float = clab.copy() vec_float = np.vstack([GMpl,Rpl,J2,J4]) - if param['ROTATION'] == 'YES': + if param['ROTATION']: vec_float = np.vstack([vec_float, Ip1, Ip2, Ip3, rotx, roty, rotz]) particle_type = "Central Body" else: @@ -350,11 +350,11 @@ def vec2xr(param, idvals, namevals, v1, v2, v3, v4, v5, v6, GMpl=None, Rpl=None, if ispl: label_float = plab.copy() vec_float = np.vstack([vec_float, GMpl]) - if param['CHK_CLOSE'] == 'YES': + if param['CHK_CLOSE']: vec_float = np.vstack([vec_float, Rpl]) - if param['RHILL_PRESENT'] == 'YES': + if param['RHILL_PRESENT']: vec_float = np.vstack([vec_float, rhill]) - if param['ROTATION'] == 'YES': + if param['ROTATION']: vec_float = np.vstack([vec_float, Ip1, Ip2, Ip3, rotx, roty, rotz]) particle_type = np.repeat("Massive Body",idvals.size) else: From c47953ac36eb1f85b10cb91a0964408eec5594b3 Mon Sep 17 00:00:00 2001 From: David Minton Date: Tue, 8 Nov 2022 20:53:58 -0500 Subject: [PATCH 10/17] Added setter and getter methods for initial conditions files --- python/swiftest/swiftest/simulation_class.py | 292 +++++++++++++++---- 1 file changed, 230 insertions(+), 62 deletions(-) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index d01da849a..1dba3a49d 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -224,7 +224,8 @@ def __init__(self, self.set_output_files(output_file_type=output_file_type, output_file_name=output_file_name, - output_format=output_format) + output_format=output_format, + verbose = False) self.set_feature(close_encounter_check=close_encounter_check, general_relativity=general_relativity, @@ -372,7 +373,7 @@ def set_simulation_time(self, return - def get_simulation_time(self, param_key: str | List[str] | None = None): + def get_simulation_time(self, arg_list: str | List[str] | None = None): """ Returns a subset of the parameter dictionary containing the current simulation time parameters. @@ -380,7 +381,7 @@ def get_simulation_time(self, param_key: str | List[str] | None = None): Parameters ---------- - param_key : str | List[str], optional + arg_list : str | List[str], optional A single string or list of strings containing the names of the simulation time parameters to extract. Default is all of: ["t0", "tstart", "tstop", "dt", "istep_out", "tstep_out", "istep_dump"] @@ -400,36 +401,37 @@ def get_simulation_time(self, param_key: str | List[str] | None = None): "istep_dump": "ISTEP_DUMP", } - units = {"T0" : self.TU_name, - "TSTART" : self.TU_name, - "TSTOP" : self.TU_name, - "DT" : self.TU_name, - "TSTEP_OUT" : self.TU_name, - "ISTEP_OUT" : "", - "ISTEP_DUMP" : ""} + units = {"t0" : self.TU_name, + "tstart" : self.TU_name, + "tstop" : self.TU_name, + "dt" : self.TU_name, + "tstep_out" : self.TU_name, + "istep_out" : "", + "istep_dump" : ""} - if "tstep_out" in param_key: + if arg_list is not None and "tstep_out" in arg_list: istep_out = self.param['ISTEP_OUT'] dt = self.param['DT'] tstep_out = istep_out * dt else: tstep_out = None - if param_key is None: - param_list = valid_var.keys() - elif type(param_key) is str: - param_list = [param_key] + if arg_list is None: + arg_list = valid_var.keys() + elif type(arg_list) is str: + arg_list = [arg_list] else: - param_list = [k for k in param_key if k in set(valid_var.keys())] + arg_list = [k for k in arg_list if k in set(valid_var.keys())] - time_dict = {valid_var[k]: self.param[valid_var[k]] for k in param_list} + time_dict = {valid_var[k]: self.param[valid_var[k]] for k in arg_list} if self.verbose: - print("Simulation feature parameters:") - for key, val in time_dict.items(): - print(f"{key:<16} {val} {units[key]}") + print("\nSimulation time parameters:") + for arg in arg_list: + key = valid_var[arg] + print(f"{arg:<32} {time_dict[key]} {units[arg]}") if tstep_out is not None: - print(f"{'TSTEP_OUT':<16} {tstep_out} {units['TSTEP_OUT']}") + print(f"{'tstep_out':<32} {tstep_out} {units['tstep_out']}") return time_dict @@ -541,7 +543,7 @@ def set_feature(self, return - def get_feature(self,feature: str | List[str] | None = None): + def get_feature(self, arg_list: str | List[str] | None = None): """ Returns a subset of the parameter dictionary containing the current value of the feature boolean values. @@ -549,7 +551,7 @@ def get_feature(self,feature: str | List[str] | None = None): Parameters ---------- - feature : str | List[str], optional + arg_list: str | List[str], optional A single string or list of strings containing the names of the features to extract. Default is all of: ["close_encounter_check", "general_relativity", "fragmentation", "rotation", "compute_conservation_values"] @@ -560,45 +562,48 @@ def get_feature(self,feature: str | List[str] | None = None): """ - valid_features = {"close_encounter_check": "CHK_CLOSE", - "general_relativity": "GR", - "fragmentation": "FRAGMENTATION", - "rotation": "ROTATION", - "compute_conservation_values": "ENERGY", - "extra_force": "EXTRA_FORCE", - "big_discard": "BIG_DISCARD", - "rhill_present": "RHILL_PRESENT", - "restart" : "RESTART" - } - - if feature is None: - feature_list = valid_features.keys() - elif type(feature) is str: - feature_list = [feature] + valid_var = {"close_encounter_check": "CHK_CLOSE", + "general_relativity": "GR", + "fragmentation": "FRAGMENTATION", + "rotation": "ROTATION", + "compute_conservation_values": "ENERGY", + "extra_force": "EXTRA_FORCE", + "big_discard": "BIG_DISCARD", + "rhill_present": "RHILL_PRESENT", + "restart" : "RESTART" + } + + if arg_list is None: + arg_list = valid_var.keys() + elif type(arg_list) is str: + arg_list = [arg_list] else: - # Only allow features to be checked if they are valid. Otherwise ignore. - feature_list = [feat for feat in feature if feat in set(valid_features.keys())] + # Only allow arg_lists to be checked if they are valid. Otherwise ignore. + arg_list = [k for k in arg_list if k in set(valid_var.keys())] - # Extract the feature dictionary - feature_dict = {valid_features[feat]:self.param[valid_features[feat]] for feat in feature_list} + # Extract the arg_list dictionary + feature_dict = {valid_var[feat]:self.param[valid_var[feat]] for feat in arg_list} if self.verbose: - print("Simulation feature parameters:") - for key,val in feature_dict.items(): - print(f"{key:<16} {val}") + print("\nSimulation feature parameters:") + for arg in arg_list: + key = valid_var[arg] + print(f"{arg:<32} {feature_dict[key]}") return feature_dict def set_init_cond_files(self, - init_cond_file_type: Literal["NETCDF_DOUBLE", "NETCDF_FLOAT", "ASCII"], - init_cond_file_name: str | os.PathLike | Dict[str, str] | Dict[str, os.PathLike] | None, - init_cond_format: Literal["EL", "XV"]): + init_cond_file_type: Literal["NETCDF_DOUBLE", "NETCDF_FLOAT", "ASCII"] | None = None, + init_cond_file_name: str | os.PathLike | Dict[str, str] | Dict[str, os.PathLike] | None = None, + init_cond_format: Literal["EL", "XV"] | None = None, + verbose: bool | None = None + ): """ Sets the initial condition file parameters in the parameters dictionary. Parameters ---------- - init_cond_file_type : {"NETCDF_DOUBLE", "NETCDF_FLOAT", "ASCII"} + init_cond_file_type : {"NETCDF_DOUBLE", "NETCDF_FLOAT", "ASCII"}, optional The file type containing initial conditions for the simulation: * NETCDF_DOUBLE: A single initial conditions input file in NetCDF file format of type NETCDF_DOUBLE * NETCDF_FLOAT: A single initial conditions input file in NetCDF file format of type NETCDF_FLOAT @@ -620,6 +625,8 @@ def set_init_cond_files(self, Indicates whether the input initial conditions are given as orbital elements or cartesian position and velocity vectors. > *Note:* If `codename` is "Swift" or "Swifter", EL initial conditions are converted to XV. + verbose: bool, optional + If passed, it will override the Simulation object's verbose flag Returns ------- @@ -627,6 +634,12 @@ def set_init_cond_files(self, """ + update_list = ["init_cond_file_name"] + if init_cond_file_type is not None: + update_list.append("init_cond_file_type") + if init_cond_format is not None: + update_list.append("init_cond_format") + def ascii_file_input_error_msg(codename): print(f"Error in set_init_cond_files: init_cond_file_name must be a dictionary of the form: ") print('{') @@ -678,19 +691,94 @@ def ascii_file_input_error_msg(codename): else: self.param["NC_IN"] = init_cond_file_name + + if verbose is None: + verbose = self.verbose + if verbose: + if len(update_list) > 0: + init_cond_file_dict = self.get_init_cond_files(update_list) + return + + def get_init_cond_files(self, arg_list: str | List[str] | None = None): + """ + + Returns a subset of the parameter dictionary containing the current initial condition file parameters + If the verbose option is set in the Simulation object, then it will also print the values. + + Parameters + ---------- + arg_list : str | List[str], optional + A single string or list of strings containing the names of the simulation time parameters to extract. + Default is all of: + ["init_cond_file_type", "init_cond_file_name", "init_cond_format"] + + Returns + ------- + init_cond_file_dict : dict + A dictionary containing the requested parameters + + """ + + valid_var = {"init_cond_file_type": "IN_TYPE", + "init_cond_format": "IN_FORM", + "init_cond_file_name" : "NC_IN", + "init_cond_file_name (cb)" : "CB_IN", + "init_cond_file_name (pl)" : "PL_IN", + "init_cond_file_name (tp)" : "TP_IN", + } + + + if arg_list is None: + arg_list = list(valid_var.keys()) + elif type(arg_list) is str: + arg_list = [arg_list] + else: + arg_list = [k for k in arg_list if k in set(valid_var.keys())] + + if "init_cond_file_name" in arg_list: + if self.param["IN_TYPE"] == "ASCII": + arg_list.remove("init_cond_file_name") + if "init_cond_file_name (cb)" not in arg_list: + arg_list.append("init_cond_file_name (cb)") + if "init_cond_file_name (pl)" not in arg_list: + arg_list.append("init_cond_file_name (pl)") + if "init_cond_file_name (tp)" not in arg_list: + arg_list.append("init_cond_file_name (tp)") + else: + if "init_cond_file_name (cb)" in arg_list: + arg_list.remove("init_cond_file_name (cb)") + if "init_cond_file_name (pl)" in arg_list: + arg_list.remove("init_cond_file_name (pl)") + if "init_cond_file_name (tp)" in arg_list: + arg_list.remove("init_cond_file_name (tp)") + + init_cond_file_dict = {valid_var[k]: self.param[valid_var[k]] for k in arg_list} + + if self.verbose: + print("\nInitial condition file parameters:") + for arg in arg_list: + key = valid_var[arg] + print(f"{arg:<32} {init_cond_file_dict[key]}") + + return init_cond_file_dict + + + + def set_output_files(self, - output_file_type: Literal[ "NETCDF_DOUBLE", "NETCDF_FLOAT", "REAL4", "REAL8", "XDR4", "XDR8"], - output_file_name: os.PathLike | str | None, - output_format: Literal["XV", "XVEL"] + output_file_type: Literal[ "NETCDF_DOUBLE", "NETCDF_FLOAT", "REAL4", "REAL8", "XDR4", "XDR8"] | None = None, + output_file_name: os.PathLike | str | None = None, + output_format: Literal["XV", "XVEL"] | None = None, + verbose: bool | None = None ): """ Sets the output file parameters in the parameters dictionary. Parameters ---------- - output_file_type : {"NETCDF_DOUBLE", "NETCDF_FLOAT","REAL4","REAL8","XDR4","XDR8"}, default "NETCDF_DOUBLE" + output_file_type : {"NETCDF_DOUBLE", "NETCDF_FLOAT","REAL4","REAL8","XDR4","XDR8"}, optional The file type for the outputs of the simulation. Compatible file types depend on the `codename` argument. * Swiftest: Only "NETCDF_DOUBLE" or "NETCDF_FLOAT" supported. * Swifter: Only "REAL4","REAL8","XDR4" or "XDR8" supported. @@ -699,25 +787,46 @@ def set_output_files(self, Name of output file to generate. If not supplied, then one of the default file names are used, depending on the value passed to `output_file_type`. If one of the NetCDF types are used, the default is "bin.nc". Otherwise, the default is "bin.dat". - output_format : {"XV","XVEL"}, default "XVEL" + output_format : {"XV","XVEL"}, optional Specifies the format for the data saved to the output file. If "XV" then cartesian position and velocity vectors for all bodies are stored. If "XVEL" then the orbital elements are also stored. + verbose: bool, optional + If passed, it will override the Simulation object's verbose flag Returns ------- Sets the appropriate initial conditions file parameters inside the self.param dictionary. """ + update_list = [] + if output_file_type is not None: + update_list.append("output_file_type") + if output_file_name is not None: + update_list.append("output_file_name") + if output_format is not None: + update_list.append("output_format") if self.codename == "Swiftest": - if output_file_type not in ["NETCDF_DOUBLE", "NETCDF_FLOAT"]: + if output_file_type is None: + output_file_type = self.param.pop("OUT_TYPE", None) + if output_file_type is None: + output_file_type = "NETCDF_DOUBLE" + elif output_file_type not in ["NETCDF_DOUBLE", "NETCDF_FLOAT"]: print(f"{output_file_type} is not compatible with Swiftest. Setting to NETCDF_DOUBLE") output_file_type = "NETCDF_DOUBLE" elif self.codename == "Swifter": - if output_file_type not in ["REAL4","REAL8","XDR4","XDR8"]: + if output_file_type is None: + output_file_type = self.param.pop("OUT_TYPE", None) + if output_file_type is None: + output_file_type = "REAL8" + elif output_file_type not in ["REAL4","REAL8","XDR4","XDR8"]: print(f"{output_file_type} is not compatible with Swifter. Setting to REAL8") output_file_type = "REAL8" elif self.codename == "Swift": + if output_file_type is None: + output_file_type = self.param.pop("OUT_TYPE", None) + if output_file_type is None: + output_file_type = "REAL4" if output_file_type not in ["REAL4"]: print(f"{output_file_type} is not compatible with Swift. Setting to REAL4") output_file_type = "REAL4" @@ -734,10 +843,65 @@ def set_output_files(self, if output_format != "XV" and self.codename != "Swiftest": print(f"{output_format} is not compatible with {self.codename}. Setting to XV") output_format = "XV" - self.param['OUT_FORM'] = output_format + self.param["OUT_FORM"] = output_format + + if self.restart: + self.param["OUT_STAT"] = "APPEND" + else: + self.param["OUT_STAT"] = "REPLACE" + + if verbose is None: + verbose = self.verbose + if verbose: + if len(update_list) > 0: + output_file_dict = self.get_output_files(update_list) return + + def get_output_files(self, arg_list: str | List[str] | None = None): + """ + + Returns a subset of the parameter dictionary containing the current output file parameters + If the verbose option is set in the Simulation object, then it will also print the values. + + Parameters + ---------- + arg_list : str | List[str], optional + A single string or list of strings containing the names of the simulation time parameters to extract. + Default is all of: + ["output_file_type", "output_file_name", "output_format"] + + Returns + ------- + output_file_dict : dict + A dictionary containing the requested parameters + + """ + + valid_var = {"output_file_type": "OUT_TYPE", + "output_file_name": "BIN_OUT", + "output_format": "OUT_FORM", + } + + if arg_list is None: + arg_list = valid_var.keys() + elif type(arg_list) is str: + arg_list = [arg_list] + else: + arg_list = [k for k in arg_list if k in set(valid_var.keys())] + + output_file_dict = {valid_var[k]: self.param[valid_var[k]] for k in arg_list} + + if self.verbose: + print("\nOutput file parameters:") + for arg in arg_list: + key = valid_var[arg] + print(f"{arg:<32} {output_file_dict[key]}") + + return output_file_dict + + def set_unit_system(self, MU: str | None = None, DU: str | None = None, @@ -1238,7 +1402,7 @@ def save(self, param_file, framenum=-1, codename="Swiftest"): """ if codename == "Swiftest": - io.swiftest_xr2infile(ds=self.ds, param=self.param, in_type=self.param['IN_TYPE'], framenum=framenum,infile_name=self.param['NC_IN']) + io.swiftest_xr2infile(ds=self.ds, param=self.param, in_type=self.param['IN_TYPE'], framenum=framenum) self.write_param(param_file) elif codename == "Swifter": if self.codename == "Swiftest": @@ -1280,25 +1444,29 @@ def initial_conditions_from_bin(self, framenum=-1, new_param=None, new_param_fil if codename != "Swiftest": self.save(new_param_file, framenum, codename) return + if new_param is None: new_param = self.param.copy() if codename == "Swiftest": if restart: new_param['T0'] = self.ds.time.values[framenum] - if self.param['OUT_TYPE'] == 'NETCDF_DOUBLE' or self.param['OUT_TYPE'] == 'REAL8': + if self.param['OUT_TYPE'] == 'NETCDF_DOUBLE': new_param['IN_TYPE'] = 'NETCDF_DOUBLE' - elif self.param['OUT_TYPE'] == 'NETCDF_FLOAT' or self.param['OUT_TYPE'] == 'REAL4': + elif self.param['OUT_TYPE'] == 'NETCDF_FLOAT': new_param['IN_TYPE'] = 'NETCDF_FLOAT' else: print(f"{self.param['OUT_TYPE']} is an invalid OUT_TYPE file") return + if self.param['BIN_OUT'] != new_param['BIN_OUT'] and restart: - print(f"Restart run with new output file. Copying {self.param['BIN_OUT']} to {new_param['BIN_OUT']}") - shutil.copy2(self.param['BIN_OUT'],new_param['BIN_OUT']) + print(f"Restart run with new output file. Copying {self.param['BIN_OUT']} to {new_param['BIN_OUT']}") + shutil.copy2(self.param['BIN_OUT'],new_param['BIN_OUT']) + new_param['IN_FORM'] = 'XV' if restart: new_param['OUT_STAT'] = 'APPEND' + new_param['FIRSTKICK'] = 'T' new_param['NC_IN'] = new_initial_conditions_file new_param.pop('PL_IN', None) From b4c6604c0b6405bc4ba40b1c3224439843a841ac Mon Sep 17 00:00:00 2001 From: David Minton Date: Wed, 9 Nov 2022 08:35:38 -0500 Subject: [PATCH 11/17] Improved the unit system setter and added a unit system getter. Added kwargs to all the setters so I can write a main set_parameters setter that can take any possible parameter as an argument. --- python/swiftest/swiftest/simulation_class.py | 166 +++++++++++++++++-- 1 file changed, 149 insertions(+), 17 deletions(-) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 1dba3a49d..f538136ed 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -23,7 +23,8 @@ from typing import ( Literal, Dict, - List + List, + Any ) @@ -48,6 +49,9 @@ def __init__(self, MU2KG: float | None = None, DU2M: float | None = None, TU2S: float | None = None, + MU_name: str | None = None, + DU_name: str | None = None, + TU_name: str | None = None, rmin: float = constants.RSun / constants.AU2M, rmax: float = 10000.0, close_encounter_check: bool = True, @@ -140,6 +144,15 @@ def __init__(self, TU2S : float, optional The conversion factor to multiply by the time unit that would convert it to seconds. Setting this overrides TU + MU_name : str, optional + The name of the mass unit. When setting one of the standard units via `MU` a name will be + automatically set for the unit, so this argument will override the automatic name. + DU_name : str, optional + The name of the distance unit. When setting one of the standard units via `DU` a name will be + automatically set for the unit, so this argument will override the automatic name. + TU_name : str, optional + The name of the time unit. When setting one of the standard units via `TU` a name will be + automatically set for the unit, so this argument will override the automatic name. rmin : float, default value is the radius of the Sun in the unit system defined by the unit input arguments. Minimum distance of the simulation (CHK_QMIN, CHK_RMIN, CHK_QMIN_RANGE[0]) rmax : float, default value is 10000 AU in the unit system defined by the unit input arguments. @@ -212,15 +225,18 @@ def __init__(self, self.verbose = verbose self.restart = restart - self.set_distance_range(rmin=rmin, rmax=rmax) + self.set_distance_range(rmin=rmin, rmax=rmax, verbose = False) self.set_unit_system(MU=MU, DU=DU, TU=TU, MU2KG=MU2KG, DU2M=DU2M, TU2S=TU2S, - recompute_values=True) + MU_name=MU_name, DU_name=DU_name, TU_name=TU_name, + recompute_values=True, + verbose = False) self.set_init_cond_files(init_cond_file_type=init_cond_file_type, init_cond_file_name=init_cond_file_name, - init_cond_format=init_cond_format) + init_cond_format=init_cond_format, + verbose = False) self.set_output_files(output_file_type=output_file_type, output_file_name=output_file_name, @@ -264,6 +280,7 @@ def set_simulation_time(self, tstep_out : float | None = None, istep_dump : int | None = None, verbose : bool | None = None, + **kwargs: Any ): """ @@ -287,6 +304,9 @@ def set_simulation_time(self, `istep_out` (or the equivalent value determined by `tstep_out`) verbose: bool, optional If passed, it will override the Simulation object's verbose flag + **kwargs + A dictionary of additional keyword argument. This allows this method to be called by the more general + set_parameters method, which takes all possible Simulation parameters as arguments, so these are ignored. Returns ------- @@ -447,6 +467,7 @@ def set_feature(self, restart: bool | None = None, tides: bool | None = None, verbose: bool | None = None, + **kwargs: Any ): """ Turns on or off various features of a simulation. @@ -486,6 +507,9 @@ def set_feature(self, when the run is executed. verbose: bool, optional If passed, it will override the Simulation object's verbose flag + **kwargs + A dictionary of additional keyword argument. This allows this method to be called by the more general + set_parameters method, which takes all possible Simulation parameters as arguments, so these are ignored. Returns ------- @@ -596,7 +620,8 @@ def set_init_cond_files(self, init_cond_file_type: Literal["NETCDF_DOUBLE", "NETCDF_FLOAT", "ASCII"] | None = None, init_cond_file_name: str | os.PathLike | Dict[str, str] | Dict[str, os.PathLike] | None = None, init_cond_format: Literal["EL", "XV"] | None = None, - verbose: bool | None = None + verbose: bool | None = None, + **kwargs: Any ): """ Sets the initial condition file parameters in the parameters dictionary. @@ -627,6 +652,9 @@ def set_init_cond_files(self, > *Note:* If `codename` is "Swift" or "Swifter", EL initial conditions are converted to XV. verbose: bool, optional If passed, it will override the Simulation object's verbose flag + **kwargs + A dictionary of additional keyword argument. This allows this method to be called by the more general + set_parameters method, which takes all possible Simulation parameters as arguments, so these are ignored. Returns ------- @@ -771,7 +799,8 @@ def set_output_files(self, output_file_type: Literal[ "NETCDF_DOUBLE", "NETCDF_FLOAT", "REAL4", "REAL8", "XDR4", "XDR8"] | None = None, output_file_name: os.PathLike | str | None = None, output_format: Literal["XV", "XVEL"] | None = None, - verbose: bool | None = None + verbose: bool | None = None, + **kwargs: Any ): """ Sets the output file parameters in the parameters dictionary. @@ -792,6 +821,9 @@ def set_output_files(self, vectors for all bodies are stored. If "XVEL" then the orbital elements are also stored. verbose: bool, optional If passed, it will override the Simulation object's verbose flag + **kwargs + A dictionary of additional keyword argument. This allows this method to be called by the more general + set_parameters method, which takes all possible Simulation parameters as arguments, so these are ignored. Returns ------- @@ -909,7 +941,12 @@ def set_unit_system(self, MU2KG: float | None = None, DU2M: float | None = None, TU2S: float | None = None, - recompute_values: bool = False): + MU_name: str | None = None, + DU_name: str | None = None, + TU_name: str | None = None, + recompute_values: bool = True, + verbose: bool | None = None, + **kwargs: Any): """ Setter for setting the unit conversion between one of the standard sets. @@ -951,11 +988,25 @@ def set_unit_system(self, TU2S : float, optional The conversion factor to multiply by the time unit that would convert it to seconds. Setting this overrides TU - recompute_values : bool, default False + MU_name : str, optional + The name of the mass unit. When setting one of the standard units via `MU` a name will be + automatically set for the unit, so this argument will override the automatic name. + DU_name : str, optional + The name of the distance unit. When setting one of the standard units via `DU` a name will be + automatically set for the unit, so this argument will override the automatic name. + TU_name : str, optional + The name of the time unit. When setting one of the standard units via `TU` a name will be + automatically set for the unit, so this argument will override the automatic name. + recompute_values : bool, default True Recompute all values into the new unit system. >*Note:* This is a destructive operation, however if not executed then the values contained in the parameter > file and input/output data files computed previously may not be consistent with the new unit conversion > factors. + verbose: bool, optional + If passed, it will override the Simulation object's verbose flag + **kwargs + A dictionary of additional keyword argument. This allows this method to be called by the more general + set_parameters method, which takes all possible Simulation parameters as arguments, so these are ignored. Returns ---------- @@ -967,6 +1018,14 @@ def set_unit_system(self, DU2M_old = None TU2S_old = None + update_list = [] + if MU is not None or MU2KG is not None: + update_list.append("MU") + if DU is not None or DU2M is not None: + update_list.append("DU") + if TU is not None or TU2S is not None: + update_list.append("TU") + if MU2KG is not None or MU is not None: MU2KG_old = self.param.pop('MU2KG',None) if MU2KG is not None: @@ -1033,23 +1092,90 @@ def set_unit_system(self, self.param['TU2S'] = constants.YR2S self.TU_name = "y" + + if MU_name is not None: + self.MU_name = MU_name + if DU_name is not None: + self.DU_name = DU_name + if TU_name is not None: + self.TU_name = TU_name + self.VU_name = f"{self.DU_name}/{self.TU_name}" self.GU = constants.GC * self.param['TU2S']**2 * self.param['MU2KG'] / self.param['DU2M']**3 if recompute_values: self.update_param_units(MU2KG_old, DU2M_old, TU2S_old) - if self.verbose: - if self.MU_name is None or self.DU_name is None or self.TU_name is None: - print(f"Custom units set.") - print(f"MU2KG: {self.param['MU2KG']}") - print(f"DU2M : {self.param['DU2M']}") - print(f"TU2S : {self.param['TU2S']}") - else: - print(f"Units set to {self.MU_name}-{self.DU_name}-{self.TU_name}") + if verbose is None: + verbose = self.verbose + + if verbose: + unit_dict = self.get_unit_system(update_list) return + def get_unit_system(self, arg_list: str | List[str] | None = None): + """ + + Returns a subset of the parameter dictionary containing the current simulation unit system. + If the verbose option is set in the Simulation object, then it will also print the values. + + Parameters + ---------- + arg_list : str | List[str], optional + A single string or list of strings containing the names of the simulation unit system + Default is all of: + ["MU", "DU", "TU"] + + Returns + ------- + time_dict : dict + A dictionary containing the requested parameters + + """ + + valid_var = { + "MU": "MU2KG", + "DU": "DU2M", + "TU": "TU2S", + } + + if self.MU_name is None: + MU_name = "mass unit" + else: + MU_name = self.MU_name + if self.DU_name is None: + DU_name = "distance unit" + else: + DU_name = self.DU_name + if self.TU_name is None: + TU_name = "time unit" + else: + TU_name = self.TU_name + + units = { + "MU" : f"kg / {MU_name}", + "DU" : f"m / {DU_name}", + "TU" : f"s / {TU_name}" + } + + + if arg_list is None: + arg_list = valid_var.keys() + elif type(arg_list) is str: + arg_list = [arg_list] + else: + arg_list = [k for k in arg_list if k in set(valid_var.keys())] + + unit_dict = {valid_var[k]: self.param[valid_var[k]] for k in arg_list} + + if self.verbose: + for arg in arg_list: + key = valid_var[arg] + print(f"{arg:<32} {unit_dict[key]} {units[arg]}") + + return unit_dict + def update_param_units(self,MU2KG_old,DU2M_old,TU2S_old): """ Updates the values of parameters that have units when the units have changed. @@ -1096,7 +1222,10 @@ def update_param_units(self,MU2KG_old,DU2M_old,TU2S_old): return - def set_distance_range(self,rmin=None,rmax=None): + def set_distance_range(self, + rmin: float | None = None, + rmax: float | None = None, + **kwargs: Any): """ Sets the minimum and maximum distances of the simulation. @@ -1106,6 +1235,9 @@ def set_distance_range(self,rmin=None,rmax=None): Minimum distance of the simulation (CHK_QMIN, CHK_RMIN, CHK_QMIN_RANGE[0]) rmax : float Maximum distance of the simulation (CHK_RMAX, CHK_QMIN_RANGE[1]) + **kwargs + A dictionary of additional keyword argument. This allows this method to be called by the more general + set_parameters method, which takes all possible Simulation parameters as arguments, so these are ignored. Returns ------- From 99c76197ab89e463bbbfabc5c38cd1347837fb71 Mon Sep 17 00:00:00 2001 From: David Minton Date: Wed, 9 Nov 2022 10:57:11 -0500 Subject: [PATCH 12/17] Improved the getters and setters for the initial conditions files --- python/swiftest/swiftest/simulation_class.py | 236 +++++++++++++------ 1 file changed, 161 insertions(+), 75 deletions(-) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index f538136ed..1c2cb3fca 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -211,12 +211,6 @@ def __init__(self, self.ds = xr.Dataset() self.param = { '! VERSION': f"Swiftest parameter input", - 'T0': 0.0, - 'TSTART' : 0.0, - 'TSTOP': 0.0, - 'DT': 0.0, - 'ISTEP_OUT': 1, - 'ISTEP_DUMP': 1, 'CHK_QMIN_COORD': "HELIO", 'INTERACTION_LOOPS': interaction_loops, 'ENCOUNTER_CHECK': encounter_check_loops @@ -436,18 +430,10 @@ def get_simulation_time(self, arg_list: str | List[str] | None = None): else: tstep_out = None - if arg_list is None: - arg_list = valid_var.keys() - elif type(arg_list) is str: - arg_list = [arg_list] - else: - arg_list = [k for k in arg_list if k in set(valid_var.keys())] - - time_dict = {valid_var[k]: self.param[valid_var[k]] for k in arg_list} + valid_arg, time_dict = self._get_valid_arg_list(arg_list, valid_var) if self.verbose: - print("\nSimulation time parameters:") - for arg in arg_list: + for arg in valid_arg: key = valid_var[arg] print(f"{arg:<32} {time_dict[key]} {units[arg]}") if tstep_out is not None: @@ -597,20 +583,10 @@ def get_feature(self, arg_list: str | List[str] | None = None): "restart" : "RESTART" } - if arg_list is None: - arg_list = valid_var.keys() - elif type(arg_list) is str: - arg_list = [arg_list] - else: - # Only allow arg_lists to be checked if they are valid. Otherwise ignore. - arg_list = [k for k in arg_list if k in set(valid_var.keys())] - - # Extract the arg_list dictionary - feature_dict = {valid_var[feat]:self.param[valid_var[feat]] for feat in arg_list} + valid_arg, feature_dict = self._get_valid_arg_list(arg_list, valid_var) if self.verbose: - print("\nSimulation feature parameters:") - for arg in arg_list: + for arg in valid_arg: key = valid_var[arg] print(f"{arg:<32} {feature_dict[key]}") @@ -662,7 +638,9 @@ def set_init_cond_files(self, """ - update_list = ["init_cond_file_name"] + update_list = [] + if init_cond_file_name is not None: + update_list.append("init_cond_file_name") if init_cond_file_type is not None: update_list.append("init_cond_file_type") if init_cond_format is not None: @@ -678,6 +656,18 @@ def ascii_file_input_error_msg(codename): print('}') return + if init_cond_format is None: + if "IN_FORM" in self.param: + init_cond_format = self.param['IN_FORM'] + else: + init_cond_format = "EL" + + if init_cond_file_type is None: + if "IN_TYPE" in self.param: + init_cond_file_type = self.param['IN_TYPE'] + else: + init_cond_file_type = "NETCDF_DOUBLE" + if self.codename == "Swiftest": init_cond_keys = ["CB", "PL", "TP"] else: @@ -689,8 +679,18 @@ def ascii_file_input_error_msg(codename): print(f"{init_cond_format} is not supported by {self.codename}. Using XV instead") init_cond_format = "XV" - self.param["IN_TYPE"] = init_cond_file_type - self.param["IN_FORM"] = init_cond_format + + valid_formats={"EL", "XV"} + if init_cond_format not in valid_formats: + print(f"{init_cond_format} is not a valid input format") + else: + self.param['IN_FORM'] = init_cond_format + + valid_types = {"NETCDF_DOUBLE", "NETCDF_FLOAT", "ASCII"} + if init_cond_file_type not in valid_types: + print(f"{init_cond_file_type} is not a valid input type") + else: + self.param['IN_TYPE'] = init_cond_file_type if init_cond_file_type == "ASCII": if init_cond_file_name is None: @@ -752,41 +752,48 @@ def get_init_cond_files(self, arg_list: str | List[str] | None = None): valid_var = {"init_cond_file_type": "IN_TYPE", "init_cond_format": "IN_FORM", "init_cond_file_name" : "NC_IN", - "init_cond_file_name (cb)" : "CB_IN", - "init_cond_file_name (pl)" : "PL_IN", - "init_cond_file_name (tp)" : "TP_IN", + "init_cond_file_name['CB']" : "CB_IN", + "init_cond_file_name['PL']" : "PL_IN", + "init_cond_file_name['TP']" : "TP_IN", } + three_file_args = ["init_cond_file_name['CB']", + "init_cond_file_name['PL']", + "init_cond_file_name['TP']"] + if self.codename == "Swifter": + three_file_args.remove("init_cond_file_name['CB']") + + # We have to figure out which initial conditions file model we are using (1 vs. 3 files) if arg_list is None: - arg_list = list(valid_var.keys()) + valid_arg = None + else: + valid_arg = arg_list.copy() + + if valid_arg is None: + valid_arg = list(valid_var.keys()) elif type(arg_list) is str: - arg_list = [arg_list] + valid_arg = [arg_list] else: - arg_list = [k for k in arg_list if k in set(valid_var.keys())] + # Only allow arg_lists to be checked if they are valid. Otherwise ignore. + valid_arg = [k for k in arg_list if k in list(valid_var.keys())] - if "init_cond_file_name" in arg_list: + # Figure out which input file model we need to use + if "init_cond_file_name" in valid_arg: if self.param["IN_TYPE"] == "ASCII": - arg_list.remove("init_cond_file_name") - if "init_cond_file_name (cb)" not in arg_list: - arg_list.append("init_cond_file_name (cb)") - if "init_cond_file_name (pl)" not in arg_list: - arg_list.append("init_cond_file_name (pl)") - if "init_cond_file_name (tp)" not in arg_list: - arg_list.append("init_cond_file_name (tp)") + valid_arg.remove("init_cond_file_name") + for key in three_file_args: + if key not in valid_arg: + valid_arg.append(key) else: - if "init_cond_file_name (cb)" in arg_list: - arg_list.remove("init_cond_file_name (cb)") - if "init_cond_file_name (pl)" in arg_list: - arg_list.remove("init_cond_file_name (pl)") - if "init_cond_file_name (tp)" in arg_list: - arg_list.remove("init_cond_file_name (tp)") + for key in three_file_args: + if key in valid_arg: + valid_arg.remove(key) - init_cond_file_dict = {valid_var[k]: self.param[valid_var[k]] for k in arg_list} + valid_arg, init_cond_file_dict = self._get_valid_arg_list(valid_arg, valid_var) if self.verbose: - print("\nInitial condition file parameters:") - for arg in arg_list: + for arg in valid_arg: key = valid_var[arg] print(f"{arg:<32} {init_cond_file_dict[key]}") @@ -794,7 +801,6 @@ def get_init_cond_files(self, arg_list: str | List[str] | None = None): - def set_output_files(self, output_file_type: Literal[ "NETCDF_DOUBLE", "NETCDF_FLOAT", "REAL4", "REAL8", "XDR4", "XDR8"] | None = None, output_file_name: os.PathLike | str | None = None, @@ -916,18 +922,10 @@ def get_output_files(self, arg_list: str | List[str] | None = None): "output_format": "OUT_FORM", } - if arg_list is None: - arg_list = valid_var.keys() - elif type(arg_list) is str: - arg_list = [arg_list] - else: - arg_list = [k for k in arg_list if k in set(valid_var.keys())] - - output_file_dict = {valid_var[k]: self.param[valid_var[k]] for k in arg_list} + valid_arg, output_file_dict = self._get_valid_arg_list(arg_list, valid_var) if self.verbose: - print("\nOutput file parameters:") - for arg in arg_list: + for arg in valid_arg: key = valid_var[arg] print(f"{arg:<32} {output_file_dict[key]}") @@ -1159,18 +1157,10 @@ def get_unit_system(self, arg_list: str | List[str] | None = None): "TU" : f"s / {TU_name}" } - - if arg_list is None: - arg_list = valid_var.keys() - elif type(arg_list) is str: - arg_list = [arg_list] - else: - arg_list = [k for k in arg_list if k in set(valid_var.keys())] - - unit_dict = {valid_var[k]: self.param[valid_var[k]] for k in arg_list} + valid_arg, unit_dict = self._get_valid_arg_list(arg_list, valid_var) if self.verbose: - for arg in arg_list: + for arg in valid_arg: key = valid_var[arg] print(f"{arg:<32} {unit_dict[key]} {units[arg]}") @@ -1262,6 +1252,102 @@ def set_distance_range(self, return + def _get_valid_arg_list(self, arg_list: str | List[str] | None = None, valid_var: Dict | None = None): + """ + Internal function for getters that extracts subset of arguments that is contained in the dictionary of valid + argument/parameter variable pairs. + + Parameters + ---------- + arg_list : str | List[str], optional + A single string or list of strings containing the Simulation argument. If none are supplied, + then it will create the arg_list out of all keys in the valid_var dictionary. + valid_var : valid_var: Dict + A dictionary where the key is the argument name and the value is the equivalent key in the Simulation + parameter dictionary (i.e. the left-hand column of a param.in file) + + Returns + ------- + valid_arg : [str] + The list of valid arguments that match the keys in valid_var + param : dict + A dictionary that is the subset of the Simulation parameter dictionary corresponding to the arguments listed + in arg_list. + + """ + + + if arg_list is None: + valid_arg = None + else: + valid_arg = arg_list.copy() + + if valid_arg is None: + valid_arg = list(valid_var.keys()) + elif type(arg_list) is str: + valid_arg = [arg_list] + else: + # Only allow arg_lists to be checked if they are valid. Otherwise ignore. + valid_arg = [k for k in arg_list if k in list(valid_var.keys())] + + # Extract the arg_list dictionary + param = {valid_var[feat]:self.param[valid_var[feat]] for feat in valid_arg} + + return valid_arg, param + + def get_distance_range(self, arg_list: str | List[str] | None = None): + """ + + Returns a subset of the parameter dictionary containing the current values of the distance range parameters. + If the verbose option is set in the Simulation object, then it will also print the values. + + Parameters + ---------- + arg_list: str | List[str], optional + A single string or list of strings containing the names of the features to extract. Default is all of: + ["rmin", "rmax"] + + Returns + ------- + range_dict : dict + A dictionary containing the requested parameters. + + """ + + valid_var = {"rmin": "CHK_RMIN", + "rmax": "CHK_RMAX", + "qmin": "CHK_QMIN", + "qminR" : "CHK_QMIN_RANGE" + } + + units = {"rmin" : self.DU_name, + "rmax" : self.DU_name, + "qmin" : self.DU_name, + "qminR" : self.DU_name, + } + + if "rmin" in arg_list: + arg_list.append("qmin") + if "rmax" in arg_list or "rmin" in arg_list: + arg_list.append("qminR") + + valid_arg, range_dict = self._get_valid_arg_list(arg_list, valid_var) + + if self.verbose: + if "rmin" in valid_arg: + key = valid_arg["rmin"] + print(f"{'rmin':<32} {range_dict[key]} {units['rmin']}") + key = valid_arg["qmin"] + print(f"{'':<32} {range_dict[key]} {units['qmin']}") + if "rmax" in valid_arg: + key = valid_arg["rmax"] + print(f"{'rmax':<32} {range_dict[key]} {units['rmax']}") + if "rmax" in valid_arg or "rmin" in valid_arg: + key = valid_arg["qminR"] + print(f"{'':<32} {range_dict[key]} {units['qminR']}") + + + return range_dict def add(self, plname, date=date.today().isoformat(), idval=None): From 41c61b1aad2a1171c29db59818ab085a7089ff3c Mon Sep 17 00:00:00 2001 From: David Minton Date: Wed, 9 Nov 2022 11:06:01 -0500 Subject: [PATCH 13/17] Prints the tstep_out value whenever istep_out is given as an argument to get_simulation_time --- python/swiftest/swiftest/simulation_class.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 1c2cb3fca..c2c557be2 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -423,7 +423,7 @@ def get_simulation_time(self, arg_list: str | List[str] | None = None): "istep_out" : "", "istep_dump" : ""} - if arg_list is not None and "tstep_out" in arg_list: + if arg_list is None or "tstep_out" in arg_list or "istep_out" in arg_list: istep_out = self.param['ISTEP_OUT'] dt = self.param['DT'] tstep_out = istep_out * dt From bd9bfad2a3b7026e609105a74faa56326da6aa14 Mon Sep 17 00:00:00 2001 From: David Minton Date: Wed, 9 Nov 2022 11:16:03 -0500 Subject: [PATCH 14/17] Fixed unit conversion bug and added simulation time parameters to Simulation class argument list --- python/swiftest/swiftest/simulation_class.py | 24 ++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index c2c557be2..174d5e782 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -1,5 +1,4 @@ """ - self.param['BIN_OUT'] = binpath Copyright 2022 - David Minton, Carlisle Wishard, Jennifer Pouplin, Jake Elliott, & Dana Singh This file is part of Swiftest. Swiftest is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License @@ -36,6 +35,13 @@ def __init__(self, codename: Literal["Swiftest", "Swifter", "Swift"] = "Swiftest", param_file: os.PathLike | str ="param.in", read_param: bool = False, + t0: float = 0.0, + tstart: float = 0.0, + tstop: float = 1.0, + dt: float = 0.002, + istep_out: int = 50, + tstep_out: float | None = None, + istep_dump: int = 50, init_cond_file_type: Literal["NETCDF_DOUBLE", "NETCDF_FLOAT", "ASCII"] = "NETCDF_DOUBLE", init_cond_file_name: str | os.PathLike | Dict[str, str] | Dict[str, os.PathLike] | None = None, init_cond_format: Literal["EL", "XV"] = "EL", @@ -248,6 +254,16 @@ def __init__(self, restart=restart, verbose = False) + self.set_simulation_time(t0=t0, + tstart=tstart, + tstop=tstop, + dt=dt, + tstep_out=tstep_out, + istep_out=istep_out, + istep_dump=istep_dump, + verbose = False + ) + # If the parameter file is in a different location than the current working directory, we will need # to use it to properly open bin files self.sim_dir = os.path.dirname(os.path.realpath(param_file)) @@ -1190,12 +1206,12 @@ def update_param_units(self,MU2KG_old,DU2M_old,TU2S_old): for k in mass_keys: if k in self.param: print(f"param['{k}']: {self.param[k]}") - self.param[k] *= self.param['MU2KG'] / MU2KG_old + self.param[k] *= MU2KG_old / self.param['MU2KG'] if DU2M_old is not None: for k in distance_keys: if k in self.param: - self.param[k] *= self.param['DU2M'] / DU2M_old + self.param[k] *= DU2M_old / self.param['DU2M'] CHK_QMIN_RANGE = self.param.pop('CHK_QMIN_RANGE', None) if CHK_QMIN_RANGE is not None: @@ -1207,7 +1223,7 @@ def update_param_units(self,MU2KG_old,DU2M_old,TU2S_old): if TU2S_old is not None: for k in time_keys: if k in self.param: - self.param[k] *= self.param['TU2S'] / TU2S_old + self.param[k] *= TU2S_old / self.param['TU2S'] return From 50f626187326682f9925eb4dc22654514a9376ff Mon Sep 17 00:00:00 2001 From: David Minton Date: Wed, 9 Nov 2022 11:46:41 -0500 Subject: [PATCH 15/17] Improved the handling of unset parameters and removed the arbitrary default values for tstop, dt, istep_out, and istep_dump. --- python/swiftest/swiftest/simulation_class.py | 89 ++++++++++++-------- 1 file changed, 53 insertions(+), 36 deletions(-) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 174d5e782..917165bd1 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -26,7 +26,6 @@ Any ) - class Simulation: """ This is a class that defines the basic Swift/Swifter/Swiftest simulation object @@ -37,11 +36,11 @@ def __init__(self, read_param: bool = False, t0: float = 0.0, tstart: float = 0.0, - tstop: float = 1.0, - dt: float = 0.002, - istep_out: int = 50, + tstop: float | None = None, + dt: float | None = None, + istep_out: int | None = None, tstep_out: float | None = None, - istep_dump: int = 50, + istep_dump: int | None = None, init_cond_file_type: Literal["NETCDF_DOUBLE", "NETCDF_FLOAT", "ASCII"] = "NETCDF_DOUBLE", init_cond_file_name: str | os.PathLike | Dict[str, str] | Dict[str, os.PathLike] | None = None, init_cond_format: Literal["EL", "XV"] = "EL", @@ -86,6 +85,22 @@ def __init__(self, a new parameter file using the arguments passed to Simulation. > *Note:* If set to true, the parameters defined in the input file will override any passed into the > arguments of Simulation. + t0 : float, default 0.0 + The reference time for the start of the simulation. Defaults is 0.0 + tstart : float, default 0.0 + The start time for a restarted simulation. For a new simulation, tstart will be set to t0 automatically. + tstop : float, optional + The stopping time for a simulation. `tstop` must be greater than `tstart`. + dt : float, optional + The step size of the simulation. `dt` must be less than or equal to `tstop-dstart`. + istep_out : int, optional + The number of time steps between outputs to file. *Note*: only `istep_out` or `toutput` can be set. + tstep_out : float, optional + The approximate time between when outputs are written to file. Passing this computes + `istep_out = floor(tstep_out/dt)`. *Note*: only `istep_out` or `toutput` can be set. + istep_dump : int, optional + The anumber of time steps between outputs to dump file. If not set, this will be set to the value of + `istep_out` (or the equivalent value determined by `tstep_out`) init_cond_file_type : {"NETCDF_DOUBLE", "NETCDF_FLOAT", "ASCII"}, default "NETCDF_DOUBLE" The file type containing initial conditions for the simulation: * NETCDF_DOUBLE: A single initial conditions input file in NetCDF file format of type NETCDF_DOUBLE @@ -346,54 +361,54 @@ def set_simulation_time(self, if tstop is None: tstop = self.param.pop("TSTOP", None) - if tstop is None: - print("Error! tstop is not set!") - return else: update_list.append("tstop") - if tstop <= tstart: - print("Error! tstop must be greater than tstart.") - return + if tstop is not None: + if tstop <= tstart: + print("Error! tstop must be greater than tstart.") + return - self.param['TSTOP'] = tstop + if tstop is not None: + self.param['TSTOP'] = tstop if dt is None: dt = self.param.pop("DT", None) - if dt is None: - print("Error! dt is not set!") - return else: update_list.append("dt") - if dt > (tstop - tstart): - print("Error! dt must be smaller than tstop-tstart") - print(f"Setting dt = {tstop-tstart} instead of {dt}") - dt = tstop - tstart + if dt is not None and tstop is not None: + if dt > (tstop - tstart): + print("Error! dt must be smaller than tstop-tstart") + print(f"Setting dt = {tstop-tstart} instead of {dt}") + dt = tstop - tstart - self.param['DT'] = dt + if dt is not None: + self.param['DT'] = dt if istep_out is None and tstep_out is None: istep_out = self.param.pop("ISTEP_OUT", None) - if istep_out is None: - print("Error! istep_out is not set") - return if istep_out is not None and tstep_out is not None: print("Error! istep_out and tstep_out cannot both be set") return - if tstep_out is not None: + if tstep_out is not None and dt is not None: istep_out = int(np.floor(tstep_out / dt)) - self.param['ISTEP_OUT'] = istep_out + if istep_out is not None: + self.param['ISTEP_OUT'] = istep_out + update_list.append("istep_out") if istep_dump is None: - istep_dump = istep_out + istep_dump = self.param.pop("ISTEP_DUMP",None) + if istep_dump is None: + istep_dump = istep_out - self.param['ISTEP_DUMP'] = istep_dump + if istep_dump is not None: + self.param['ISTEP_DUMP'] = istep_dump + update_list.append("istep_dump") - update_list.extend(["istep_out", "tstep_out", "istep_dump"]) if verbose is None: verbose = self.verbose @@ -439,19 +454,22 @@ def get_simulation_time(self, arg_list: str | List[str] | None = None): "istep_out" : "", "istep_dump" : ""} + tstep_out = None if arg_list is None or "tstep_out" in arg_list or "istep_out" in arg_list: - istep_out = self.param['ISTEP_OUT'] - dt = self.param['DT'] - tstep_out = istep_out * dt - else: - tstep_out = None + if "ISTEP_OUT" in self.param and "DT" in self.param: + istep_out = self.param['ISTEP_OUT'] + dt = self.param['DT'] + tstep_out = istep_out * dt valid_arg, time_dict = self._get_valid_arg_list(arg_list, valid_var) if self.verbose: for arg in valid_arg: key = valid_var[arg] - print(f"{arg:<32} {time_dict[key]} {units[arg]}") + if key in time_dict: + print(f"{arg:<32} {time_dict[key]} {units[arg]}") + else: + print(f"{arg:<32} NOT SET") if tstep_out is not None: print(f"{'tstep_out':<32} {tstep_out} {units['tstep_out']}") @@ -1292,7 +1310,6 @@ def _get_valid_arg_list(self, arg_list: str | List[str] | None = None, valid_var """ - if arg_list is None: valid_arg = None else: @@ -1307,7 +1324,7 @@ def _get_valid_arg_list(self, arg_list: str | List[str] | None = None, valid_var valid_arg = [k for k in arg_list if k in list(valid_var.keys())] # Extract the arg_list dictionary - param = {valid_var[feat]:self.param[valid_var[feat]] for feat in valid_arg} + param = {valid_var[arg]:self.param[valid_var[arg]] for arg in valid_arg if valid_var[arg] in self.param} return valid_arg, param From 99f8c8b7a627062c4b86ddfebbb141ef3fbe2e47 Mon Sep 17 00:00:00 2001 From: David Minton Date: Wed, 9 Nov 2022 11:52:09 -0500 Subject: [PATCH 16/17] Improved display of units in get_unit_system --- python/swiftest/swiftest/simulation_class.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 917165bd1..21e090747 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -1115,7 +1115,7 @@ def set_unit_system(self, self.TU_name = "y" elif TU.upper() == "DAY" or TU.upper() == "D" or TU.upper() == "JD" or TU.upper() == "DAYS": self.param['TU2S'] = constants.JD2S - self.TU_name = "Day" + self.TU_name = "d" elif TU.upper() == "S" or TU.upper() == "SECONDS" or TU.upper() == "SEC": self.param['TU2S'] = 1.0 self.TU_name = "s" @@ -1185,7 +1185,12 @@ def get_unit_system(self, arg_list: str | List[str] | None = None): else: TU_name = self.TU_name - units = { + units1 = { + "MU" : MU_name, + "DU" : DU_name, + "TU" : TU_name + } + units2 = { "MU" : f"kg / {MU_name}", "DU" : f"m / {DU_name}", "TU" : f"s / {TU_name}" @@ -1196,7 +1201,7 @@ def get_unit_system(self, arg_list: str | List[str] | None = None): if self.verbose: for arg in valid_arg: key = valid_var[arg] - print(f"{arg:<32} {unit_dict[key]} {units[arg]}") + print(f"{arg}: {units1[arg]:<28} {unit_dict[key]} {units2[arg]}") return unit_dict From 1b962e262681385d8dc4430ae8d844b3ffa1b8c7 Mon Sep 17 00:00:00 2001 From: David Minton Date: Wed, 9 Nov 2022 14:44:45 -0500 Subject: [PATCH 17/17] Fixed typos in getter methods and improved the handling of rmin and rmax reporting in the get_distance_range method --- python/swiftest/swiftest/simulation_class.py | 47 ++++++++++++-------- 1 file changed, 29 insertions(+), 18 deletions(-) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 21e090747..b72542aaa 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -475,6 +475,26 @@ def get_simulation_time(self, arg_list: str | List[str] | None = None): return time_dict + def set_parameters(self, **kwargs): + """ + Setter for all possible parameters. Calls each of the specialized setters using keyword arguments + Parameters + ---------- + **kwargs : [TODO: write this documentation] + + Returns + ------- + param : A dictionary of all Simulation parameters that changed + + """ + self.set_unit_system(**kwargs) + self.set_distance_range(**kwargs) + self.set_feature(**kwargs) + self.set_init_cond_files(**kwargs) + self.set_output_files(**kwargs) + self.set_simulation_time(**kwargs) + + def set_feature(self, close_encounter_check: bool | None = None, general_relativity: bool | None = None, @@ -800,16 +820,10 @@ def get_init_cond_files(self, arg_list: str | List[str] | None = None): # We have to figure out which initial conditions file model we are using (1 vs. 3 files) if arg_list is None: - valid_arg = None - else: - valid_arg = arg_list.copy() - - if valid_arg is None: valid_arg = list(valid_var.keys()) elif type(arg_list) is str: valid_arg = [arg_list] else: - # Only allow arg_lists to be checked if they are valid. Otherwise ignore. valid_arg = [k for k in arg_list if k in list(valid_var.keys())] # Figure out which input file model we need to use @@ -1364,26 +1378,23 @@ def get_distance_range(self, arg_list: str | List[str] | None = None): "qminR" : self.DU_name, } - if "rmin" in arg_list: - arg_list.append("qmin") - if "rmax" in arg_list or "rmin" in arg_list: - arg_list.append("qminR") + if type(arg_list) is str: + arg_list = [arg_list] + if arg_list is not None: + if "rmin" in arg_list: + arg_list.append("qmin") + if "rmax" in arg_list or "rmin" in arg_list: + arg_list.append("qminR") valid_arg, range_dict = self._get_valid_arg_list(arg_list, valid_var) if self.verbose: if "rmin" in valid_arg: - key = valid_arg["rmin"] + key = valid_var["rmin"] print(f"{'rmin':<32} {range_dict[key]} {units['rmin']}") - key = valid_arg["qmin"] - print(f"{'':<32} {range_dict[key]} {units['qmin']}") if "rmax" in valid_arg: - key = valid_arg["rmax"] + key = valid_var["rmax"] print(f"{'rmax':<32} {range_dict[key]} {units['rmax']}") - if "rmax" in valid_arg or "rmin" in valid_arg: - key = valid_arg["qminR"] - print(f"{'':<32} {range_dict[key]} {units['qminR']}") - return range_dict