diff --git a/examples/Basic_Simulation/initial_conditions.py b/examples/Basic_Simulation/initial_conditions.py index 640330f1f..c9d0823af 100644 --- a/examples/Basic_Simulation/initial_conditions.py +++ b/examples/Basic_Simulation/initial_conditions.py @@ -30,48 +30,14 @@ from numpy.random import default_rng # Initialize the simulation object as a variable -sim = swiftest.Simulation() +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) # 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['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 +61,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/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: diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index 4678b4784..28303f8f8 100644 --- a/python/swiftest/swiftest/io.py +++ b/python/swiftest/swiftest/io.py @@ -17,9 +17,102 @@ import tempfile import re -newfeaturelist = ("FRAGMENTATION", "ROTATION", "TIDES", "ENERGY", "GR", "YARKOVSKY", "YORP", "IN_FORM") +# 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"] +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): """ Converts a Fortran-generated ASCII string of a real value into a numpy float type. Handles cases where double precision @@ -27,7 +120,7 @@ def real2float(realstr): Parameters ---------- - realstr : string + realstr : str Fortran-generated ASCII string of a real value. Returns @@ -80,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']) @@ -89,22 +183,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() - 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() - 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 @@ -140,7 +227,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", @@ -148,9 +235,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", } @@ -180,10 +267,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: @@ -345,15 +431,24 @@ 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: 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 @@ -484,12 +579,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') @@ -502,7 +597,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') @@ -589,14 +684,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: @@ -620,17 +715,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: @@ -680,14 +775,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]) @@ -1019,14 +1114,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']) @@ -1043,7 +1138,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() @@ -1052,11 +1147,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) @@ -1066,7 +1161,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() @@ -1115,7 +1210,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'] @@ -1131,11 +1226,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) @@ -1182,10 +1277,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: @@ -1213,9 +1308,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'] @@ -1254,17 +1349,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 == '': @@ -1310,11 +1405,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()] @@ -1434,12 +1529,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]) @@ -1609,7 +1704,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) @@ -1698,7 +1793,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) @@ -1770,7 +1865,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 5ad945d69..b72542aaa 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 @@ -9,6 +8,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,61 +19,276 @@ import numpy as np import os import shutil +from typing import ( + Literal, + Dict, + List, + Any +) 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, + t0: float = 0.0, + tstart: float = 0.0, + tstop: float | None = None, + dt: float | None = None, + istep_out: int | None = None, + tstep_out: float | None = None, + 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", + 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, + 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, + 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, + restart: 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. + 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 + * 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 + 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. + 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. + 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. + 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. + * "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", - 'T0': "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", - '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': "", - 'PARTICLE_OUT' : "", - 'BIG_DISCARD': "NO", - 'CHK_CLOSE': "YES", - 'RHILL_PRESENT': "YES", - 'FRAGMENTATION': "NO", - 'ROTATION': "NO", - 'TIDES': "NO", - 'ENERGY': "NO", - '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_unit_system() + self.restart = restart + + 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, + 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, + verbose = False) + + self.set_output_files(output_file_type=output_file_type, + output_file_name=output_file_name, + output_format=output_format, + verbose = False) + + 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, + 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)) - 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() @@ -81,8 +296,703 @@ def __init__(self, codename="Swiftest", param_file="param.in", readbin=False, ve 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, + **kwargs: Any + ): + """ + + 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 + **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 + ------- + 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) + else: + update_list.append("tstop") + + if tstop is not None: + if tstop <= tstart: + print("Error! tstop must be greater than tstart.") + return + + if tstop is not None: + self.param['TSTOP'] = tstop + + if dt is None: + dt = self.param.pop("DT", None) + else: + update_list.append("dt") + + 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 + + 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 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 and dt is not None: + istep_out = int(np.floor(tstep_out / dt)) + + if istep_out is not None: + self.param['ISTEP_OUT'] = istep_out + update_list.append("istep_out") + + if istep_dump is None: + istep_dump = self.param.pop("ISTEP_DUMP",None) + if istep_dump is None: + istep_dump = istep_out + + if istep_dump is not None: + self.param['ISTEP_DUMP'] = istep_dump + update_list.append("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, arg_list: 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 + ---------- + 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"] + + 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" : ""} + + tstep_out = None + if arg_list is None or "tstep_out" in arg_list or "istep_out" in arg_list: + 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] + 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']}") + + 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, + 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, + **kwargs: Any + ): + """ + 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) + 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 + **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 + ------- + 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 restart is not None: + self.param["RESTART"] = restart + update_list.append("restart") + + self.param["TIDES"] = False + + if verbose is None: + verbose = self.verbose + if verbose: + if len(update_list) > 0: + feature_dict = self.get_feature(update_list) + return + + + 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. + 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: + ["close_encounter_check", "general_relativity", "fragmentation", "rotation", "compute_conservation_values"] + + Returns + ------- + feature_dict : dict + A dictionary containing the requested features. + + """ + + 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" + } + + valid_arg, feature_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} {feature_dict[key]}") + + return feature_dict + + 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, + **kwargs: Any + ): + """ + Sets the initial condition file parameters in the parameters dictionary. + + Parameters + ---------- + 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 + * 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. + 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 + ------- + Sets the appropriate initial conditions file parameters inside the self.param dictionary. + + """ + + 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: + 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('{') + 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 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: + 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" + + + 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: + # 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 + + + 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", + } + + 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: + valid_arg = list(valid_var.keys()) + elif type(arg_list) is str: + valid_arg = [arg_list] + else: + 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 + if "init_cond_file_name" in valid_arg: + if self.param["IN_TYPE"] == "ASCII": + valid_arg.remove("init_cond_file_name") + for key in three_file_args: + if key not in valid_arg: + valid_arg.append(key) + else: + for key in three_file_args: + if key in valid_arg: + valid_arg.remove(key) + + valid_arg, init_cond_file_dict = self._get_valid_arg_list(valid_arg, valid_var) + + if self.verbose: + for arg in valid_arg: + 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"] | None = None, + output_file_name: os.PathLike | str | None = None, + output_format: Literal["XV", "XVEL"] | None = None, + verbose: bool | None = None, + **kwargs: Any + ): + """ + Sets the output file parameters in the parameters dictionary. + + Parameters + ---------- + 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. + * 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"}, 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 + **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 + ------- + 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 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 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" + + 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 + + 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 - def set_unit_system(self,MU="Msun",DU="AU",TU="YR",MU2KG=None,DU2M=None,TU2S=None): + """ + + valid_var = {"output_file_type": "OUT_TYPE", + "output_file_name": "BIN_OUT", + "output_format": "OUT_FORM", + } + + valid_arg, output_file_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} {output_file_dict[key]}") + + return output_file_dict + + + 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, + 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. @@ -96,116 +1006,399 @@ 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, 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 + 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. - MU2KG : float, default `None` - The conversion factor to multiply by the mass unit that would convert it to kilogram. Setting this overrides MU + Returns + ---------- + 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 + """ - DU2M : float, default `None` - The conversion factor to multiply by the distance unit that would convert it to meter. Setting this overrides DU + MU2KG_old = None + DU2M_old = None + TU2S_old = None - TU2S : float, default `None` - The conversion factor to multiply by the time unit that would convert it to seconds. Setting this overrides TU + 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: + 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" + 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: + 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: + 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 = "d" + 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" + + + 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 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 + + 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}" + } + + valid_arg, unit_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}: {units1[arg]:<28} {unit_dict[key]} {units2[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. + + Parameters ---------- - Sets the values of MU2KG, DU2M, and TU2S in the param dictionary to the appropriate units. Also computes the gravitational constant, GU + 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] *= MU2KG_old / self.param['MU2KG'] + + if DU2M_old is not None: + for k in distance_keys: + if k in self.param: + 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: + 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] *= TU2S_old / self.param['TU2S'] + + return + + + def set_distance_range(self, + rmin: float | None = None, + rmax: float | None = None, + **kwargs: Any): """ + 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]) + **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 + ------- + Sets the appropriate param dictionary values. - if MU2KG is not None: - self.param['MU2KG'] = MU2KG - self.MU_name = None + """ + CHK_QMIN_RANGE = self.param.pop('CHK_QMIN_RANGE',None) + if CHK_QMIN_RANGE is None: + CHK_QMIN_RANGE = [-1, -1] 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" - else: - print(f"{MU} not a recognized unit system. Using MSun as a default.") - self.param['MU2KG'] = constants.MSun - self.MU_name = "MSun" + 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]}" - if DU2M is not None: - self.param['DU2M'] = DU2M - self.DU_name = None + 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: - 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" + valid_arg = arg_list.copy() - if TU2S is not None: - self.param['TU2S'] = TU2S - self.TU_name = None + if valid_arg is None: + valid_arg = list(valid_var.keys()) + elif type(arg_list) is str: + valid_arg = [arg_list] else: - if TU.upper() == "YR": - self.param['TU2S'] = constants.YR2S - self.TU_name = "y" - elif TU.upper() == "DAY" or TU.upper() == "D" or TU.upper() == "JD": - self.param['TU2S'] = constants.JD2S - self.TU_name = "Day" - elif TU.upper() == "S": - 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" + # 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())] - 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 + # Extract the arg_list dictionary + 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 + + 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 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 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 "rmin" in valid_arg: + key = valid_var["rmin"] + print(f"{'rmin':<32} {range_dict[key]} {units['rmin']}") + if "rmax" in valid_arg: + key = valid_var["rmax"] + print(f"{'rmax':<32} {range_dict[key]} {units['rmax']}") + + return range_dict + - return - def add(self, plname, date=date.today().isoformat(), idval=None): """ Adds a solar system body to an existing simulation DataSet. @@ -476,7 +1669,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": @@ -518,25 +1711,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)