From 3b441e7b1f581abd0011c7f10b087fb5f18dfbae Mon Sep 17 00:00:00 2001 From: David Minton Date: Tue, 10 Aug 2021 13:06:54 -0400 Subject: [PATCH 1/2] Added back the swiftest_particle_2xr function to io.py. Also added two spaces between function definitions --- python/swiftest/swiftest/io.py | 37 ++++++++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) 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 From 1fa1832eca0c30eaac14460de59d108ab250389f Mon Sep 17 00:00:00 2001 From: David Minton Date: Tue, 10 Aug 2021 13:27:12 -0400 Subject: [PATCH 2/2] Added back the particle info reader --- .../symba_energy_momentum/collision_movie.py | 2 +- python/swiftest/swiftest/io.py | 99 ++++++++++++++----- python/swiftest/swiftest/simulation_class.py | 1 + 3 files changed, 75 insertions(+), 27 deletions(-) diff --git a/examples/symba_energy_momentum/collision_movie.py b/examples/symba_energy_momentum/collision_movie.py index ec4741895..85a020183 100755 --- a/examples/symba_energy_momentum/collision_movie.py +++ b/examples/symba_energy_momentum/collision_movie.py @@ -260,7 +260,7 @@ def data_stream(self, frame=0): vx = np.nan_to_num(vx, copy=False) vy = np.nan_to_num(vy, copy=False) radmarker = np.nan_to_num(radmarker, copy=False) - GMass = np.nan_to_num(Mass, copy=False) + GMass = np.nan_to_num(GMass, copy=False) Radius = np.nan_to_num(Radius, copy=False) rotx = np.nan_to_num(rotx, copy=False) roty = np.nan_to_num(roty, copy=False) diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index 037bc0806..0492a05f8 100644 --- a/python/swiftest/swiftest/io.py +++ b/python/swiftest/swiftest/io.py @@ -73,6 +73,8 @@ def read_swiftest_param(param_file_name, param): param['CHK_CLOSE'] = param['CHK_CLOSE'].upper() param['RHILL_PRESENT'] = param['RHILL_PRESENT'].upper() param['FRAGMENTATION'] = param['FRAGMENTATION'].upper() + if param['FRAGMENTATION'] == 'YES' and param['PARTICLE_OUT'] == '': + param['PARTICLE_OUT'] = 'particle.dat' param['ROTATION'] = param['ROTATION'].upper() param['TIDES'] = param['TIDES'].upper() param['ENERGY'] = param['ENERGY'].upper() @@ -612,26 +614,29 @@ def swiftest2xr(param): cb = [] pl = [] tp = [] - with FortranFile(param['BIN_OUT'], 'r') as f: - for t, cbid, cvec, clab, \ - npl, plid, pvec, plab, \ - ntp, tpid, tvec, tlab in swiftest_stream(f, param): - # Prepare frames by adding an extra axis for the time coordinate - cbframe = np.expand_dims(cvec, axis=0) - plframe = np.expand_dims(pvec, axis=0) - tpframe = np.expand_dims(tvec, axis=0) + try: + with FortranFile(param['BIN_OUT'], 'r') as f: + for t, cbid, cvec, clab, \ + npl, plid, pvec, plab, \ + ntp, tpid, tvec, tlab in swiftest_stream(f, param): + # Prepare frames by adding an extra axis for the time coordinate + cbframe = np.expand_dims(cvec, axis=0) + plframe = np.expand_dims(pvec, axis=0) + tpframe = np.expand_dims(tvec, axis=0) + + # Create xarray DataArrays out of each body type + cbxr = xr.DataArray(cbframe, dims=dims, coords={'time': t, 'id': cbid, 'vec': clab}) + plxr = xr.DataArray(plframe, dims=dims, coords={'time': t, 'id': plid, 'vec': plab}) + tpxr = xr.DataArray(tpframe, dims=dims, coords={'time': t, 'id': tpid, 'vec': tlab}) + + cb.append(cbxr) + pl.append(plxr) + tp.append(tpxr) + sys.stdout.write('\r' + f"Reading in time {t[0]:.3e}") + sys.stdout.flush() + except IOError: + print(f"Error encountered reading in {param['BIN_OUT']}") - # Create xarray DataArrays out of each body type - cbxr = xr.DataArray(cbframe, dims=dims, coords={'time': t, 'id': cbid, 'vec': clab}) - plxr = xr.DataArray(plframe, dims=dims, coords={'time': t, 'id': plid, 'vec': plab}) - tpxr = xr.DataArray(tpframe, dims=dims, coords={'time': t, 'id': tpid, 'vec': tlab}) - - cb.append(cbxr) - pl.append(plxr) - tp.append(tpxr) - sys.stdout.write('\r' + f"Reading in time {t[0]:.3e}") - sys.stdout.flush() - cbda = xr.concat(cb, dim='time') plda = xr.concat(pl, dim='time') tpda = xr.concat(tp, dim='time') @@ -642,22 +647,59 @@ def swiftest2xr(param): print('\nCreating Dataset') ds = xr.combine_by_coords([cbds, plds, tpds]) print(f"Successfully converted {ds.sizes['time']} output frames.") + if param['PARTICLE_OUT'] != "": + ds = swiftest_particle_2xr(ds, param) + return ds +def swiftest_particle_stream(f): + """ + Reads in a Swiftest particle.dat file and returns a single frame of particle data as a datastream + + Parameters + ---------- + f : file object + param : dict + + Yields + ------- + plid : int + ID of massive bodie + origin_type : string + The origin type for the body (Initial conditions, disruption, supercatastrophic, hit and run, etc) + origin_xh : float array + The origin heliocentric position vector + origin_vh : float array + The origin heliocentric velocity vector + """ + while True: # Loop until you read the end of file + try: + # Read multi-line header + plid = f.read_ints() # Try first part of the header + except: + break + origin_rec = f.read_record(np.dtype('a32'), np.dtype((' ") + swiftest_param['! VERSION'] = "Swiftest parameter file converted from Swifter" return swiftest_param diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index bea59ae5f..fc5075ab9 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -38,6 +38,7 @@ def __init__(self, codename="Swiftest", param_file=""): 'DU2M': constants.AU2M, 'EXTRA_FORCE': "NO", 'DISCARD_OUT': "discard.out", + 'PARTICLE_OUT' : "", 'BIG_DISCARD': "NO", 'CHK_CLOSE': "YES", 'RHILL_PRESENT': "YES",