Wrote a new swiftestio method that can create a Dataset from JPL/Hori…

…zons data. Updated inititial conditions generator to use it in one example. Only the Swiftest version has been completed
MintonGroup · Jun 17, 2021 · 6658513 · 6658513
1 parent 4f1069d
commit 6658513
Show file tree

Hide file tree

Showing 12 changed files with 525 additions and 267 deletions.
diff --git a/examples/rmvs_swifter_comparison/1pl_1tp_encounter/.idea/.gitignore b/examples/rmvs_swifter_comparison/1pl_1tp_encounter/.idea/.gitignore
diff --git a/examples/rmvs_swifter_comparison/9pl_18tp_encounters/cb.swiftest.in b/examples/rmvs_swifter_comparison/9pl_18tp_encounters/cb.swiftest.in
@@ -1,4 +1,4 @@
-0.00029591220819207775568
-0.0046504672609621575315
-4.753580694812735539e-12
--2.2473967953572828617e-18
+0.0002959122081920778
+0.004650467260962157
+4.7535806948127355e-12
+-2.2473967953572827e-18
diff --git a/examples/rmvs_swifter_comparison/9pl_18tp_encounters/cb_test.in b/examples/rmvs_swifter_comparison/9pl_18tp_encounters/cb_test.in
@@ -0,0 +1,4 @@
+0.0002959122081920778
+0.004650467260962157
+4.7535806948127355e-12
+-2.2473967953572827e-18
diff --git a/examples/rmvs_swifter_comparison/9pl_18tp_encounters/config.swiftest.in b/examples/rmvs_swifter_comparison/9pl_18tp_encounters/config.swiftest.in
@@ -1,6 +1,6 @@
 ! Swiftest input file generated using init_cond.py
 T0             0 
-TSTOP          80.0  
+TSTOP          365.25
 DT             1.0
 CB_IN          cb.swiftest.in
 PL_IN          pl.swiftest.in

diff --git a/examples/rmvs_swifter_comparison/9pl_18tp_encounters/init_cond.py b/examples/rmvs_swifter_comparison/9pl_18tp_encounters/init_cond.py
@@ -3,11 +3,25 @@
 To use the script, modify the variables just after the  "if __name__ == '__main__':" line
 """
 import numpy as np
-from astroquery.jplhorizons import Horizons
-import astropy.constants as const 
 import swiftestio as swio
-from scipy.io import FortranFile
+import astropy.constants as const
 import sys
+import xarray as xr
+
+# Swiftest configuration file
+# Both codes use the same tp input file
+tpin = "tp.in"
+
+swifter_input  = "param.swifter.in"
+swifter_pl     = "pl.swifter.in"
+swifter_bin    = "bin.swifter.dat"
+swifter_enc    = "enc.swifter.dat"
+
+swiftest_input = "config.swiftest.in"
+swiftest_pl    = "pl.swiftest.in"
+swiftest_cb    = "cb.swiftest.in"
+swiftest_bin   = "bin.swiftest.dat"
+swiftest_enc   = "enc.swiftest.dat"
 
 #Values from JPL Horizons
 AU2M = np.longdouble(const.au.value)
@@ -24,70 +38,6 @@
 GU       = np.longdouble(GC / (DU2M**3 / (MU2KG * TU2S**2)))
 GMSun = np.longdouble(GMSunSI / (DU2M**3 / TU2S**2))
 
-# Solar oblatenes values: From Mecheri et al. (2004), using Corbard (b) 2002 values (Table II)
-J2 = np.longdouble(2.198e-7) * (Rsun / DU2M)**2
-J4 = np.longdouble(-4.805e-9) * (Rsun / DU2M)**4
-
-npl = 9
-ntp = 2 * npl
-
-# Planet ids
-plid = {
-   '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
-Rpl = {
-   '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)
-}
-
-THIRDLONG = np.longdouble(1.0) / np.longdouble(3.0)
-
-# Both codes use the same tp input file
-tpin = "tp.in"
-
-swifter_input  = "param.swifter.in"
-swifter_pl     = "pl.swifter.in"
-swifter_bin    = "bin.swifter.dat"
-swifter_enc    = "enc.swifter.dat"
-
-swiftest_input = "config.swiftest.in"
-swiftest_pl    = "pl.swiftest.in"
-swiftest_cb    = "cb.swiftest.in"
-swiftest_bin   = "bin.swiftest.dat"
-swiftest_enc   = "enc.swiftest.dat"
-
-# Simple initial conditions of a circular planet with one test particle in a close encounter state
-# Simulation start, stop, and output cadence times
 t_0	  = 0 # simulation start time
 deltaT	= 1.00 * JD / TU2S   # simulation step size
 end_sim = year / TU2S # simulation end time
@@ -102,112 +52,6 @@
 tend = '2021-06-16'
 tstep = '1d'
 
-#######################################################
-# Start generating initial conditions for the planets #
-#######################################################
-pldata = {}
-tpdata = {}
-plvec = {}
-tpvec = {}
-Rhill = {}
-
-for key,val in plid.items():
-   pldata[key] = Horizons(id=val, id_type='majorbody',location='@sun',
-            epochs={'start': tstart, 'stop': tend,
-            'step': tstep})
-   plvec[key] = np.array([
-         pldata[key].vectors()['x'][0],
-         pldata[key].vectors()['y'][0],
-         pldata[key].vectors()['z'][0],
-         pldata[key].vectors()['vx'][0],
-         pldata[key].vectors()['vy'][0],
-         pldata[key].vectors()['vz'][0]
-         ])
-   Rhill[key] = pldata[key].elements()['a'][0] * (3 * MSun_over_Mpl[key])**(-1.0 / 3.0)
-
-# For each planet, we will initialize a pair of test particles. One on its way in, and one on its way out. We will also initialize two additional particles that don't encounter anything
-tpnames = range(101, 101 + ntp)
-tpxv1 = np.empty((6))
-tpxv2 = np.empty((6))
-for idx,key in enumerate(plid):
-    rstart = 2 * Rpl[key] / DU2M # Start the test particles at a multiple of the planet radius away
-    vstart = 1.5 * np.sqrt(2 * GMSun / MSun_over_Mpl[key] / rstart)  # Start the test particle velocities at a multiple of the escape speed
-    xvstart = np.array([rstart / np.sqrt(2.0), rstart / np.sqrt(2.0), 0.0, vstart, 0.0, 0.0])
-    # The positions and velocities of each pair of test particles will be in reference to a planet
-    tpxv1 = plvec[key] + xvstart
-    tpxv2 = plvec[key] - xvstart
-    tpvec[tpnames[idx]] = tpxv1
-    tpvec[tpnames[idx + npl]] = tpxv2
-
-# TP file
-tpfile = open(tpin, 'w')
-print(ntp, file=tpfile)
-for tpid, tp in tpvec.items():
-    print(tpid, file=tpfile)
-    print(f'{tp[0]} {tp[1]} {tp[2]}', file=tpfile)
-    print(f'{tp[3]} {tp[4]} {tp[5]}', file=tpfile)
-tpfile.close()
-
-# Swifter PL file
-plfile = open(swifter_pl, 'w')
-print(npl + 1, file=plfile)
-print(1,GMSun,file=plfile)
-print('0.0 0.0 0.0',file=plfile)
-print('0.0 0.0 0.0',file=plfile)
-for key, pl in plvec.items():
-    print(f'{int(plid[key])+1} {GMSun / MSun_over_Mpl[key]} {Rhill[key]} ! {key}', file=plfile)
-    print(f'{Rpl[key] / DU2M}', file=plfile)
-    print(f'{pl[0]} {pl[1]} {pl[2]}', file=plfile)
-    print(f'{pl[3]} {pl[4]} {pl[5]}', file=plfile)
-plfile.close()
-
-# Swiftest Central body and pl file
-cbfile = open(swiftest_cb, 'w')
-print(GMSun, file=cbfile)
-print(rmin, file=cbfile)
-print(J2, file=cbfile)
-print(J4, file=cbfile)
-cbfile.close()
-plfile = open(swiftest_pl, 'w')
-print(npl, file=plfile)
-for key, pl in plvec.items():
-    print(f'{int(plid[key]) + 1} {GMSun / MSun_over_Mpl[key]} ! {key}', file=plfile)
-    print(f'{Rpl[key] / DU2M}', file=plfile)
-    print(f'{pl[0]} {pl[1]} {pl[2]}', file=plfile)
-    print(f'{pl[3]} {pl[4]} {pl[5]}', file=plfile)
-plfile.close()
-
-# Swifter parameter file
-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         {tpin}')
-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      XV')
-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')
-sys.stdout = sys.__stdout__
-
-# Swiftest configuration file
 sys.stdout = open(swiftest_input, "w")
 print(f'! Swiftest input file generated using init_cond.py')
 print(f'T0             {t_0} ')
@@ -238,6 +82,68 @@
 print(f'MU2KG          {MU2KG}')
 print(f'DU2M           {DU2M}')
 print(f'TU2S           {TU2S}')
+sys.stdout = sys.__stdout__
+config = swio.read_swiftest_config(swiftest_input)
+ds = swio.solar_system_pl(config, tstart)
+
+ntp = 18
+dims = ['time', 'id', 'vec']
+tp = []
+t = np.array([0.0])
+clab, plab, tlab = swio.make_swiftest_labels(config)
+
+# For each planet, we will initialize a pair of test particles. One on its way in, and one on its way out. We will also initialize two additional particles that don't encounter anything
+tpnames = np.arange(101, 101 + ntp)
+tpxv1 = np.empty((6))
+tpxv2 = np.empty((6))
+cb = ds.sel(id=0)
+pl = ds.where(ds.id > 0, drop=True)
+
+p1 = []
+p2 = []
+p3 = []
+p4 = []
+p5 = []
+p6 = []
+
+for i in pl.id:
+    pli = pl.sel(id=i)
+    rstart = 2 * np.double(pli['Radius'])  # Start the test particles at a multiple of the planet radius away
+    vstart = 1.5 * np.sqrt(2 * np.double(pli['Mass'])  / rstart)  # Start the test particle velocities at a multiple of the escape speed
+    xvstart = np.array([rstart / np.sqrt(2.0), rstart / np.sqrt(2.0), 0.0, vstart, 0.0, 0.0])
+    # The positions and velocities of each pair of test particles will be in reference to a planet
+    plvec = np.array([np.double(pli['px']),
+                      np.double(pli['py']),
+                      np.double(pli['pz']),
+                      np.double(pli['vx']),
+                      np.double(pli['vy']),
+                      np.double(pli['vz'])])
+    tpxv1 = plvec + xvstart
+    tpxv2 = plvec - xvstart
+    p1.append(tpxv1[0])
+    p1.append(tpxv2[0])
+    p2.append(tpxv1[1])
+    p2.append(tpxv2[1])
+    p3.append(tpxv1[2])
+    p3.append(tpxv2[2])
+    p4.append(tpxv1[3])
+    p4.append(tpxv2[3])
+    p5.append(tpxv1[4])
+    p5.append(tpxv2[4])
+    p6.append(tpxv1[5])
+    p6.append(tpxv2[5])
+
+tvec = np.vstack([p1, p2, p3, p4, p5, p6])
+tpframe = np.expand_dims(tvec.T, axis=0)
+tpxr = xr.DataArray(tpframe, dims = dims, coords = {'time' : t, 'id' : tpnames, 'vec' : tlab})
+
+tp = [tpxr]
+tpda = xr.concat(tp,dim='time')
+tpds = tpda.to_dataset(dim = 'vec')
+
+ds = xr.combine_by_coords([ds, tpds])
+swio.swiftest_xr2_infile(ds, config)
+
 
 
 

diff --git a/examples/rmvs_swifter_comparison/9pl_18tp_encounters/pl.swiftest.in b/examples/rmvs_swifter_comparison/9pl_18tp_encounters/pl.swiftest.in
@@ -1,37 +1,37 @@
 9
-2 4.9125474498983625e-11 ! mercury
-1.6306381826061646e-05
-0.008059842448018334 -0.4616051037329109 -0.03846017738329229
-0.02248719132054853 0.001934639213990692 -0.001904656977422976
-3 7.243452483873647e-10 ! venus
-4.0453784346544176e-05
--0.5115875215389065 0.5030818749037324 0.03642547299277956
--0.01425515725454357 -0.01452868630179309 0.0006232072038298823
-4 8.997011382166019e-10 ! earthmoon
-4.25875607065041e-05
--0.1090020607540907 -1.009893805009766 4.823302918632528e-05
-0.01682491922568941 -0.001910549762056979 3.992660742687128e-08
-5 9.549535102761465e-11 ! mars
-2.2657408050928896e-05
--1.342897929331636 0.9778655112682739 0.05343398538723887
--0.007712315645393206 -0.01011917844182223 -2.287744801261131e-05
-6 2.8253459086313547e-07 ! jupiter
-0.0004673261703049093
-3.923184193414315 -3.168419770483168 -0.0746147877972047
-0.004655552638985802 0.006232623300954468 -0.0001300429201057457
-7 8.459715183006416e-08 ! saturn
-0.00038925687730393614
-6.185794462795267 -7.804174837804826 -0.110498432926239
-0.004066833203985018 0.003458637040736611 -0.0002219310939327014
-8 1.2920249163736674e-08 ! uranus
-0.00016953449859497232
-14.9290976575471 12.92949673572929 -0.1454099139559955
--0.002599557960646664 0.002795888198858545 4.391864857782088e-05
-9 1.5243589003230834e-08 ! neptune
-0.00016458790412449367
-29.54416169025338 -4.716921603714237 -0.5838030174427992
-0.0004792636209523189 0.00312573757291745 -7.53264045199501e-05
-10 2.1919422829042796e-12 ! plutocharon
-7.943294877391593e-06
-14.54448346259197 -31.05223519593471 -0.8828000265625595
-0.002923077617691739 0.0006625916902153526 -0.0009142553677224461
+1 4.9125474498983625056e-11
+2439400.0
+0.008059842448018333591 -0.46160510373291091524 -0.038460177383292291908
+0.02248719132054852949 0.0019346392139906921279 -0.0019046569774229759606
+2 7.243452483873647106e-10
+6051800.0
+-0.51158752153890652004 0.5030818749037323512 0.036425472992779560355
+-0.01425515725454356988 -0.014528686301793089943 0.00062320720382988232425
+3 8.997011382166018993e-10
+6371008.4000000003725
+-0.109002060754090704386 -1.0098938050097661101 4.8233029186325282966e-05
+0.016824919225689409317 -0.0019105497620569790936 3.9926607426871282392e-08
+4 9.549535102761465872e-11
+3389500.0
+-1.3428979293316360977 0.97786551126827392366 0.053433985387238869258
+-0.007712315645393206092 -0.0101191784418222296525 -2.2877448012611311785e-05
+5 2.8253459086313549713e-07
+69911000.0
+3.923184193414314791 -3.1684197704831680298 -0.07461478779720470689
+0.0046555526389858022113 0.006232623300954467766 -0.00013004292010574569365
+6 8.45971518300641563e-08
+58232000.0
+6.1857944627952665684 -7.804174837804826126 -0.11049843292623899582
+0.0040668332039850179674 0.0034586370407366113193 -0.00022193109393270141328
+7 1.2920249163736673984e-08
+25362000.0
+14.929097657547099942 12.9294967357292893695 -0.14540991395599550673
+-0.0025995579606466640267 0.0027958881988585450113 4.391864857782088156e-05
+8 1.5243589003230834746e-08
+24622000.0
+29.544161690253378794 -4.7169216037142369657 -0.58380301744279916587
+0.00047926362095231893815 0.00312573757291745008 -7.532640451995010825e-05
+9 2.1919422829042797324e-12
+1188300.0
+14.544483462591969669 -31.052235195934709822 -0.88280002656255951443
+0.0029230776176917390448 0.0006625916902153525834 -0.0009142553677224461557