Merge branch 'Fragmentation' of https://github.itap.purdue.edu/Minton…

…Group/swiftest into Fragmentation

python/swiftestio/swiftestio.py

@@ -2,16 +2,22 @@
 import pandas as pd
 from scipy.io import FortranFile
 import xarray as xr
-import matplotlib.pyplot as plt
-import sys, time
-
-
-
-#I/O Routines for reading in Swifter and Swiftest parameter and binary data files
-
-
-GC = 6.6743E-11
-einstinC = 299792458.0
+from astroquery.jplhorizons import Horizons
+import astropy.constants as const
+import datetime
+import sys
+
+# Constants in SI units
+AU2M = np.longdouble(const.au.value)
+GMSunSI = np.longdouble(const.GM_sun.value)
+RSun = np.longdouble(const.R_sun.value)
+GC = np.longdouble(const.G.value)
+JD2S = 86400
+year = np.longdouble(365.25 * JD2S)
+einsteinC = np.longdouble(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)
 
 def read_swifter_param(inparfile):
    """
@@ -289,6 +295,38 @@ def swifter_stream(f, param):
       yield t, npl, plid, pvec.T, plab, \
            ntp, tpid, tvec.T, tlab
 
+def make_swiftest_labels(config):
+   tlab = []
+   if config['OUT_FORM'] == 'XV':
+      tlab.append('px')
+      tlab.append('py')
+      tlab.append('pz')
+      tlab.append('vx')
+      tlab.append('vy')
+      tlab.append('vz')
+   elif config['OUT_FORM'] == 'EL':
+      tlab.append('a')
+      tlab.append('e')
+      tlab.append('inc')
+      tlab.append('capom')
+      tlab.append('omega')
+      tlab.append('capm')
+   plab = tlab.copy()
+   plab.append('mass')
+   plab.append('radius')
+   if config['ROTATION'] == 'YES':
+      plab.append('Ip_x')
+      plab.append('Ip_y')
+      plab.append('Ip_z')
+      plab.append('rot_x')
+      plab.append('rot_y')
+      plab.append('rot_z')
+   if config['TIDES'] == 'YES':
+      plab.append('k2')
+      plab.append('Q')
+   return plab, tlab
+
+
 def swiftest_stream(f, config):
    """
    Reads in a Swiftest bin.dat file and returns a single frame of data as a datastream
@@ -354,31 +392,7 @@ def swiftest_stream(f, config):
          t5 = f.read_reals(np.float64)
          t6 = f.read_reals(np.float64)
 
-      tlab = []
-      if config['OUT_FORM'] == 'XV':
-         tlab.append('px')
-         tlab.append('py')
-         tlab.append('pz')
-         tlab.append('vx')
-         tlab.append('vy')
-         tlab.append('vz')
-      elif config['OUT_FORM'] == 'EL':
-         tlab.append('a')
-         tlab.append('e')
-         tlab.append('inc')
-         tlab.append('capom')
-         tlab.append('omega')
-         tlab.append('capm')
-      plab = tlab.copy()
-      plab.append('mass')
-      plab.append('radius')
-      if config['ROTATION'] == 'YES':
-         plab.append('rot_x')
-         plab.append('rot_y')
-         plab.append('rot_z')
-         plab.append('Ip_1')
-         plab.append('Ip_2')
-         plab.append('Ip_3')
+      plab, tlab = make_swiftest_labels(config)
 
       if npl > 0:
          if config['ROTATION'] == 'YES':
@@ -516,10 +530,237 @@ def swiftest_particle_2xr(ds, config):
    infoxr = vecda.to_dataset(dim='vec')
    infoxr['origin_type'] = typeda
 
-   print('\nAdding particle info Dataset')
+   print('\nAdding particle info to Dataset')
    ds = xr.merge([ds, infoxr])
    return ds
 
+
+def solar_system_pl(config, ephemerides_start_date):
+   """
+   Initializes a Swiftest dataset containing the major planets of the Solar System at a particular data from JPL/Horizons
+
+   Parameters
+   ----------
+   config : dict
+       Swiftest Configuration parameters. This method uses the unit conversion factors to convert from JPL's AU-day system into the system specified in the config file
+   ephemerides_start_date : string
+       Date to use when obtaining the ephemerides in the format YYYY-MM-DD
+
+   Returns
+   -------
+   xarray dataset
+   """
+   # Planet ids
+   planetid = {
+      'mercury': '1',
+      'venus': '2',
+      'earthmoon': '3',
+      'mars': '4',
+      'jupiter': '5',
+      'saturn': '6',
+      'uranus': '7',
+      'neptune': '8',
+      'plutocharon': '9'
+   }
+
+   # Planet Msun/M ratio
+   MSun_over_Mpl = {
+      'mercury': np.longdouble(6023600.0),
+      'venus': np.longdouble(408523.71),
+      'earthmoon': np.longdouble(328900.56),
+      'mars': np.longdouble(3098708.),
+      'jupiter': np.longdouble(1047.3486),
+      'saturn': np.longdouble(3497.898),
+      'uranus': np.longdouble(22902.98),
+      'neptune': np.longdouble(19412.24),
+      'plutocharon': np.longdouble(1.35e8)
+   }
+
+   # Planet radii in meters
+   planetradius = {
+      'mercury': np.longdouble(2439.4e3),
+      'venus': np.longdouble(6051.8e3),
+      'earthmoon': np.longdouble(6371.0084e3),  # Earth only for radius
+      'mars': np.longdouble(3389.50e3),
+      'jupiter': np.longdouble(69911e3),
+      'saturn': np.longdouble(58232.0e3),
+      'uranus': np.longdouble(25362.e3),
+      'neptune': np.longdouble(24622.e3),
+      'plutocharon': np.longdouble(1188.3e3)
+   }
+
+   # Unit conversion factors
+   DCONV = AU2M / config['DU2M']
+   VCONV = (AU2M / JD2S) / (config['DU2M'] / config['TU2S'])
+   THIRDLONG = np.longdouble(1.0) / np.longdouble(3.0)
+
+   # Central body value vectors
+   GMcb = np.array([GMSunSI * config['TU2S'] ** 2 / config['DU2M'] ** 3])
+   Rcb = np.array([RSun / config['DU2M']])
+   J2RP2 = np.array([J2Sun * (RSun / config['DU2M']) ** 2])
+   J4RP4 = np.array([J4Sun * (RSun / config['DU2M']) ** 4])
+   cbid = np.array([0])
+   cvec = np.vstack([GMcb, Rcb, J2RP2, J4RP4])
+
+   # Horizons date time internal variables
+   tstart = datetime.date.fromisoformat(ephemerides_start_date)
+   tstep = datetime.timedelta(days=1)
+   tend = tstart + tstep
+   ephemerides_end_date = tend.isoformat()
+   ephemerides_step = '1d'
+
+   pldata = {}
+   p1 = []
+   p2 = []
+   p3 = []
+   p4 = []
+   p5 = []
+   p6 = []
+   Rhill = []
+   Rpl = []
+   GMpl = []
+
+   # Fetch solar system ephemerides from Horizons
+   for key, val in planetid.items():
+      pldata[key] = Horizons(id=val, id_type='majorbody', location='@sun',
+                             epochs={'start': ephemerides_start_date, 'stop': ephemerides_end_date,
+                                     'step': ephemerides_step})
+      if config['OUT_FORM'] == 'XV':
+         p1.append(pldata[key].vectors()['x'][0] * DCONV)
+         p2.append(pldata[key].vectors()['y'][0] * DCONV)
+         p3.append(pldata[key].vectors()['z'][0] * DCONV)
+         p4.append(pldata[key].vectors()['vx'][0] * VCONV)
+         p5.append(pldata[key].vectors()['vy'][0] * VCONV)
+         p6.append(pldata[key].vectors()['vz'][0] * VCONV)
+      elif config['OUT_FORM'] == 'EL':
+         p1.append(pldata[key].elements()['a'][0] * DCONV)
+         p2.append(pldata[key].elements()['e'][0])
+         p3.append(pldata[key].elements()['inc'][0] * np.pi / 180.0)
+         p4.append(pldata[key].elements()['Omega'][0] * np.pi / 180.0)
+         p5.append(pldata[key].elements()['w'][0] * np.pi / 180.0)
+         p6.append(pldata[key].elements()['M'][0] * np.pi / 180.0)
+      Rhill.append(pldata[key].elements()['a'][0] * (3 * MSun_over_Mpl[key]) ** (-THIRDLONG))
+      Rpl.append(planetradius[key] * DCONV)
+      GMpl.append(GMcb[0] / MSun_over_Mpl[key])
+   # Generate planet value vectors
+   plid = np.fromiter(planetid.values(), dtype=int)
+   pvec = np.vstack([p1, p2, p3, p4, p5, p6, GMpl, Rpl])
+
+   dims = ['time', 'id', 'vec']
+   cb = []
+   pl = []
+   tp = []
+   t = np.array([0.0])
+
+   clab, plab, tlab = make_swiftest_labels(config)
+
+   # Prepare frames by adding an extra axis for the time coordinate
+   cbframe = np.expand_dims(cvec.T, axis=0)
+   plframe = np.expand_dims(pvec.T, axis=0)
+
+   # Create xarray DataArrays out of each body type
+   cbxr = xr.DataArray(cbframe, dims=dims, coords={'time': t, 'id': cbid, 'vec': clab})
+   plxr = xr.DataArray(plframe, dims=dims, coords={'time': t, 'id': plid, 'vec': plab})
+
+   cb.append(cbxr)
+   pl.append(plxr)
+
+   cbda = xr.concat(cb, dim='time')
+   plda = xr.concat(pl, dim='time')
+
+   cbds = cbda.to_dataset(dim='vec')
+   plds = plda.to_dataset(dim='vec')
+   ds = xr.combine_by_coords([cbds, plds])
+   return ds
+
+
+def swiftest_xr2_infile(ds, config, framenum=-1):
+   """
+   Writes a set of Swiftest input files from a single frame of a Swiftest xarray dataset
+
+   Parameters
+   ----------
+   ds : xarray dataset
+       Dataset containing Swiftest n-body data in XV format
+   framenum : int
+       Time frame to use to generate the initial conditions. If this argument is not passed, the default is to use the last frame in the dataset.
+   config : dict
+       Swiftest Configuration parameters. This method uses the names of the cb, pl, and tp files from the configuration
+
+   Returns
+   -------
+   A set of three input files for a Swiftest run
+   """
+   frame = ds.isel(time=framenum)
+   pl = frame.where(np.invert(np.isnan(frame['Mass'])), drop=True)
+   tp = frame.where(np.isnan(frame['Mass']), drop=True).drop_vars(['Mass', 'Radius'])
+   tp = tp.where(np.invert(np.isnan(tp['px'])), drop=True)
+
+   GMSun = np.double(pl['Mass'].isel(id=0))
+   RSun = np.double(pl['Radius'].isel(id=0))
+
+   if config['IN_TYPE'] == 'ASCII':
+      # Swiftest PL file
+      plfile = open(config['PL_IN'], 'w')
+      print(pl.id.count().values, file=plfile)
+      for i in pl.id:
+         pli = pl.sel(id=i)
+         print(i.values, pli['Mass'].values, file=plfile)
+         if i != 0:
+            print(pli['Radius'].values, file=plfile)
+         print(pli['px'].values, pli['py'].values, pli['pz'].values, file=plfile)
+         print(pli['vx'].values, pli['vy'].values, pli['vz'].values, file=plfile)
+      if config['ROTATION'] == 'YES':
+         print(pli['Ip_x'].values, pli['Ip_y'].values, pli['Ip_z'].values, file=plfile)
+         print(pli['rot_x'].values, pli['rot_y'].values, pli['rot_z'].values, file=plfile)
+
+         plfile.close()
+
+      # TP file
+      tpfile = open(config['TP_IN'], 'w')
+      print(tp.id.count().values, file=tpfile)
+      for i in tp.id:
+         tpi = tp.sel(id=i)
+         print(i.values, file=tpfile)
+         print(tpi['px'].values, tpi['py'].values, tpi['pz'].values, file=tpfile)
+         print(tpi['vx'].values, tpi['vy'].values, tpi['vz'].values, file=tpfile)
+      tpfile.close()
+   elif config['IN_TYPE'] == 'REAL8':
+      # Now make Swiftest files
+      cbfile = FortranFile(swiftest_cb, 'w')
+      Msun = np.double(1.0)
+      cbfile.write_record(np.double(GMSun))
+      cbfile.write_record(np.double(rmin))
+      cbfile.write_record(np.double(J2))
+      cbfile.write_record(np.double(J4))
+      cbfile.close()
+
+      plfile = FortranFile(swiftest_pl, 'w')
+      plfile.write_record(npl)
+
+      plfile.write_record(plid)
+      plfile.write_record(p_pl[0])
+      plfile.write_record(p_pl[1])
+      plfile.write_record(p_pl[2])
+      plfile.write_record(v_pl[0])
+      plfile.write_record(v_pl[1])
+      plfile.write_record(v_pl[2])
+      plfile.write_record(mass)
+      plfile.write_record(radius)
+      plfile.close()
+      tpfile = FortranFile(swiftest_tp, 'w')
+      ntp = 1
+      tpfile.write_record(ntp)
+      tpfile.write_record(tpid)
+      tpfile.write_record(p_tp[0])
+      tpfile.write_record(p_tp[1])
+      tpfile.write_record(p_tp[2])
+      tpfile.write_record(v_tp[0])
+      tpfile.write_record(v_tp[1])
+      tpfile.write_record(v_tp[2])
+   else:
+      print(f"{config['IN_TYPE']} is an unknown file type")
+
 if __name__ == '__main__':
 
    workingdir = '/Users/daminton/git/swiftest/examples/symba_mars_disk/'