diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index 702f05fcd..224f78ed9 100644 --- a/python/swiftest/swiftest/io.py +++ b/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) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 9e44f95f2..832894790 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/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 \ No newline at end of file + return frame \ No newline at end of file