Merge branch 'debug' of https://github.itap.purdue.edu/MintonGroup/sw…

…iftest into debug

examples/symba_energy_momentum/collision_movie.py

@@ -260,7 +260,7 @@ def data_stream(self, frame=0):
             vx = np.nan_to_num(vx, copy=False)
             vy = np.nan_to_num(vy, copy=False)
             radmarker = np.nan_to_num(radmarker, copy=False)
-            GMass = np.nan_to_num(Mass, copy=False)
+            GMass = np.nan_to_num(GMass, copy=False)
             Radius = np.nan_to_num(Radius, copy=False)
             rotx = np.nan_to_num(rotx, copy=False)
             roty = np.nan_to_num(roty, copy=False)

python/swiftest/swiftest/io.py

@@ -24,6 +24,7 @@ def real2float(realstr):
     """
     return float(realstr.replace('d', 'E').replace('D', 'E'))
 
+
 def read_swiftest_param(param_file_name, param):
     """
     Reads in a Swiftest param.in file and saves it as a dictionary
@@ -72,6 +73,8 @@ def read_swiftest_param(param_file_name, param):
         param['CHK_CLOSE'] = param['CHK_CLOSE'].upper()
         param['RHILL_PRESENT'] = param['RHILL_PRESENT'].upper()
         param['FRAGMENTATION'] = param['FRAGMENTATION'].upper()
+        if param['FRAGMENTATION'] == 'YES' and param['PARTICLE_OUT'] == '':
+            param['PARTICLE_OUT'] = 'particle.dat'
         param['ROTATION'] = param['ROTATION'].upper()
         param['TIDES'] = param['TIDES'].upper()
         param['ENERGY'] = param['ENERGY'].upper()
@@ -82,6 +85,7 @@ def read_swiftest_param(param_file_name, param):
         print(f"{param_file_name} not found.")
     return param
 
+
 def read_swifter_param(param_file_name):
     """
     Reads in a Swifter param.in file and saves it as a dictionary
@@ -165,6 +169,7 @@ def read_swifter_param(param_file_name):
 
     return param
 
+
 def read_swift_param(param_file_name, startfile="swift.in"):
     """
     Reads in a Swift param.in file and saves it as a dictionary
@@ -251,6 +256,7 @@ def read_swift_param(param_file_name, startfile="swift.in"):
 
     return param
 
+
 def write_swift_param(param, param_file_name):
     outfile = open(param_file_name, 'w')
     print(param['T0'], param['TSTOP'], param['DT'], file=outfile)
@@ -262,6 +268,7 @@ def write_swift_param(param, param_file_name):
     outfile.close()
     return
 
+
 def write_labeled_param(param, param_file_name):
     outfile = open(param_file_name, 'w')
     keylist = ['! VERSION',
@@ -300,6 +307,7 @@ def write_labeled_param(param, param_file_name):
     outfile.close()
     return
 
+
 def swifter_stream(f, param):
     """
     Reads in a Swifter bin.dat file and returns a single frame of data as a datastream
@@ -544,6 +552,7 @@ def swiftest_stream(f, param):
               npl, plid, pvec.T, plab, \
               ntp, tpid, tvec.T, tlab
 
+
 def swifter2xr(param):
     """
     Converts a Swifter binary data file into an xarray DataSet.
@@ -586,6 +595,7 @@ def swifter2xr(param):
         print(f"Successfully converted {ds.sizes['time']} output frames.")
     return ds
 
+
 def swiftest2xr(param):
     """
     Converts a Swiftest binary data file into an xarray DataSet.
@@ -604,26 +614,29 @@ def swiftest2xr(param):
     cb = []
     pl = []
     tp = []
-    with FortranFile(param['BIN_OUT'], 'r') as f:
-        for t, cbid, cvec, clab, \
-            npl, plid, pvec, plab, \
-            ntp, tpid, tvec, tlab in swiftest_stream(f, param):
-            # Prepare frames by adding an extra axis for the time coordinate
-            cbframe = np.expand_dims(cvec, axis=0)
-            plframe = np.expand_dims(pvec, axis=0)
-            tpframe = np.expand_dims(tvec, axis=0)
+    try:
+        with FortranFile(param['BIN_OUT'], 'r') as f:
+            for t, cbid, cvec, clab, \
+                npl, plid, pvec, plab, \
+                ntp, tpid, tvec, tlab in swiftest_stream(f, param):
+                # Prepare frames by adding an extra axis for the time coordinate
+                cbframe = np.expand_dims(cvec, axis=0)
+                plframe = np.expand_dims(pvec, axis=0)
+                tpframe = np.expand_dims(tvec, 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})
+                tpxr = xr.DataArray(tpframe, dims=dims, coords={'time': t, 'id': tpid, 'vec': tlab})
+
+                cb.append(cbxr)
+                pl.append(plxr)
+                tp.append(tpxr)
+                sys.stdout.write('\r' + f"Reading in time {t[0]:.3e}")
+                sys.stdout.flush()
+    except IOError:
+        print(f"Error encountered reading in {param['BIN_OUT']}")
 
-            # 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})
-            tpxr = xr.DataArray(tpframe, dims=dims, coords={'time': t, 'id': tpid, 'vec': tlab})
-
-            cb.append(cbxr)
-            pl.append(plxr)
-            tp.append(tpxr)
-            sys.stdout.write('\r' + f"Reading in time {t[0]:.3e}")
-            sys.stdout.flush()
-
     cbda = xr.concat(cb, dim='time')
     plda = xr.concat(pl, dim='time')
     tpda = xr.concat(tp, dim='time')
@@ -634,8 +647,75 @@ def swiftest2xr(param):
     print('\nCreating Dataset')
     ds = xr.combine_by_coords([cbds, plds, tpds])
     print(f"Successfully converted {ds.sizes['time']} output frames.")
+    if param['PARTICLE_OUT'] != "":
+       ds = swiftest_particle_2xr(ds, param)
+
     return ds
 
+
+def swiftest_particle_stream(f):
+   """
+   Reads in a Swiftest particle.dat file and returns a single frame of particle data as a datastream
+
+   Parameters
+   ----------
+   f : file object
+   param : dict
+
+   Yields
+   -------
+   plid : int
+      ID of massive bodie
+   origin_type : string
+      The origin type for the body (Initial conditions, disruption, supercatastrophic, hit and run, etc)
+   origin_xh : float array
+      The origin heliocentric position vector
+   origin_vh : float array
+      The origin heliocentric velocity vector
+   """
+   while True:  # Loop until you read the end of file
+      try:
+         # Read multi-line header
+         plid = f.read_ints()  # Try first part of the header
+      except:
+         break
+      origin_rec = f.read_record(np.dtype('a32'), np.dtype(('<f8', (7))))
+      origin_type = np.char.strip(str(origin_rec[0], encoding='utf-8'))
+      origin_vec = origin_rec[1]
+      yield plid, origin_type, origin_vec
+
+
+def swiftest_particle_2xr(ds, param):
+   """Reads in the Swiftest SyMBA-generated PARTICLE_OUT  and converts it to an xarray Dataset"""
+   veclab = ['time_origin', 'px_origin', 'py_origin', 'pz_origin', 'vx_origin', 'vy_origin', 'vz_origin']
+   id_list = []
+   origin_type_list = []
+   origin_vec_list = []
+
+   try:
+      with FortranFile(param['PARTICLE_OUT'], 'r') as f:
+          for id, origin_type, origin_vec in swiftest_particle_stream(f):
+             id_list.append(id)
+             origin_type_list.append(origin_type)
+             origin_vec_list.append(origin_vec)
+   except IOError:
+       print(f"Error reading in {param['PARTICLE_OUT']} ")
+
+   id_list =  np.asarray(id_list)[:,0]
+   origin_type_list = np.asarray(origin_type_list)
+   origin_vec_list = np.vstack(origin_vec_list)
+
+   typeda = xr.DataArray(origin_type_list, dims=['id'], coords={'id' : id_list})
+   vecda = xr.DataArray(origin_vec_list, dims=['id', 'vec'], coords={'id' : id_list, 'vec' : veclab})
+
+   infoxr = vecda.to_dataset(dim='vec')
+   infoxr['origin_type'] = typeda
+
+   print('\nAdding particle info to Dataset')
+   ds = xr.merge([ds, infoxr])
+   return ds
+
+
 def swiftest_xr2infile(ds, param, framenum=-1):
     """
     Writes a set of Swiftest input files from a single frame of a Swiftest xarray dataset
@@ -1234,6 +1314,10 @@ def swifter2swiftest(swifter_param, plname="", tpname="", cbname="", conversion_
         if swiftest_param['DISCARD_OUT'] == '':
             swiftest_param['DISCARD_OUT'] = "discard.out"
 
+    swiftest_param['PARTICLE_OUT'] = conversion_questions.get('PARTICLE_OUT', '')
+    if not swiftest_param['PARTICLE_OUT']:
+        swiftest_param['PARTICLE_OUT'] = input("PARTICLE_OUT: Particle info file name (Only used by SyMBA): []> ")
+
     swiftest_param['! VERSION'] = "Swiftest parameter file converted from Swifter"
     return swiftest_param
 

python/swiftest/swiftest/simulation_class.py

@@ -38,6 +38,7 @@ def __init__(self, codename="Swiftest", param_file=""):
             'DU2M': constants.AU2M,
             'EXTRA_FORCE': "NO",
             'DISCARD_OUT': "discard.out",
+            'PARTICLE_OUT' : "",
             'BIG_DISCARD': "NO",
             'CHK_CLOSE': "YES",
             'RHILL_PRESENT': "YES",