Merge remote-tracking branch 'origin/debug' into debug

python/swiftest/swiftest/constants.py

@@ -13,15 +13,18 @@
 import astropy.constants as const
 
 # Constants in SI units
-GC = np.longdouble(const.G.value)
-AU2M = np.longdouble(const.au.value)
-GMSunSI = np.longdouble(const.GM_sun.value)
-MSun = np.longdouble(const.M_sun.value)
-RSun = np.longdouble(const.R_sun.value)
+GC = const.G.value
+AU2M = const.au.value
+GMSun = const.GM_sun.value
+MSun = const.M_sun.value
+RSun = const.R_sun.value
+MEarth = const.M_earth.value
+REarth = const.R_earth.value
+GMEarth = const.GM_earth.value
 JD2S = 86400
-YR2S = np.longdouble(365.25 * JD2S)
-einsteinC = np.longdouble(299792458.0)
+YR2S = 365.25 * JD2S
+einsteinC = 299792458.0
 # Solar oblatenes values: From Mecheri et al. (2004), using Corbard (b) 2002 values (Table II)
-J2Sun = np.longdouble(2.198e-7)
-J4Sun = np.longdouble(-4.805e-9)
+J2Sun = 2.198e-7
+J4Sun = -4.805e-9
 

python/swiftest/swiftest/init_cond.py

@@ -120,7 +120,7 @@ def solar_system_horizons(plname, idval, param, ephemerides_start_date, ds):
     THIRDLONG = np.longdouble(1.0) / np.longdouble(3.0)
 
     # Central body value vectors
-    GMcb = np.array([swiftest.GMSunSI * param['TU2S'] ** 2 / param['DU2M'] ** 3])
+    GMcb = np.array([swiftest.GMSun * param['TU2S'] ** 2 / param['DU2M'] ** 3])
     Rcb = np.array([swiftest.RSun / param['DU2M']])
     J2RP2 = np.array([swiftest.J2Sun * (swiftest.RSun / param['DU2M']) ** 2])
     J4RP4 = np.array([swiftest.J4Sun * (swiftest.RSun / param['DU2M']) ** 4])

python/swiftest/swiftest/simulation_class.py

@@ -1,4 +1,5 @@
 """
+                    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 
@@ -23,7 +24,7 @@ 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=True, verbose=True):
+    def __init__(self, codename="Swiftest", param_file="param.in", readbin=False, verbose=True):
         self.ds = xr.Dataset()
         self.param = {
             '! VERSION': f"Swiftest parameter input",
@@ -49,9 +50,6 @@ def __init__(self, codename="Swiftest", param_file="param.in", readbin=True, ver
             'CHK_QMIN_COORD': "HELIO",
             'CHK_QMIN_RANGE': "-1.0 -1.0",
             'ENC_OUT': "",
-            'MU2KG': constants.MSun,
-            'TU2S': constants.JD2S,
-            'DU2M': constants.AU2M,
             'EXTRA_FORCE': "NO",
             'DISCARD_OUT': "",
             'PARTICLE_OUT' : "",
@@ -68,16 +66,144 @@ def __init__(self, codename="Swiftest", param_file="param.in", readbin=True, ver
         }
         self.codename = codename
         self.verbose = verbose
-        if param_file != "" :
-            dir_path = os.path.dirname(os.path.realpath(param_file))
+        self.set_unit_system()
+
+        # 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:
-                binpath = os.path.join(dir_path,self.param['BIN_OUT'])
-                if os.path.exists(binpath):
-                    self.param['BIN_OUT'] = binpath
-                    self.bin2xr()
-                else:
-                    print(f"BIN_OUT file {self.param['BIN_OUT']} not found.")
+        if readbin:
+            binpath = os.path.join(self.sim_dir,self.param['BIN_OUT'])
+            if os.path.exists(binpath):
+                self.bin2xr()
+            else:
+                print(f"BIN_OUT file {binpath} not found.")
+        return
+
+
+    def set_unit_system(self,MU="Msun",DU="AU",TU="YR",MU2KG=None,DU2M=None,TU2S=None):
+        """
+        Setter for setting the unit conversion between one of the standard sets.
+
+        The units can be set one of two ways:
+        1) The user can supply string values to the arguments MU, DU, and TU to select between common systems
+        2) The user can supply float values to the arguments MU2KG, DU2M, and TU2S to manually set the conversion
+           factor between the desired unit and the SI unit (kg-m-s).
+
+        The two sets of arguments are mutually exclusive. Any values passed to MU2KG, DU2M, or TU2S will override any
+        specified in MU, DU, or TU, respectively. The default system is Msun-AU-YR. MU, DU, and TU are case-insenstive
+
+        Parameters
+        ----------
+        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, default "AU"
+            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"
+            The time unit system to use. Case-insensitive valid options are:
+            "YR"     : Year
+            "DAY"    : Julian day
+            "d"      : Julian day
+            "JD"     : Julian day
+            "s"      : second
+
+        MU2KG : float, default `None`
+            The conversion factor to multiply by the mass unit that would convert it to kilogram. Setting this overrides MU
+
+        DU2M  : float, default `None`
+            The conversion factor to multiply by the distance unit that would convert it to meter. Setting this overrides DU
+
+        TU2S  : float, default `None`
+            The conversion factor to multiply by the time unit that would convert it to seconds. Setting this overrides TU
+
+        Returns
+        ----------
+        Sets the values of MU2KG, DU2M, and TU2S in the param dictionary to the appropriate units. Also computes the gravitational constant, GU
+        """
+
+        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:
+            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:
+            self.param['TU2S'] = TU2S
+            self.TU_name = None
+        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"
+
+        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
+
+        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}")
+
         return
 
     def add(self, plname, date=date.today().isoformat(), idval=None):
@@ -273,11 +399,17 @@ def bin2xr(self):
         -------
             self.ds : xarray dataset
         """
+
+        # Make a temporary copy of the parameter dictionary so we can supply the absolute path of the binary file
+        # This is done to handle cases where the method is called from a different working directory than the simulation
+        # results
+        param_tmp = self.param.copy()
+        param_tmp['BIN_OUT'] = os.path.join(self.dir_path,self.param['BIN_OUT'])
         if self.codename == "Swiftest":
-            self.ds = io.swiftest2xr(self.param, verbose=self.verbose)
+            self.ds = io.swiftest2xr(param_tmp, verbose=self.verbose)
             if self.verbose: print('Swiftest simulation data stored as xarray DataSet .ds')
         elif self.codename == "Swifter":
-            self.ds = io.swifter2xr(self.param, verbose=self.verbose)
+            self.ds = io.swifter2xr(param_tmp, verbose=self.verbose)
             if self.verbose: print('Swifter simulation data stored as xarray DataSet .ds')
         elif self.codename == "Swift":
             print("Reading Swift simulation data is not implemented yet")