mwtkrayer
/
suspendtools
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

first production version for filtered plotting

This commit is contained in:
Michael Krayer 2021-06-01 12:16:10 +02:00
parent 8dd8161537
commit 3366816adf
3 changed files with 759 additions and 378 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

343
field.py
View File

@ -4,19 +4,74 @@ class Field3d:
assert len(origin)==3, "'origin' must be of length 3" assert len(origin)==3, "'origin' must be of length 3"
assert len(spacing)==3, "'spacing' must be of length 3" assert len(spacing)==3, "'spacing' must be of length 3"
self.data = numpy.array(data) self.data = numpy.array(data)
self.numpoints = self.data.shape
self.origin = tuple([float(x) for x in origin]) self.origin = tuple([float(x) for x in origin])
self.spacing = tuple([float(x) for x in spacing]) self.spacing = tuple([float(x) for x in spacing])
self.eps_collapse = 1e-12 self.eps_collapse = 1e-12
self._dim = None
return return
def __str__(self): def __str__(self):
str = 'Field3d with\n' str = 'Field3d with\n'
str+= ' dimension: {}, {}, {}\n'.format(*self.numpoints) str+= ' dimension: {}, {}, {}\n'.format(*self.dim())
str+= ' origin: {:G}, {:G}, {:G}\n'.format(*self.origin) str+= ' origin: {:G}, {:G}, {:G}\n'.format(*self.origin)
str+= ' spacing: {:G}, {:G}, {:G}'.format(*self.spacing) str+= ' spacing: {:G}, {:G}, {:G}\n'.format(*self.spacing)
str+= ' datatype: {}'.format(self.dtype())
return str return str
def __add__(self,other):
if isinstance(other,Field3d):
assert self.has_same_grid(other), "Grid mismatch."
return Field3d(self.data+other.data,self.origin,self.spacing)
else:
return Field3d(self.data+other,self.origin,self.spacing)
def __sub__(self,other):
if isinstance(other,Field3d):
assert self.has_same_grid(other), "Grid mismatch."
return Field3d(self.data-other.data,self.origin,self.spacing)
else:
return Field3d(self.data-other,self.origin,self.spacing)
def __mul__(self,other):
if isinstance(other,Field3d):
assert self.has_same_grid(other), "Grid mismatch."
return Field3d(self.data*other.data,self.origin,self.spacing)
else:
return Field3d(self.data*other,self.origin,self.spacing)
def __radd__(self,other):
return Field3d(other+self.data,self.origin,self.spacing)
def __rmul__(self,other):
return Field3d(other*self.data,self.origin,self.spacing)
def __pow__(self,other):
return Field3d(self.data**other,self.origin,self.spacing)
def __iadd__(self,other):
if isinstance(other,Field3d):
assert self.has_same_grid(other), "Grid mismatch."
self.data += other.data
else:
self.data += other
return self
def __isub__(self,other):
if isinstance(other,Field3d):
assert self.has_same_grid(other), "Grid mismatch."
self.data -= other.data
else:
self.data -= other
return self
def __imul__(self,other):
if isinstance(other,Field3d):
assert self.has_same_grid(other), "Grid mismatch."
self.data *= other.data
else:
self.data *= other
return self
# TBD: this should return another Field3d object # TBD: this should return another Field3d object
# def __getitem__(self,val): # def __getitem__(self,val):
# assert isinstance(val,tuple) and len(val)==3, "Field3d must be indexed by [ii,jj,kk]." # assert isinstance(val,tuple) and len(val)==3, "Field3d must be indexed by [ii,jj,kk]."
@ -45,6 +100,30 @@ class Field3d:
assert (chunk['nzl']+2*chunk['ighost'])==nz, "Invalid chunk data: nzl != chunk['data'].shape[2]" 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)) return cls(chunk['data'],origin=(xo,yo,zo),spacing=(dx,dy,dz))
@classmethod
def allocate(cls,dim,origin,spacing,fill=None,dtype=numpy.float64,pseudo=False):
'''Allocates an empty field, or a field filled with 'fill'.'''
assert isinstance(dim,(tuple,list,numpy.ndarray)) and len(dim)==3,\
"'dim' must be a tuple/list of length 3."
assert isinstance(origin,(tuple,list,numpy.ndarray)) and len(origin)==3,\
"'origin' must be a tuple/list of length 3."
assert isinstance(spacing,(tuple,list,numpy.ndarray)) and len(spacing)==3,\
"'spacing' must be a tuple/list of length 3."
if pseudo:
data = numpy.empty((0,0,0))
r = cls(data,origin,spacing)
r._dim = dim
return r
else:
if fill is None:
data = numpy.empty(dim,dtype=dtype)
else:
data = numpy.full(dim,fill,dtype=dtype)
return cls(data,origin,spacing)
def copy(self):
return Field3d(self.data.copy(),self.origin,self.spacing)
def insert_subfield(self,subfield): def insert_subfield(self,subfield):
assert all([abs(subfield.spacing[ii]-self.spacing[ii])<self.eps_collapse assert all([abs(subfield.spacing[ii]-self.spacing[ii])<self.eps_collapse
for ii in range(3)]), "spacing differs." for ii in range(3)]), "spacing differs."
@ -54,16 +133,16 @@ class Field3d:
assert all(self.is_within_bounds(subfield.endpoint(),axis=None)), "subfield endpoint is out-of-bounds." assert all(self.is_within_bounds(subfield.endpoint(),axis=None)), "subfield endpoint is out-of-bounds."
#ib,jb,kb = [int(round((subfield.origin[ii]-self.origin[ii])/self.spacing[ii])) for ii in range(3)] #ib,jb,kb = [int(round((subfield.origin[ii]-self.origin[ii])/self.spacing[ii])) for ii in range(3)]
ib,jb,kb = self.nearest_gridpoint(subfield.origin,axis=None) ib,jb,kb = self.nearest_gridpoint(subfield.origin,axis=None)
nx,ny,nz = subfield.numpoints nx,ny,nz = subfield.dim()
self.data[ib:ib+nx,jb:jb+ny,kb:kb+nz] = subfield.data[:,:,:] self.data[ib:ib+nx,jb:jb+ny,kb:kb+nz] = subfield.data[:,:,:]
return return
def extract_subfield(self,idx_origin,numpoints,stride=(1,1,1)): def extract_subfield(self,idx_origin,dim,stride=(1,1,1)):
assert all(idx_origin[ii]>=0 and idx_origin[ii]<self.numpoints[ii] for ii in range(3)),\ assert all(idx_origin[ii]>=0 and idx_origin[ii]<self.dim(axis=ii) for ii in range(3)),\
"'origin' is out-of-bounds." "'origin' is out-of-bounds."
assert all(idx_origin[ii]+stride[ii]*(numpoints[ii]-1)<self.numpoints[ii] for ii in range(3)),\ assert all(idx_origin[ii]+stride[ii]*(dim[ii]-1)<self.dim(axis=ii) for ii in range(3)),\
"endpoint is out-of-bounds." "endpoint is out-of-bounds."
sl = tuple(slice(idx_origin[ii],idx_origin[ii]+stride[ii]*numpoints[ii],stride[ii]) for ii in range(3)) sl = tuple(slice(idx_origin[ii],idx_origin[ii]+stride[ii]*dim[ii],stride[ii]) for ii in range(3))
origin = self.coordinate(idx_origin) origin = self.coordinate(idx_origin)
spacing = tuple(self.spacing[ii]*stride[ii] for ii in range(3)) spacing = tuple(self.spacing[ii]*stride[ii] for ii in range(3))
data = self.data[sl].copy() data = self.data[sl].copy()
@ -74,7 +153,7 @@ class Field3d:
assert len(idx)==3, "If 'axis' is None, 'idx' must be a tuple/list of length 3." assert len(idx)==3, "If 'axis' is None, 'idx' must be a tuple/list of length 3."
return tuple(self.coordinate(idx[ii],axis=ii) for ii in range(3)) return tuple(self.coordinate(idx[ii],axis=ii) for ii in range(3))
assert axis<3, "'axis' must be one of 0,1,2." assert axis<3, "'axis' must be one of 0,1,2."
assert idx<self.numpoints[axis], "'idx' is out-of-bounds." assert idx<self.dim(axis=axis), "'idx' is out-of-bounds."
return self.origin[axis]+idx*self.spacing[axis] return self.origin[axis]+idx*self.spacing[axis]
def nearest_gridpoint(self,coord,axis=None,lower=False): def nearest_gridpoint(self,coord,axis=None,lower=False):
@ -103,61 +182,84 @@ class Field3d:
return tuple(self.is_within_bounds(coord[ii],axis=ii) for ii in range(3)) return tuple(self.is_within_bounds(coord[ii],axis=ii) for ii in range(3))
assert axis<3, "'axis' must be one of 0,1,2." assert axis<3, "'axis' must be one of 0,1,2."
idx_nearest = self.nearest_gridpoint(coord,axis=axis) idx_nearest = self.nearest_gridpoint(coord,axis=axis)
if idx_nearest>0 and idx_nearest<self.numpoints[axis]-1: if idx_nearest>0 and idx_nearest<self.dim(axis=axis)-1:
return True return True
dist_nearest = self.distance_to_nearest_gridpoint(coord,axis=axis) dist_nearest = self.distance_to_nearest_gridpoint(coord,axis=axis)
if (idx_nearest==0 or idx_nearest==self.numpoints[axis]-1) and abs(dist_nearest)<self.eps_collapse: if (idx_nearest==0 or idx_nearest==self.dim(axis=axis)-1) and abs(dist_nearest)<self.eps_collapse:
return True return True
else: else:
return False return False
def has_same_grid(self,other):
if not self.has_same_origin(other): return False
if not self.has_same_spacing(other): return False
if not self.has_same_origin(other): return False
return True
def has_same_origin(self,other,axis=None):
if axis is None:
return all([self.has_same_origin(other,axis=ii) for ii in range(3)])
origin = other.origin if isinstance(other,Field3d) else other
return abs(origin[axis]-self.origin[axis])<self.eps_collapse
def has_same_spacing(self,other,axis=None):
if axis is None:
return all([self.has_same_spacing(other,axis=ii) for ii in range(3)])
spacing = other.spacing if isinstance(other,Field3d) else other
return abs(spacing[axis]-self.spacing[axis])<self.eps_collapse
def has_same_dim(self,other,axis=None):
if axis is None:
return all([self.has_same_dim(other,axis=ii) for ii in range(3)])
dim = other.dim(axis=axis) if isinstance(other,Field3d) else other
return dim==self.dim(axis=axis)
def dim(self,axis=None): def dim(self,axis=None):
if axis is None: if axis is None:
return tuple(self.dim(axis=ii) for ii in range(3)) return tuple(self.dim(axis=ii) for ii in range(3))
assert axis<3, "'axis' must be one of 0,1,2." assert axis<3, "'axis' must be one of 0,1,2."
return self.numpoints[axis] if self._dim is not None:
return self._dim[axis]
else:
return self.data.shape[axis]
def endpoint(self,axis=None): def endpoint(self,axis=None):
if axis is None: if axis is None:
return tuple(self.endpoint(axis=ii) for ii in range(3)) return tuple(self.endpoint(axis=ii) for ii in range(3))
assert axis<3, "'axis' must be one of 0,1,2." assert axis<3, "'axis' must be one of 0,1,2."
return self.origin[axis]+(self.numpoints[axis]-1)*self.spacing[axis] return self.origin[axis]+(self.dim(axis=axis)-1)*self.spacing[axis]
def derivative(self,axis,only_keep_interior=False,add_border='before',preserve_origin=False): def dtype(self):
if axis is None: return self.data.dtype
return tuple(self.derivative(axis,preserve_grid=preserve_grid,padding=padding,numpad=numpad)
for axis in range(0,3)) def convert_dtype(self,dtype):
self.data = self.data.astype(dtype,copy=False)
return
def derivative(self,axis,only_keep_interior=False,shift_origin='before'):
'''Computes derivative wrt to direction 'axis' with 2nd order finite differences
centered between the origin grid points'''
from scipy import ndimage
assert axis<3, "'axis' must be one of 0,1,2." assert axis<3, "'axis' must be one of 0,1,2."
if not preserve_origin: origin = list(self.origin)
origin = list(self.origin) assert shift_origin in ('before','after'), "'shift_origin' must be one of {'before','after'}."
numpoints = list(self.numpoints) if shift_origin=='left':
origin[axis] += 0.5*self.spacing[axis] corr_orig = 0
numpoints[axis] -= 1 origin[axis] -= 0.5*self.spacing[axis]
if not only_keep_interior:
assert add_border in ('before','after'), "'add_border' must be one of {'before','after'}."
origin,numpoints,sl_out = padding(origin,self.spacing,numpoints,add_border,1)
else:
sl_out = tuple(slice(None),slice(None),slice(None))
data = numpy.zeros(numpoints,dtype=self.data.dtype)
# 2nd order FD with h = dx/2
slhi = 3*[slice(None)]
sllo = 3*[slice(None)]
slhi[axis] = slice(1,None)
sllo[axis] = slice(0,-1)
slhi = tuple(slhi)
sllo = tuple(sllo)
data[sl_out] = (1./self.spacing[axis])*(self.data[slhi]-self.data[sllo])
else: else:
origin,numpoints,sl_out = padding(self.origin,self.spacing,self.numpoints,padding,numpad) corr_orig = -1
data = numpy.zeros(numpoints,dtype=self.data.dtype) origin[axis] += 0.5*self.spacing[axis]
# 2nd order FD with h = dx, one-sided at boundaries data = ndimage.correlate1d(self.data,numpy.array([-1.,1.])/self.spacing[axis],axis=axis,
# https://numpy.org/doc/stable/reference/generated/numpy.gradient.html mode='constant',cval=numpy.nan,origin=corr_orig)
data[sl_out] = numpy.gradient(self.data,self.spacing[axis],axis=axis,edge_order=2) if only_keep_interior:
if only_keep_interior: sl = 3*[slice(None)]
sl_out = 3*[slice(None)] if shift_origin=='before':
sl_out[axis] = slice(1,-1) sl[axis] = slice(-1,None)
data = data[tuple(sl_out)]
origin[axis] += self.spacing[axis] origin[axis] += self.spacing[axis]
else:
sl[axis] = slice(0,-1)
data = data[tuple(sl)]
return Field3d(data,origin,self.spacing) return Field3d(data,origin,self.spacing)
def gradient(self,axis,preserve_origin=False,only_keep_interior=False,add_border='before'): def gradient(self,axis,preserve_origin=False,only_keep_interior=False,add_border='before'):
@ -167,9 +269,9 @@ class Field3d:
def laplacian(self,only_keep_interior=False): def laplacian(self,only_keep_interior=False):
'''Computes the Laplacian of a field.''' '''Computes the Laplacian of a field.'''
from scipy import ndimage from scipy import ndimage
data = ndimage.correlate1d(self.data,[1.,-2.,1.]/self.spacing[0],axis=0,mode='constant',cval=numpy.nan,origin=0) data = ndimage.correlate1d(self.data,numpy.array([1.,-2.,1.])/self.spacing[0]**2,axis=0,mode='constant',cval=numpy.nan,origin=0)
data += ndimage.correlate1d(self.data,[1.,-2.,1.]/self.spacing[1],axis=1,mode='constant',cval=numpy.nan,origin=0) data += ndimage.correlate1d(self.data,numpy.array([1.,-2.,1.])/self.spacing[1]**2,axis=1,mode='constant',cval=numpy.nan,origin=0)
data += ndimage.correlate1d(self.data,[1.,-2.,1.]/self.spacing[2],axis=2,mode='constant',cval=numpy.nan,origin=0) data += ndimage.correlate1d(self.data,numpy.array([1.,-2.,1.])/self.spacing[2]**2,axis=2,mode='constant',cval=numpy.nan,origin=0)
origin = list(self.origin) origin = list(self.origin)
if only_keep_interior: if only_keep_interior:
data = data[1:-1,1:-1,1:-1] data = data[1:-1,1:-1,1:-1]
@ -191,7 +293,7 @@ class Field3d:
if integrate_axis[axis]: if integrate_axis[axis]:
axes.append(axis) axes.append(axis)
if average: if average:
weight *= self.numpoints[axis] weight *= self.dim(axis=axis)
else: else:
weight *= self.spacing[axis] weight *= self.spacing[axis]
axes = tuple(axes) axes = tuple(axes)
@ -240,49 +342,89 @@ class Field3d:
radius.append(int(truncate*sigma_img[ii]+0.5)) radius.append(int(truncate*sigma_img[ii]+0.5))
return tuple(radius) return tuple(radius)
def shift_origin(self,rel_shift): # def shift_origin(self,rel_shift):
raise NotImplementedError("Routine has not been verified yet.") # raise NotImplementedError("Routine has not been verified yet.")
#TBD: verify this routine # #TBD: verify this routine
# assert isinstance(rel_shift,(tuple,list,numpy.ndarray)) and len(rel_shift)==3,\
# "'shift' must be tuple/list with length 3."
# assert all([rel_shift[ii]>-1.0 and rel_shift[ii]<1.0 for ii in range(3)]),\
# "'shift' must be in (-1.0,1.0)."
# #data = numpy.full(self.dim,numpy.nan)
# data = self.data.copy()
# origin = list(self.origin)
# for axis in range(3):
# if abs(rel_shift[axis])<self.eps_collapse:
# continue
# origin[axis] += rel_shift[axis]*self.spacing[axis]
# sl_hi = 3*[slice(None)]
# sl_lo = 3*[slice(None)]
# sl_out = 3*[slice(None)]
# sl_clr = 3*[slice(None)]
# sl_hi[axis] = slice(1,None)
# sl_lo[axis] = slice(0,-1)
# if rel_shift[ii]<0.0:
# w = 1.0+rel_shift[axis]
# sl_out[axis] = slice(1,None)
# sl_clr[axis] = slice(0,1)
# else:
# w = rel_shift[axis]
# sl_out[axis] = slice(0,-1)
# sl_clr[axis] = slice(-1,None)
# sl_hi = tuple(sl_hi)
# sl_lo = tuple(sl_lo)
# sl_out = tuple(sl_out)
# sl_clr = tuple(sl_clr)
# data[sl_out] = (1.0-w)*data[sl_lo] + w*data[sl_hi]
# data[sl_clr] = numpy.nan
# return Field3d(data,origin,self.spacing)
def shift_origin(self,rel_shift,only_keep_interior=False):
'''Shifts the origin of a field by multiple of spacing.'''
from scipy import ndimage
assert isinstance(rel_shift,(tuple,list,numpy.ndarray)) and len(rel_shift)==3,\ assert isinstance(rel_shift,(tuple,list,numpy.ndarray)) and len(rel_shift)==3,\
"'shift' must be tuple/list with length 3." "'shift' must be tuple/list with length 3."
assert all([rel_shift[ii]>-1.0 and rel_shift[ii]<1.0 for ii in range(3)]),\ assert all([rel_shift[ii]>=-1.0 and rel_shift[ii]<=1.0 for ii in range(3)]),\
"'shift' must be in (-1.0,1.0)." "'shift' must be in (-1.0,1.0)."
#data = numpy.full(self.numpoints,numpy.nan) data = self.data.copy()
data = self.data.copy()
origin = list(self.origin) origin = list(self.origin)
sl = 3*[slice(None)]
for axis in range(3): for axis in range(3):
if abs(rel_shift[axis])<self.eps_collapse: if abs(rel_shift[axis])<self.eps_collapse:
continue continue
origin[axis] += rel_shift[axis]*self.spacing[axis] elif rel_shift[axis]>0:
sl_hi = 3*[slice(None)] weights = (1.0-rel_shift[axis],rel_shift[axis])
sl_lo = 3*[slice(None)] data = ndimage.correlate1d(data,weights,axis=axis,mode='constant',cval=numpy.nan,origin=-1)
sl_out = 3*[slice(None)] origin[axis] += rel_shift[axis]*self.spacing[axis]
sl_clr = 3*[slice(None)] if only_keep_interior:
sl_hi[axis] = slice(1,None) sl[axis] = slice(0,-1)
sl_lo[axis] = slice(0,-1)
if rel_shift[ii]<0.0:
w = 1.0+rel_shift[axis]
sl_out[axis] = slice(1,None)
sl_clr[axis] = slice(0,1)
else: else:
w = rel_shift[axis] weights = (-rel_shift[axis],1.0+rel_shift[axis])
sl_out[axis] = slice(0,-1) data = ndimage.correlate1d(data,weights,axis=axis,mode='constant',cval=numpy.nan,origin=0)
sl_clr[axis] = slice(-1,None) origin[axis] += rel_shift[axis]*self.spacing[axis]
sl_hi = tuple(sl_hi) if only_keep_interior:
sl_lo = tuple(sl_lo) sl[axis] = slice(1,None)
sl_out = tuple(sl_out) origin[axis] += self.spacing[axis]
sl_clr = tuple(sl_clr) if only_keep_interior:
data[sl_out] = (1.0-w)*data[sl_lo] + w*data[sl_hi] data = data[sl]
data[sl_clr] = numpy.nan
return Field3d(data,origin,self.spacing) return Field3d(data,origin,self.spacing)
def change_grid(self,origin,spacing,numpoints,padding=None,numpad=1): def relative_shift(self,field):
'''Compute the relative shift (in terms of spacing) to shift self onto field.'''
assert self.has_same_spacing(field), "spacing differs."
rel_shift = [0.0,0.0,0.0]
for axis in range(3):
dist = field.origin[axis]-self.origin[axis]
if abs(dist)>self.eps_collapse:
rel_shift[axis] = dist/self.spacing[axis]
return tuple(rel_shift)
def change_grid(self,origin,spacing,dim,padding=None,numpad=1):
assert all([origin[ii]>=self.origin[ii] for ii in range(0,3)]), "New origin is out of bounds." 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)] endpoint = [origin[ii]+(dim[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." assert all([endpoint[ii]<=self.endpoint(ii) for ii in range(0,3)]), "New end point is out of bounds."
# Allocate (possibly padded array) # Allocate (possibly padded array)
origin_pad,numpoints_pad,sl_out = padding(origin,spacing,numpoints,padding,numpad): origin_pad,dim_pad,sl_out = padding(origin,spacing,dim,padding,numpad)
data = numpy.zeros(numpoints_pad,dtype=self.data.dtype) data = numpy.zeros(dim_pad,dtype=self.data.dtype)
# Trilinear interpolation # Trilinear interpolation
if numpy.allclose(spacing,self.spacing): if numpy.allclose(spacing,self.spacing):
# spacing is the same: we can construct universal weights for the stencil # spacing is the same: we can construct universal weights for the stencil
@ -295,14 +437,14 @@ class Field3d:
for kk in range(0,2): for kk in range(0,2):
if c[ii,jj,kk]>self.eps_collapse: if c[ii,jj,kk]>self.eps_collapse:
data[sl_out] += c[ii,jj,kk]*self.data[ data[sl_out] += c[ii,jj,kk]*self.data[
i0+ii:i0+ii+numpoints[0], i0+ii:i0+ii+dim[0],
j0+jj:j0+jj+numpoints[1], j0+jj:j0+jj+dim[1],
k0+kk:k0+kk+numpoints[2]] k0+kk:k0+kk+dim[2]]
else: else:
data_ref = data[sl_out] data_ref = data[sl_out]
for ii in range(0,numpoints[0]): for ii in range(0,dim[0]):
for jj in range(0,numpoints[1]): for jj in range(0,dim[1]):
for kk in range(0,numpoints[2]): for kk in range(0,dim[2]):
coord = ( coord = (
origin[0]+ii*spacing[0], origin[0]+ii*spacing[0],
origin[1]+jj*spacing[1], origin[1]+jj*spacing[1],
@ -326,36 +468,36 @@ class Field3d:
return val return val
@staticmethod @staticmethod
def padding(origin,spacing,numpoints,padding,numpad): def padding(origin,spacing,dim,padding,numpad):
if isinstance(numpad,int): if isinstance(numpad,int):
numpad = numpy.fill(3,numpad,dtype=int) numpad = numpy.fill(3,numpad,dtype=int)
else: else:
numpad = numpy.array(numpad,dtype=int) numpad = numpy.array(numpad,dtype=int)
assert len(numpad)==3, "'numpad' must be either an integer or tuple/list of length 3." assert len(numpad)==3, "'numpad' must be either an integer or tuple/list of length 3."
origin_pad = numpy.array(origin) origin_pad = numpy.array(origin)
numpoints_pad = numpy.array(numpoints) dim_pad = numpy.array(dim)
sl_out = [slice(None),slice(None),slice(None)] sl_out = [slice(None),slice(None),slice(None)]
if padding is not None: if padding is not None:
if padding=='before': if padding=='before':
numpoints_pad += numpad dim_pad += numpad
origin_pad -= numpad*spacing origin_pad -= numpad*spacing
for axis in range(3): for axis in range(3):
sl_out[axis] = slice(numpad[axis],None) sl_out[axis] = slice(numpad[axis],None)
elif padding=='after': elif padding=='after':
numpoints_pad += numpad dim_pad += numpad
for axis in range(3): for axis in range(3):
sl_out[axis] = slice(0,-numpad[axis]) sl_out[axis] = slice(0,-numpad[axis])
elif padding=='both': elif padding=='both':
numpoints_pad += 2*numpad dim_pad += 2*numpad
origin_pad -= numpad*spacing origin_pad -= numpad*spacing
for axis in range(3): for axis in range(3):
sl_out[axis] = slice(numpad[axis],-numpad[axis]) sl_out[axis] = slice(numpad[axis],-numpad[axis])
else: else:
raise ValueError("'padding' must either be None or one of {'before','after','both'}.") raise ValueError("'padding' must either be None or one of {'before','after','both'}.")
sl_out = tuple(sl_out) sl_out = tuple(sl_out)
origin_pad = tuple(origin_pad) origin_pad = tuple(origin_pad)
numpoints_pad = tuple(numpoints_pad) dim_pad = tuple(dim_pad)
return (origin_pad,numpoints_pad,sl_out) return (origin_pad,dim_pad,sl_out)
def weights_trilinear(self,rel_dist): def weights_trilinear(self,rel_dist):
assert len(rel_dist)==3, "len(rel_dist) must be 3." assert len(rel_dist)==3, "len(rel_dist) must be 3."
@ -408,6 +550,19 @@ class Field3d:
"'origin' must be tuple of length 3." "'origin' must be tuple of length 3."
return self.to_vtk(deep=deep).slice(normal=normal,origin=origin) return self.to_vtk(deep=deep).slice(normal=normal,origin=origin)
def gaussian_filter_umean_channel(array,spacing,sigma,truncate=4.0):
'''Applies a Gaussian filter to a numpy array of dimension (1,ny,1) which
contains the mean streamwise velocity of the channel (or possibly sth else).
Since yu[0] = c and yu[-1] = d, we can use scipy's mirror settings and don't
need ghost cells.'''
from scipy import ndimage
assert array.ndim==3, "Expected an array with three dimensions/axes."
assert array.shape[0]==1 and array.shape[1]>1 and array.shape[2]==1,\
"Expected an array with shape (1,ny,1)."
sigma_img = sigma/spacing
array = ndimage.gaussian_filter1d(array,sigma_img,axis=1,truncate=truncate,mode='mirror')
return array
class ChunkIterator: class ChunkIterator:
'''Iterates through all chunks. 'snapshot' must be an instance '''Iterates through all chunks. 'snapshot' must be an instance
of a class which returns a Field3d from the method call of a class which returns a Field3d from the method call

683
parallel.py
View File

@ -1,22 +1,43 @@
import numpy as np
class PPP: class PPP:
"""Parallel Python Postprocessor for suspend""" """Parallel Python Postprocessor for suspend"""
def __init__(self,comm,func_load,num_ghost,precision,origin,spacing,periodicity,bounds,proc,chunks_per_proc): def __init__(self,comm,func_load,
'''Constructor: except for comm, only rank 0 needs initialized data.''' num_ghost,chunks_per_proc,origin,spacing,periodicity,bounds,
self.comm = comm proc_grid,nxp,nyp,nzp,
self.rank = comm.Get_rank() proc_grid_ext,nxp_ext,nyp_ext,nzp_ext,
self.func_load = func_load nghbr,field,symmetries):
self.init_settings(num_ghost,precision) '''Class constructor: it is expected that all communcation is done in the corresponding
self.init_domain(origin,spacing,periodicity,bounds) class methods, such that every rank already holds the required information.'''
self.init_procgrid(proc,chunks_per_proc) # Initialize object with passed arguments
self.field = {} self.comm = comm
self.symmetries = {} self.func_load = func_load
return self.num_ghost = num_ghost
self.chunks_per_proc = chunks_per_proc
self.origin = origin
self.spacing = spacing
self.periodicity = periodicity
self.bounds = bounds
self.proc_grid = proc_grid
self.nxp,self.nyp,self.nzp = nxp,nyp,nzp
self.proc_grid_ext = proc_grid_ext
self.nxp_ext,self.nyp_ext,self.nzp_ext = nxp_ext,nyp_ext,nzp_ext
self.nghbr = nghbr
self.field = field
self.symmetries = symmetries
# Initialize some deduced variables
self.rank = comm.Get_rank()
self.nproc = nxp*nyp*nzp
self.ip,self.jp,self.kp = self.position_from_rank(self.rank,external=False)
self.nghx,self.nghy,self.nghz = num_ghost
self.xperiodic,self.yperiodic,self.zperiodic = periodicity
return
@classmethod @classmethod
def from_snapshot(cls,snap,chunks_per_proc=(1,1,1),num_ghost=(1,1,1),precision='float64'): def from_snapshot(cls,snap,chunks_per_proc=(1,1,1),num_ghost=(1,1,1)):
from mpi4py import MPI from mpi4py import MPI
comm = MPI.COMM_WORLD comm = MPI.COMM_WORLD
rank = comm.Get_rank() rank = comm.Get_rank()
# Let rank 0 read the processor grid as well as some info and then broadcast it
if rank==0: if rank==0:
proc = snap.procgrid() proc = snap.procgrid()
origin = snap.origin() origin = snap.origin()
@ -24,182 +45,53 @@ class PPP:
periodicity = snap.periodicity() periodicity = snap.periodicity()
bounds = snap.bounds() bounds = snap.bounds()
else: else:
proc = None proc,origin,spacing,periodicity,bounds = 5*[None]
origin = None proc = comm.bcast(proc ,root=0)
spacing = None origin = comm.bcast(origin ,root=0)
periodicity = None spacing = comm.bcast(spacing ,root=0)
bounds = None periodicity = comm.bcast(periodicity,root=0)
bounds = comm.bcast(bounds ,root=0)
# Make sure all processors got the same arguments
chunks_per_proc = comm.bcast(chunks_per_proc,root=0)
num_ghost = comm.bcast(num_ghost,root=0)
# Set data loading function
func_load = snap.field_chunk func_load = snap.field_chunk
return cls(comm,func_load,num_ghost,precision,origin,spacing,periodicity,bounds,proc,chunks_per_proc) # Construct new processor grid
(proc_grid,nxp,nyp,nzp,
proc_grid_ext,nxp_ext,nyp_ext,nzp_ext,
nghbr) = PPP._construct_procgrid(comm,proc,chunks_per_proc,num_ghost,periodicity)
# Initially, no field is loaded
field = {}
symmetries = {}
return cls(comm,func_load,
num_ghost,chunks_per_proc,origin,spacing,periodicity,bounds,
proc_grid,nxp,nyp,nzp,
proc_grid_ext,nxp_ext,nyp_ext,nzp_ext,
nghbr,field,symmetries)
def init_settings(self,num_ghost,precision): @classmethod
'''Initializes PPP settings for all processors.''' def from_state(cls,filename):
# TBD: some assertions import pickle
self.num_ghost = self.comm.bcast(num_ghost,root=0) from mpi4py import MPI
self.precision = self.comm.bcast(precision,root=0) comm = MPI.COMM_WORLD
# Some shortcuts rank = comm.Get_rank()
self.nghx,self.nghy,self.nghz = self.num_ghost fin = filename+'.{:05d}.pickle'.format(rank)
with open(fin,"rb") as f:
payload = pickle.load(f)
(num_ghost,chunks_per_proc,origin,spacing,periodicity,bounds,
proc_grid,nxp,nyp,nzp,nghbr,field,symmetries) = payload
assert nxp*nyp*nzp==comm.Get_size(), "The loaded state requires {} processors, but "\
"we are currently running with {}.".format(nxp*nyp*nzp,comm.Get_size())
func_load = None
proc_grid_ext = None
nxp_ext,nyp_ext,nzp_ext = 3*[None]
return cls(comm,func_load,
num_ghost,chunks_per_proc,origin,spacing,periodicity,bounds,
proc_grid,nxp,nyp,nzp,
proc_grid_ext,nxp_ext,nyp_ext,nzp_ext,
nghbr,field,symmetries)
def init_symmetries(self,key,mirror=(True,True)): def load_field(self,key,io_limit=None,barrier=False,dtype=np.float64):
'''Sets the symmetries for ghost cells behind the wall'''
# Example: wall in y
# No-slip boundary (no mirror) free-slip boundary (mirror)
# u -> -u u -> u
# v -> -v v -> -v
# w -> -w w -> w
# p -> p p -> p
import numpy as np
self.symmetries[key] = np.zeros((3,3,3),dtype='i')
if not self.xperiodic:
if key=='u':
self.symmetries[key][0,1,1] = -1
self.symmetries[key][2,1,1] = -1
else:
self.symmetries[key][0,1,1] = 1 if mirror[0] else -1
self.symmetries[key][2,1,1] = 1 if mirror[1] else -1
if not self.yperiodic:
if key=='v':
self.symmetries[key][1,0,1] = -1
self.symmetries[key][1,2,1] = -1
else:
self.symmetries[key][1,0,1] = 1 if mirror[0] else -1
self.symmetries[key][1,2,1] = 1 if mirror[1] else -1
if not self.zperiodic:
if key=='w':
self.symmetries[key][1,1,0] = -1
self.symmetries[key][1,1,2] = -1
else:
self.symmetries[key][1,1,0] = 1 if mirror[0] else -1
self.symmetries[key][1,1,2] = 1 if mirror[1] else -1
def init_domain(self,origin,spacing,periodicity,bounds):
'''Sets up domain information for all processors'''
# TBD: some assertions
self.origin = self.comm.bcast(origin,root=0)
self.spacing = self.comm.bcast(spacing,root=0)
self.periodicity = self.comm.bcast(periodicity,root=0)
self.bounds = self.comm.bcast(bounds,root=0)
# Some shortcuts
self.xperiodic,self.yperiodic,self.zperiodic = self.periodicity
return
def init_procgrid(self,proc,chunks_per_proc):
# Note: requires nghx, xperiodic to be set
'''Read input processor grid, compute processor grid for workers'''
import numpy as np
self.chunks_per_proc = self.comm.bcast(chunks_per_proc,root=0)
self.nxcpp = chunks_per_proc[0]
self.nycpp = chunks_per_proc[1]
self.nzcpp = chunks_per_proc[2]
if self.rank==0:
# Assert proc and add it to class
assert all(k in proc for k in ('u','v','w','p','s')), "'proc' must be a dictionary with "\
"keys 'u','v','w','p','s'"
for key in proc:
assert len(proc[key])==6, "Entries of 'proc' must have length of 6."
proc_grid_ext = proc
# Initialize chunks per processor
# Verify that this processor grid can be redistributed onto the requested processor layout
nxp_ext = len(proc_grid_ext['u'][0])
nyp_ext = len(proc_grid_ext['u'][2])
nzp_ext = len(proc_grid_ext['u'][4])
nproc_ext = nxp_ext*nyp_ext*nzp_ext
#
assert nxp_ext%self.nxcpp==0, "Number of processors must be divisible by the number "\
"of processors per process. (nxp_ext={}, nxcpp={})".format(
nxp_ext,self.nxcpp)
assert nyp_ext%self.nycpp==0, "Number of processors must be divisible by the number "\
"of processors per process. (nyp_ext={}, nycpp={})".format(
nyp_ext,self.nycpp)
assert nzp_ext%self.nzcpp==0, "Number of processors must be divisible by the number "\
"of processors per process. (nzp_ext={}, nzcpp={})".format(
nzp_ext,self.nzcpp)
# Determine the new processor layout and verify total number of MPI processes
nxp = nxp_ext//self.nxcpp
nyp = nyp_ext//self.nycpp
nzp = nzp_ext//self.nzcpp
nproc = nxp*nyp*nzp
assert nproc==self.comm.Get_size(), "Number of MPI processes does not match the requested "\
"processor layout. (MPI procs: {}, required procs: {})".format(
self.comm.Get_size(),nproc)
# Construct internal processor grid
proc_grid = {}
for key in proc_grid_ext:
proc_grid[key] = [None]*6
proc_grid[key][0] = proc_grid_ext[key][0][0::self.nxcpp]
proc_grid[key][1] = proc_grid_ext[key][1][self.nxcpp-1::self.nxcpp]
proc_grid[key][2] = proc_grid_ext[key][2][0::self.nycpp]
proc_grid[key][3] = proc_grid_ext[key][3][self.nycpp-1::self.nycpp]
proc_grid[key][4] = proc_grid_ext[key][4][0::self.nzcpp]
proc_grid[key][5] = proc_grid_ext[key][5][self.nzcpp-1::self.nzcpp]
else:
proc_grid_ext = None
proc_grid = None
nxp_ext,nyp_ext,nzp_ext,nproc_ext = None,None,None,None
nxp,nyp,nzp,nproc = None,None,None,None
# Broadcast the data
self.proc_grid_ext = self.comm.bcast(proc_grid_ext,root=0)
self.proc_grid = self.comm.bcast(proc_grid,root=0)
self.nxp_ext,self.nyp_ext,\
self.nzp_ext,self.nproc_ext = self.comm.bcast((nxp_ext,nyp_ext,nzp_ext,nproc_ext),root=0)
self.nxp,self.nyp,\
self.nzp,self.nproc = self.comm.bcast((nxp,nyp,nzp,nproc),root=0)
# Get position in processor grid
self.ip,self.jp,self.kp = self.position_from_rank(self.rank,external=False)
# Determine local grid indices and size
for key in self.proc_grid:
nxl = self.proc_grid[key][1][self.ip]-self.proc_grid[key][0][self.ip]+1
nyl = self.proc_grid[key][3][self.jp]-self.proc_grid[key][2][self.jp]+1
nzl = self.proc_grid[key][5][self.kp]-self.proc_grid[key][4][self.kp]+1
# Verify that local grid size is not smaller than ghost cell size
assert (nxl>=self.nghx and nyl>=self.nghy and nzl>=self.nghz),\
"Local grid size must be greater than number "\
"of ghost cells in each direction!"
# Initialize neighbor array
nghbr = np.empty((3,3,3),dtype='int')
# wrap-around x
ipl = self.ip-1
if ipl<0:
if self.xperiodic: ipl = self.nxp-1
else: ipl = -1
ipr = self.ip+1
if ipr>self.nxp-1:
if self.xperiodic: ipr = 0
else: ipr = -1
inghbr = (ipl,self.ip,ipr)
# wrap-around y
jpl = self.jp-1
if jpl<0:
if self.yperiodic: jpl = self.nyp-1
else: jpl = -1
jpr = self.jp+1
if jpr>self.nyp-1:
if self.yperiodic: jpr = 0
else: jpr = -1
jnghbr = (jpl,self.jp,jpr)
# wrap-around z
kpl = self.kp-1
if kpl<0:
if self.zperiodic: kpl = self.nzp-1
else: kpl = -1
kpr = self.kp+1
if kpr>self.nzp-1:
if self.zperiodic: kpr = 0
else: kpr = -1
knghbr = (kpl,self.kp,kpr)
# Construct array of neighbors
for ip in range(3):
for jp in range(3):
for kp in range(3):
# Assign rank to neighbor array
if inghbr[ip]<0 or jnghbr[jp]<0 or knghbr[kp]<0:
nghbr[ip,jp,kp] = -1
else:
nghbr[ip,jp,kp] = self.rank_from_position(inghbr[ip],jnghbr[jp],knghbr[kp],external=False)
# Save neighbors as class variable
self.nghbr = nghbr
def load_field(self,key,io_limit=None,barrier=False):
'''Loads the required chunks from file''' '''Loads the required chunks from file'''
from .field import Field3d from .field import Field3d
import numpy as np import numpy as np
@ -226,7 +118,7 @@ class PPP:
data = np.empty( data = np.empty(
(nxl+2*self.nghx, (nxl+2*self.nghx,
nyl+2*self.nghy, nyl+2*self.nghy,
nzl+2*self.nghz),dtype=self.precision) nzl+2*self.nghz),dtype=dtype)
# Compute origin of subfield # Compute origin of subfield
origin = (self.origin[key][0]+(ib-1-self.nghx)*self.spacing[0], origin = (self.origin[key][0]+(ib-1-self.nghx)*self.spacing[0],
self.origin[key][1]+(jb-1-self.nghy)*self.spacing[1], self.origin[key][1]+(jb-1-self.nghy)*self.spacing[1],
@ -257,34 +149,32 @@ class PPP:
# Exchange ghost cells # Exchange ghost cells
self.exchange_ghost_cells(key) self.exchange_ghost_cells(key)
# Initialize symmetries and impose BC # Initialize symmetries and impose BC
self.init_symmetries(key) self.init_symmetry(key)
self.impose_boundary_conditions(key) self.impose_boundary_conditions(key)
# Syncronize processes if requested # Syncronize processes if requested
if barrier: self.comm.Barrier() if barrier: self.comm.Barrier()
def differentiate(self,key,axis,key_out=None): def differentiate(self,key,axis,key_out=None,on_pressure_grid=False):
assert axis<3, "'axis' must be one of 0,1,2." assert axis<3, "'axis' must be one of 0,1,2."
if key_out is None: if key_out is None:
key_out = key+('x','y','z')[axis] key_out = 'd'+key+('/dx','/dy','/dz')[axis]
origin = list(self.origin) origin = list(self.origin[key])
shifting_state = self.shifting_state(key,axis=axis) shifting_state = self.shifting_state(key,axis=axis)
if shifting_state==-1: if shifting_state==-1:
padding = 'after' shift_origin = 'after'
origin[axis] += 0.5*self.spacing[axis] origin[axis] += 0.5*self.spacing[axis]
elif shifting_state==0: elif shifting_state==0:
padding = 'before' shift_origin = 'before'
origin[axis] -= 0.5*self.spacing[axis] origin[axis] -= 0.5*self.spacing[axis]
elif shifting_state==1: elif shifting_state==1:
padding = 'before' shift_origin = 'before'
origin[axis] -= 0.5*self.spacing[axis] origin[axis] -= 0.5*self.spacing[axis]
else: else:
raise ValueError("Invalid shifting state.") raise ValueError("Invalid shifting state.")
self.field[key_out] = self.field[key].derivative(axis,padding=padding) self.field[key_out] = self.field[key].derivative(axis,only_keep_interior=False,shift_origin=shift_origin)
self.origin[key_out] = tuple(origin) self.origin[key_out] = tuple(origin)
self.spacing[key_out] = self.spacing[key].copy() self.symmetries[key_out] = self.symmetries[key].copy()
self.symmetries[key_out] = self.symmetries[key].copy() self.proc_grid[key_out] = self.proc_grid[key].copy()
self.proc_grid[key_out] = self.proc_grid[key].copy()
self.proc_grid_ext[key_out] = self.proc_grid_ext[key].copy()
if axis==0: if axis==0:
self.symmetries[key_out][0,1,1] = -self.symmetries[key_out][0,1,1] self.symmetries[key_out][0,1,1] = -self.symmetries[key_out][0,1,1]
self.symmetries[key_out][2,1,1] = -self.symmetries[key_out][2,1,1] self.symmetries[key_out][2,1,1] = -self.symmetries[key_out][2,1,1]
@ -297,15 +187,90 @@ class PPP:
# Exchange ghost cells and set boundary conditions # Exchange ghost cells and set boundary conditions
self.exchange_ghost_cells(key_out) self.exchange_ghost_cells(key_out)
self.impose_boundary_conditions(key_out) self.impose_boundary_conditions(key_out)
# Shift to pressure grid if desired
if on_pressure_grid:
self.shift_to_pressure_grid(key_out,key_out=key_out)
def allocate(self,key_grid,fill=None,pseudo=False,dtype=np.float64):
from .field import Field3d
ib = self.proc_grid[key_grid][0][self.ip]
ie = self.proc_grid[key_grid][1][self.ip]
jb = self.proc_grid[key_grid][2][self.jp]
je = self.proc_grid[key_grid][3][self.jp]
kb = self.proc_grid[key_grid][4][self.kp]
ke = self.proc_grid[key_grid][5][self.kp]
origin = (self.origin[key_grid][0]+(ib-1-self.nghx)*self.spacing[0],
self.origin[key_grid][1]+(jb-1-self.nghy)*self.spacing[1],
self.origin[key_grid][2]+(kb-1-self.nghz)*self.spacing[2])
dim = (ie-ib+1+2*self.nghx,
je-jb+1+2*self.nghy,
ke-kb+1+2*self.nghz)
return Field3d.allocate(dim,origin,self.spacing,fill=fill,pseudo=pseudo,dtype=dtype)
def qcriterion(self,key_out=None,from_pressure=False,keep_derivatives=False):
'''Computes Q-criterion'''
from .field import Field3d
if key_out is None:
key_out = 'Q'
if from_pressure:
assert 'p' in self.field, "Pressure field is not loaded."
self.copy('p',key_out,skip_data=True)
self.field[key_out] = 0.5*self.field['p'].laplacian(only_keep_interior=False)
else:
assert 'u' in self.field, "Velocity field is not loaded: u"
assert 'v' in self.field, "Velocity field is not loaded: v"
assert 'w' in self.field, "Velocity field is not loaded: w"
self.copy('p',key_out,skip_data=True,skip_symmetries=True)
self.init_symmetry(key_out)
self.field[key_out] = self.allocate('p',fill=0.0,dtype=self.field['u'].dtype())
keyx = ('x','y','z')
keyu = ('u','v','w')
for ii in range(3):
key_tmp = 'd'+keyu[ii]+'/d'+keyx[ii]
parprint(key_tmp)
self.differentiate(keyu[ii],axis=ii,key_out=key_tmp,on_pressure_grid=True)
self.field[key_out] += self.field[key_tmp]**2
if not keep_derivatives:
self.delete(key_tmp)
#
iip1 = (ii+1)%3
key_tmp1 = 'd'+keyu[ii]+'/d'+keyx[iip1]
key_tmp2 = 'd'+keyu[iip1]+'/d'+keyx[ii]
parprint(key_tmp1,key_tmp2)
self.differentiate(keyu[ii],axis=iip1,key_out=key_tmp1,on_pressure_grid=True)
self.differentiate(keyu[iip1],axis=ii,key_out=key_tmp2,on_pressure_grid=True)
self.field[key_out] += 2.0*self.field[key_tmp1]*self.field[key_tmp2]
if not keep_derivatives:
self.delete(key_tmp1)
self.delete(key_tmp2)
self.field[key_out] *= -0.5
# Exchange ghost cells and set boundary conditions
self.exchange_ghost_cells(key_out)
self.impose_boundary_conditions(key_out)
return
#def vorticity(self,axis=None,key_out=None,keep_derivatives=False,on_pressure_grid=True):
def multiply(self,key1,key2,key_out=None,on_pressure_grid=True):
if key_out is None:
key_out = key1
if not self.field[key1].has_same_grid(self.field[key2]) or on_pressure_grid:
self.shift_to_pressure_grid(key1,key_out)
self.shift_to_pressure_grid(key2,'tmp')
self.field[key_out] *= self.field['tmp']
self.delete('tmp')
else:
self.copy(key1,key_out)
self.field[key_out] *= self.field[key2]
return
def gaussian_filter(self,key,sigma,truncate=4.0,key_out=None,iterate=False): def gaussian_filter(self,key,sigma,truncate=4.0,key_out=None,iterate=False):
'''Applies a gaussian filter to a field as in-place operation. Sigma is the std of the filter in terms of grid width.''' '''Applies a gaussian filter to a field as in-place operation. Sigma is the std of the filter in terms of grid width.'''
import numpy as np import numpy as np
if key_out is None: if key_out is None:
key_out = key key_out = key
else: if key!=key_out:
self.origin[key_out] = self.origin[key].copy() self.copy(key,key_out,skip_data=True)
self.spacing[key_out] = self.spacing[key].copy()
# Compute radius of Gaussian filter # Compute radius of Gaussian filter
radius = self.field[key].gaussian_filter_radius(sigma,truncate=truncate) radius = self.field[key].gaussian_filter_radius(sigma,truncate=truncate)
if not iterate: if not iterate:
@ -324,23 +289,27 @@ class PPP:
assert all([radius[ii]>0 if sigma[ii]>0.0 else True for ii in range(3)]),\ assert all([radius[ii]>0 if sigma[ii]>0.0 else True for ii in range(3)]),\
"Iterative procedure leads to invalid stencil radius: "\ "Iterative procedure leads to invalid stencil radius: "\
"increase number of ghost cells. {}".format(radius) "increase number of ghost cells. {}".format(radius)
print('Iterations: {}, stencil radius: {}'.format(niter,radius)) parprint('Gaussian filter: iterations={}, stencil radius={}'.format(niter,radius))
for iiter in range(niter): for iiter in range(niter):
parprint('Iter #{:d}'.format(iiter))
# Filter field: if key_out is None, perform operation inplace # Filter field: if key_out is None, perform operation inplace
self.field[key_out] = self.field[key].gaussian_filter(sigma, self.field[key_out] = self.field[key].gaussian_filter(sigma,
truncate=truncate,border='constant',const_val=0.0) truncate=truncate,only_keep_interior=False)
# Exchange ghost cells and set boundary conditions # Exchange ghost cells and set boundary conditions
self.exchange_ghost_cells(key_out) self.exchange_ghost_cells(key_out)
self.impose_boundary_conditions(key_out) self.impose_boundary_conditions(key_out)
# Iterate inplace from now on # Iterate inplace from now on
key = key_out key = key_out
def broadcast(self,key,arg,operation): def broadcast(self,key,operation,arg,key_out=None):
'''Broadcasts an inplace operation involving a scalar or matrix on '''Broadcasts an operation involving a scalar or matrix on
the entire grid. If 'arg' is a matrix, it must be three-dimensional the entire grid. If 'arg' is a matrix, it must be three-dimensional
and its axes must be singular or of length nx/ny/nz.''' and its axes must be singular or of length nx/ny/nz.'''
import numpy as np import numpy as np
import operator import operator
# Broadcast argument
arg = self.comm.bcast(arg,root=0)
# Select operator
if operation in ('add','+'): if operation in ('add','+'):
op = operator.iadd op = operator.iadd
elif operation in ('subtract','sub','-'): elif operation in ('subtract','sub','-'):
@ -358,7 +327,7 @@ class PPP:
for axis in range(3): for axis in range(3):
if arg.shape[axis]==1: if arg.shape[axis]==1:
continue continue
elif arg.shape[axis]==self.numpoints(key,axis=axis): elif arg.shape[axis]==self.dim(key,axis=axis):
pos = self.position_from_rank(self.rank,external=False)[axis] pos = self.position_from_rank(self.rank,external=False)[axis]
sl_arg[axis] = slice( sl_arg[axis] = slice(
self.proc_grid[key][2*axis][pos]-1, self.proc_grid[key][2*axis][pos]-1,
@ -366,16 +335,24 @@ class PPP:
else: else:
raise ValueError("'arg' must either be singular or match global "\ raise ValueError("'arg' must either be singular or match global "\
"grid dimension. (axis={}: got {:d}, expected {:d}".format( "grid dimension. (axis={}: got {:d}, expected {:d}".format(
axis,arg.shape[axis],self.numpoints(key,axis=axis))) axis,arg.shape[axis],self.dim(key,axis=axis)))
# Only operate on interior and communcate ghosts later # Only operate on interior and communcate ghosts later
sl_int = tuple(slice(self.num_ghost[ii],-self.num_ghost[ii]) for ii in range(3)) sl_int = tuple(slice(self.num_ghost[ii],-self.num_ghost[ii]) for ii in range(3))
sl_arg = tuple(sl_arg) sl_arg = tuple(sl_arg)
op(self.field[key].data[sl_int],arg[sl_arg]) if key_out is None:
op(self.field[key].data[sl_int],arg[sl_arg])
else:
self.copy(key,key_out)
op(self.field[key_out].data[sl_int],arg[sl_arg])
# Exchange ghost cells and set boundary conditions # Exchange ghost cells and set boundary conditions
self.exchange_ghost_cells(key) self.exchange_ghost_cells(key)
self.impose_boundary_conditions(key) self.impose_boundary_conditions(key)
elif isinstance(arg,(int,float)): elif isinstance(arg,(int,float)):
op(self.field[key].data,arg) if key_out is None:
op(self.field[key].data,arg)
else:
self.copy(key,key_out)
op(self.field[key_out].data,arg)
return return
def integrate(self,key,integrate_axis,ufunc=None,ignore_nan=False,average=False): def integrate(self,key,integrate_axis,ufunc=None,ignore_nan=False,average=False):
@ -396,7 +373,7 @@ class PPP:
weights_local = np.array(weights_local) weights_local = np.array(weights_local)
# Simple implementation: all processor communicate directly to root # Simple implementation: all processor communicate directly to root
if self.rank==0: if self.rank==0:
dim_final = tuple(1 if integrate_axis[axis] else self.numpoints(key,axis=axis) for axis in range(3)) dim_final = tuple(1 if integrate_axis[axis] else self.dim(key,axis=axis) for axis in range(3))
# Allocate a nice spot of memory # Allocate a nice spot of memory
integ = np.zeros(dim_final,dtype=integ_local.dtype) integ = np.zeros(dim_final,dtype=integ_local.dtype)
if average: if average:
@ -432,47 +409,169 @@ class PPP:
self.comm.send(weights_local,dest=0,tag=2) self.comm.send(weights_local,dest=0,tag=2)
return None return None
def shift_to_pressure_grid(self,key,key_out=None):
from .field import Field3d
if key_out is None:
if key in ('u','v','w'):
raise ValueError("In-place shift of {'u','v','w'} is not permitted.")
key_out = key
# If pressure is loaded, use field information, else create a
# fake Field3d object so that we can use its methods.
field_ref = self.allocate('p',pseudo=True)
# Compute the relative shift and shift it
rel_shift = self.field[key].relative_shift(field_ref)
self.field[key_out] = self.field[key].shift_origin(rel_shift)
# Match number of gridpoints
sl = 3*[slice(None)]
for axis in range(3):
if not self.field[key_out].has_same_dim(field_ref,axis=axis):
sl[axis] = slice(0,field_ref.dim(axis=axis))
self.field[key_out].data = self.field[key_out].data[sl]
# Copy metadata of pressure grid
self.copy('p',key_out,skip_data=True,skip_symmetries=True)
self.symmetries[key_out] = self.symmetries[key].copy()
# Exchange ghost cells and set boundary conditions
self.exchange_ghost_cells(key_out)
self.impose_boundary_conditions(key_out)
def init_symmetry(self,key,mirror=(True,True)):
'''Sets the symmetries for ghost cells behind the wall'''
# Example: wall in y
# No-slip boundary (no mirror) free-slip boundary (mirror)
# u -> -u u -> u
# v -> -v v -> -v
# w -> -w w -> w
# p -> p p -> p
import numpy as np
self.symmetries[key] = np.zeros((3,3,3),dtype='i')
if not self.xperiodic:
if key=='u':
self.symmetries[key][0,1,1] = -1
self.symmetries[key][2,1,1] = -1
else:
self.symmetries[key][0,1,1] = 1 if mirror[0] else -1
self.symmetries[key][2,1,1] = 1 if mirror[1] else -1
if not self.yperiodic:
if key=='v':
self.symmetries[key][1,0,1] = -1
self.symmetries[key][1,2,1] = -1
else:
self.symmetries[key][1,0,1] = 1 if mirror[0] else -1
self.symmetries[key][1,2,1] = 1 if mirror[1] else -1
if not self.zperiodic:
if key=='w':
self.symmetries[key][1,1,0] = -1
self.symmetries[key][1,1,2] = -1
else:
self.symmetries[key][1,1,0] = 1 if mirror[0] else -1
self.symmetries[key][1,1,2] = 1 if mirror[1] else -1
def copy(self,key,key_out,skip_data=False,skip_symmetries=False):
'''Copies a field to a new key.'''
if key==key_out:
return
self.origin[key_out] = tuple(self.origin[key])
self.proc_grid[key_out] = self.proc_grid[key].copy()
if not skip_symmetries:
self.symmetries[key_out] = self.symmetries[key].copy()
if not skip_data:
self.field[key_out] = self.field[key].copy()
return
def delete(self,key):
'''Removes a field from memory.'''
assert key in self.field, "Field {} cannot be deleted, because it "\
"does not exist.".format(key)
import gc
del self.field[key]
del self.symmetries[key]
if key not in ('u','v','w','p','s'):
del self.origin[key]
del self.proc_grid[key]
gc.collect()
return
def merge(self,key,stride=(1,1,1)):
'''Returns the complete (possibly downsampled) field at rank 0.'''
return
def save_state(self,filename):
import pickle
fout = filename+'.{:05d}.pickle'.format(self.rank)
payload = (
self.num_ghost,
self.chunks_per_proc,
self.origin,
self.spacing,
self.periodicity,
self.bounds,
self.proc_grid,
self.nxp,
self.nyp,
self.nzp,
self.nghbr,
self.field,
self.symmetries)
with open(fout,"wb") as f:
pickle.dump(payload,f)
def to_vtk(self,key,no_ghost=False):
'''Returns the field (only its interior + some ghost cells for plotting)
as a vtk mesh.'''
return self._subfield_for_vtk(key,no_ghost).to_vtk()
def vtk_contour(self,key,val): def vtk_contour(self,key,val):
'''Compute isocontour for chunks.''' '''Compute isocontour for chunks.'''
if any([self.num_ghost[ii]>1 for ii in range(3)]): return self._subfield_for_vtk(key,False).vtk_contour(val)
idx_origin = tuple(self.num_ghost[ii]-1 for ii in range(3))
numpoints = self.chunk_size(axis=None)
return self.field[key].extract_subfield(
idx_origin,numpoints).vtk_contour(val)
else:
return self.field[key].vtk_contour(val)
return
def vtk_slice(self,key,normal,origin): def vtk_slice(self,key,normal,origin):
'''Extracts a plane from field.''' '''Extracts a plane from field.'''
if any([self.num_ghost[ii]>1 for ii in range(3)]): return self._subfield_for_vtk(key,False).vtk_slice(normal,origin)
idx_origin = tuple(self.num_ghost[ii]-1 for ii in range(3))
numpoints = tuple(self.field[key].numpoints[ii]-2*(self.num_ghost[ii]-1) def _subfield_for_vtk(self,key,no_ghost):
for ii in range(3)) num_ghost = np.array(self.num_ghost)
return self.field[key].extract_subfield( idx_origin = num_ghost
idx_origin,numpoints).vtk_slice(normal,origin) dim = np.array(self.field[key].dim(axis=None))
else: dim -= 2*num_ghost
return self.field[key].vtk_slice(normal,origin) if not no_ghost:
return idx_origin -= 1
dim += 1
if self.ip==self.nxp-1: dim[0]+=1
if self.jp==self.nyp-1: dim[1]+=1
if self.kp==self.nzp-1: dim[2]+=1
idx_origin = tuple(idx_origin)
dim = tuple(dim)
return self.field[key].extract_subfield(idx_origin,dim)
def rank_from_position(self,ip,jp,kp,external=False): def rank_from_position(self,ip,jp,kp,external=False):
if external: if external:
nyp,nzp = self.nyp_ext,self.nzp_ext nyp,nzp = self.nyp_ext,self.nzp_ext
else: else:
nyp,nzp = self.nyp,self.nzp nyp,nzp = self.nyp,self.nzp
return ip*nyp*nzp+jp*nzp+kp return self._rank_from_position(ip,jp,kp,nyp,nzp)
def position_from_rank(self,rank,external=False): def position_from_rank(self,rank,external=False):
if external: if external:
nyp,nzp = self.nyp_ext,self.nzp_ext nyp,nzp = self.nyp_ext,self.nzp_ext
else: else:
nyp,nzp = self.nyp,self.nzp nyp,nzp = self.nyp,self.nzp
return self._position_from_rank(rank,nyp,nzp)
@staticmethod
def _rank_from_position(ip,jp,kp,nyp,nzp):
return ip*nyp*nzp+jp*nzp+kp
@staticmethod
def _position_from_rank(rank,nyp,nzp):
ip = rank//(nyp*nzp) ip = rank//(nyp*nzp)
jp = (rank//nzp)%nyp jp = (rank//nzp)%nyp
kp = rank%nzp kp = rank%nzp
return (ip,jp,kp) return (ip,jp,kp)
def shifting_state(self,key,axis=None): def shifting_state(self,key,axis=None):
'''Returns the relative shift of first global grid point to domain
boundaries as either 1 (= +0.5h), 0 (= 0h) or -1 (= -0.5h).'''
if axis is None: if axis is None:
return tuple(self.shifting_state(key,axis=ii) for ii in range(3)) return tuple(self.shifting_state(key,axis=ii) for ii in range(3))
assert axis<3, "'axis' must be one of 0,1,2." assert axis<3, "'axis' must be one of 0,1,2."
@ -483,13 +582,13 @@ class PPP:
if axis is None: if axis is None:
return tuple(self.chunk_size(key,axis=ii) for ii in range(3)) return tuple(self.chunk_size(key,axis=ii) for ii in range(3))
assert axis<3, "'axis' must be one of 0,1,2." assert axis<3, "'axis' must be one of 0,1,2."
return self.field[key].numpoints[axis]-2*self.num_ghost[axis] return self.field[key].dim(axis=axis)-2*self.num_ghost[axis]
def numpoints(self,key,axis=None): def dim(self,key,axis=None):
'''Returns the total number of gridpoints across all processors '''Returns the total number of gridpoints across all processors
without ghost cells.''' without ghost cells.'''
if axis is None: if axis is None:
return tuple(self.numpoints(key,axis=ii) for ii in range(3)) return tuple(self.dim(key,axis=ii) for ii in range(3))
assert axis<3, "'axis' must be one of 0,1,2." assert axis<3, "'axis' must be one of 0,1,2."
return self.proc_grid[key][2*axis+1][-1] return self.proc_grid[key][2*axis+1][-1]
@ -600,7 +699,7 @@ class PPP:
recvbuf = np.zeros( recvbuf = np.zeros(
(ii_dst.stop-ii_dst.start, (ii_dst.stop-ii_dst.start,
jj_dst.stop-jj_dst.start, jj_dst.stop-jj_dst.start,
kk_dst.stop-kk_dst.start),dtype=self.precision) kk_dst.stop-kk_dst.start),dtype=self.field[key].data.dtype)
reqrecv = self.comm.Irecv(recvbuf,source=rank_src,tag=tag) reqrecv = self.comm.Irecv(recvbuf,source=rank_src,tag=tag)
# [recv] wait for data to be received # [recv] wait for data to be received
if reqrecv is not None: if reqrecv is not None:
@ -625,7 +724,7 @@ class PPP:
for axis in range(3): for axis in range(3):
if positionBd[axis]==-1: # lower boundary if positionBd[axis]==-1: # lower boundary
# index of first point within the domain # index of first point within the domain
idx = self.field[key].nearest_gridpoint(self.bounds[2*axis],axis=axis,lower=True)+1 idx = self.field[key].nearest_gridpoint(self.bounds[2*axis],axis=axis,lower=True)+1
# distance first point to wall: should be either 0 or dx/2 # distance first point to wall: should be either 0 or dx/2
dist = np.abs(self.field[key].coordinate(idx,axis=axis)-self.bounds[2*axis]) dist = np.abs(self.field[key].coordinate(idx,axis=axis)-self.bounds[2*axis])
if dist>=0.25*self.spacing[axis]: # no point on boundary if dist>=0.25*self.spacing[axis]: # no point on boundary
@ -637,19 +736,123 @@ class PPP:
break break
elif positionBd[axis]==1: # upper boundary elif positionBd[axis]==1: # upper boundary
# index of last point within the domain # index of last point within the domain
idx = self.field[key].nearest_gridpoint(self.bounds[2*axis+1],axis=axis,lower=True) idx = self.field[key].nearest_gridpoint(self.bounds[2*axis+1],axis=axis,lower=True)
# distance last point to wall: should be either 0 or -dx/2 # distance last point to wall: should be either 0 or -dx/2
dist = np.abs(self.field[key].coordinate(idx,axis=axis)-self.bounds[2*axis+1]) dist = np.abs(self.field[key].coordinate(idx,axis=axis)-self.bounds[2*axis+1])
if dist>=0.25*self.spacing[axis]: # no point on boundary if dist>=0.25*self.spacing[axis]: # no point on boundary
sl_dst[axis] = slice(idx+1,self.field[key].numpoints[axis],1) sl_dst[axis] = slice(idx+1,self.field[key].dim(axis=axis),1)
sl_src[axis] = slice(idx,2*idx+1-self.field[key].numpoints[axis],-1) sl_src[axis] = slice(idx,2*idx+1-self.field[key].dim(axis=axis),-1)
else: # point on boundary else: # point on boundary
sl_dst[axis] = slice(idx+1,self.field[key].numpoints[axis],1) sl_dst[axis] = slice(idx+1,self.field[key].dim(axis=axis),1)
sl_src[axis] = slice(idx-1,2*idx-self.field[key].numpoints[axis],-1) sl_src[axis] = slice(idx-1,2*idx-self.field[key].dim(axis=axis),-1)
break break
# print(sl_src,sl_dst,self.field[key].dim())
self.field[key].data[tuple(sl_dst)] = self.symmetries[key][inghbr,jnghbr,knghbr]*\ self.field[key].data[tuple(sl_dst)] = self.symmetries[key][inghbr,jnghbr,knghbr]*\
self.field[key].data[tuple(sl_src)] self.field[key].data[tuple(sl_src)]
@staticmethod
def _construct_procgrid(comm,proc,chunks_per_proc,num_ghost,periodicity):
import numpy as np
rank = comm.Get_rank()
nxcpp,nycpp,nzcpp = chunks_per_proc
nghx,nghy,nghz = num_ghost
xperiodic,yperiodic,zperiodic = periodicity
# Assert proc
assert all(k in proc for k in ('u','v','w','p','s')), "'proc' must be a dictionary with "\
"keys 'u','v','w','p','s'"
for key in proc:
assert len(proc[key])==6, "Entries of 'proc' must have length of 6."
proc_grid_ext = proc
# Initialize chunks per processor
# Verify that this processor grid can be redistributed onto the requested processor layout
nxp_ext = len(proc_grid_ext['u'][0])
nyp_ext = len(proc_grid_ext['u'][2])
nzp_ext = len(proc_grid_ext['u'][4])
nproc_ext = nxp_ext*nyp_ext*nzp_ext
#
assert nxp_ext%nxcpp==0, "Number of processors must be divisible by the number "\
"of processors per process. (nxp_ext={}, nxcpp={})".format(
nxp_ext,nxcpp)
assert nyp_ext%nycpp==0, "Number of processors must be divisible by the number "\
"of processors per process. (nyp_ext={}, nycpp={})".format(
nyp_ext,nycpp)
assert nzp_ext%nzcpp==0, "Number of processors must be divisible by the number "\
"of processors per process. (nzp_ext={}, nzcpp={})".format(
nzp_ext,nzcpp)
# Determine the new processor layout and verify total number of MPI processes
nxp = nxp_ext//nxcpp
nyp = nyp_ext//nycpp
nzp = nzp_ext//nzcpp
nproc = nxp*nyp*nzp
assert nproc==comm.Get_size(), "Number of MPI processes does not match the requested "\
"processor layout. (MPI procs: {}, required procs: {})".format(
comm.Get_size(),nproc)
# Construct internal processor grid
proc_grid = {}
for key in proc_grid_ext:
proc_grid[key] = [None]*6
proc_grid[key][0] = proc_grid_ext[key][0][0::nxcpp]
proc_grid[key][1] = proc_grid_ext[key][1][nxcpp-1::nxcpp]
proc_grid[key][2] = proc_grid_ext[key][2][0::nycpp]
proc_grid[key][3] = proc_grid_ext[key][3][nycpp-1::nycpp]
proc_grid[key][4] = proc_grid_ext[key][4][0::nzcpp]
proc_grid[key][5] = proc_grid_ext[key][5][nzcpp-1::nzcpp]
# Get position in processor grid
ip,jp,kp = PPP._position_from_rank(rank,nyp,nzp)
# Determine local grid indices and size
for key in proc_grid:
nxl = proc_grid[key][1][ip]-proc_grid[key][0][ip]+1
nyl = proc_grid[key][3][jp]-proc_grid[key][2][jp]+1
nzl = proc_grid[key][5][kp]-proc_grid[key][4][kp]+1
# Verify that local grid size is not smaller than ghost cell size
assert (nxl>=nghx and nyl>=nghy and nzl>=nghz),\
"Local grid size must be greater than number "\
"of ghost cells in each direction!"
# Initialize neighbor array
nghbr = np.empty((3,3,3),dtype='int')
# wrap-around x
ipl = ip-1
if ipl<0:
if xperiodic: ipl = nxp-1
else: ipl = -1
ipr = ip+1
if ipr>nxp-1:
if xperiodic: ipr = 0
else: ipr = -1
inghbr = (ipl,ip,ipr)
# wrap-around y
jpl = jp-1
if jpl<0:
if yperiodic: jpl = nyp-1
else: jpl = -1
jpr = jp+1
if jpr>nyp-1:
if yperiodic: jpr = 0
else: jpr = -1
jnghbr = (jpl,jp,jpr)
# wrap-around z
kpl = kp-1
if kpl<0:
if zperiodic: kpl = nzp-1
else: kpl = -1
kpr = kp+1
if kpr>nzp-1:
if zperiodic: kpr = 0
else: kpr = -1
knghbr = (kpl,kp,kpr)
# Construct array of neighbors
for ip in range(3):
for jp in range(3):
for kp in range(3):
# Assign rank to neighbor array
if inghbr[ip]<0 or jnghbr[jp]<0 or knghbr[kp]<0:
nghbr[ip,jp,kp] = -1
else:
nghbr[ip,jp,kp] = PPP._rank_from_position(inghbr[ip],jnghbr[jp],knghbr[kp],nyp,nzp)
return (proc_grid,nxp,nyp,nzp,
proc_grid_ext,nxp_ext,nyp_ext,nzp_ext,
nghbr)
def _baton_wait(self,batch_size,tag=420): def _baton_wait(self,batch_size,tag=420):
'''Block execution until an empty message from rank-batch_wait '''Block execution until an empty message from rank-batch_wait
is received (issued by _baton_pass)''' is received (issued by _baton_pass)'''

111
particle.py
View File

@ -1,37 +1,36 @@
import numpy as np
class Particles: class Particles:
def __init__(self,num,time,attr,period): def __init__(self,num,time,attr,period):
'''Initialize a Particles object. '''Initialize a Particles object.
num number of particles num number of particles
time time value time time value
attr dictionary of numpy arrays with length equal to num attr dictionary of np arrays with length equal to num
period the domain periods period the domain periods
''' '''
import numpy
assert 'id' in attr, "Need attribute 'id' to initialize Particles" assert 'id' in attr, "Need attribute 'id' to initialize Particles"
assert 'x' in attr, "Need attribute 'x' to initialize Particles" assert 'x' in attr, "Need attribute 'x' to initialize Particles"
assert 'y' in attr, "Need attribute 'y' to initialize Particles" assert 'y' in attr, "Need attribute 'y' to initialize Particles"
assert 'z' in attr, "Need attribute 'z' to initialize Particles" assert 'z' in attr, "Need attribute 'z' to initialize Particles"
assert isinstance(period,(tuple,list)) and len(period)==3, "'period' must be tuple/list with length 3" assert isinstance(period,(tuple,list)) and len(period)==3, "'period' must be tuple/list with length 3"
self.num = num self.num = num
if isinstance(time,(tuple,list,numpy.ndarray)): if isinstance(time,(tuple,list,np.ndarray)):
assert len(time)==1, "'time' must be a scalar value." assert len(time)==1, "'time' must be a scalar value."
time = time[0] time = time[0]
self.time = time self.time = time
self.attr = {} self.attr = {}
sortidx = attr['id'].argsort() sortidx = attr['id'].argsort()
for key in attr: for key in attr:
assert attr[key].ndim==1, "Attributes must be a one-dimensional numpy array." assert attr[key].ndim==1, "Attributes must be a one-dimensional np array."
assert attr[key].shape[0]==num, "Attributes must be of length 'num'." assert attr[key].shape[0]==num, "Attributes must be of length 'num'."
if key=='id': if key=='id':
self.add_attribute(key,attr[key][sortidx].astype('int')) self.add_attribute(key,attr[key][sortidx].astype('int'))
else: else:
self.add_attribute(key,attr[key][sortidx]) self.add_attribute(key,attr[key][sortidx])
self.period = tuple(period) self.period = tuple(period)
self.frame_velocity = numpy.zeros(3) self.frame_velocity = np.zeros(3)
return return
@classmethod @classmethod
def from_array(cls,pp,col,time,period,select_col=None): def from_array(cls,pp,col,time,period,select_col=None):
import numpy
assert 'id' in col, "Need column 'id' to initialize Particles" assert 'id' in col, "Need column 'id' to initialize Particles"
assert 'x' in col, "Need column 'x' to initialize Particles" assert 'x' in col, "Need column 'x' to initialize Particles"
assert 'y' in col, "Need column 'y' to initialize Particles" assert 'y' in col, "Need column 'y' to initialize Particles"
@ -39,7 +38,7 @@ class Particles:
assert pp.ndim==2 or pp.shape[2]==1, "'pp' cannot be a time series." assert pp.ndim==2 or pp.shape[2]==1, "'pp' cannot be a time series."
num = pp.shape[1] num = pp.shape[1]
if pp.ndim==2: if pp.ndim==2:
pp = numpy.expand_dims(pp,axis=2) pp = np.expand_dims(pp,axis=2)
if select_col is not None: if select_col is not None:
for scol in select_col: for scol in select_col:
assert scol in col, "Selected column '{}' not found in 'col'.".format(scol) assert scol in col, "Selected column '{}' not found in 'col'.".format(scol)
@ -78,21 +77,19 @@ class Particles:
attr[key] = self.attr[key].copy() attr[key] = self.attr[key].copy()
return Particles(self.num,self.time,attr,self.period) return Particles(self.num,self.time,attr,self.period)
def add_attribute(self,key,val): def add_attribute(self,key,val):
import numpy if isinstance(val,(tuple,list,np.ndarray)):
if isinstance(val,(tuple,list,numpy.ndarray)):
assert len(val)==self.num and val.ndim==1, "Invalid 'val'." assert len(val)==self.num and val.ndim==1, "Invalid 'val'."
self.attr[key] = numpy.array(val) self.attr[key] = np.array(val)
else: else:
self.attr[key] = numpy.full(self.num,val) self.attr[key] = np.full(self.num,val)
return return
def del_attribute(self,key): def del_attribute(self,key):
del self.attr[key] del self.attr[key]
def get_attribute(self,key): def get_attribute(self,key):
return self.attr[key] return self.attr[key]
def get_position(self,axis=None): def get_position(self,axis=None):
import numpy
if axis is None: if axis is None:
return numpy.vstack([self.attr[key].copy() for key in ('x','y','z')]) return np.vstack([self.attr[key].copy() for key in ('x','y','z')])
else: else:
assert axis<3, "'axis' must be smaller than 3." assert axis<3, "'axis' must be smaller than 3."
key = ('x','y','z')[axis] key = ('x','y','z')[axis]
@ -113,7 +110,8 @@ class Particles:
valdiff = val-self.frame_velocity[axis] valdiff = val-self.frame_velocity[axis]
self.translate(-self.time*valdiff,axis=axis) self.translate(-self.time*valdiff,axis=axis)
key = ('u','v','w')[axis] key = ('u','v','w')[axis]
self.attr[key] -= valdiff if self.has_attribute(key):
self.attr[key] -= valdiff
self.frame_velocity[axis] = val self.frame_velocity[axis] = val
return return
def enforce_periodicity(self,axis=None): def enforce_periodicity(self,axis=None):
@ -127,7 +125,6 @@ class Particles:
return return
def to_vtk(self,deep=False): def to_vtk(self,deep=False):
import pyvista as pv import pyvista as pv
import numpy as np
position = np.vstack([self.attr[key] for key in ('x','y','z')]).transpose() position = np.vstack([self.attr[key] for key in ('x','y','z')]).transpose()
mesh = pv.PolyData(position,deep=deep) mesh = pv.PolyData(position,deep=deep)
for key in self.attr: for key in self.attr:
@ -143,12 +140,11 @@ class Particles:
class Trajectories: class Trajectories:
def __init__(self,part,unraveled=False,copy_particles=False): def __init__(self,part,unraveled=False,copy_particles=False):
import numpy
self.numtime = 1 self.numtime = 1
self.numpart = part.num self.numpart = part.num
self.copy_particles = copy_particles self.copy_particles = copy_particles
# Attributes are first stored as a list of view to instances of # Attributes are first stored as a list of view to instances of
# Particles. Later they will be concatenated into numpy array # Particles. Later they will be concatenated into np array
# and the attributes of Particles will be views. # and the attributes of Particles will be views.
self.attr = {} self.attr = {}
self.part = [] self.part = []
@ -161,7 +157,8 @@ class Trajectories:
self.attr[key] = [self.part[0].attr[key].view()] self.attr[key] = [self.part[0].attr[key].view()]
self.period = self.part[0].period self.period = self.part[0].period
self.unraveled = unraveled self.unraveled = unraveled
self.frame_velocity = numpy.zeros(3) self.frame_velocity = np.zeros(3)
def __str__(self): def __str__(self):
str = 'Trajectories with\n' str = 'Trajectories with\n'
str+= ' time steps: {:d}\n'.format(self.numtime) str+= ' time steps: {:d}\n'.format(self.numtime)
@ -173,6 +170,7 @@ class Trajectories:
str+= ' frame velocity: {}\n'.format(self.frame_velocity) str+= ' frame velocity: {}\n'.format(self.frame_velocity)
str+= ' unraveled: {}'.format(self.unraveled) str+= ' unraveled: {}'.format(self.unraveled)
return str return str
def __iadd__(self,other): def __iadd__(self,other):
if isinstance(other,Trajectories): if isinstance(other,Trajectories):
self.add_trajectories(other) self.add_trajectories(other)
@ -181,6 +179,7 @@ class Trajectories:
else: else:
raise TypeError('Cannot add type {}'.format(type(other))) raise TypeError('Cannot add type {}'.format(type(other)))
return self return self
def __getitem__(self,val): def __getitem__(self,val):
assert isinstance(val,tuple) and len(val)==2, "Trajectories must be indexed by [part,time]." assert isinstance(val,tuple) and len(val)==2, "Trajectories must be indexed by [part,time]."
sl = [] sl = []
@ -194,6 +193,7 @@ class Trajectories:
else: else:
raise TypeError("Trajectories can only be sliced by slice objects or integers.") raise TypeError("Trajectories can only be sliced by slice objects or integers.")
return self._slice(slice_part=sl[0],slice_time=sl[1]) return self._slice(slice_part=sl[0],slice_time=sl[1])
@classmethod @classmethod
def from_mat(cls,file,unraveled=False,select_col=None): def from_mat(cls,file,unraveled=False,select_col=None):
from .helper import load_mat from .helper import load_mat
@ -206,6 +206,7 @@ class Trajectories:
for key in col: for key in col:
col[key]-=1 col[key]-=1
return cls.from_array(pp,col,time,period,unraveled=unraveled,select_col=select_col) return cls.from_array(pp,col,time,period,unraveled=unraveled,select_col=select_col)
@classmethod @classmethod
def from_array(cls,pp,col,time,period,unraveled=False,select_col=None): def from_array(cls,pp,col,time,period,unraveled=False,select_col=None):
ntime = len(time) ntime = len(time)
@ -213,6 +214,7 @@ class Trajectories:
for ii in range(1,ntime): for ii in range(1,ntime):
traj += Particles.from_array(pp[:,:,ii],col,time[ii],period,select_col=select_col) traj += Particles.from_array(pp[:,:,ii],col,time[ii],period,select_col=select_col)
return traj return traj
def _slice(self,slice_part=slice(None),slice_time=slice(None)): def _slice(self,slice_part=slice(None),slice_time=slice(None)):
assert isinstance(slice_part,slice), "'slice_part' must be a slice object." assert isinstance(slice_part,slice), "'slice_part' must be a slice object."
assert isinstance(slice_time,slice), "'slice_time' must be a slice object." assert isinstance(slice_time,slice), "'slice_time' must be a slice object."
@ -225,8 +227,8 @@ class Trajectories:
traj += Trajectories(part[slice_part]) traj += Trajectories(part[slice_part])
traj.frame_velocity = self.frame_velocity traj.frame_velocity = self.frame_velocity
return traj return traj
def add_particles(self,part): def add_particles(self,part):
import numpy
# Verify part # Verify part
assert part.time>self.part[-1].time, "Time steps must be added in monotonically increasing order." assert part.time>self.part[-1].time, "Time steps must be added in monotonically increasing order."
assert part.num==self.numpart, "Number of particles is different from previous time steps." assert part.num==self.numpart, "Number of particles is different from previous time steps."
@ -246,6 +248,7 @@ class Trajectories:
self.attr[key].append(self.part[-1].attr[key].view()) self.attr[key].append(self.part[-1].attr[key].view())
self.numtime += 1 self.numtime += 1
return return
def add_trajectories(self,traj): def add_trajectories(self,traj):
# Verify traj # Verify traj
assert traj.part[0].time>self.part[-1].time, "Time steps must be added in monotonically increasing order." assert traj.part[0].time>self.part[-1].time, "Time steps must be added in monotonically increasing order."
@ -267,26 +270,26 @@ class Trajectories:
self.attr[key].append(self.part[-traj.numtime+ii].attr[key].view()) self.attr[key].append(self.part[-traj.numtime+ii].attr[key].view())
self.numtime += traj.numtime self.numtime += traj.numtime
return return
def get_time(self): def get_time(self):
import numpy time = np.zeros(self.numtime)
time = numpy.zeros(self.numtime)
for ii,part in enumerate(self.part): for ii,part in enumerate(self.part):
time[ii] = part.time time[ii] = part.time
return time return time
def get_trajectories(self,slice_part=slice(None),slice_time=slice(None)): def get_trajectories(self,slice_part=slice(None),slice_time=slice(None)):
'''Get (x,y,z) trajectories in numpy array of dimension (3,npart,ntime).''' '''Get (x,y,z) trajectories in np array of dimension (3,npart,ntime).'''
import numpy
assert isinstance(slice_part,slice), "'slice_part' must be a slice." assert isinstance(slice_part,slice), "'slice_part' must be a slice."
assert isinstance(slice_time,slice), "'slice_time' must be a slice." assert isinstance(slice_time,slice), "'slice_time' must be a slice."
self._make_data_array() self._make_data_array()
return numpy.stack([ return np.stack([
self.attr['x'][slice_part,slice_time], self.attr['x'][slice_part,slice_time],
self.attr['y'][slice_part,slice_time], self.attr['y'][slice_part,slice_time],
self.attr['z'][slice_part,slice_time]],axis=0) self.attr['z'][slice_part,slice_time]],axis=0)
def get_trajectories_segmented(self,slice_part=slice(None),slice_time=slice(None),restore_ravel_state=True): def get_trajectories_segmented(self,slice_part=slice(None),slice_time=slice(None),restore_ravel_state=True):
'''Get (x,y,z) segments of trajectories as a tuple (len: npart) of '''Get (x,y,z) segments of trajectories as a tuple (len: npart) of
lists (len: nsegments) of numpy arrays (shape: 3 X ntime of the segment)''' lists (len: nsegments) of np arrays (shape: 3 X ntime of the segment)'''
import numpy
assert isinstance(slice_part,slice), "'slice_part' must be a slice." assert isinstance(slice_part,slice), "'slice_part' must be a slice."
assert isinstance(slice_time,slice), "'slice_time' must be a slice." assert isinstance(slice_time,slice), "'slice_time' must be a slice."
was_unraveled = self.unraveled was_unraveled = self.unraveled
@ -296,21 +299,21 @@ class Trajectories:
ntime = xyzpath.shape[2] ntime = xyzpath.shape[2]
# Initialize output and helper arrays # Initialize output and helper arrays
out = tuple([] for ii in range(npart)) out = tuple([] for ii in range(npart))
lastJump = numpy.zeros(npart,dtype=numpy.uint) lastJump = np.zeros(npart,dtype=np.uint)
lastPeriod = numpy.zeros((3,npart),dtype=numpy.int) lastPeriod = np.zeros((3,npart),dtype=np.int)
for axis in range(3): for axis in range(3):
if self.period[axis] is not None: if self.period[axis] is not None:
lastPeriod[axis,:] = numpy.floor_divide(xyzpath[axis,:,0],self.period[axis]) lastPeriod[axis,:] = np.floor_divide(xyzpath[axis,:,0],self.period[axis])
# Construct output tuple # Construct output tuple
for itime in range(1,ntime+1): for itime in range(1,ntime+1):
thisPeriod = numpy.zeros((3,npart),dtype=int) thisPeriod = np.zeros((3,npart),dtype=int)
if itime==ntime: if itime==ntime:
hasJumped = numpy.ones(npart,dtype=bool) hasJumped = np.ones(npart,dtype=bool)
else: else:
for axis in range(3): for axis in range(3):
if self.period[axis] is not None: if self.period[axis] is not None:
thisPeriod[axis,:] = numpy.floor_divide(xyzpath[axis,:,itime],self.period[axis]) thisPeriod[axis,:] = np.floor_divide(xyzpath[axis,:,itime],self.period[axis])
hasJumped = numpy.any(thisPeriod!=lastPeriod,axis=0) hasJumped = np.any(thisPeriod!=lastPeriod,axis=0)
for ipart in range(npart): for ipart in range(npart):
if hasJumped[ipart]: if hasJumped[ipart]:
sl = slice(lastJump[ipart],itime) sl = slice(lastJump[ipart],itime)
@ -325,8 +328,20 @@ class Trajectories:
if restore_ravel_state and not was_unraveled: if restore_ravel_state and not was_unraveled:
self.enforce_periodicity() self.enforce_periodicity()
return out return out
def slice_time(self,time_beg=None,time_end=None): def slice_time(self,time_beg=None,time_end=None):
return time = self.get_time()
if time_beg is not None:
idx_beg = np.searchsorted(time,time_beg,side='left')
else:
idx_beg = None
if time_end is not None:
idx_end = np.searchsorted(time,time_end,side='left')
else:
idx_end = None
sl = slice(idx_beg,idx_end+1)
return self._slice(slice_time=sl)
def set_frame_velocity(self,val,axis=0): def set_frame_velocity(self,val,axis=0):
assert axis<3, "'axis' must be smaller than 3." assert axis<3, "'axis' must be smaller than 3."
for part in self.part: for part in self.part:
@ -335,8 +350,8 @@ class Trajectories:
if not self.unraveled: if not self.unraveled:
self.enforce_periodicity() self.enforce_periodicity()
return return
def unravel(self): def unravel(self):
import numpy
if self.unraveled: return if self.unraveled: return
#raise NotImplementedError('Implemented but untested!') #raise NotImplementedError('Implemented but untested!')
self._make_data_array() self._make_data_array()
@ -345,48 +360,56 @@ class Trajectories:
key = ('x','y','z')[axis] key = ('x','y','z')[axis]
#posdiff = (self.part[ii].attr[key]-self.part[ii-1].attr[key]).squeeze() #posdiff = (self.part[ii].attr[key]-self.part[ii-1].attr[key]).squeeze()
posdiff = self.attr[key][:,1:]-self.attr[key][:,0:-1] posdiff = self.attr[key][:,1:]-self.attr[key][:,0:-1]
coeff = -numpy.sign(posdiff)*(numpy.abs(posdiff)>0.5*self.period[axis]) coeff = -np.sign(posdiff)*(np.abs(posdiff)>0.5*self.period[axis])
coeff = numpy.cumsum(coeff,axis=1) coeff = np.cumsum(coeff,axis=1)
self.attr[key][:,1:] += coeff*self.period[axis] self.attr[key][:,1:] += coeff*self.period[axis]
self.unraveled = True self.unraveled = True
return return
def enforce_periodicity(self): def enforce_periodicity(self):
for part in self.part: for part in self.part:
part.enforce_periodicity(axis=None) part.enforce_periodicity(axis=None)
self.unraveled = False self.unraveled = False
return return
def to_vtk(self):
import pyvista as pv
mesh = pv.PolyData()
for part in self.get_trajectories_segmented():
for seg in part:
mesh += pv.helpers.lines_from_points(seg.transpose())
return mesh
def _make_data_array(self): def _make_data_array(self):
'''Transforms self.attr[key] to a numpy array of dimension npart X ntime '''Transforms self.attr[key] to a np array of dimension npart X ntime
and updates self.part[itime].attr[key] to views on self.attr[key][:,itime].''' and updates self.part[itime].attr[key] to views on self.attr[key][:,itime].'''
import numpy
#print('DEBUG: _make_data_array') #print('DEBUG: _make_data_array')
if not self._is_data_list: if not self._is_data_list:
return return
#print([x.shape for x in self.attr['x']]) #print([x.shape for x in self.attr['x']])
for key in self.attr: for key in self.attr:
self.attr[key] = numpy.stack(self.attr[key],axis=1) self.attr[key] = np.stack(self.attr[key],axis=1)
for ii in range(0,self.numtime): for ii in range(0,self.numtime):
self.part[ii].attr[key] = self.attr[key][:,ii] self.part[ii].attr[key] = self.attr[key][:,ii]
#print(self.attr['x'].shape) #print(self.attr['x'].shape)
self._is_data_list = False self._is_data_list = False
return return
def _make_data_list(self): def _make_data_list(self):
'''Transforms self.attr[key] to a list of length ntime of numpy arrays of '''Transforms self.attr[key] to a list of length ntime of np arrays of
dimension npart and updates self.part[itime].attr[key] to views on self.attr[key][itime].''' dimension npart and updates self.part[itime].attr[key] to views on self.attr[key][itime].'''
import numpy
#print('DEBUG: _make_data_list') #print('DEBUG: _make_data_list')
if self._is_data_list: if self._is_data_list:
return return
#print(self.attr['x'].shape) #print(self.attr['x'].shape)
for key in self.attr: for key in self.attr:
self.attr[key] = [numpy.squeeze(x) for x in numpy.split(self.attr[key],self.numtime,axis=1)] self.attr[key] = [np.squeeze(x) for x in np.split(self.attr[key],self.numtime,axis=1)]
for ii in range(0,self.numtime): for ii in range(0,self.numtime):
self.part[ii].attr[key] = self.attr[key][ii].view() self.part[ii].attr[key] = self.attr[key][ii].view()
#print([x.shape for x in self.attr['x']]) #print([x.shape for x in self.attr['x']])
self._is_data_list = True self._is_data_list = True
return return
def sort(pp,col): def sort(pp,col):
ncol,npart,ntime = pp.shape ncol,npart,ntime = pp.shape
assert('id' in col) assert('id' in col)