196 lines
9.1 KiB
Python
196 lines
9.1 KiB
Python
import numpy
|
|
class Field3d:
|
|
def __init__(self,data,origin,spacing,bounds,period):
|
|
assert len(origin)==3, "'origin' must be of length 3"
|
|
assert len(spacing)==3, "'spacing' must be of length 3"
|
|
assert len(period)==3, "'period' must be of length 3"
|
|
self.data = numpy.array(data)
|
|
self.data.setflags(write=False)
|
|
self.numpoints = self.data.shape
|
|
self.origin = tuple([float(x) for x in origin])
|
|
self.spacing = tuple([float(x) for x in spacing])
|
|
self.period = tuple(period)
|
|
return
|
|
|
|
def __str__(self):
|
|
str = 'Field3d with\n'
|
|
str+= ' dimension: {}, {}, {}\n'.format(*self.numpoints)
|
|
str+= ' origin: {:G}, {:G}, {:G}\n'.format(*self.origin)
|
|
str+= ' spacing: {:G}, {:G}, {:G}\n'.format(*self.spacing)
|
|
str+= ' period: {}, {}, {}'.format(*self.period)
|
|
return str
|
|
|
|
@classmethod
|
|
def from_chunk(cls,chunk,gridg,period=(None,None,None)):
|
|
'''Initialize Field3d from chunk data and global grid.'''
|
|
xg,yg,zg = gridg
|
|
ib,jb,kb = chunk['ibeg']-1, chunk['jbeg']-1, chunk['kbeg']-1
|
|
dx,dy,dz = xg[1]-xg[0], yg[1]-yg[0], zg[1]-zg[0]
|
|
xo,yo,zo = xg[ib]-chunk['ighost']*dx, yg[jb]-chunk['ighost']*dy, zg[kb]-chunk['ighost']*dz
|
|
nx,ny,nz = chunk['data'].shape
|
|
assert (chunk['nxl']+2*chunk['ighost'])==nx, "Invalid chunk data: nxl != chunk['data'].shape[0]"
|
|
assert (chunk['nyl']+2*chunk['ighost'])==ny, "Invalid chunk data: nyl != chunk['data'].shape[1]"
|
|
assert (chunk['nzl']+2*chunk['ighost'])==nz, "Invalid chunk data: nzl != chunk['data'].shape[2]"
|
|
return cls(chunk['data'],origin=(xo,yo,zo),spacing=(dx,dy,dz),period=period)
|
|
|
|
def dim(self,axis=None):
|
|
if axis is None:
|
|
return tuple(self.dim(axis=ii) for ii in range(0,3))
|
|
assert axis<3, "'axis' must be one of 0,1,2."
|
|
return self.numpoints[axis]
|
|
|
|
def endpoint(self,axis):
|
|
assert axis<3, "'axis' must be one of 0,1,2."
|
|
return self.origin[axis]+(self.numpoints[axis]-1)*self.spacing[axis]
|
|
|
|
def derivative(self,axis,keep_origin=False):
|
|
if not keep_origin:
|
|
# 2nd order FD with h = dx/2
|
|
origin = list(self.origin)
|
|
if self.period[axis] is not None:
|
|
slhi = [slice(None),slice(None),slice(None)]
|
|
sllo = [slice(None),slice(None),slice(None)]
|
|
slmi = [slice(None),slice(None),slice(None)]
|
|
if numpy.isclose(self.origin[axis],-self.spacing[axis]):
|
|
slhi[axis] = slice(1,-1)
|
|
sllo[axis] = slice(0,-2)
|
|
origin[axis] += 0.5*self.spacing[axis]
|
|
else:
|
|
slhi[axis] = slice(2,None)
|
|
sllo[axis] = slice(1,-1)
|
|
origin[axis] -= 0.5*self.spacing[axis]
|
|
slmi[axis] = slice(1,-1)
|
|
slhi = tuple(slhi)
|
|
sllo = tuple(sllo)
|
|
slmi = tuple(slmi)
|
|
data = numpy.zeros_like(self.data)
|
|
data[slmi] = (1./self.spacing[axis])*(self.data[slhi]-self.data[sllo])
|
|
Field3d.__update_ghostcells_periodic(data,axis)
|
|
else:
|
|
slhi = [slice(None),slice(None),slice(None)]
|
|
sllo = [slice(None),slice(None),slice(None)]
|
|
slhi[axis] = slice(1,None)
|
|
sllo[axis] = slice(0,-1)
|
|
slhi = tuple(slhi)
|
|
sllo = tuple(sllo)
|
|
data = (1./self.spacing[axis])*(self.data[slhi]-self.data[sllo])
|
|
origin[axis] += 0.5*self.spacing[axis]
|
|
else:
|
|
# 2nd order FD with h = dx, one-sided at boundaries
|
|
# https://numpy.org/doc/stable/reference/generated/numpy.gradient.html
|
|
origin = self.origin
|
|
if self.period[axis] is not None:
|
|
data = numpy.gradient(self.data,self.spacing[axis],axis=axis,edge_order=1)
|
|
Field3d.__update_ghostcells_periodic(data,axis)
|
|
else:
|
|
data = numpy.gradient(self.data,self.spacing[axis],axis=axis,edge_order=2)
|
|
return Field3d(data,origin,self.spacing,self.period)
|
|
|
|
def gradient(self,axis,keep_origin=False):
|
|
return [self.derivative(axis,keep_origin=keep_origin) for axis in range(0,3)]
|
|
|
|
def change_grid(self,origin,spacing,numpoints,EPS_COLLAPSE=1e-12):
|
|
# TBD: periodicity
|
|
assert all([origin[ii]>=self.origin[ii] for ii in range(0,3)]), "New origin is out of bounds."
|
|
endpoint = [origin[ii]+(numpoints[ii]-1)*spacing[ii] for ii in range(0,3)]
|
|
assert all([endpoint[ii]<=self.endpoint(ii) for ii in range(0,3)]), "New end point is out of bounds."
|
|
data = numpy.zeros(numpoints)
|
|
if numpy.allclose(spacing,self.spacing):
|
|
# spacing is the same: we can construct universal weights for the stencil
|
|
i0,j0,k0 = [numpy.floor((origin[ii]-self.origin[ii])/self.spacing[ii]).astype('int') for ii in range(0,3)]
|
|
cx,cy,cz = [numpy.remainder(origin[ii]-self.origin[ii],self.spacing[ii]) for ii in range(0,3)]
|
|
c = weights_trilinear((cx,cy,cz))
|
|
for ii in range(0,2):
|
|
for jj in range(0,2):
|
|
for kk in range(0,2):
|
|
if c[ii,jj,kk]>EPS_COLLAPSE:
|
|
data += c[ii,jj,kk]*self.data[
|
|
i0+ii:i0+ii+numpoints[0],
|
|
j0+jj:j0+jj+numpoints[1],
|
|
k0+kk:k0+kk+numpoints[2]]
|
|
else:
|
|
for ii in range(0,numpoints[0]):
|
|
for jj in range(0,numpoints[1]):
|
|
for kk in range(0,numpoints[2]):
|
|
coordinate = (
|
|
origin[0]+ii*spacing[0],
|
|
origin[1]+jj*spacing[1],
|
|
origin[2]+kk*spacing[2])
|
|
data[ii,jj,kk] = self.interpolate(coordinate)
|
|
return Field3d(data,origin,spacing,self.period)
|
|
|
|
def interpolate(self,coordinate,EPS_COLLAPSE=1e-12):
|
|
assert all([coordinate[ii]>=self.origin[ii] for ii in range(0,3)]), "'coordinate' is out of bounds."
|
|
assert all([coordinate[ii]<=self.endpoint(ii) for ii in range(0,3)]), "'coordinate' is out of bounds."
|
|
i0,j0,k0 = [numpy.floor((coordinate[ii]-self.origin[ii])/self.spacing[ii]).astype('int') for ii in range(0,3)]
|
|
cx,cy,cz = [numpy.remainder(coordinate[ii]-self.origin[ii],self.spacing[ii]) for ii in range(0,3)]
|
|
c = weights_trilinear((cx,cy,cz))
|
|
val = 0.0
|
|
for ii in range(0,2):
|
|
for jj in range(0,2):
|
|
for kk in range(0,2):
|
|
if c[ii,jj,kk]>EPS_COLLAPSE:
|
|
val += c[ii,jj,kk]*self.data[i0+ii,j0+jj,k0+kk]
|
|
return val
|
|
|
|
def move_origin_periodic(self,shift,axis):
|
|
'''Moves origin of grid in periodic direction by shift*spacing.'''
|
|
assert self.period[axis], "Origin can only be moved in periodic directions."
|
|
self.data = numpy.roll(self.data,shift,axis=axis)
|
|
self.data.setflags(write=False)
|
|
origin = list(self.origin)
|
|
origin[axis] += shift*self.spacing[axis]
|
|
self.origin = tuple(origin)
|
|
return
|
|
@staticmethod
|
|
def __update_ghostcells_periodic(data,axis):
|
|
slsrc = [slice(None),slice(None),slice(None)]
|
|
sldst = [slice(None),slice(None),slice(None)]
|
|
slsrc[axis] = slice(1,2)
|
|
sldst[axis] = slice(-1,None)
|
|
data[sldst] = data[slsrc]
|
|
slsrc = [slice(None),slice(None),slice(None)]
|
|
sldst = [slice(None),slice(None),slice(None)]
|
|
slsrc[axis] = slice(-2,-1)
|
|
sldst[axis] = slice(0,1)
|
|
data[sldst] = data[slsrc]
|
|
return
|
|
|
|
class ChunkIterator:
|
|
'''Iterates through all chunks. 'snapshot' must be an instance
|
|
of a class which returns a Field3d from the method call
|
|
snapshot.field_chunk(rank,key,keep_ghost=keep_ghost).
|
|
One example implementation is UCFSnapshot from suspendtools.ucf.'''
|
|
def __init__(self,snapshot,key,keep_ghost=True):
|
|
self.snapshot = snapshot
|
|
self.key = key
|
|
self.keep_ghost = keep_ghost
|
|
self.iter_rank = 0
|
|
def __iter__(self):
|
|
self.iter_rank = 0
|
|
return self
|
|
def __next__(self):
|
|
if self.iter_rank<self.snapshot.nproc():
|
|
field = self.snapshot.field_chunk(
|
|
self.iter_rank,self.key,keep_ghost=self.keep_ghost)
|
|
self.iter_rank += 1
|
|
return field
|
|
else:
|
|
raise StopIteration
|
|
|
|
def weights_trilinear(rel_dist):
|
|
assert len(rel_dist)==3, "len(rel_dist) must be 3."
|
|
assert all([rel_dist[ii]>=0.0 for ii in range(0,3)]), "'rel_dist' must be >=0"
|
|
assert all([rel_dist[ii]<=1.0 for ii in range(0,3)]), "'rel_dist' must be <=1"
|
|
cx,cy,cz = rel_dist
|
|
c = numpy.zeros((2,2,2))
|
|
c[0,0,0] = 1.0-(cx+cy+cz)+(cx*cy+cx*cz+cy*cz)-(cx*cy*cz)
|
|
c[1,0,0] = cx-(cx*cy+cx*cz)+(cx*cy*cz)
|
|
c[0,1,0] = cy-(cx*cy+cy*cz)+(cx*cy*cz)
|
|
c[0,0,1] = cz-(cx*cz+cy*cz)+(cx*cy*cz)
|
|
c[1,1,0] = (cx*cy)-(cx*cy*cz)
|
|
c[1,0,1] = (cx*cz)-(cx*cy*cz)
|
|
c[0,1,1] = (cy*cz)-(cx*cy*cz)
|
|
c[1,1,1] = (cx*cy*cz)
|
|
return c
|