diff --git a/examples/helio_gr_test/init_cond.py b/examples/helio_gr_test/init_cond.py new file mode 100755 index 000000000..9e89dc358 --- /dev/null +++ b/examples/helio_gr_test/init_cond.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python3 +import swiftest + +sim = swiftest.Simulation() +sim.param['PL_IN'] = "pl.swiftest.in" +sim.param['TP_IN'] = "tp.swiftest.in" +sim.param['CB_IN'] = "cb.swiftest.in" +sim.param['BIN_OUT'] = "bin.swiftest.nc" + +sim.param['MU2KG'] = swiftest.MSun +sim.param['TU2S'] = swiftest.YR2S +sim.param['DU2M'] = swiftest.AU2M +sim.param['T0'] = 0.0 +sim.param['DT'] = 0.25 * swiftest.JD2S / swiftest.YR2S +sim.param['TSTOP'] = 1000.0 +sim.param['ISTEP_OUT'] = 1461 +sim.param['ISTEP_DUMP'] = 1461 +sim.param['CHK_QMIN_COORD'] = "HELIO" +sim.param['CHK_QMIN'] = swiftest.RSun / swiftest.AU2M +sim.param['CHK_QMIN_RANGE'] = f"{swiftest.RSun / swiftest.AU2M} 1000.0" +sim.param['CHK_RMIN'] = swiftest.RSun / swiftest.AU2M +sim.param['CHK_RMAX'] = 1000.0 +sim.param['CHK_EJECT'] = 1000.0 +sim.param['OUT_STAT'] = "UNKNOWN" +sim.param['IN_FORM'] = "EL" +sim.param['OUT_FORM'] = "XVEL" +sim.param['OUT_TYPE'] = "NETCDF_DOUBLE" +sim.param['RHILL_PRESENT'] = "YES" +sim.param['GR'] = 'YES' + +bodyid = { + "Sun": 0, + "Mercury": 1, + "Venus": 2, + "Earth": 3, + "Mars": 4, + "Jupiter": 5, + "Saturn": 6, + "Uranus": 7, + "Neptune": 8, +} + +for name, id in bodyid.items(): + sim.add(name, idval=id, date="2027-04-30") + +sim.save("param.swiftest.in") + + diff --git a/examples/whm_gr_test/init_cond.py b/examples/whm_gr_test/init_cond.py new file mode 100644 index 000000000..7904eb100 --- /dev/null +++ b/examples/whm_gr_test/init_cond.py @@ -0,0 +1,226 @@ +import numpy as np +import sys +from astroquery.jplhorizons import Horizons +import astropy.constants as const + +#Values from JPL Horizons +AU2M = const.au.value +GMSunSI = const.GM_sun.value +Rsun = const.R_sun.value +GC = const.G.value +JD = 86400 +year = 365.25 * JD +c = 299792458.0 +MSun_over_Mpl = [6023600.0, + 408523.71, + 328900.56, + 3098708., + 1047.3486, + 3497.898, + 22902.98, + 19412.24, + 1.35e8] + +MU2KG = GMSunSI / GC #Conversion from mass unit to kg +DU2M = AU2M #Conversion from radius unit to centimeters +TU2S = year #Conversion from time unit to seconds +GU = GC / (DU2M**3 / (MU2KG * TU2S**2)) + +GMSun = GMSunSI / (DU2M**3 / TU2S**2) + +t_print = 10.e0 * year / TU2S #output interval to print results +deltaT = 0.25 * JD / TU2S #timestep simulation +end_sim = 1.0e3 * year / TU2S + t_print #end time + +# Solar oblatenes values: From Mecheri et al. (2004), using Corbard (b) 2002 values (Table II) +J2 = 2.198e-7 * (Rsun / DU2M)**2 +J4 = -4.805e-9 * (Rsun / DU2M)**4 + +tstart = '2021-01-28' +tend = '2021-01-29' +tstep = '1d' +planetid = { + 'mercury' : '1', + 'venus' : '2', + 'earthmoon' : '3', + 'mars' : '4', + 'jupiter' : '5', + 'saturn' : '6', + 'uranus' : '7', + 'neptune' : '8', + 'plutocharon' : '9' +} +npl = 9 + +#Planet Msun/M ratio +MSun_over_Mpl = { + 'mercury' : 6023600.0, + 'venus' : 408523.71, + 'earthmoon' : 328900.56, + 'mars' : 3098708., + 'jupiter' : 1047.3486, + 'saturn' : 3497.898, + 'uranus' : 22902.98, + 'neptune' : 19412.24, + 'plutocharon' : 1.35e8 +} + +#Planet radii in meters +Rpl = { + 'mercury' : 2439.4e3, + 'venus' : 6051.8e3, + 'earthmoon' : 6371.0084e3, # Earth only for radius + 'mars' : 3389.50e3, + 'jupiter' : 69911e3, + 'saturn' : 58232.0e3, + 'uranus' : 25362.e3, + 'neptune' : 24622.e3, + 'plutocharon' : 1188.3e3 +} + +pdata = {} +plvec = {} +Rhill = {} + +for key,val in planetid.items(): + pdata[key] = Horizons(id=val, id_type='majorbody',location='@sun', + epochs={'start': tstart, 'stop': tend, + 'step': tstep}) + plvec[key] = np.array([pdata[key].vectors()['x'][0], + pdata[key].vectors()['y'][0], + pdata[key].vectors()['z'][0], + pdata[key].vectors()['vx'][0], + pdata[key].vectors()['vy'][0], + pdata[key].vectors()['vz'][0] + ]) + Rhill[key] = pdata[key].elements()['a'][0] * (3 * MSun_over_Mpl[key])**(-1.0 / 3.0) + + +if __name__ == '__main__': + # Convert from AU-day to AU-year just because I find it easier to keep track of the sim progress + for plid in plvec: + plvec[plid][3:] *= year / JD + + # Names of all output files + swifter_input = "param.swifter.in" + swifter_pl = "pl.swifter.in" + swifter_tp = "tp.swifter.in" + swifter_bin = "bin.swifter.dat" + swifter_enc = "enc.swifter.dat" + + swiftest_input = "param.swiftest.in" + swiftest_pl = "pl.swiftest.in" + swiftest_tp = "tp.swiftest.in" + swiftest_cb = "cb.swiftest.in" + swiftest_bin = "bin.swiftest.dat" + swiftest_enc = "enc.swiftest.dat" + + # Simulation start, stop, and output cadence times + t_0 = 0 # simulation start time + end_sim = 1000.0e0 * year / TU2S # simulation end time + deltaT = 0.25 * JD / TU2S # simulation step size + t_print = 1.0 * year / TU2S #output interval to print results + + iout = int(np.ceil(t_print / deltaT)) + rmin = Rsun / DU2M + rmax = 1000.0 + + #Make Swifter files + plfile = open(swifter_pl, 'w') + print(f'{npl+1} ! Planet input file generated using init_cond.py using JPL Horizons data for the major planets (and Pluto) for epoch {tstart}' ,file=plfile) + print(f'1 {GMSun}',file=plfile) + print(f'0.0 0.0 0.0',file=plfile) + print(f'0.0 0.0 0.0',file=plfile) + for i, plid in enumerate(plvec): + print(f'{i + 2} {GMSun / MSun_over_Mpl[plid]} {Rhill[plid]}', file=plfile) + print(f'{Rpl[plid] / DU2M}', file=plfile) + print(f'{plvec[plid][0]} {plvec[plid][1]} {plvec[plid][2]}', file=plfile) + print(f'{plvec[plid][3]} {plvec[plid][4]} {plvec[plid][5]}', file=plfile) + plfile.close() + + tpfile = open(swifter_tp, 'w') + print(0,file=tpfile) + tpfile.close() + + sys.stdout = open(swifter_input, "w") + print(f'! Swifter input file generated using init_cond.py') + print(f'T0 {t_0} ') + print(f'TSTOP {end_sim}') + print(f'DT {deltaT}') + print(f'PL_IN {swifter_pl}') + print(f'TP_IN {swifter_tp}') + print(f'IN_TYPE ASCII') + print(f'ISTEP_OUT {iout:d}') + print(f'ISTEP_DUMP {iout:d}') + print(f'BIN_OUT {swifter_bin}') + print(f'OUT_TYPE REAL8') + print(f'OUT_FORM EL') + print(f'OUT_STAT NEW') + print(f'J2 {J2}') + print(f'J4 {J4}') + print(f'CHK_CLOSE yes') + print(f'CHK_RMIN {rmin}') + print(f'CHK_RMAX {rmax}') + print(f'CHK_EJECT {rmax}') + print(f'CHK_QMIN {rmin}') + print(f'CHK_QMIN_COORD HELIO') + print(f'CHK_QMIN_RANGE {rmin} {rmax}') + print(f'ENC_OUT {swifter_enc}') + print(f'EXTRA_FORCE no') + print(f'BIG_DISCARD no') + print(f'RHILL_PRESENT yes') + print(f'C {c / (DU2M / TU2S)}') + + #Now make Swiftest files + cbfile = open(swiftest_cb, 'w') + print(f'{1.0}',file=cbfile) + print(f'{rmin}',file=cbfile) + print(f'{J2}',file=cbfile) + print(f'{J4}',file=cbfile) + + plfile = open(swiftest_pl, 'w') + print(npl,file=plfile) + + for i, plid in enumerate(plvec): + print(f'{i + 2} {1.0 / MSun_over_Mpl[plid]}', file=plfile) + print(f'{Rpl[plid] / DU2M}', file=plfile) + print(f'{plvec[plid][0]} {plvec[plid][1]} {plvec[plid][2]}', file=plfile) + print(f'{plvec[plid][3]} {plvec[plid][4]} {plvec[plid][5]}', file=plfile) + plfile.close() + tpfile = open(swiftest_tp, 'w') + print(0,file=tpfile) + tpfile.close() + + sys.stdout = open(swiftest_input, "w") + print(f'! Swiftest input file generated using init_cond.py') + print(f'T0 {t_0} ') + print(f'TSTOP {end_sim}') + print(f'DT {deltaT}') + print(f'CB_IN {swiftest_cb}') + print(f'PL_IN {swiftest_pl}') + print(f'TP_IN {swiftest_tp}') + print(f'IN_TYPE ASCII') + print(f'ISTEP_OUT {iout:d}') + print(f'ISTEP_DUMP {iout:d}') + print(f'BIN_OUT {swiftest_bin}') + print(f'OUT_TYPE REAL8') + print(f'OUT_FORM EL') + print(f'OUT_STAT REPLACE') + print(f'CHK_CLOSE yes') + print(f'CHK_RMIN {rmin}') + print(f'CHK_RMAX {rmax}') + print(f'CHK_EJECT {rmax}') + print(f'CHK_QMIN {rmin}') + print(f'CHK_QMIN_COORD HELIO') + print(f'CHK_QMIN_RANGE {rmin} {rmax}') + print(f'ENC_OUT {swiftest_enc}') + print(f'EXTRA_FORCE no') + print(f'BIG_DISCARD no') + print(f'ROTATION no') + print(f'GR yes') + print(f'MU2KG {MU2KG}') + print(f'DU2M {DU2M}') + print(f'TU2S {TU2S}') + + + sys.stdout = sys.__stdout__ diff --git a/examples/whm_gr_test/swiftest_relativity.ipynb b/examples/whm_gr_test/swiftest_relativity.ipynb new file mode 100644 index 000000000..0cd73753f --- /dev/null +++ b/examples/whm_gr_test/swiftest_relativity.ipynb @@ -0,0 +1,1050 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import swiftest\n", + "from astroquery.jplhorizons import Horizons\n", + "import datetime\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Creating the Sun as a central body\n", + "Fetching ephemerides data for Mercury from JPL/Horizons\n", + "Fetching ephemerides data for Venus from JPL/Horizons\n", + "Fetching ephemerides data for Earth from JPL/Horizons\n", + "Fetching ephemerides data for Mars from JPL/Horizons\n", + "Fetching ephemerides data for Jupiter from JPL/Horizons\n", + "Fetching ephemerides data for Saturn from JPL/Horizons\n", + "Fetching ephemerides data for Uranus from JPL/Horizons\n", + "Fetching ephemerides data for Neptune from JPL/Horizons\n", + "Writing initial conditions to file init_cond.nc\n", + "Writing parameter inputs to file /home/daminton/git_debug/swiftest/examples/whm_gr_test/param.gr.in\n" + ] + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "
<xarray.Dataset>\n",
+       "Dimensions:        (name: 9, time: 1)\n",
+       "Coordinates:\n",
+       "  * name           (name) <U32 'Sun' 'Mercury' 'Venus' ... 'Uranus' 'Neptune'\n",
+       "  * time           (time) float64 0.0\n",
+       "Data variables: (12/14)\n",
+       "    particle_type  (name) <U32 'Central Body' 'Massive Body' ... 'Massive Body'\n",
+       "    id             (name) int64 0 1 2 3 4 5 6 7 8\n",
+       "    Gmass          (time, name) float64 39.48 6.554e-06 ... 0.001724 0.002034\n",
+       "    radius         (time, name) float64 0.00465 1.631e-05 ... 0.0001646\n",
+       "    j2rp2          (time, name) float64 4.754e-12 nan nan nan ... nan nan nan\n",
+       "    j4rp4          (time, name) float64 -2.247e-18 nan nan nan ... nan nan nan\n",
+       "    ...             ...\n",
+       "    xhz            (time, name) float64 nan -0.002529 ... -0.0407 -0.7287\n",
+       "    vhx            (time, name) float64 nan -9.241 3.454 ... -1.315 -0.08755\n",
+       "    vhy            (time, name) float64 nan 7.755 6.477 ... 1.932 0.5372 1.149\n",
+       "    vhz            (time, name) float64 nan 1.481 -0.1103 ... 0.01905 -0.02168\n",
+       "    ntp            (time) int64 0\n",
+       "    npl            (time) int64 8
" + ], + "text/plain": [ + "\n", + "Dimensions: (name: 9, time: 1)\n", + "Coordinates:\n", + " * name (name) \n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "
<xarray.Dataset>\n",
+       "Dimensions:        (name: 9, time: 1)\n",
+       "Coordinates:\n",
+       "  * name           (name) <U32 'Sun' 'Mercury' 'Venus' ... 'Uranus' 'Neptune'\n",
+       "  * time           (time) float64 0.0\n",
+       "Data variables: (12/14)\n",
+       "    particle_type  (name) <U32 'Central Body' 'Massive Body' ... 'Massive Body'\n",
+       "    id             (name) int64 0 1 2 3 4 5 6 7 8\n",
+       "    Gmass          (time, name) float64 39.48 6.554e-06 ... 0.001724 0.002034\n",
+       "    radius         (time, name) float64 0.00465 1.631e-05 ... 0.0001646\n",
+       "    j2rp2          (time, name) float64 4.754e-12 nan nan nan ... nan nan nan\n",
+       "    j4rp4          (time, name) float64 -2.247e-18 nan nan nan ... nan nan nan\n",
+       "    ...             ...\n",
+       "    xhz            (time, name) float64 nan -0.002529 ... -0.0407 -0.7287\n",
+       "    vhx            (time, name) float64 nan -9.241 3.454 ... -1.315 -0.08755\n",
+       "    vhy            (time, name) float64 nan 7.755 6.477 ... 1.932 0.5372 1.149\n",
+       "    vhz            (time, name) float64 nan 1.481 -0.1103 ... 0.01905 -0.02168\n",
+       "    ntp            (time) int64 0\n",
+       "    npl            (time) int64 8
" + ], + "text/plain": [ + "\n", + "Dimensions: (name: 9, time: 1)\n", + "Coordinates:\n", + " * name (name) 0, drop=True) pl = pl.where(np.invert(np.isnan(pl['Gmass'])), drop=True).drop_vars(['j2rp2', 'j4rp4']) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 7c17c3d32..c6fe5e8da 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -68,6 +68,8 @@ def __init__(self,read_param: bool = True, **kwargs: Any): integrator : {"symba","rmvs","whm","helio"}, default "symba" Name of the n-body integrator that will be used when executing a run. Parameter input file equivalent: None + read_param : bool, default False + Read the parameter file given by `param_file`. param_file : str, path-like, or file-lke, default "param.in" Name of the parameter input file that will be passed to the integrator. Parameter input file equivalent: None @@ -291,7 +293,7 @@ def __init__(self,read_param: bool = True, **kwargs: Any): # Set the location of the parameter input file param_file = kwargs.pop("param_file",self.param_file) - read_param = kwargs.pop("read_param",True) + read_param = kwargs.pop("read_param",False) self.set_parameter(verbose=False,param_file=param_file) #----------------------------------------------------------------- @@ -301,9 +303,8 @@ def __init__(self,read_param: bool = True, **kwargs: Any): # If the user asks to read in an old parameter file, override any default parameters with values from the file # If the file doesn't exist, flag it for now so we know to create it if read_param: - if os.path.exists(self.param_file): - self.read_param(self.param_file, codename=self.codename, verbose=self.verbose) - param_file_found = True + #good_param is self.read_param() + if self.read_param(): # We will add the parameter file to the kwarg list. This will keep the set_parameter method from # overriding everything with defaults when there are no arguments passed to Simulation() kwargs['param_file'] = self.param_file @@ -743,7 +744,7 @@ def get_parameter(self, **kwargs): return param_dict def set_integrator(self, - codename: Literal["swiftest", "swifter", "swift"] | None = None, + codename: Literal["Swiftest", "Swifter", "Swift"] | None = None, integrator: Literal["symba","rmvs","whm","helio"] | None = None, mtiny: float | None = None, gmtiny: float | None = None, @@ -2316,58 +2317,82 @@ def _combine_and_fix_dsnew(self,dsnew): Updated Dataset with ntp, npl values and types fixed. """ + if "id" not in self.data.dims: + if len(np.unique(dsnew['name'])) == len(dsnew['name']): + dsnew = dsnew.swap_dims({"id" : "name"}) + dsnew = dsnew.reset_coords("id") + else: + warnings.warn("Non-unique names detected for bodies. The Dataset will be dimensioned by integer id instead of name.") + warnings.warn("Consider using unique names instead.") + + if self.param['OUT_TYPE'] == "NETCDF_DOUBLE": + dsnew = io.fix_types(dsnew, ftype=np.float64) + elif self.param['OUT_TYPE'] == "NETCDF_FLOAT": + dsnew = io.fix_types(dsnew, ftype=np.float32) self.data = xr.combine_by_coords([self.data, dsnew]) def get_nvals(ds): + if "name" in ds.dims: + count_dim = "name" + elif "id" in ds.dims: + count_dim = "id" if "Gmass" in ds: - ds['ntp'] = ds['id'].where(np.isnan(ds['Gmass'])).count(dim="id") - ds['npl'] = ds['id'].where(~(np.isnan(ds['Gmass']))).count(dim="id") - 1 + ds['ntp'] = ds[count_dim].where(np.isnan(ds['Gmass'])).count(dim=count_dim) + ds['npl'] = ds[count_dim].where(~(np.isnan(ds['Gmass']))).count(dim=count_dim) - 1 else: - ds['ntp'] = ds['id'].count(dim="id") + ds['ntp'] = ds[count_dim].count(dim=count_dim) ds['npl'] = xr.full_like(ds['ntp'],0) return ds dsnew = get_nvals(dsnew) self.data = get_nvals(self.data) - if self.param['OUT_TYPE'] == "NETCDF_DOUBLE": - dsnew = io.fix_types(dsnew, ftype=np.float64) - self.data = io.fix_types(self.data, ftype=np.float64) - elif self.param['OUT_TYPE'] == "NETCDF_FLOAT": - dsnew = io.fix_types(dsnew, ftype=np.float32) - self.data = io.fix_types(self.data, ftype=np.float32) - return dsnew - def read_param(self, param_file, codename="Swiftest", verbose=True): + def read_param(self, + param_file : os.PathLike | str | None = None, + codename: Literal["Swiftest", "Swifter", "Swift"] | None = None, + verbose: bool | None = None): """ Reads in an input parameter file and stores the values in the param dictionary. Parameters ---------- - param_file : string + param_file : str or path-like, default is the value of the Simulation object's internal `param_file`. File name of the input parameter file - codename : string + codename : {"Swiftest", "Swifter", "Swift"}, default is the value of the Simulation object's internal`codename` Type of parameter file, either "Swift", "Swifter", or "Swiftest" + verbose : bool, default is the value of the Simulation object's internal `verbose` + 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. Returns ------- - + True if the parameter file exists and is read correctly. False otherwise. """ + if param_file is None: + param_file = self.param_file + + if coename is None: + codename = self.codename + + if verbose is None: + verbose = self.verbose + + if not os.path.exists(param_file): + return False + if codename == "Swiftest": - param_old = self.param.copy() - self.param = io.read_swiftest_param(param_file, param_old, verbose=verbose) - self.codename = "Swiftest" + self.param = io.read_swiftest_param(param_file, param, verbose=verbose) elif codename == "Swifter": self.param = io.read_swifter_param(param_file, verbose=verbose) - self.codename = "Swifter" elif codename == "Swift": self.param = io.read_swift_param(param_file, verbose=verbose) - self.codename = "Swift" else: warnings.warn(f'{codename} is not a recognized code name. Valid options are "Swiftest", "Swifter", or "Swift".') - self.codename = "Unknown" - return + return False + + return True def write_param(self, codename: Literal["Swiftest", "Swifter", "Swift"] | None = None,