Fixed issues of indexing in the NetCDF code, and started process of c…

…leaning up the strings from NetCDF reads
MintonGroup · Aug 27, 2021 · 0ad8959 · 0ad8959
1 parent be5e5d2
commit 0ad8959
Show file tree

Hide file tree

Showing 12 changed files with 2,073 additions and 44 deletions.
diff --git a/examples/symba_mars_disk/testnetcdf.ipynb b/examples/symba_mars_disk/testnetcdf.ipynb
diff --git a/examples/whm_gr_test/Untitled.ipynb b/examples/whm_gr_test/Untitled.ipynb
diff --git a/examples/whm_gr_test/cb.swiftest.in b/examples/whm_gr_test/cb.swiftest.in
@@ -1,4 +1,4 @@
-0
+Sun
 39.476926408897626
 0.004650467260962157
 4.7535806948127355e-12

diff --git a/examples/whm_gr_test/init_cond.py b/examples/whm_gr_test/init_cond.py
@@ -13,9 +13,13 @@
 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['TSTOP'] = 1000.0
+#sim.param['ISTEP_OUT']  = 1461
+#sim.param['ISTEP_DUMP'] = 1461
+sim.param['TSTOP'] = 2*sim.param['DT']
+sim.param['ISTEP_OUT']  = 1
+sim.param['ISTEP_DUMP'] = 1
+
 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"

diff --git a/examples/whm_gr_test/param.swifter.in b/examples/whm_gr_test/param.swifter.in
@@ -1,9 +1,9 @@
 ! VERSION        Swifter parameter file converted from Swiftest
 T0               0.0
-TSTOP            1000.0
+TSTOP            0.0013689253935660506
 DT               0.0006844626967830253
-ISTEP_OUT        1461
-ISTEP_DUMP       1461
+ISTEP_OUT        1
+ISTEP_DUMP       1
 OUT_FORM         EL
 OUT_TYPE         REAL8
 OUT_STAT         UNKNOWN

diff --git a/examples/whm_gr_test/param.swiftest.in b/examples/whm_gr_test/param.swiftest.in
@@ -1,9 +1,9 @@
 ! VERSION        Swiftest parameter input
 T0               0.0
-TSTOP            1000.0
+TSTOP            0.0013689253935660506
 DT               0.0006844626967830253
-ISTEP_OUT        1461
-ISTEP_DUMP       1461
+ISTEP_OUT        1
+ISTEP_DUMP       1
 OUT_FORM         XVEL
 OUT_TYPE         NETCDF_DOUBLE
 OUT_STAT         UNKNOWN

diff --git a/examples/whm_gr_test/pl.swifter.in b/examples/whm_gr_test/pl.swifter.in
@@ -2,35 +2,35 @@
 1 39.476926408897626
 0.0 0.0 0.0
 0.0 0.0 0.0
-2 6.553709809565314146e-06 0.0014751261514880139061
+2 6.5537098095653139645e-06 0.001475126151488013514
 1.6306381826061645943e-05
 -0.29510017042975300594 -0.34346884022084378518 -0.000997917547895216684
 5.709185893355676925 -6.2220892824670268354 -1.0321515701207669188
-3 9.6633133995815381836e-05 0.006759061578633710828
+3 9.663313399581537916e-05 0.006759061578633709628
 4.0453784346544178454e-05
 -0.21096294566811030213 -0.6945147553261317164 0.0026420019351886940041
 7.0187176732445560167 -2.1783670470162428854 -0.43491289164169576724
-4 0.000120026935827952456416 0.010044949983771724966
+4 0.000120026935827952453094 0.010044949983771722159
 4.25875607065040958e-05
 0.905522725638602366 -0.44829515638659761523 1.8122901980659660508e-05
 2.685845700565266851 5.6080546137104218133 -0.00027833838024725729542
-5 1.2739802010675941808e-05 0.0072465915674003790445
+5 1.2739802010675941456e-05 0.0072465915674003775008
 2.265740805092889601e-05
 -1.6511945936824949932 0.1180803833522415941 0.042978148735422203042
 -0.17443328105136805607 -4.661619009141641736 -0.09341597039948347882
-6 0.03769225108898567778 0.3552713110772063853
+6 0.037692251088985676735 0.3552713110772063034
 0.00046732617030490929307
 4.2381319671740662614 -2.694827110197309139 -0.08362807329786287047
 1.44742540330701551 2.4581907268113588696 -0.042593445938391914576
-7 0.01128589982009127331 0.4376635990332856823
+7 0.011285899820091272997 0.43766359903328559694
 0.00038925687730393611812
 6.4776155764849425722 -7.5454781609219372385 -0.12660625214421539209
 1.4343694668413992401 1.3251694474665614901 -0.08010594537316981756
-8 0.001723658947826773068 0.46966224198242572768
+8 0.0017236589478267730203 0.4696622419824256353
 0.00016953449859497231466
 14.737783583010530819 13.132284780084109599 -0.14218874866247160904
 -0.96494157148906816704 1.0080364706941240677 0.016191990849809560611
-9 0.0020336100526728302882 0.7814394516095526881
+9 0.0020336100526728302319 0.781439451609552476
 0.000164587904124493665
 29.578253698940308425 -4.488584904681241383 -0.58928426126360722304
 0.16609282485651713797 1.143247554888599065 -0.027336661118935745503
diff --git a/examples/whm_gr_test/pl.swiftest.in b/examples/whm_gr_test/pl.swiftest.in
@@ -1,33 +1,33 @@
 8
-2 6.553709809565314146e-06 0.0014751261514880139061
+Mercury 6.5537098095653139645e-06 0.001475126151488013514
 1.6306381826061645943e-05
 0.38709861919270799335 0.20562987442219879397 7.0036598799530471737
 48.303764444546942514 29.187012577257871015 139.02452846395490837
-3 9.6633133995815381836e-05 0.006759061578633710828
+Venus 9.663313399581537916e-05 0.006759061578633709628
 4.0453784346544178454e-05
 0.72332350338494522113 0.0067851993472706276234 3.3945100118236060105
 76.62172575937908903 55.11451498626085055 120.69543184874230235
-4 0.000120026935827952456416 0.010044949983771724966
+Earth 0.000120026935827952453094 0.010044949983771722159
 4.25875607065040958e-05
 1.000022414803547921 0.016679693167334301573 0.002750719340522077977
 175.59912721852418827 287.3487256951211748 232.20474506920808722
-5 1.2739802010675941808e-05 0.0072465915674003790445
+Mars 1.2739802010675941456e-05 0.0072465915674003775008
 2.265740805092889601e-05
 1.5237056817307590428 0.09335454089002033495 1.8479086301002540793
 49.490790672135332784 286.7030449733272235 203.56009534652309867
-6 0.03769225108898567778 0.3552713110772063853
+Jupiter 0.037692251088985676735 0.3552713110772063034
 0.00046732617030490929307
 5.203524963998765074 0.048518619089771883313 1.3035691332389880426
 100.516740776557597314 273.38311024861741316 317.57998688455870706
-7 0.01128589982009127331 0.4376635990332856823
+Saturn 0.011285899820091272997 0.43766359903328559694
 0.00038925687730393611812
 9.581861578191695372 0.05220296447788015659 2.4862598063103709123
 113.59526938704850352 335.69019804854252698 225.44748806008931297
-8 0.001723658947826773068 0.46966224198242572768
+Uranus 0.0017236589478267730203 0.4696622419824256353
 0.00016953449859497231466
 19.23638216159032055 0.04433184777155944195 0.7703424893861580136
 74.09557761028084144 95.84993853316392176 235.82773290830229485
-9 0.0020336100526728302882 0.7814394516095526881
+Neptune 0.0020336100526728302319 0.781439451609552476
 0.000164587904124493665
 30.289653279202511271 0.013458737599393380546 1.769000156955224945
 131.74519418988560915 245.79890439350270981 334.51418242279709148
diff --git a/python/swiftest/swiftest/init_cond.py b/python/swiftest/swiftest/init_cond.py
@@ -171,6 +171,7 @@ def solar_system_horizons(plname, idval, param, ephemerides_start_date, ds):
         print("Creating the Sun as a central body")
         cbframe = np.expand_dims(cvec.T, axis=0)
         cbxr = xr.DataArray(cbframe, dims=dims, coords={'time': t, 'id': cbid, 'vec': clab})
+        cbxr = cbxr.assign_coords(name=('id', [key]))
         cbds = cbxr.to_dataset(dim='vec')
         ds = xr.combine_by_coords([ds, cbds])
     else: # Fetch solar system ephemerides from Horizons
@@ -270,12 +271,13 @@ def solar_system_horizons(plname, idval, param, ephemerides_start_date, ds):
         # Prepare frames by adding an extra axis for the time coordinate
         plframe = np.expand_dims(pvec.T, axis=0)
         plxr = xr.DataArray(plframe, dims=dims, coords={'time': t, 'id': plid, 'vec': plab})
+        plxr = plxr.assign_coords(name=('id', [plname]))
         plds = plxr.to_dataset(dim='vec')
         ds = xr.combine_by_coords([ds, plds])
 
     return ds
 
-def vec2xr(param, idvals, v1, v2, v3, v4, v5, v6, GMpl=None, Rpl=None, rhill=None, Ip1=None, Ip2=None, Ip3=None, rotx=None, roty=None, rotz=None, t=0.0):
+def vec2xr(param, names, idvals, v1, v2, v3, v4, v5, v6, GMpl=None, Rpl=None, rhill=None, Ip1=None, Ip2=None, Ip3=None, rotx=None, roty=None, rotz=None, t=0.0):
 
     if param['ROTATION'] == 'YES':
         if Ip1 is None:
@@ -317,6 +319,8 @@ def vec2xr(param, idvals, v1, v2, v3, v4, v5, v6, GMpl=None, Rpl=None, rhill=Non
         bodyxr = xr.DataArray(bodyframe, dims=dims, coords={'time': [t], 'id': idvals, 'vec': plab})
     else:
         bodyxr = xr.DataArray(bodyframe, dims=dims, coords={'time': [t], 'id': idvals, 'vec': tlab})
+
+    bodyxr = bodyxr.assign_coords(name=('id', names))
 
     ds = bodyxr.to_dataset(dim='vec')
     return ds
diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py
@@ -696,6 +696,7 @@ def swiftest2xr(param):
     elif ((param['OUT_TYPE'] == 'NETCDF_DOUBLE') or (param['OUT_TYPE'] == 'NETCDF_FLOAT')):
         print('\nCreating Dataset')
         ds = xr.open_dataset(param['BIN_OUT'])
+        ds = clean_string_values(param, ds)
     else:
         print(f"Error encountered. OUT_TYPE {param['OUT_TYPE']} not recognized.")
         return None
@@ -797,6 +798,7 @@ def swiftest_xr2infile(ds, param, framenum=-1):
     RSun = np.double(cb['radius'])
     J2 = np.double(cb['J_2'])
     J4 = np.double(cb['J_4'])
+    cbname = cb['name'].values[0]
     if param['ROTATION'] == 'YES':
         Ip1cb = np.double(cb['Ip1'])
         Ip2cb = np.double(cb['Ip2'])
@@ -809,7 +811,7 @@ def swiftest_xr2infile(ds, param, framenum=-1):
     if param['IN_TYPE'] == 'ASCII':
         # Swiftest Central body file
         cbfile = open(param['CB_IN'], 'w')
-        print(0, file=cbfile)
+        print(cbname, file=cbfile)
         print(GMSun, file=cbfile)
         print(RSun, file=cbfile)
         print(J2, file=cbfile)
@@ -824,9 +826,9 @@ def swiftest_xr2infile(ds, param, framenum=-1):
         for i in pl.id:
             pli = pl.sel(id=i)
             if param['RHILL_PRESENT'] == 'YES':
-               print(i.values, pli['Gmass'].values, pli['rhill'].values, file=plfile)
+               print(pli['name'].values, pli['Gmass'].values, pli['rhill'].values, file=plfile)
             else:
-               print(i.values, pli['Gmass'].values, file=plfile)
+               print(pli['name'].values, pli['Gmass'].values, file=plfile)
             print(pli['radius'].values, file=plfile)
             if param['IN_FORM'] == 'XV':
                 print(pli['xhx'].values, pli['xhy'].values, pli['xhz'].values, file=plfile)
@@ -846,7 +848,7 @@ def swiftest_xr2infile(ds, param, framenum=-1):
         print(tp.id.count().values, file=tpfile)
         for i in tp.id:
             tpi = tp.sel(id=i)
-            print(i.values, file=tpfile)
+            print(tpi['name'].values, file=tpfile)
             if param['IN_FORM'] == 'XV':
                 print(tpi['xhx'].values, tpi['xhy'].values, tpi['xhz'].values, file=tpfile)
                 print(tpi['vhx'].values, tpi['vhy'].values, tpi['vhz'].values, file=tpfile)

diff --git a/src/io/io.f90 b/src/io/io.f90
@@ -903,8 +903,8 @@ module subroutine io_read_in_cb(self, param)
       integer(I4B)            :: ierr, idold
 
       if (param%in_type == 'ASCII') then
-         self%id = 1
-         param%maxid = 1
+         self%id = 0
+         param%maxid = 0
          open(unit = iu, file = param%incbfile, status = 'old', form = 'FORMATTED', err = 667, iomsg = errmsg)
          read(iu, *, err = 667, iomsg = errmsg) self%name
          read(iu, *, err = 667, iomsg = errmsg) self%Gmass
@@ -1070,11 +1070,13 @@ module function io_read_frame_body(self, iu, param) result(ierr)
                   read(iu, err = 667, iomsg = errmsg) pl%rot(3, :)
                end if
                if (param%ltides) then
-                  read(iu, err = 667, iomsg = errmsg) pl%k2(1:n)
-                  read(iu, err = 667, iomsg = errmsg) pl%Q(1:n)
+                  read(iu, err = 667, iomsg = errmsg) pl%k2(:)
+                  read(iu, err = 667, iomsg = errmsg) pl%Q(:)
                end if
             end select
 
+            param%maxid = max(param%maxid, maxval(self%id(1:n)))
+
          case (ASCII_TYPE)
             do i = 1, n
                select type(self)

diff --git a/src/netcdf/netcdf.f90 b/src/netcdf/netcdf.f90
@@ -214,7 +214,7 @@ module subroutine netcdf_write_frame_base(self, iu, param)
       class(netcdf_parameters),   intent(inout) :: iu     !! Parameters used to identify a particular NetCDF dataset
       class(swiftest_parameters), intent(in)    :: param  !! Current run configuration parameters
       ! Internals
-      integer(I4B)                              :: i, j, id, tslot, strlen, idslot
+      integer(I4B)                              :: i, j, tslot, strlen, idslot
       integer(I4B), dimension(:), allocatable   :: ind
       character(len=:), allocatable             :: name
 
@@ -232,10 +232,10 @@ module subroutine netcdf_write_frame_base(self, iu, param)
             do i = 1, n
                j = ind(i)
                idslot = self%id(j) + 1
-               call check( nf90_put_var(iu%ncid, iu%id_varid, id, start=[idslot]) )
+               call check( nf90_put_var(iu%ncid, iu%id_varid, self%id(j), start=[idslot]) )
                name = trim(adjustl(self%name(j)))
                strlen = len(name)
-               call check( nf90_put_var(iu%ncid, iu%name_varid, name, start=[id, 1], count=[1, strlen]) )
+               call check( nf90_put_var(iu%ncid, iu%name_varid, name, start=[idslot, 1], count=[1, strlen]) )
                if ((param%out_form == XV) .or. (param%out_form == XVEL)) then
                   call check( nf90_put_var(iu%ncid, iu%xhx_varid, self%xh(1, j), start=[tslot, idslot]) )
                   call check( nf90_put_var(iu%ncid, iu%xhy_varid, self%xh(2, j), start=[tslot, idslot]) )
@@ -277,11 +277,11 @@ module subroutine netcdf_write_frame_base(self, iu, param)
             end do
          end associate
       class is (swiftest_cb)
-         id = self%id
+         idslot = self%id + 1
          call check( nf90_put_var(iu%ncid, iu%id_varid, id, start=[idslot]) )
          name = trim(adjustl(self%name))
          strlen = len(name)
-         call check( nf90_put_var(iu%ncid, iu%name_varid, name, start=[id, 1], count=[1, strlen]) )
+         call check( nf90_put_var(iu%ncid, iu%name_varid, name, start=[idslot, 1], count=[1, strlen]) )
          call check( nf90_put_var(iu%ncid, iu%Gmass_varid, self%Gmass, start=[tslot, idslot]) )
          call check( nf90_put_var(iu%ncid, iu%radius_varid, self%radius, start=[tslot, idslot]) )
          if (param%lrotation) then