ucftools/python/ucf/tools.py

183 lines
6.1 KiB
Python

def read_grid(file,verbosity=False,debug=False):
from .base import UCF
obj = UCF(file=file,verbosity=verbosity,debug=debug)
output = []
for iset in range(0,obj.NumDataset):
(data,params) = obj.readSet(step=1,dset=iset+1)
nx = params[0]
ny = params[1]
nz = params[2]
output.append((data[0:nx],data[nx:nx+ny],data[nx+ny:nx+ny+nz]))
#if obj.UCFVersion<2:
if obj.NumDataset<5:
output.extend(output[-1:])
obj.close()
return output
def read_procgrid(file,verbosity=False,debug=False):
from .base import UCF
obj = UCF(file=file,verbosity=verbosity,debug=debug)
output = []
for iset in range(0,obj.NumDataset):
(data,params) = obj.readSet(step=1,dset=iset+1)
nxp = params[0]
nyp = params[1]
nzp = params[2]
output.append((
data[0:nxp], # ibeg
data[nxp:2*nxp], # iend
data[2*nxp:2*nxp+nyp], # jbeg
data[2*nxp+nyp:2*nxp+2*nyp], # jend
data[2*nxp+2*nyp:2*nxp+2*nyp+nzp], # kbeg
data[2*nxp+2*nyp+nzp:2*nxp+2*nyp*2*nzp] # kend
))
#if obj.UCFVersion<2:
if obj.NumDataset<5:
output.extend(output[-1:])
obj.close()
return output
def read_chunk(file,step=1,dset=-1,keep_ghost=True,verbosity=False,debug=False):
from .base import UCF
obj = UCF(file=file,verbosity=verbosity,debug=debug)
if not isinstance(dset,list):
if dset==-1:
dset = range(1,obj.NumDataset+1) # fix that maybe later (this is maximum over all timesteps)
else:
dset = [dset]
output = []
for ii in dset:
tmp = dict()
(data,params) = obj.readSet(step=step,dset=ii)
tmp['ighost'] = params[0]
tmp['ibeg'] = params[1]
tmp['jbeg'] = params[2]
tmp['kbeg'] = params[3]
tmp['nxl'] = params[4]
tmp['nyl'] = params[5]
tmp['nzl'] = params[6]
tmp['nx'] = params[7]
tmp['ny'] = params[8]
tmp['nz'] = params[9]
tmp['data'] = data.reshape((tmp['nxl']+2*tmp['ighost'],
tmp['nyl']+2*tmp['ighost'],
tmp['nzl']+2*tmp['ighost']),
order='F')
tmp['rank'] = obj.IORank[0]
tmp['rankijk']= obj.IORank[1:]
if not keep_ghost and ighost:
tmp['data'] = tmp['data'][1:-1,1:-1,1:-1]
tmp['ghost'] = 0
output.append(tmp)
obj.close()
return output
def read_particles(file,step=-1,verbosity=False,debug=False):
from .base import UCF
import numpy
obj = UCF(file=file,verbosity=verbosity,debug=debug)
if not isinstance(step,list):
if step==-1:
step = range(1,obj.NumTimestep+1)
else:
step = [step]
# The output will be the following:
# 1) numpy array with dimension (ncol,np,ntime)
# 2) dictionary which specifies the columns
# We read the data step by step in a list, which is then converted to a 3D array
pp = []
for ii in step:
(data,params) = obj.readSet(step=ii,dset=1)
npart = params[0]
ncol = params[1]
ncol_rank = params[2]
ncol_hybrid = params[3]
ncol_dem = params[4]
ncol_scalar = params[5]
nscal = ncol_scalar//2
pp.append(data.reshape((ncol,npart),order='F'))
# Close UCF obeject
obj.close()
# Convert list of 2D arrays to 3D array
pp = numpy.stack(pp,axis=2)
# Create the dictionary
col = colmap_from_flags(ncol_rank,ncol_hybrid,ncol_dem,nscal)
# Return result
return (pp,col)
def colmap_from_flags(irank,ihybrid,idem,iscal):
'''Creates a dictionary which specifies the columns of a particle array.'''
col = {}
ioffset = 0
if irank>0:
col['rank'] = ioffset; ioffset+=1
if ihybrid>0:
col['id'] = ioffset; ioffset+=1
col['x'] = ioffset; ioffset+=1
col['y'] = ioffset; ioffset+=1
col['z'] = ioffset; ioffset+=1
col['r'] = ioffset; ioffset+=1
col['rho']= ioffset; ioffset+=1
col['ax'] = ioffset; ioffset+=1
col['ay'] = ioffset; ioffset+=1
col['az'] = ioffset; ioffset+=1
col['u'] = ioffset; ioffset+=1
col['v'] = ioffset; ioffset+=1
col['w'] = ioffset; ioffset+=1
col['ox'] = ioffset; ioffset+=1
col['oy'] = ioffset; ioffset+=1
col['oz'] = ioffset; ioffset+=1
col['fx'] = ioffset; ioffset+=1
col['fy'] = ioffset; ioffset+=1
col['fz'] = ioffset; ioffset+=1
col['tx'] = ioffset; ioffset+=1
col['ty'] = ioffset; ioffset+=1
col['tz'] = ioffset; ioffset+=1
if idem>0:
col['fxc'] = ioffset; ioffset+=1
col['fyc'] = ioffset; ioffset+=1
col['fzc'] = ioffset; ioffset+=1
col['txc'] = ioffset; ioffset+=1
col['tyc'] = ioffset; ioffset+=1
col['tzc'] = ioffset; ioffset+=1
if iscal>0:
for ii in range(0,iscal):
col['s'+str(ii)] = ioffset; ioffset+=1
col['q'+str(ii)] = ioffset; ioffset+=1
return col
def grid_chunk(grid_global,chunk):
import numpy
xg,yg,zg = grid_global
# Shift indices so that they start from zero
ib = chunk['ibeg']-1
jb = chunk['jbeg']-1
kb = chunk['kbeg']-1
nxl = chunk['nxl']
nyl = chunk['nyl']
nzl = chunk['nzl']
ighost = chunk['ighost']
if ighost:
nxg = len(xg)
nyg = len(yg)
nzg = len(zg)
dx = xg[2]-xg[1]
dy = yg[2]-yg[1]
dz = zg[2]-zg[1]
xl = numpy.zeros(nxl+2)
yl = numpy.zeros(nyl+2)
zl = numpy.zeros(nzl+2)
xl[0] = xg[ib]-dx
xl[1:-1] = xg[ib:ib+nxl]
xl[-1] = xg[ib+nxl-1]+dx
yl[0] = yg[jb]-dy
yl[1:-1] = yg[jb:jb+nyl]
yl[-1] = yg[jb+nyl-1]+dy
zl[0] = zg[kb]-dz
zl[1:-1] = zg[kb:kb+nzl]
zl[-1] = zg[kb+nzl-1]+dz
else:
xl = xg[ibeg:ib+nxl-1]
yl = yg[jbeg:jb+nyl-1]
zl = zg[kbeg:kb+nzl-1]
return (xl,yl,zl)