From 09a700e161a86ba72a5797c8bb65bbc151007701 Mon Sep 17 00:00:00 2001 From: David Minton Date: Thu, 8 Sep 2022 14:53:22 -0400 Subject: [PATCH 1/6] Added a new method that can write a set of initial conditions file extracted from an existing simulation dataset. --- python/swiftest/swiftest/simulation_class.py | 37 ++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 8242f8088..b88920d1a 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,40 @@ def save(self, param_file, framenum=-1, codename="Swiftest"): else: print(f'Saving to {codename} not supported') + 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", new_bin_out_file=None, 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 new_bin_out_file is None: + new_param['BIN_OUT'] = os.path.basename(self.param['BIN_OUT']) + else: + if restart: + shutil.copy2(self.param['BIN_OUT'],new_bin_out_file) + new_param['BIN_OUT'] = new_bin_out_file + 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) + io.swiftest_xr2infile(self.ds, new_param) + self.write_param(new_param_file, param=new_param) + return \ No newline at end of file From cce040fa1f980bebbb5ea30cedb2367fd0631c4d Mon Sep 17 00:00:00 2001 From: David Minton Date: Thu, 8 Sep 2022 15:18:39 -0400 Subject: [PATCH 2/6] Added helpful messages when using initial_conditions_from_bin method --- python/swiftest/swiftest/simulation_class.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index b88920d1a..9e44f95f2 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -263,6 +263,7 @@ def initial_conditions_from_bin(self, framenum=-1, new_param=None, new_param_fil new_param['BIN_OUT'] = os.path.basename(self.param['BIN_OUT']) else: if restart: + print(f"Restart run with new output file. Copying {self.param['BIN_OUT']} to {new_bin_out_file}") shutil.copy2(self.param['BIN_OUT'],new_bin_out_file) new_param['BIN_OUT'] = new_bin_out_file new_param['IN_FORM'] = 'XV' @@ -273,7 +274,9 @@ def initial_conditions_from_bin(self, framenum=-1, new_param=None, new_param_fil 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']}") 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 From 67a6173f5320fe5d39e073e75bb6bef92cc68d3a Mon Sep 17 00:00:00 2001 From: David Minton Date: Wed, 14 Sep 2022 20:09:59 -0400 Subject: [PATCH 3/6] Updated the python package to return the output frame when saving initial conditions from a bin file --- python/swiftest/swiftest/io.py | 20 ++++++++++---------- python/swiftest/swiftest/simulation_class.py | 4 ++-- 2 files changed, 12 insertions(+), 12 deletions(-) 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 From f4818800e711c8f00e5b77f06fc2788557f611ef Mon Sep 17 00:00:00 2001 From: David Minton Date: Wed, 14 Sep 2022 20:42:34 -0400 Subject: [PATCH 4/6] changed how active particles are selected --- python/swiftest/swiftest/io.py | 22 ++++------------------ 1 file changed, 4 insertions(+), 18 deletions(-) diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index 224f78ed9..a83c07f06 100644 --- a/python/swiftest/swiftest/io.py +++ b/python/swiftest/swiftest/io.py @@ -875,29 +875,15 @@ def select_active_from_frame(ds, param, framenum=-1): # Select the input time frame from the Dataset frame = ds.isel(time=[framenum]) - - # 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()) + 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)) - # 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),dask='parallelized') + 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),dask='parallelized') - 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.sel(id=iactive.values) return frame From 88b2b57b7b32ac4b437b292cbb9a4edc35d59936 Mon Sep 17 00:00:00 2001 From: David A Minton Date: Wed, 14 Sep 2022 20:57:33 -0400 Subject: [PATCH 5/6] Fixed bug in which the last frame was always selected when generating initial conditions --- python/swiftest/swiftest/io.py | 5 ++--- python/swiftest/swiftest/simulation_class.py | 4 ++-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index a83c07f06..ee262a6a6 100644 --- a/python/swiftest/swiftest/io.py +++ b/python/swiftest/swiftest/io.py @@ -883,7 +883,7 @@ def select_active_from_frame(ds, param, framenum=-1): iactive = iframe['id'].where((~np.isnan(iframe['Gmass'])) | (~np.isnan(iframe['xhx'])), drop=True).id else: iactive = iframe['id'].where((~np.isnan(iframe['Gmass'])) | (~np.isnan(iframe['a'])), drop = True).id - frame = frame.sel(id=iactive.values) + frame = frame.isel(id=iactive.values) return frame @@ -904,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": @@ -1685,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 832894790..fb6f01192 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']}") - frame = io.swiftest_xr2infile(self.ds, new_param) + 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 \ No newline at end of file + return frame From 4c501a1d8eb351ab4ddf896a25c03dbdb11c4a62 Mon Sep 17 00:00:00 2001 From: David Minton Date: Thu, 15 Sep 2022 10:24:09 -0400 Subject: [PATCH 6/6] Cleaned up the way that BIN_OUT is handled when extracting initial conditions. Instead of having a specific new_bin_out argument, just read it in from new_param if it gets passed in --- python/swiftest/swiftest/simulation_class.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 832894790..b73106a48 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -242,7 +242,7 @@ def save(self, param_file, framenum=-1, codename="Swiftest"): 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", new_bin_out_file=None, restart=False, codename="Swiftest"): + 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 @@ -259,13 +259,9 @@ def initial_conditions_from_bin(self, framenum=-1, new_param=None, new_param_fil else: print(f"{self.param['OUT_TYPE']} is an invalid OUT_TYPE file") return - if new_bin_out_file is None: - new_param['BIN_OUT'] = os.path.basename(self.param['BIN_OUT']) - else: - if restart: - print(f"Restart run with new output file. Copying {self.param['BIN_OUT']} to {new_bin_out_file}") - shutil.copy2(self.param['BIN_OUT'],new_bin_out_file) - new_param['BIN_OUT'] = new_bin_out_file + 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'