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

Commit

Permalink
Enabled non-unique character coordinate name to bodies. Tested in REA…
Browse files Browse the repository at this point in the history 
…L8 mode, but does not yet work in NetCDF
  • Loading branch information
@daminton
daminton committed Aug 27, 2021
1 parent e9dc3bf commit be5e5d2
Show file tree
Hide file tree
Showing 5 changed files with 112 additions and 77 deletions.
36 changes: 29 additions & 7 deletions python/swiftest/swiftest/io.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -468,6 +468,7 @@ def swiftest_stream(f, param):
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
Expand All @@ -478,6 +479,9 @@ def swiftest_stream(f, param):
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)
Expand All @@ -494,6 +498,11 @@ def swiftest_stream(f, param):
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)
Expand Down Expand Up @@ -523,6 +532,11 @@ def swiftest_stream(f, param):
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)
Expand Down Expand Up @@ -550,6 +564,7 @@ def swiftest_stream(f, param):
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':
Expand All @@ -560,6 +575,7 @@ def swiftest_stream(f, param):
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'] == 'YES':
if npl > 0:
Expand All @@ -572,9 +588,9 @@ def swiftest_stream(f, param):
cvec = np.vstack([cvec, k2cb, Qcb])
if npl > 0:
pvec = np.vstack([pvec, k2pl, Qpl])
yield t, cbid, cvec.T, clab, \
npl, plid, pvec.T, plab, \
ntp, tpid, tvec.T, tlab
yield t, cbid, cbnames, cvec.T, clab, \
npl, plid, plnames, pvec.T, plab, \
ntp, tpid, tpnames, tvec.T, tlab


def swifter2xr(param):
Expand Down Expand Up @@ -638,24 +654,30 @@ def swiftest2xr(param):
cb = []
pl = []
tp = []
cbn = None
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):
for t, cbid, cbnames, cvec, clab, \
npl, plid, plnames, pvec, plab, \
ntp, tpid, tpnames, 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})
cbxr = cbxr.assign_coords(name=("id", cbnames))
plxr = xr.DataArray(plframe, dims=dims, coords={'time': t, 'id': plid, 'vec': plab})
plxr = plxr.assign_coords(name=("id", plnames))
tpxr = xr.DataArray(tpframe, dims=dims, coords={'time': t, 'id': tpid, 'vec': tlab})

tpxr = tpxr.assign_coords(name=("id", tpnames))

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:
Expand Down
8 changes: 5 additions & 3 deletions src/io/io.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1036,7 +1036,7 @@ module function io_read_frame_body(self, iu, param) result(ierr)
select case(param%in_type)
case (REAL4_TYPE, REAL8_TYPE)
read(iu, err = 667, iomsg = errmsg) self%id(:)
!read(iu, err = 667, iomsg = errmsg) self%name(:)
read(iu, err = 667, iomsg = errmsg) self%name(:)

select case (param%in_form)
case (XV)
Expand Down Expand Up @@ -1158,7 +1158,7 @@ module function io_read_frame_cb(self, iu, param) result(ierr)
character(len=STRMAX) :: errmsg

read(iu, err = 667, iomsg = errmsg) self%id
!read(iu, err = 667, iomsg = errmsg) self%name
read(iu, err = 667, iomsg = errmsg) self%name
read(iu, err = 667, iomsg = errmsg) self%Gmass
self%mass = self%Gmass / param%GU
read(iu, err = 667, iomsg = errmsg) self%radius
Expand Down Expand Up @@ -1473,6 +1473,7 @@ module subroutine io_write_encounter(self, pl, encbody, param)
call util_exit(FAILURE)
end subroutine io_write_encounter


module subroutine io_write_frame_body(self, iu, param)
!! author: David A. Minton
!!
Expand All @@ -1492,6 +1493,7 @@ module subroutine io_write_frame_body(self, iu, param)
associate(n => self%nbody)
if (n == 0) return
write(iu, err = 667, iomsg = errmsg) self%id(1:n)
write(iu, err = 667, iomsg = errmsg) self%name(1:n)
if ((param%out_form == XV) .or. (param%out_form == XVEL)) then
write(iu, err = 667, iomsg = errmsg) self%xh(1, 1:n)
write(iu, err = 667, iomsg = errmsg) self%xh(2, 1:n)
Expand Down Expand Up @@ -1551,8 +1553,8 @@ module subroutine io_write_frame_cb(self, iu, param)
character(len=STRMAX) :: errmsg

associate(cb => self)
!write(iu, err = 667, iomsg = errmsg) cb%name
write(iu, err = 667, iomsg = errmsg) cb%id
write(iu, err = 667, iomsg = errmsg) cb%name
write(iu, err = 667, iomsg = errmsg) cb%Gmass
write(iu, err = 667, iomsg = errmsg) cb%radius
write(iu, err = 667, iomsg = errmsg) cb%j2rp2
Expand Down
53 changes: 27 additions & 26 deletions src/modules/swiftest_classes.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -86,10 +86,11 @@ module swiftest_classes
type :: netcdf_parameters
integer(I4B) :: out_type !! NetCDF output type (will be assigned either NF90_DOUBLE or NF90_FLOAT, depending on the user parameter)
integer(I4B) :: ncid !! NetCDF ID for the output file
integer(I4B) :: dimids(2) !! Dimensions of the NetCDF file
integer(I4B) :: dimids(3) !! Dimensions of the NetCDF file
integer(I4B) :: time_dimid !! NetCDF ID for the time dimension
integer(I4B) :: id_dimid !! NetCDF ID for the particle id dimension
integer(I4B) :: str_dimid !! NetCDF ID for the character string dimension
integer(I4B) :: time_varid !! NetCDF ID for the time variable
integer(I4B) :: id_dimid !! NetCDF ID for the particle name dimension
integer(I4B) :: id_varid !! NetCDF ID for the particle name variable
integer(I4B) :: name_varid !! NetCDF ID for the semimajor axis variable
integer(I4B) :: npl_varid !! NetCDF ID for the number of active massive bodies variable
Expand Down Expand Up @@ -156,7 +157,7 @@ module swiftest_classes
!********************************************************************************************************************************
!> A concrete lass for the central body in a Swiftest simulation
type, abstract, extends(swiftest_base) :: swiftest_cb
character(len=STRMAX) :: name !! Non-unique name
character(len=NAMELEN) :: name !! Non-unique name
integer(I4B) :: id = 0 !! External identifier (unique)
real(DP) :: mass = 0.0_DP !! Central body mass (units MU)
real(DP) :: Gmass = 0.0_DP !! Central mass gravitational term G * mass (units GU * MU)
Expand Down Expand Up @@ -193,29 +194,29 @@ module swiftest_classes
!> An abstract class for a generic collection of Swiftest bodies
type, abstract, extends(swiftest_base) :: swiftest_body
!! Superclass that defines the generic elements of a Swiftest particle
logical :: lfirst = .true. !! Run the current step as a first
integer(I4B) :: nbody = 0 !! Number of bodies
character(len=STRMAX), dimension(:), allocatable :: name !! Non-unique name
integer(I4B), dimension(:), allocatable :: id !! External identifier (unique)
integer(I4B), dimension(:), allocatable :: status !! An integrator-specific status indicator
logical, dimension(:), allocatable :: ldiscard !! Body should be discarded
logical, dimension(:), allocatable :: lmask !! Logical mask used to select a subset of bodies when performing certain operations (drift, kick, accel, etc.)
real(DP), dimension(:), allocatable :: mu !! G * (Mcb + [m])
real(DP), dimension(:,:), allocatable :: xh !! Swiftestcentric position
real(DP), dimension(:,:), allocatable :: vh !! Swiftestcentric velocity
real(DP), dimension(:,:), allocatable :: xb !! Barycentric position
real(DP), dimension(:,:), allocatable :: vb !! Barycentric velocity
real(DP), dimension(:,:), allocatable :: ah !! Total heliocentric acceleration
real(DP), dimension(:,:), allocatable :: aobl !! Barycentric accelerations of bodies due to central body oblatenes
real(DP), dimension(:,:), allocatable :: atide !! Tanngential component of acceleration of bodies due to tides
real(DP), dimension(:,:), allocatable :: agr !! Acceleration due to post-Newtonian correction
real(DP), dimension(:), allocatable :: ir3h !! Inverse heliocentric radius term (1/rh**3)
real(DP), dimension(:), allocatable :: a !! Semimajor axis (pericentric distance for a parabolic orbit)
real(DP), dimension(:), allocatable :: e !! Eccentricity
real(DP), dimension(:), allocatable :: inc !! Inclination
real(DP), dimension(:), allocatable :: capom !! Longitude of ascending node
real(DP), dimension(:), allocatable :: omega !! Argument of pericenter
real(DP), dimension(:), allocatable :: capm !! Mean anomaly
logical :: lfirst = .true. !! Run the current step as a first
integer(I4B) :: nbody = 0 !! Number of bodies
character(len=NAMELEN), dimension(:), allocatable :: name !! Non-unique name
integer(I4B), dimension(:), allocatable :: id !! External identifier (unique)
integer(I4B), dimension(:), allocatable :: status !! An integrator-specific status indicator
logical, dimension(:), allocatable :: ldiscard !! Body should be discarded
logical, dimension(:), allocatable :: lmask !! Logical mask used to select a subset of bodies when performing certain operations (drift, kick, accel, etc.)
real(DP), dimension(:), allocatable :: mu !! G * (Mcb + [m])
real(DP), dimension(:,:), allocatable :: xh !! Swiftestcentric position
real(DP), dimension(:,:), allocatable :: vh !! Swiftestcentric velocity
real(DP), dimension(:,:), allocatable :: xb !! Barycentric position
real(DP), dimension(:,:), allocatable :: vb !! Barycentric velocity
real(DP), dimension(:,:), allocatable :: ah !! Total heliocentric acceleration
real(DP), dimension(:,:), allocatable :: aobl !! Barycentric accelerations of bodies due to central body oblatenes
real(DP), dimension(:,:), allocatable :: atide !! Tanngential component of acceleration of bodies due to tides
real(DP), dimension(:,:), allocatable :: agr !! Acceleration due to post-Newtonian correction
real(DP), dimension(:), allocatable :: ir3h !! Inverse heliocentric radius term (1/rh**3)
real(DP), dimension(:), allocatable :: a !! Semimajor axis (pericentric distance for a parabolic orbit)
real(DP), dimension(:), allocatable :: e !! Eccentricity
real(DP), dimension(:), allocatable :: inc !! Inclination
real(DP), dimension(:), allocatable :: capom !! Longitude of ascending node
real(DP), dimension(:), allocatable :: omega !! Argument of pericenter
real(DP), dimension(:), allocatable :: capm !! Mean anomaly
!! Note to developers: If you add components to this class, be sure to update methods and subroutines that traverse the
!! component list, such as setup_body and util_spill
contains
Expand Down
2 changes: 2 additions & 0 deletions src/modules/swiftest_globals.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ module swiftest_globals
integer(I4B), parameter :: RINGMOONS = 9

integer(I4B), parameter :: STRMAX = 512 !! Maximum size of character strings
integer(I4B), parameter :: NAMELEN = 32 !! Maximum size of name strings

character(*), parameter :: ASCII_TYPE = 'ASCII' !! Symbolic name for ASCII file type
character(*), parameter :: REAL4_TYPE = 'REAL4' !! Symbolic name for binary file type REAL4
Expand Down Expand Up @@ -130,6 +131,7 @@ module swiftest_globals
character(*), parameter :: NETCDF_OUTFILE = 'bin.nc' !! Default output file name
character(*), parameter :: TIME_DIMNAME = "time" !! NetCDF name of the time dimension
character(*), parameter :: ID_DIMNAME = "id" !! NetCDF name of the particle id dimension
character(*), parameter :: STR_DIMNAME = "str" !! NetCDF name of the particle id dimension
character(*), parameter :: NAME_VARNAME = "name" !! NetCDF name of the particle name variable
character(*), parameter :: NPL_VARNAME = "npl" !! NetCDF name of the number of active massive bodies variable
character(*), parameter :: NTP_VARNAME = "ntp" !! NetCDF name of the number of active test particles variable
Expand Down
Loading

0 comments on commit be5e5d2

Please sign in to comment.