mirror of
https://github.com/JHUAPL/kaipy.git
synced 2026-01-09 14:28:02 -05:00
Chimp docstings plus reduce toc depth
This commit is contained in:
wiltbemj
@@ -2,7 +2,7 @@ kaipy
|
||||
=====
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 4
|
||||
:maxdepth: 2
|
||||
|
||||
cdaweb_utils
|
||||
chimp
|
||||
|
||||
@@ -3,92 +3,184 @@ import numpy as np
|
||||
import h5py
|
||||
|
||||
#Return time at step n
|
||||
def tStep(fname,nStp=0):
|
||||
with h5py.File(fname,'r') as hf:
|
||||
gID = "Step#%d"%(nStp)
|
||||
def tStep(fname, nStp=0):
|
||||
"""
|
||||
Get the time attribute for a given step in an HDF5 file.
|
||||
|
||||
Args:
|
||||
fname (str): The path to the HDF5 file.
|
||||
nStp (int): The step number. Default is 0.
|
||||
|
||||
Returns:
|
||||
float: The time attribute for the specified step.
|
||||
"""
|
||||
with h5py.File(fname, 'r') as hf:
|
||||
gID = "Step#%d" % (nStp)
|
||||
grp = hf.get(gID)
|
||||
t = grp.attrs.get("time")
|
||||
return t
|
||||
|
||||
|
||||
#Count number of timesteps in an h5p file
|
||||
def cntSteps(fname):
|
||||
with h5py.File(fname,'r') as hf:
|
||||
"""
|
||||
Count the number of steps in an H5 file.
|
||||
|
||||
Args:
|
||||
fname (str): The path to the H5 file.
|
||||
|
||||
Returns:
|
||||
int: The number of steps in the H5 file.
|
||||
"""
|
||||
with h5py.File(fname, 'r') as hf:
|
||||
grps = hf.values()
|
||||
grpNames = [str(grp.name) for grp in grps]
|
||||
#Steps = [stp if "/Step#" in stp for stp in grpNames]
|
||||
Steps = [stp for stp in grpNames if "/Step#" in stp]
|
||||
nSteps = len(Steps)
|
||||
|
||||
return nSteps
|
||||
|
||||
|
||||
#Count number of particles in an h5p file
|
||||
def cntTPs(fname):
|
||||
with h5py.File(fname,'r') as hf:
|
||||
"""
|
||||
Count the number of time points in the given H5 file.
|
||||
|
||||
Args:
|
||||
fname (str): The path to the H5 file.
|
||||
|
||||
Returns:
|
||||
int: The number of time points in the H5 file.
|
||||
"""
|
||||
with h5py.File(fname, 'r') as hf:
|
||||
grp = hf.get("Step#0")
|
||||
ids = (grp.get("id")[()])
|
||||
ids = grp.get("id")[()]
|
||||
Np = ids.shape[0]
|
||||
return Np
|
||||
|
||||
|
||||
def bndTPs(fname):
|
||||
with h5py.File(fname,'r') as hf:
|
||||
"""
|
||||
Get the number of particles, minimum ID, and maximum ID from an HDF5 file.
|
||||
|
||||
Args:
|
||||
fname (str): The path to the HDF5 file.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing the number of particles (Np), the minimum ID (nS), and the maximum ID (nE).
|
||||
"""
|
||||
with h5py.File(fname, 'r') as hf:
|
||||
grp = hf.get("Step#0")
|
||||
ids = (grp.get("id")[()])
|
||||
ids = grp.get("id")[()]
|
||||
Np = ids.shape[0]
|
||||
nS = ids.min()
|
||||
nE = ids.max()
|
||||
return Np,nS,nE
|
||||
return Np, nS, nE
|
||||
|
||||
|
||||
#Find array index for a given particle ID (ie if block doesn't start at 1)
|
||||
def locPID(fname,pid):
|
||||
with h5py.File(fname,'r') as hf:
|
||||
def locPID(fname, pid):
|
||||
"""
|
||||
Find the location of a particle with a given ID in an HDF5 file.
|
||||
|
||||
Args:
|
||||
fname (str): The path to the HDF5 file.
|
||||
pid (int): The ID of the particle to locate.
|
||||
|
||||
Returns:
|
||||
int or None: The index of the particle in the file, or None if the particle is not found.
|
||||
"""
|
||||
with h5py.File(fname, 'r') as hf:
|
||||
grp = hf.get("Step#0")
|
||||
ids = grp.get("id")[()]
|
||||
isP = (ids == pid)
|
||||
loc = isP.argmax()
|
||||
if (ids[loc] != pid):
|
||||
print("Didn't find particle %d ..."%(pid))
|
||||
print("Didn't find particle %d ..." % (pid))
|
||||
loc = None
|
||||
quit()
|
||||
return loc
|
||||
|
||||
|
||||
|
||||
#Given an h5part file, create a time series for a single particle w/ ID = pid
|
||||
def getH5pid(fname,vId,pid):
|
||||
p0 = locPID(fname,pid) #Find particle in array
|
||||
Nt = cntSteps(fname) #Find number of slices
|
||||
def getH5pid(fname, vId, pid):
|
||||
"""
|
||||
Retrieve the time and values of a specific particle from an H5 file.
|
||||
|
||||
Args:
|
||||
fname (str): The path to the H5 file.
|
||||
vId (str): The ID of the values to retrieve.
|
||||
pid (int): The ID of the particle to retrieve.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing two arrays - the time array and the values array.
|
||||
"""
|
||||
p0 = locPID(fname, pid) # Find particle in array
|
||||
Nt = cntSteps(fname) # Find number of slices
|
||||
|
||||
V = np.zeros(Nt)
|
||||
t = np.zeros(Nt)
|
||||
with h5py.File(fname,'r') as hf:
|
||||
with h5py.File(fname, 'r') as hf:
|
||||
for n in range(Nt):
|
||||
#Create gId
|
||||
gId = "Step#%d"%(n)
|
||||
# Create gId
|
||||
gId = "Step#%d" % (n)
|
||||
grp = hf.get(gId)
|
||||
t[n] = grp.attrs.get("time")
|
||||
V[n] = (grp.get(vId)[()])[p0]
|
||||
return t,V
|
||||
return t, V
|
||||
|
||||
|
||||
|
||||
#Given an h5part file, create a time series from an input string
|
||||
def getH5p(fname,vId,Mask=None):
|
||||
def getH5p(fname, vId, Mask=None):
|
||||
"""
|
||||
Retrieve time and voltage data from an H5p file.
|
||||
|
||||
Args:
|
||||
fname (str): The file path of the H5p file.
|
||||
vId (str): The identifier of the voltage data to retrieve.
|
||||
Mask (ndarray, optional): An optional mask to apply to the voltage data.
|
||||
|
||||
Returns:
|
||||
ndarray: An array of time values.
|
||||
ndarray: An array of voltage values.
|
||||
|
||||
If Mask is not provided, the function returns the entire voltage data.
|
||||
If Mask is provided, the function returns the voltage data with the applied mask.
|
||||
"""
|
||||
Nt = cntSteps(fname)
|
||||
Np = cntTPs(fname)
|
||||
t = np.zeros(Nt)
|
||||
V = np.zeros((Nt,Np))
|
||||
with h5py.File(fname,'r') as hf:
|
||||
V = np.zeros((Nt, Np))
|
||||
with h5py.File(fname, 'r') as hf:
|
||||
for n in range(Nt):
|
||||
#Create gId
|
||||
gId = "Step#%d"%(n)
|
||||
# Create gId
|
||||
gId = "Step#%d" % (n)
|
||||
grp = hf.get(gId)
|
||||
t[n] = grp.attrs.get("time")
|
||||
V[n,:] = grp.get(vId)[()]
|
||||
if (Mask is None):
|
||||
return t,V
|
||||
|
||||
V[n, :] = grp.get(vId)[()]
|
||||
if Mask is None:
|
||||
return t, V
|
||||
else:
|
||||
return t,V[:,Mask]
|
||||
return t, V[:, Mask]
|
||||
|
||||
|
||||
#Given an h5p file and step, get one slice of data
|
||||
def getH5pT(fname,vID="isIn",nStp=0,cutIn=False):
|
||||
with h5py.File(fname,'r') as hf:
|
||||
gID = "Step#%d"%(nStp)
|
||||
def getH5pT(fname, vID="isIn", nStp=0, cutIn=False):
|
||||
"""
|
||||
Retrieve a specific dataset from an H5p file.
|
||||
|
||||
Args:
|
||||
fname (str): The path to the H5p file.
|
||||
vID (str): The name of the dataset to retrieve. Default is "isIn".
|
||||
nStp (int): The step number. Default is 0.
|
||||
cutIn (bool): Whether to cut in the dataset. Default is False.
|
||||
|
||||
Returns:
|
||||
The retrieved dataset.
|
||||
"""
|
||||
with h5py.File(fname, 'r') as hf:
|
||||
gID = "Step#%d" % (nStp)
|
||||
V = hf.get(gID).get(vID)[()]
|
||||
|
||||
return V
|
||||
|
||||
return V
|
||||
@@ -9,33 +9,46 @@ import h5py
|
||||
#xyz = |Nt,Np,3|
|
||||
#V = |Nt,Np,Nv|
|
||||
#vIDs = |Nv|
|
||||
def AddTPs(xyz,V=None,vIDs=None,fOut="tps.h5",nStp=0):
|
||||
gID = "Step#%d"%(nStp)
|
||||
oH5 = h5py.File(fOut,'w')
|
||||
Nt,Np,Nd = xyz.shape
|
||||
if (vIDs is not None):
|
||||
def AddTPs(xyz, V=None, vIDs=None, fOut="tps.h5", nStp=0):
|
||||
"""
|
||||
Add time points (TPs) to an HDF5 file.
|
||||
|
||||
Args:
|
||||
xyz (numpy.ndarray): The coordinates of the TPs. Shape (Nt, Np, Nd).
|
||||
V (numpy.ndarray, optional): The values associated with the TPs. Shape (Nt, Np, Nv).
|
||||
vIDs (list of str, optional): The names of the value datasets.
|
||||
fOut (str, optional): The output file path.
|
||||
nStp (int, optional): The step number.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
gID = "Step#%d" % (nStp)
|
||||
oH5 = h5py.File(fOut, 'w')
|
||||
Nt, Np, Nd = xyz.shape
|
||||
if vIDs is not None:
|
||||
Nv = len(vIDs)
|
||||
else:
|
||||
Nv = 0
|
||||
|
||||
#Create step data
|
||||
# Create step data
|
||||
oH5.create_group(gID)
|
||||
oH5[gID].attrs.create("time",nStp)
|
||||
oH5[gID].attrs.create("time", nStp)
|
||||
|
||||
for m in range(Np):
|
||||
lID = "Line#%d"%(m)
|
||||
lID = "Line#%d" % (m)
|
||||
|
||||
#Add connectivity
|
||||
ijL = np.zeros((Nt-1,2),dtype=int)
|
||||
ijL[:,0] = np.arange(Nt-1)
|
||||
ijL[:,1] = np.arange(1,Nt)
|
||||
#Save line data
|
||||
# Add connectivity
|
||||
ijL = np.zeros((Nt - 1, 2), dtype=int)
|
||||
ijL[:, 0] = np.arange(Nt - 1)
|
||||
ijL[:, 1] = np.arange(1, Nt)
|
||||
# Save line data
|
||||
oH5[gID].create_group(lID)
|
||||
oH5[gID][lID].attrs.create("Np",Nt)
|
||||
oH5[gID][lID].create_dataset("xyz",data=np.single(xyz[:,m,:]))
|
||||
oH5[gID][lID].create_dataset("LCon",data=ijL)
|
||||
oH5[gID][lID].attrs.create("Np", Nt)
|
||||
oH5[gID][lID].create_dataset("xyz", data=np.single(xyz[:, m, :]))
|
||||
oH5[gID][lID].create_dataset("LCon", data=ijL)
|
||||
|
||||
for i in range(Nv):
|
||||
oH5[gID][lID].create_dataset(vIDs[i],data=np.single(V[:,m,i]))
|
||||
oH5[gID][lID].create_dataset(vIDs[i], data=np.single(V[:, m, i]))
|
||||
|
||||
oH5.close()
|
||||
oH5.close()
|
||||
|
||||
@@ -5,25 +5,41 @@ import h5py
|
||||
import kaipy.kaiViz as kv
|
||||
|
||||
#Get grid from K-Cyl
|
||||
def getGrid(fIn,do4D=False):
|
||||
with h5py.File(fIn,'r') as hf:
|
||||
def getGrid(fIn, do4D=False):
|
||||
"""
|
||||
Retrieve grid data from an HDF5 file.
|
||||
|
||||
Args:
|
||||
fIn (str): The path to the HDF5 file.
|
||||
do4D (bool, optional): Flag indicating whether to retrieve 4D data. Default is False.
|
||||
|
||||
Returns:
|
||||
xx (ndarray): X-coordinate values of the grid.
|
||||
yy (ndarray): Y-coordinate values of the grid.
|
||||
Ki (ndarray): Original 3D data.
|
||||
Kc (ndarray): Averaged 3D data.
|
||||
Ai (ndarray, optional): Original 4D data. Only returned if do4D is True.
|
||||
Ac (ndarray, optional): Averaged 4D data. Only returned if do4D is True.
|
||||
"""
|
||||
with h5py.File(fIn, 'r') as hf:
|
||||
X3 = hf["X"][()].T
|
||||
Y3 = hf["Y"][()].T
|
||||
Z3 = hf["Z"][()].T
|
||||
if (do4D):
|
||||
if do4D:
|
||||
Ai = hf["A"][()].T
|
||||
|
||||
xx = X3[:,:,0]
|
||||
yy = Y3[:,:,0]
|
||||
Zi = Z3[0,0,:]
|
||||
Ki = 10**Zi
|
||||
Kc = 0.5*(Ki[0:-1] + Ki[1:])
|
||||
if (do4D):
|
||||
Ac = 0.5*(Ai[0:-1] + Ai[1:])
|
||||
return xx,yy,Ki,Kc,Ai,Ac
|
||||
else:
|
||||
return xx,yy,Ki,Kc
|
||||
|
||||
xx = X3[:, :, 0]
|
||||
yy = Y3[:, :, 0]
|
||||
Zi = Z3[0, 0, :]
|
||||
Ki = 10 ** Zi
|
||||
Kc = 0.5 * (Ki[0:-1] + Ki[1:])
|
||||
if do4D:
|
||||
Ac = 0.5 * (Ai[0:-1] + Ai[1:])
|
||||
return xx, yy, Ki, Kc, Ai, Ac
|
||||
else:
|
||||
return xx, yy, Ki, Kc
|
||||
|
||||
|
||||
def getSlc(fIn,nStp=0,vID="jPSD",doWrap=False):
|
||||
gID = "Step#%d"%(nStp)
|
||||
with h5py.File(fIn,'r') as hf:
|
||||
@@ -37,50 +53,85 @@ def getSlc(fIn,nStp=0,vID="jPSD",doWrap=False):
|
||||
#doAsym : Px/Pz-1
|
||||
#!doAsym: Px/(Px+Pz)
|
||||
|
||||
def PIso(fIn,nStp=0,pCut=1.0e-3,doAsym=False):
|
||||
Pxy = getSlc(fIn,nStp,"Pxy")
|
||||
Pz = getSlc(fIn,nStp,"Pz" )
|
||||
Pk = 2*Pxy+Pz
|
||||
Nx,Ny = Pz.shape
|
||||
pR = np.zeros((Nx,Ny))
|
||||
def PIso(fIn, nStp=0, pCut=1.0e-3, doAsym=False):
|
||||
"""
|
||||
Calculate the isotropy parameter for a given input file.
|
||||
|
||||
Args:
|
||||
fIn (str): The input file path.
|
||||
nStp (int, optional): The number of steps. Default is 0.
|
||||
pCut (float, optional): The cutoff value. Default is 1.0e-3.
|
||||
doAsym (bool, optional): Flag to calculate asymmetric isotropy parameter. Default is False.
|
||||
|
||||
Returns:
|
||||
pR (numpy.ndarray): The calculated isotropy parameter.
|
||||
|
||||
"""
|
||||
Pxy = getSlc(fIn, nStp, "Pxy")
|
||||
Pz = getSlc(fIn, nStp, "Pz")
|
||||
Pk = 2 * Pxy + Pz
|
||||
Nx, Ny = Pz.shape
|
||||
pR = np.zeros((Nx, Ny))
|
||||
|
||||
for i in range(Nx):
|
||||
for j in range(Ny):
|
||||
|
||||
if (Pk[i,j]>pCut and Pz[i,j]>pCut):
|
||||
if (doAsym):
|
||||
pR[i,j] = Pxy[i,j]/Pz[i,j] - 1.0
|
||||
if Pk[i, j] > pCut and Pz[i, j] > pCut:
|
||||
if doAsym:
|
||||
pR[i, j] = Pxy[i, j] / Pz[i, j] - 1.0
|
||||
else:
|
||||
pR[i,j] = Pxy[i,j]/(Pxy[i,j]+Pz[i,j])
|
||||
pR[i, j] = Pxy[i, j] / (Pxy[i, j] + Pz[i, j])
|
||||
else:
|
||||
if (doAsym):
|
||||
pR[i,j] = 0.0
|
||||
if doAsym:
|
||||
pR[i, j] = 0.0
|
||||
else:
|
||||
pR[i,j] = np.nan
|
||||
pR[i, j] = np.nan
|
||||
return pR
|
||||
|
||||
|
||||
#Equatorial grids (option for wrapping for contours)
|
||||
def getEQGrid(fIn,doCenter=False,doWrap=False):
|
||||
if (doWrap):
|
||||
def getEQGrid(fIn, doCenter=False, doWrap=False):
|
||||
"""
|
||||
Get the equidistant grid coordinates from an HDF5 file.
|
||||
|
||||
Args:
|
||||
fIn (str): The path to the HDF5 file.
|
||||
doCenter (bool): Flag indicating whether to center the grid coordinates. Default is False.
|
||||
doWrap (bool): Flag indicating whether to wrap the grid coordinates. Default is False.
|
||||
|
||||
Returns:
|
||||
If doCenter is False:
|
||||
xx (ndarray): The x-coordinates of the grid.
|
||||
yy (ndarray): The y-coordinates of the grid.
|
||||
If doCenter is True and doWrap is False:
|
||||
xxc (ndarray): The centered x-coordinates of the grid.
|
||||
yyc (ndarray): The centered y-coordinates of the grid.
|
||||
If doCenter and doWrap are both True:
|
||||
xxc (ndarray): The wrapped and centered x-coordinates of the grid.
|
||||
yyc (ndarray): The wrapped and centered y-coordinates of the grid.
|
||||
"""
|
||||
if doWrap:
|
||||
doCenter = True
|
||||
with h5py.File(fIn,'r') as hf:
|
||||
|
||||
with h5py.File(fIn, 'r') as hf:
|
||||
xx = hf["X"][()].T
|
||||
yy = hf["Y"][()].T
|
||||
|
||||
if (not doCenter):
|
||||
return xx,yy
|
||||
if not doCenter:
|
||||
return xx, yy
|
||||
|
||||
Ngi,Ngj = xx.shape
|
||||
Ni = Ngi-1; Nj = Ngj-1
|
||||
xxc = np.zeros((Ni,Nj))
|
||||
yyc = np.zeros((Ni,Nj))
|
||||
Ngi, Ngj = xx.shape
|
||||
Ni = Ngi - 1
|
||||
Nj = Ngj - 1
|
||||
xxc = np.zeros((Ni, Nj))
|
||||
yyc = np.zeros((Ni, Nj))
|
||||
|
||||
xxc = 0.25*( xx[0:Ngi-1,0:Ngj-1] + xx[1:Ngi,0:Ngj-1] + xx[0:Ngi-1,1:Ngj] + xx[1:Ngi,1:Ngj])
|
||||
yyc = 0.25*( yy[0:Ngi-1,0:Ngj-1] + yy[1:Ngi,0:Ngj-1] + yy[0:Ngi-1,1:Ngj] + yy[1:Ngi,1:Ngj])
|
||||
xxc = 0.25 * (xx[0:Ngi - 1, 0:Ngj - 1] + xx[1:Ngi, 0:Ngj - 1] + xx[0:Ngi - 1, 1:Ngj] + xx[1:Ngi, 1:Ngj])
|
||||
yyc = 0.25 * (yy[0:Ngi - 1, 0:Ngj - 1] + yy[1:Ngi, 0:Ngj - 1] + yy[0:Ngi - 1, 1:Ngj] + yy[1:Ngi, 1:Ngj])
|
||||
|
||||
if (not doWrap):
|
||||
return xxc,yyc
|
||||
if not doWrap:
|
||||
return xxc, yyc
|
||||
else:
|
||||
return kv.reWrap(xxc),kv.reWrap(yyc)
|
||||
return kv.reWrap(xxc), kv.reWrap(yyc)
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user