Fixed some issues with the 9pl9tp test case Swifter initial conditions generator

MintonGroup · Jun 29, 2021 · fb340b1 · fb340b1
1 parent 658c916
commit fb340b1
Show file tree

Hide file tree

Showing 9 changed files with 109 additions and 89 deletions.
diff --git a/examples/rmvs_swifter_comparison/9pl_18tp_encounters/cb_test.in b/examples/rmvs_swifter_comparison/9pl_18tp_encounters/cb_test.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
@@ -1,14 +1,9 @@
-"""
-For testing RMVS, the code generates clones of test particles based on one that is fated to impact Mercury.
-To use the script, modify the variables just after the  "if __name__ == '__main__':" line
-"""
 import numpy as np
-import swiftestio as swio
+import swiftest.swiftestio as swio
 import astropy.constants as const
 import sys
 import xarray as xr
 
-# Swiftest paramuration file
 # Both codes use the same tp input file
 tpin = "tp.in"
 
@@ -47,11 +42,6 @@
 rmin = Rsun / DU2M
 rmax = 1000.0
 
-# Dates to fetch planet ephemerides from JPL Horizons
-tstart = '2021-06-15'
-tend = '2021-06-16'
-tstep = '1d'
-
 sys.stdout = open(swiftest_input, "w")
 print(f'! Swiftest input file generated using init_cond.py')
 print(f'T0             {t_0} ')
@@ -84,7 +74,13 @@
 print(f'TU2S           {TU2S}')
 sys.stdout = sys.__stdout__
 param = swio.read_swiftest_param(swiftest_input)
+
+# Dates to fetch planet ephemerides from JPL Horizons
+tstart = '2021-06-15'
 ds = swio.solar_system_pl(param, tstart)
+cb = ds.sel(id=0)
+pl = ds.where(ds.id > 0, drop=True)
+npl = pl.id.size
 
 ntp = 18
 dims = ['time', 'id', 'vec']
@@ -96,8 +92,6 @@
 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 = []
@@ -144,8 +138,47 @@
 ds = xr.combine_by_coords([ds, tpds])
 swio.swiftest_xr2_infile(ds, param)
 
-
-
-
-
+# Swifter PL file
+plfile = open(swifter_pl, 'w')
+print(npl + 1, file=plfile)
+print(0,GMSun,file=plfile)
+print('0.0 0.0 0.0',file=plfile)
+print('0.0 0.0 0.0',file=plfile)
+for i in pl.id:
+    pli = pl.sel(id=i)
+    print(f"{int(i)} {pli['Mass'].values[0]} {pli['Rhill'].values[0]}", file=plfile)
+    print(f"{pli['Radius'].values[0]}", file=plfile)
+    print(f"{pli['px'].values[0]} {pli['py'].values[0]} {pli['pz'].values[0]}", file=plfile)
+    print(f"{pli['vx'].values[0]} {pli['vy'].values[0]} {pli['vz'].values[0]}", 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      UNKNOWN")
+print(f"J2            {param['J2']}")
+print(f"J4            {param['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__
 
diff --git a/examples/rmvs_swifter_comparison/9pl_18tp_encounters/param.swifter.in b/examples/rmvs_swifter_comparison/9pl_18tp_encounters/param.swifter.in
@@ -1,6 +1,6 @@
 ! Swifter input file generated using init_cond.py
 T0            0 
-TSTOP         80.0  
+TSTOP         365.25
 DT            1.0
 PL_IN         pl.swifter.in
 TP_IN         tp.in
@@ -10,9 +10,9 @@ ISTEP_DUMP    1
 BIN_OUT       bin.swifter.dat
 OUT_TYPE      REAL8
 OUT_FORM      XV
-OUT_STAT      NEW
-J2            4.7535806948127355e-12
-J4            -2.2473967953572827e-18
+OUT_STAT      UNKNOWN
+J2            0.0
+J4            0.0
 CHK_CLOSE     yes
 CHK_RMIN      0.004650467260962157
 CHK_RMAX      1000.0

diff --git a/examples/rmvs_swifter_comparison/9pl_18tp_encounters/pl.swifter.in b/examples/rmvs_swifter_comparison/9pl_18tp_encounters/pl.swifter.in
@@ -1,40 +1,40 @@
 10
-1 0.00029591220819207775568
+0 0.00029591220819207776388
 0.0 0.0 0.0
 0.0 0.0 0.0
-2 4.9125474498983625e-11 0.0014751258227142052 ! mercury
-1.6306381826061646e-05
+1 4.9125474498983625e-11 0.0014751258227142052
+2439400.0
 0.008059842448018334 -0.4616051037329109 -0.03846017738329229
 0.02248719132054853 0.001934639213990692 -0.001904656977422976
-3 7.243452483873647e-10 0.006759134232034942 ! venus
-4.0453784346544176e-05
+2 7.243452483873647e-10 0.006759134232034941
+6051800.0
 -0.5115875215389065 0.5030818749037324 0.03642547299277956
 -0.01425515725454357 -0.01452868630179309 0.0006232072038298823
-4 8.997011382166019e-10 0.010044625087011913 ! earthmoon
-4.25875607065041e-05
+3 8.997011382166019e-10 0.010044625087011915
+6371008.4
 -0.1090020607540907 -1.009893805009766 4.823302918632528e-05
 0.01682491922568941 -0.001910549762056979 3.992660742687128e-08
-5 9.549535102761465e-11 0.007246789790242477 ! mars
-2.2657408050928896e-05
+4 9.549535102761465e-11 0.0072467897902424765
+3389500.0
 -1.342897929331636 0.9778655112682739 0.05343398538723887
 -0.007712315645393206 -0.01011917844182223 -2.287744801261131e-05
-6 2.8253459086313547e-07 0.3552720805286442 ! jupiter
-0.0004673261703049093
+5 2.825345908631355e-07 0.3552720805286442
+69911000.0
 3.923184193414315 -3.168419770483168 -0.0746147877972047
 0.004655552638985802 0.006232623300954468 -0.0001300429201057457
-7 8.459715183006416e-08 0.4376460836930155 ! saturn
-0.00038925687730393614
+6 8.459715183006416e-08 0.4376460836930155
+58232000.0
 6.185794462795267 -7.804174837804826 -0.110498432926239
 0.004066833203985018 0.003458637040736611 -0.0002219310939327014
-8 1.2920249163736674e-08 0.46946272948265794 ! uranus
-0.00016953449859497232
+7 1.2920249163736674e-08 0.46946272948265794
+25362000.0
 14.9290976575471 12.92949673572929 -0.1454099139559955
 -0.002599557960646664 0.002795888198858545 4.391864857782088e-05
-9 1.5243589003230834e-08 0.7811947848333599 ! neptune
-0.00016458790412449367
+8 1.5243589003230834e-08 0.78119478483336
+24622000.0
 29.54416169025338 -4.716921603714237 -0.5838030174427992
 0.0004792636209523189 0.00312573757291745 -7.53264045199501e-05
-10 2.1919422829042796e-12 0.05379680851617536 ! plutocharon
-7.943294877391593e-06
+9 2.1919422829042796e-12 0.05379680851617536
+1188300.0
 14.54448346259197 -31.05223519593471 -0.8828000265625595
 0.002923077617691739 0.0006625916902153526 -0.0009142553677224461
diff --git a/examples/rmvs_swifter_comparison/9pl_18tp_encounters/pl_test.in b/examples/rmvs_swifter_comparison/9pl_18tp_encounters/pl_test.in
diff --git a/examples/rmvs_swifter_comparison/9pl_18tp_encounters/tp.swifter.in b/examples/rmvs_swifter_comparison/9pl_18tp_encounters/tp.swifter.in
diff --git a/examples/rmvs_swifter_comparison/9pl_18tp_encounters/tp.swiftest.in b/examples/rmvs_swifter_comparison/9pl_18tp_encounters/tp.swiftest.in
diff --git a/examples/rmvs_swifter_comparison/9pl_18tp_encounters/tp_test.in b/examples/rmvs_swifter_comparison/9pl_18tp_encounters/tp_test.in
diff --git a/python/swiftest/swiftest/swiftestio.py b/python/swiftest/swiftest/swiftestio.py
@@ -715,6 +715,12 @@ def solar_system_pl(param, ephemerides_start_date):
     p4 = []
     p5 = []
     p6 = []
+    p7 = []
+    p8 = []
+    p9 = []
+    p10 = []
+    p11 = []
+    p12 = []
     Rhill = []
     Rpl = []
     GMpl = []
@@ -731,27 +737,54 @@ def solar_system_pl(param, ephemerides_start_date):
             p4.append(pldata[key].vectors()['vx'][0] * VCONV)
             p5.append(pldata[key].vectors()['vy'][0] * VCONV)
             p6.append(pldata[key].vectors()['vz'][0] * VCONV)
+            p7.append(pldata[key].elements()['a'][0] * DCONV)
+            p8.append(pldata[key].elements()['e'][0])
+            p9.append(pldata[key].elements()['incl'][0] * np.pi / 180.0)
+            p10.append(pldata[key].elements()['Omega'][0] * np.pi / 180.0)
+            p11.append(pldata[key].elements()['w'][0] * np.pi / 180.0)
+            p12.append(pldata[key].elements()['M'][0] * np.pi / 180.0)
         elif param['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)
+            p3.append(pldata[key].elements()['incl'][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)
+            p7.append(pldata[key].vectors()['x'][0] * DCONV)
+            p8.append(pldata[key].vectors()['y'][0] * DCONV)
+            p9.append(pldata[key].vectors()['z'][0] * DCONV)
+            p10.append(pldata[key].vectors()['vx'][0] * VCONV)
+            p11.append(pldata[key].vectors()['vy'][0] * VCONV)
+            p12.append(pldata[key].vectors()['vz'][0] * VCONV)
         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])
+    pvec = np.vstack([p1, p2, p3, p4, p5, p6, GMpl, Rpl, p7, p8, p9, p10, p11, p12, Rhill])
 
     dims = ['time', 'id', 'vec']
     cb = []
     pl = []
     t = np.array([0.0])
 
     clab, plab, tlab = make_swiftest_labels(param)
-
+    if param['OUT_FORM'] == 'XV':
+        plab.append('a')
+        plab.append('e')
+        plab.append('inc')
+        plab.append('capom')
+        plab.append('omega')
+        plab.append('capm')
+    elif param['OUT_FORM'] == 'EL':
+        plab.append('px')
+        plab.append('py')
+        plab.append('pz')
+        plab.append('vx')
+        plab.append('vy')
+        plab.append('vz')
+    plab.append('Rhill')
+
     # 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)