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",