diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index 702f05fcd..ee262a6a6 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): @@ -872,32 +875,15 @@ def select_active_from_frame(ds, param, framenum=-1): # Select the input time frame from the Dataset frame = ds.isel(time=[framenum]) + iframe = frame.isel(time=0) # 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()) - - 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)) + iactive = iframe['id'].where((~np.isnan(iframe['Gmass'])) | (~np.isnan(iframe['xhx'])), drop=True).id else: - active_masker = xr.apply_ufunc(is_not_nan, frame_both['a'].isel(time=0)) - active_masker = active_masker.drop_vars("time") - active_masker = xr.where(active_masker['id'] == 0, True, active_masker) - - frame_id_only = frame_id_only.where(active_masker, drop=True) - frame_both = frame_both.where(active_masker, drop=True) - - frame = frame_both.combine_first(frame_time_only).combine_first(frame_id_only) - frame = frame.transpose("time", "id") + iactive = iframe['id'].where((~np.isnan(iframe['Gmass'])) | (~np.isnan(iframe['a'])), drop = True).id + frame = frame.isel(id=iactive.values) return frame @@ -918,7 +904,6 @@ def swiftest_xr2infile(ds, param, framenum=-1): ------- A set of input files for a new Swiftest run """ - frame = select_active_from_frame(ds, param, framenum) if param['IN_TYPE'] == "NETCDF_DOUBLE" or param['IN_TYPE'] == "NETCDF_FLOAT": @@ -927,7 +912,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) @@ -1699,4 +1684,4 @@ def swifter2swift_param(swifter_param, J2=0.0, J4=0.0): swift_param['L2'] = 'T' - return swift_param \ No newline at end of file + return swift_param diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 8242f8088..740658d07 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -6,6 +6,7 @@ import xarray as xr import numpy as np import os +import shutil class Simulation: """ @@ -239,4 +240,39 @@ def save(self, param_file, framenum=-1, codename="Swiftest"): else: print(f'Saving to {codename} not supported') - return \ No newline at end of file + return + + def initial_conditions_from_bin(self, framenum=-1, new_param=None, new_param_file="param.new.in", new_initial_conditions_file="bin_in.nc", restart=False, codename="Swiftest"): + if codename != "Swiftest": + self.save(new_param_file, framenum, codename) + return + if new_param is None: + new_param = self.param.copy() + + if codename == "Swiftest": + if restart: + new_param['T0'] = self.ds.time.values[framenum] + if self.param['OUT_TYPE'] == 'NETCDF_DOUBLE' or self.param['OUT_TYPE'] == 'REAL8': + new_param['IN_TYPE'] = 'NETCDF_DOUBLE' + elif self.param['OUT_TYPE'] == 'NETCDF_FLOAT' or self.param['OUT_TYPE'] == 'REAL4': + new_param['IN_TYPE'] = 'NETCDF_FLOAT' + else: + print(f"{self.param['OUT_TYPE']} is an invalid OUT_TYPE file") + return + if self.param['BIN_OUT'] != new_param['BIN_OUT'] and restart: + print(f"Restart run with new output file. Copying {self.param['BIN_OUT']} to {new_param['BIN_OUT']}") + shutil.copy2(self.param['BIN_OUT'],new_param['BIN_OUT']) + new_param['IN_FORM'] = 'XV' + if restart: + new_param['OUT_STAT'] = 'APPEND' + new_param['FIRSTKICK'] = 'T' + new_param['NC_IN'] = new_initial_conditions_file + new_param.pop('PL_IN', None) + 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']}") + frame = io.swiftest_xr2infile(self.ds, new_param, framenum) + print(f"Saving parameter configuration file to {new_param_file}") + self.write_param(new_param_file, param=new_param) + + return frame