Skip to content
This repository was archived by the owner on Aug 28, 2024. It is now read-only.

Commit

Permalink
Cleaned up bugs that were preventing Swifter input file from being re…
Browse files Browse the repository at this point in the history 
…ad properly
  • Loading branch information
@daminton
daminton committed May 15, 2023
1 parent fb15c36 commit 1620fdf
Show file tree
Hide file tree
Showing 3 changed files with 179 additions and 376 deletions.
6 changes: 3 additions & 3 deletions python/swiftest/swiftest/init_cond.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -233,7 +233,7 @@ def vec2xr(param: Dict, **kwargs: Any):
Parameters
----------
param : dict
Swiftest paramuration parameters.
Swiftest simulation parameters.
name : str or array-like of str, optional
Name or names of Bodies. If none passed, name will be "Body<id>"
id : int or array-like of int, optional
Expand All @@ -252,9 +252,9 @@ def vec2xr(param: Dict, **kwargs: Any):
capm : float or array-like of float, optional
mean anomaly for param['IN_FORM'] == "EL"
rh : (n,3) array-like of float, optional
Position vector array. This can be used instead of passing v1, v2, and v3 sepearately for "XV" input format
Position vector array.
vh : (n,3) array-like of float, optional
Velocity vector array. This can be used instead of passing v4, v5, and v6 sepearately for "XV" input format
Velocity vector array.
Gmass : float or array-like of float, optional
G*mass values if these are massive bodies (only one of mass or Gmass can be passed)
radius : float or array-like of float, optional
Expand Down
304 changes: 51 additions & 253 deletions python/swiftest/swiftest/io.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
"YARKOVSKY",
"YORP",
"IN_FORM",
"NC_IN",
"SEED",
"INTERACTION_LOOPS",
"ENCOUNTER_CHECK",
Expand Down Expand Up @@ -483,6 +484,7 @@ def swifter_stream(f, param):
tlab : string list
Labels for the tvec data
"""

while True: # Loop until you read the end of file
try:
# Read single-line header
Expand All @@ -492,32 +494,52 @@ def swifter_stream(f, param):
t = record[0]
npl = record[1][0] - 1
ntp = record[2][0]

pvec = np.empty((6, npl))

if param['OUT_FORM'] == 'XV':
rpvec = np.empty((npl,3))
vpvec = np.empty((npl,3))
rtvec = np.empty((ntp,3))
vtvec = np.empty((ntp,3))
elif param['OUT_FORM'] == 'EL':
elempvec = np.empty((npl,6))
elemtvec = np.empty((ntp,6))
plid = np.empty(npl, dtype='int')
tvec = np.empty((6, ntp))
tpid = np.empty(ntp, dtype='int')
if npl > 0:
GMpl = np.empty(npl)
Rpl = np.empty(npl)
for i in range(npl):
# Read single-line pl frame for
record = f.read_record('<i4', '<f8', '<f8', '(6,)<f8')
if param['OUT_FORM'] == 'XV':
record = f.read_record('<i4', '<f8', '<f8', '(3,)<f8', '(3,)<f8')
elif param['OUT_FORM'] == 'EL':
record = f.read_record('<i4', '<f8', '<f8', '(6,)<f8')
plid[i] = record[0]
GMpl[i] = record[1]
Rpl[i] = record[2]
pvec[:, i] = record[3]
if param['OUT_FORM'] == 'XV':
rpvec[i,:] = record[3]
vpvec[i,:] = record[4]
elif param['OUT_FORM'] == 'EL':
elempvec[i,:] = record[3]
if ntp > 0:
for i in range(ntp):
record = f.read_record('<i4', '(6,)<f8')
if param['OUT_FORM'] == 'XV':
record = f.read_record('<i4', '(3,)<f8', '(3,)<f8')
elif param['OUT_FORM'] == 'EL':
record = f.read_record('<i4', '(6,)<f8')
tpid[i] = record[0]
tvec[:, i] = record[1]
if param['OUT_FORM'] == 'XV':
rtvec[i,:] = record[1]
vtvec[i,:] = record[2]
elif param['OUT_FORM'] == 'EL':
elemtvec[i,:] = record[1]

tlab = []
if param['OUT_FORM'] == 'XV' or param['OUT_FORM'] == 'XVEL':
tlab = ['id']
if param['OUT_FORM'] == 'XV':
tlab.append('rh')
tlab.append('vh')
if param['OUT_FORM'] == 'EL' or param['OUT_FORM'] == 'XVEL':
elif param['OUT_FORM'] == 'EL':
tlab.append('a')
tlab.append('e')
tlab.append('inc')
Expand All @@ -527,245 +549,13 @@ def swifter_stream(f, param):
plab = tlab.copy()
plab.append('Gmass')
plab.append('radius')
pvec = np.vstack([pvec, GMpl, Rpl])

yield t, npl, plid, pvec.T, plab, \
ntp, tpid, tvec.T, tlab


def make_swiftest_labels(param):
"""
Creates the lables for the variables to be included in the output file.
Parameters
----------
param : dictionary
The entries in the user parameter file
Returns
-------
clab : string list
Labels for the cvec data
plab : string list
Labels for the pvec data
tlab : string list
Labels for the tvec data
infolab_float : string list
Labels for floating point data
infolab_int :
Labels for integer data
infolab_str :
Labels for string data
"""
tlab = []
if param['OUT_FORM'] == 'XV' or param['OUT_FORM'] == 'XVEL':
tlab.append('rh')
tlab.append('vh')

if param['OUT_FORM'] == 'EL' or param['OUT_FORM'] == 'XVEL':
tlab.append('a')
tlab.append('e')
tlab.append('inc')
tlab.append('capom')
tlab.append('omega')
tlab.append('capm')
plab = tlab.copy()
plab.append('Gmass')
if param['CHK_CLOSE']:
plab.append('radius')
if param['RHILL_PRESENT']:
plab.append('rhill')
clab = ['Gmass', 'radius', 'j2rp2', 'j4rp4']
if param['ROTATION']:
clab.append('Ip')
clab.append('rot')
plab.append('Ip')
plab.append('rot')
if param['TIDES']:
clab.append('k2')
clab.append('Q')
plab.append('k2')
plab.append('Q')

infolab_float = [
"origin_time",
"origin_rh",
"origin_vh",
"discard_time",
"discard_rh",
"discard_vh",
]
infolab_int = [
"collision_id",
"discard_body_id"
]
infolab_str = [
"particle_type",
"origin_type",
"status",
]

return clab, plab, tlab, infolab_float, infolab_int, infolab_str


def swiftest_stream(f, param):
"""
Reads in a Swifter bin.dat file and returns a single frame of data as a datastream
Parameters
----------
f : file object
param : dict
Yields
-------
t : float
Time of this frame
cbid : int array
ID of central body (always returns 0)
cvec : float array
(npl,1) - vector of quantities for the massive body (Gmass, radius, J2, J4, etc)
npl : int
Number of massive bodies
plid : int array
IDs of massive bodies
pvec : float array
(npl,N) - vector of N quantities or each particle (6 of XV/EL + Gmass, radius, etc)
plab : string list
Labels for the pvec data
ntp : int
Number of test particles
tpid : int array
Ids of test particles
tvec : float array
(ntp,N) - vector of N quantities for each particle (6 of XV/EL, etc.)
tlab : string list
Labels for the tvec data
"""
NAMELEN = 32
while True: # Loop until you read the end of file
try:
# Read multi-line header
t = f.read_reals(np.float64) # Try first part of the header
except:
break
npl = f.read_ints()
ntp = f.read_ints()
iout_form = f.read_reals('c')
cbid = f.read_ints()
dtstr = f'a{NAMELEN}'
cbnames = f.read_record(dtstr)
cbnames = [np.char.strip(str(cbnames[0], encoding='utf-8'))]
Mcb = f.read_reals(np.float64)
Rcb = f.read_reals(np.float64)
J2cb = f.read_reals(np.float64)
J4cb = f.read_reals(np.float64)
if param['ROTATION']:
Ipcbx = f.read_reals(np.float64)
Ipcby = f.read_reals(np.float64)
Ipcbz = f.read_reals(np.float64)
rotcbx = f.read_reals(np.float64)
rotcby = f.read_reals(np.float64)
rotcbz = f.read_reals(np.float64)
if param['TIDES']:
k2cb = f.read_reals(np.float64)
Qcb = f.read_reals(np.float64)
if npl[0] > 0:
plid = f.read_ints()
dtstr = f'({npl[0]},)a{NAMELEN}'
names = f.read_record(np.dtype(dtstr))
plnames = []
for i in range(npl[0]):
plnames.append(np.char.strip(str(names[i], encoding='utf-8')))
p1 = f.read_reals(np.float64)
p2 = f.read_reals(np.float64)
p3 = f.read_reals(np.float64)
p4 = f.read_reals(np.float64)
p5 = f.read_reals(np.float64)
p6 = f.read_reals(np.float64)
if param['OUT_FORM'] == 'XVEL':
p7 = f.read_reals(np.float64)
p8 = f.read_reals(np.float64)
p9 = f.read_reals(np.float64)
p10 = f.read_reals(np.float64)
p11 = f.read_reals(np.float64)
p12 = f.read_reals(np.float64)
GMpl = f.read_reals(np.float64)
if param['RHILL_PRESENT']:
rhill = f.read_reals(np.float64)
Rpl = f.read_reals(np.float64)
if param['ROTATION']:
Ipplx = f.read_reals(np.float64)
Ipply = f.read_reals(np.float64)
Ipplz = f.read_reals(np.float64)
rotplx = f.read_reals(np.float64)
rotply = f.read_reals(np.float64)
rotplz = f.read_reals(np.float64)
if param['TIDES']:
k2pl = f.read_reals(np.float64)
Qpl = f.read_reals(np.float64)
if ntp[0] > 0:
tpid = f.read_ints()
dtstr = f'({ntp[0]},)a{NAMELEN}'
names = f.read_record(np.dtype(dtstr))
tpnames = []
for i in range(npl[0]):
tpnames.append(np.char.strip(str(names[i], encoding='utf-8')))
t1 = f.read_reals(np.float64)
t2 = f.read_reals(np.float64)
t3 = f.read_reals(np.float64)
t4 = f.read_reals(np.float64)
t5 = f.read_reals(np.float64)
t6 = f.read_reals(np.float64)
if param['OUT_FORM'] == 'XVEL':
t7 = f.read_reals(np.float64)
t8 = f.read_reals(np.float64)
t9 = f.read_reals(np.float64)
t10 = f.read_reals(np.float64)
t11 = f.read_reals(np.float64)
t12 = f.read_reals(np.float64)

clab, plab, tlab, infolab_float, infolab_int, infolab_str = make_swiftest_labels(param)

if npl > 0:
pvec = np.vstack([p1, p2, p3, p4, p5, p6])
if param['OUT_FORM'] == 'XVEL':
pvec = np.vstack([pvec, p7, p8, p9, p10, p11, p12])
pvec = np.vstack([pvec, GMpl, Rpl])
else:
if param['OUT_FORM'] == 'XVEL':
pvec = np.empty((14, 0))
else:
pvec = np.empty((8, 0))
plid = np.empty(0)
plnames = np.empty(0)
if ntp > 0:
tvec = np.vstack([t1, t2, t3, t4, t5, t6])
if param['OUT_FORM'] == 'XVEL':
tvec = np.vstack([tvec, t7, t8, t9, t10, t11, t12])
else:
if param['OUT_FORM'] == 'XVEL':
tvec = np.empty((12, 0))
else:
tvec = np.empty((6, 0))
tpid = np.empty(0)
tpnames = np.empty(0)
cvec = np.array([Mcb, Rcb, J2cb, J4cb])
if param['RHILL_PRESENT']:
if npl > 0:
pvec = np.vstack([pvec, rhill])
if param['ROTATION']:
cvec = np.vstack([cvec, Ipcbx, Ipcby, Ipcbz, rotcbx, rotcby, rotcbz])
if npl > 0:
pvec = np.vstack([pvec, Ipplx, Ipply, Ipplz, rotplx, rotply, rotplz])
if param['TIDES']:
cvec = np.vstack([cvec, k2cb, Qcb])
if npl > 0:
pvec = np.vstack([pvec, k2pl, Qpl])
yield t, cbid, cbnames, cvec.T, clab, \
npl, plid, plnames, pvec.T, plab, \
ntp, tpid, tpnames, tvec.T, tlab

if param['OUT_FORM'] == 'XV':
yield t, npl, [plid, rpvec, vpvec, GMpl, Rpl], plab, \
ntp, [tpid, rtvec, vtvec], tlab
elif param['OUT_FORM'] == 'EL':
yield t, npl, [plid, elempvec, GMpl, Rpl], plab, \
ntp, [tpid, elemtvec], tlab

def swifter2xr(param, verbose=True):
"""
Expand All @@ -783,9 +573,16 @@ def swifter2xr(param, verbose=True):
dims = ['time', 'id','vec']
pl = []
tp = []
ds = []
with FortranFile(param['BIN_OUT'], 'r') as f:
for t, npl, plid, pvec, plab, \
ntp, tpid, tvec, tlab in swifter_stream(f, param):
for t, npl, pvec, plab, \
ntp, tvec, tlab in swifter_stream(f, param):

pvec_args = dict(zip(plab,pvec))
plds = swiftest.init_cond.vec2xr(param,**pvec_args)
if ntp > 0:
tvec_args = dict(zip(tlab,tvec))

# Prepare frames by adding an extra axis for the time coordinate
plframe = np.expand_dims(pvec, axis=0)
tpframe = np.expand_dims(tvec, axis=0)
Expand All @@ -806,6 +603,7 @@ def swifter2xr(param, verbose=True):
tpds = tpda.to_dataset(dim='vec')
if verbose: print('\nCreating Dataset')
ds = xr.combine_by_coords([plds, tpds])
#if param['OUT_FORM'] = 'XV':
if verbose: print(f"Successfully converted {ds.sizes['time']} output frames.")
return ds

Expand Down Expand Up @@ -1848,9 +1646,9 @@ def swiftest2swifter_param(swiftest_param, J2=0.0, J4=0.0):
DU2M = swifter_param.pop("DU2M", 1.0)
TU2S = swifter_param.pop("TU2S", 1.0)
GR = swifter_param.pop("GR", None)
if GR is not None:
if GR:
swifter_param['C'] = swiftest.einsteinC * np.longdouble(TU2S) / np.longdouble(DU2M)
# if GR is not None:
# if GR:
# swifter_param['C'] = swiftest.einsteinC * np.longdouble(TU2S) / np.longdouble(DU2M)
for key in newfeaturelist:
tmp = swifter_param.pop(key, None)
if "ISTEP_DUMP" not in swifter_param:
Expand Down
Loading

0 comments on commit 1620fdf

Please sign in to comment.