Updated the python package to return the output frame when saving initial conditions from a bin file

python/swiftest/swiftest/io.py

@@ -50,8 +50,11 @@ def read_swiftest_param(param_file_name, param, verbose=True):
             for line in f.readlines():
                 fields = line.split()
                 if len(fields) > 0:
-                    for key in param:
-                        if (key == fields[0].upper()): param[key] = fields[1]
+                    if fields[0][0] != '!':
+                        key = fields[0].upper()
+                        param[key] = fields[1]
+                    #for key in param:
+                    #    if (key == fields[0].upper()): param[key] = fields[1]
                     # Special case of CHK_QMIN_RANGE requires a second input
                     if fields[0].upper() == 'CHK_QMIN_RANGE':
                         alo = real2float(fields[1])
@@ -744,7 +747,7 @@ def swiftest2xr(param, verbose=True):
 
 def xstrip(a):
     func = lambda x: np.char.strip(x)
-    return xr.apply_ufunc(func, a.str.decode(encoding='utf-8'))
+    return xr.apply_ufunc(func, a.str.decode(encoding='utf-8'),dask='parallelized')
 
 def string_converter(da):
    if da.dtype == np.dtype(object):
@@ -873,23 +876,20 @@ def select_active_from_frame(ds, param, framenum=-1):
     # Select the input time frame from the Dataset
     frame = ds.isel(time=[framenum])
 
-    # Select only the active particles at this time step
-    def is_not_nan(x):
-        return np.invert(np.isnan(x))
-
     # For NETCDF input file type, just do a dump of the system at the correct time frame
     frame_id_only = frame.drop_dims('time')
     frame_time_only = frame.drop_dims('id')
     frame_both = frame.drop_vars(frame_id_only.keys()).drop_vars(frame_time_only.keys())
 
+    # Select only the active particles at this time step
     def is_not_nan(x):
         return np.invert(np.isnan(x))
 
     # Remove the inactive particles
     if param['OUT_FORM'] == 'XV' or param['OUT_FORM'] == 'XVEL':
-        active_masker = xr.apply_ufunc(is_not_nan, frame_both['xhx'].isel(time=0))
+        active_masker = xr.apply_ufunc(is_not_nan, frame_both['xhx'].isel(time=0),dask='parallelized')
     else:
-        active_masker = xr.apply_ufunc(is_not_nan, frame_both['a'].isel(time=0))
+        active_masker = xr.apply_ufunc(is_not_nan, frame_both['a'].isel(time=0),dask='parallelized')
     active_masker = active_masker.drop_vars("time")
     active_masker = xr.where(active_masker['id'] == 0, True, active_masker)
 
@@ -927,7 +927,7 @@ def swiftest_xr2infile(ds, param, framenum=-1):
         # will rename this after reading
         frame = unclean_string_values(frame)
         frame.to_netcdf(path=param['NC_IN'])
-        return
+        return frame
 
     # All other file types need seperate files for each of the inputs
     cb = frame.where(frame.id == 0, drop=True)

python/swiftest/swiftest/simulation_class.py

@@ -275,8 +275,8 @@ def initial_conditions_from_bin(self, framenum=-1, new_param=None, new_param_fil
             new_param.pop('TP_IN', None)
             new_param.pop('CB_IN', None)
             print(f"Extracting data from dataset at time frame number {framenum} and saving it to {new_param['NC_IN']}")
-            io.swiftest_xr2infile(self.ds, new_param)
+            frame = io.swiftest_xr2infile(self.ds, new_param)
             print(f"Saving parameter configuration file to {new_param_file}")
             self.write_param(new_param_file, param=new_param)
 
-        return
+        return frame