diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index 37f3370fd..037bc0806 100644 --- a/python/swiftest/swiftest/io.py +++ b/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 @@ -82,6 +83,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 +167,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 +254,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 +266,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 +305,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 +550,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 +593,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. @@ -636,6 +644,35 @@ def swiftest2xr(param): print(f"Successfully converted {ds.sizes['time']} output frames.") return ds + +def swiftest_particle_2xr(ds, param): + """Reads in the Swiftest 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 = [] + + with FortranFile(param['PARTICLE_OUT'], 'r') as f: + for plid, origin_type, origin_vec in swiftest_particle_stream(f): + id_list.append(plid) + + origin_type_list.append(origin_type) + origin_vec_list.append(origin_vec) + + 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