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

parallel tools added.

This commit is contained in:
Michael Krayer 2021-05-27 21:37:00 +02:00
parent 77f56261be
commit 1756a7a916
3 changed files with 877 additions and 117 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

318
field.py
View File

@ -1,27 +1,39 @@
import numpy import numpy
class Field3d: class Field3d:
def __init__(self,data,origin,spacing,bounds,period): def __init__(self,data,origin,spacing):
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"
assert len(period)==3, "'period' must be of length 3"
self.data = numpy.array(data) self.data = numpy.array(data)
self.data.setflags(write=False)
self.numpoints = self.data.shape 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.period = tuple(period) self.eps_collapse = 1e-12
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.numpoints)
str+= ' origin: {:G}, {:G}, {:G}\n'.format(*self.origin) str+= ' origin: {:G}, {:G}, {:G}\n'.format(*self.origin)
str+= ' spacing: {:G}, {:G}, {:G}\n'.format(*self.spacing) str+= ' spacing: {:G}, {:G}, {:G}'.format(*self.spacing)
str+= ' period: {}, {}, {}'.format(*self.period)
return str return str
# TBD: this should return another Field3d object
# def __getitem__(self,val):
# assert isinstance(val,tuple) and len(val)==3, "Field3d must be indexed by [ii,jj,kk]."
# sl = []
# for x in val:
# if isinstance(x,int):
# lo,hi = x,x+1
# hi = hi if hi!=0 else None
# sl.append(slice(lo,hi))
# elif isinstance(x,slice):
# sl.append(x)
# else:
# raise TypeError("Trajectories can only be sliced by slice objects or integers.")
# return self.data[sl[0],sl[1],sl[2]]
@classmethod @classmethod
def from_chunk(cls,chunk,gridg,period=(None,None,None)): def from_chunk(cls,chunk,gridg):
'''Initialize Field3d from chunk data and global grid.''' '''Initialize Field3d from chunk data and global grid.'''
xg,yg,zg = gridg xg,yg,zg = gridg
ib,jb,kb = chunk['ibeg']-1, chunk['jbeg']-1, chunk['kbeg']-1 ib,jb,kb = chunk['ibeg']-1, chunk['jbeg']-1, chunk['kbeg']-1
@ -31,79 +43,174 @@ class Field3d:
assert (chunk['nxl']+2*chunk['ighost'])==nx, "Invalid chunk data: nxl != chunk['data'].shape[0]" 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['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]" 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) return cls(chunk['data'],origin=(xo,yo,zo),spacing=(dx,dy,dz))
def insert_subfield(self,subfield):
assert all([abs(subfield.spacing[ii]-self.spacing[ii])<self.eps_collapse
for ii in range(3)]), "spacing differs."
assert all([self.distance_to_nearest_gridpoint(subfield.origin[ii],axis=ii)<self.eps_collapse
for ii in range(3)]), "subfield has shifted origin."
assert all(self.is_within_bounds(subfield.origin,axis=None)), "subfield origin 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 = self.nearest_gridpoint(subfield.origin,axis=None)
nx,ny,nz = subfield.numpoints
self.data[ib:ib+nx,jb:jb+ny,kb:kb+nz] = subfield.data[:,:,:]
return
def extract_subfield(self,idx_origin,numpoints,stride=(1,1,1)):
assert all(idx_origin[ii]>=0 and idx_origin[ii]<self.numpoints[ii] for ii in range(3)),\
"'origin' is out-of-bounds."
assert all(idx_origin[ii]+stride[ii]*(numpoints[ii]-1)<self.numpoints[ii] for ii in range(3)),\
"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))
origin = self.coordinate(idx_origin)
spacing = tuple(self.spacing[ii]*stride[ii] for ii in range(3))
data = self.data[sl].copy()
return Field3d(data,origin,spacing)
def coordinate(self,idx,axis=None):
if axis is None:
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))
assert axis<3, "'axis' must be one of 0,1,2."
assert idx<self.numpoints[axis], "'idx' is out-of-bounds."
return self.origin[axis]+idx*self.spacing[axis]
def nearest_gridpoint(self,coord,axis=None,lower=False):
if axis is None:
assert len(coord)==3, "If 'axis' is None, 'coord' must be a tuple/list of length 3."
return tuple(self.nearest_gridpoint(coord[ii],axis=ii,lower=lower) for ii in range(3))
assert axis<3, "'axis' must be one of 0,1,2."
if lower:
return numpy.floor((coord+self.eps_collapse-self.origin[axis])/self.spacing[axis]).astype('int')
else:
return numpy.round((coord-self.origin[axis])/self.spacing[axis]).astype('int')
def distance_to_nearest_gridpoint(self,coord,axis=None,lower=False):
if axis is None:
assert len(coord)==3, "If 'axis' is None, 'coord' must be a tuple/list of length 3."
return tuple(self.distance_to_nearest_gridpoint(coord[ii],axis=ii,lower=lower) for ii in range(3))
assert axis<3, "'axis' must be one of 0,1,2."
val = numpy.remainder(coord+self.eps_collapse-self.origin[axis],self.spacing[axis])-self.eps_collapse
if not lower and val>0.5*self.spacing[axis]:
val = self.spacing[axis]-val
return val
def is_within_bounds(self,coord,axis=None):
if axis is None:
assert len(coord)==3, "If 'axis' is None, 'coord' must be a tuple/list of length 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."
idx_nearest = self.nearest_gridpoint(coord,axis=axis)
if idx_nearest>0 and idx_nearest<self.numpoints[axis]-1:
return True
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:
return True
else:
return False
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(0,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] return self.numpoints[axis]
def endpoint(self,axis): def endpoint(self,axis=None):
if axis is None:
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.numpoints[axis]-1)*self.spacing[axis]
def derivative(self,axis,keep_origin=False): def derivative(self,axis,preserve_grid=False,padding=None):
if not keep_origin: if not preserve_grid:
# 2nd order FD with h = dx/2
origin = list(self.origin) origin = list(self.origin)
if self.period[axis] is not None: if padding is None:
slhi = [slice(None),slice(None),slice(None)] data = numpy.zeros(self.data.shape,dtype=self.data.dtype)
sllo = [slice(None),slice(None),slice(None)] slout = [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] origin[axis] += 0.5*self.spacing[axis]
elif padding=='before':
shape = numpy.array(self.data.shape)
shape[axis] += 1
data = numpy.zeros(tuple(shape),dtype=self.data.dtype)
slout = [slice(None),slice(None),slice(None)]
slout[axis] = slice(1,None)
origin[axis] -= 0.5*self.spacing[axis]
elif padding=='after':
shape = numpy.array(self.data.shape)
shape[axis] += 1
data = numpy.zeros(tuple(shape),dtype=self.data.dtype)
slout = [slice(None),slice(None),slice(None)]
slout[axis] = slice(0,-1)
origin[axis] += 0.5*self.spacing[axis]
elif padding=='both':
shape = numpy.array(self.data.shape)
shape[axis] += 2
data = numpy.zeros(tuple(shape),dtype=self.data.dtype)
slout = [slice(None),slice(None),slice(None)]
slout[axis] = slice(1,-1)
origin[axis] -= 0.5*self.spacing[axis]
else:
raise ValueError("'padding' must be one of 'none','before','after'.")
slout = tuple(slout)
# 2nd order FD with h = dx/2
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[slout] = (1./self.spacing[axis])*(self.data[slhi]-self.data[sllo])
else: else:
# 2nd order FD with h = dx, one-sided at boundaries # 2nd order FD with h = dx, one-sided at boundaries
# https://numpy.org/doc/stable/reference/generated/numpy.gradient.html # https://numpy.org/doc/stable/reference/generated/numpy.gradient.html
origin = self.origin origin = self.origin
if self.period[axis] is not None: data = numpy.gradient(self.data,self.spacing[axis],axis=axis,edge_order=2)
data = numpy.gradient(self.data,self.spacing[axis],axis=axis,edge_order=1) return Field3d(data,origin,self.spacing)
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): def gradient(self,axis,preserve_origin=False,padding=None):
return [self.derivative(axis,keep_origin=keep_origin) for axis in range(0,3)] return [self.derivative(axis,preserve_origin=preserve_origin,padding=padding) for axis in range(0,3)]
def change_grid(self,origin,spacing,numpoints,EPS_COLLAPSE=1e-12): def gaussian_filter(self,sigma,truncate=4.0,border='constant',const_val=numpy.nan):
# TBD: periodicity '''Applies a gaussian filter: sigma is standard deviation for Gaussian kernel for each axis.'''
from scipy import ndimage
assert isinstance(sigma,(tuple,list,numpy.ndarray)) and len(sigma)==3,\
"'sigma' must be a tuple/list of length 3"
# Convert sigma from simulation length scales to grid points as required by ndimage
sigma_img = tuple(sigma[ii]/self.spacing[ii] for ii in range(3))
data = ndimage.gaussian_filter(self.data,sigma_img,truncate=truncate,mode=border,cval=const_val)
print(data.shape,self.data.shape)
print(numpy.linalg.norm(data-self.data))
print(numpy.max(data))
return Field3d(data,self.origin,self.spacing)
def gaussian_filter_radius(self,sigma,truncate=4.0):
'''Radius of Gaussian filter. Stencil width is 2*radius+1.'''
assert isinstance(sigma,(tuple,list,numpy.ndarray)) and len(sigma)==3,\
"'sigma' must be a tuple/list of length 3"
# Convert sigma from simulation length scales to grid points as required by ndimage
sigma_img = tuple(sigma[ii]/self.spacing[ii] for ii in range(3))
radius = []
for ii in range(3):
radius.append(int(truncate*sigma_img[ii]+0.5))
return tuple(radius)
def change_grid(self,origin,spacing,numpoints):
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]+(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." assert all([endpoint[ii]<=self.endpoint(ii) for ii in range(0,3)]), "New end point is out of bounds."
data = numpy.zeros(numpoints) data = numpy.zeros(numpoints)
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
i0,j0,k0 = [numpy.floor((origin[ii]-self.origin[ii])/self.spacing[ii]).astype('int') for ii in range(0,3)] i0,j0,k0 = self.nearest_gridpoint(origin,axis=None,lower=True)
cx,cy,cz = [numpy.remainder(origin[ii]-self.origin[ii],self.spacing[ii]) for ii in range(0,3)] cx,cy,cz = [self.distance_to_nearest_gridpoint(origin[ii],axis=ii,lower=True)/self.spacing[ii]
c = weights_trilinear((cx,cy,cz)) for ii in range(3)]
c = self.weights_trilinear((cx,cy,cz))
for ii in range(0,2): for ii in range(0,2):
for jj in range(0,2): for jj in range(0,2):
for kk in range(0,2): for kk in range(0,2):
if c[ii,jj,kk]>EPS_COLLAPSE: if c[ii,jj,kk]>self.eps_collapse:
data += c[ii,jj,kk]*self.data[ data += c[ii,jj,kk]*self.data[
i0+ii:i0+ii+numpoints[0], i0+ii:i0+ii+numpoints[0],
j0+jj:j0+jj+numpoints[1], j0+jj:j0+jj+numpoints[1],
@ -112,50 +219,75 @@ class Field3d:
for ii in range(0,numpoints[0]): for ii in range(0,numpoints[0]):
for jj in range(0,numpoints[1]): for jj in range(0,numpoints[1]):
for kk in range(0,numpoints[2]): for kk in range(0,numpoints[2]):
coordinate = ( coord = (
origin[0]+ii*spacing[0], origin[0]+ii*spacing[0],
origin[1]+jj*spacing[1], origin[1]+jj*spacing[1],
origin[2]+kk*spacing[2]) origin[2]+kk*spacing[2])
data[ii,jj,kk] = self.interpolate(coordinate) data[ii,jj,kk] = self.interpolate(coord)
return Field3d(data,origin,spacing,self.period) return Field3d(data,origin,spacing)
def interpolate(self,coordinate,EPS_COLLAPSE=1e-12): def interpolate(self,coord):
assert all([coordinate[ii]>=self.origin[ii] for ii in range(0,3)]), "'coordinate' is out of bounds." assert all([coord[ii]>=self.origin[ii] for ii in range(0,3)]), "'coord' is out of bounds."
assert all([coordinate[ii]<=self.endpoint(ii) for ii in range(0,3)]), "'coordinate' is out of bounds." assert all([coord[ii]<=self.endpoint(ii) for ii in range(0,3)]), "'coord' 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)] i0,j0,k0 = self.nearest_gridpoint(coord,axis=None,lower=True)
cx,cy,cz = [numpy.remainder(coordinate[ii]-self.origin[ii],self.spacing[ii]) for ii in range(0,3)] cx,cy,cz = [self.distance_to_nearest_gridpoint(coord[ii],axis=ii,lower=True)/self.spacing[ii]
c = weights_trilinear((cx,cy,cz)) for ii in range(3)]
c = self.weights_trilinear((cx,cy,cz))
val = 0.0 val = 0.0
for ii in range(0,2): for ii in range(0,2):
for jj in range(0,2): for jj in range(0,2):
for kk in range(0,2): for kk in range(0,2):
if c[ii,jj,kk]>EPS_COLLAPSE: if c[ii,jj,kk]>self.eps_collapse:
val += c[ii,jj,kk]*self.data[i0+ii,j0+jj,k0+kk] val += c[ii,jj,kk]*self.data[i0+ii,j0+jj,k0+kk]
return val return val
def move_origin_periodic(self,shift,axis): def weights_trilinear(self,rel_dist):
'''Moves origin of grid in periodic direction by shift*spacing.''' assert len(rel_dist)==3, "len(rel_dist) must be 3."
assert self.period[axis], "Origin can only be moved in periodic directions." cx,cy,cz = rel_dist
self.data = numpy.roll(self.data,shift,axis=axis) if cx<0.0 and cx>-self.eps_collapse: cx=0.0
self.data.setflags(write=False) if cy<0.0 and cy>-self.eps_collapse: cy=0.0
origin = list(self.origin) if cz<0.0 and cz>-self.eps_collapse: cz=0.0
origin[axis] += shift*self.spacing[axis] if cx>1.0 and cx<1.0+self.eps_collapse: cx=1.0
self.origin = tuple(origin) if cy>1.0 and cy<1.0+self.eps_collapse: cy=1.0
return if cz>1.0 and cz<1.0+self.eps_collapse: cz=1.0
@staticmethod assert cx>=0.0 and cy>=0.0 and cz>=0.0, "'rel_dist' must be >=0"
def __update_ghostcells_periodic(data,axis): assert cx<=1.0 and cy<=1.0 and cz<=1.0, "'rel_dist' must be <=1"
slsrc = [slice(None),slice(None),slice(None)] c = numpy.zeros((2,2,2))
sldst = [slice(None),slice(None),slice(None)] c[0,0,0] = 1.0-(cx+cy+cz)+(cx*cy+cx*cz+cy*cz)-(cx*cy*cz)
slsrc[axis] = slice(1,2) c[1,0,0] = cx-(cx*cy+cx*cz)+(cx*cy*cz)
sldst[axis] = slice(-1,None) c[0,1,0] = cy-(cx*cy+cy*cz)+(cx*cy*cz)
data[sldst] = data[slsrc] c[0,0,1] = cz-(cx*cz+cy*cz)+(cx*cy*cz)
slsrc = [slice(None),slice(None),slice(None)] c[1,1,0] = (cx*cy)-(cx*cy*cz)
sldst = [slice(None),slice(None),slice(None)] c[1,0,1] = (cx*cz)-(cx*cy*cz)
slsrc[axis] = slice(-2,-1) c[0,1,1] = (cy*cz)-(cx*cy*cz)
sldst[axis] = slice(0,1) c[1,1,1] = (cx*cy*cz)
data[sldst] = data[slsrc] return c
def set_writable(self,flag):
self.data.setflags(write=flag)
return return
def to_vtk(self):
import pyvista as pv
mesh = pv.UniformGrid()
mesh.dimensions = self.dim(axis=None)
mesh.origin = self.origin
mesh.spacing = self.spacing
mesh.point_arrays['data'] = self.data.flatten(order='F')
return mesh
def vtk_contour(self,val):
if not isinstance(val,(tuple,list)):
val = [val]
return self.to_vtk().contour(val)
def vtk_slice(self,normal,origin):
assert (normal in ('x','y','z') or (isinstance(normal,(tuple,list))
and len(normal)==3)), "'normal' must be 'x','y','z' or tuple of length 3."
assert isinstance(origin,(tuple,list)) and len(origin)==3,\
"'origin' must be tuple of length 3."
return self.to_vtk().slice(normal=normal,origin=origin)
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
@ -177,19 +309,3 @@ class ChunkIterator:
return field return field
else: else:
raise StopIteration 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

624
parallel.py Normal file
View File

@ -0,0 +1,624 @@
class PPP:
"""Parallel Python Postprocessor for suspend"""
def __init__(self,comm,func_load,num_ghost,precision,origin,spacing,periodicity,bounds,proc,chunks_per_proc):
'''Constructor: except for comm, only rank 0 needs initialized data.'''
self.comm = comm
self.rank = comm.Get_rank()
self.func_load = func_load
self.init_settings(num_ghost,precision)
self.init_domain(origin,spacing,periodicity,bounds)
self.init_procgrid(proc,chunks_per_proc)
self.field = {}
self.symmetries = {}
return
@classmethod
def from_snapshot(cls,snap,chunks_per_proc=(1,1,1),num_ghost=(1,1,1),precision='float64'):
from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
if rank==0:
proc = snap.procgrid()
origin = snap.origin()
spacing = snap.spacing()
periodicity = snap.periodicity()
bounds = snap.bounds()
else:
proc = None
origin = None
spacing = None
periodicity = None
bounds = None
func_load = snap.field_chunk
return cls(comm,func_load,num_ghost,precision,origin,spacing,periodicity,bounds,proc,chunks_per_proc)
def init_settings(self,num_ghost,precision):
'''Initializes PPP settings for all processors.'''
# TBD: some assertions
self.num_ghost = self.comm.bcast(num_ghost,root=0)
self.precision = self.comm.bcast(precision,root=0)
# Some shortcuts
self.nghx,self.nghy,self.nghz = self.num_ghost
def init_symmetries(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 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
self.chunk_bounds = {}
self.chunk_size = {}
for key in self.proc_grid:
self.chunk_bounds[key] = [None]*6
self.chunk_bounds[key][0] = self.proc_grid[key][0][self.ip]
self.chunk_bounds[key][1] = self.proc_grid[key][1][self.ip]
self.chunk_bounds[key][2] = self.proc_grid[key][2][self.jp]
self.chunk_bounds[key][3] = self.proc_grid[key][3][self.jp]
self.chunk_bounds[key][4] = self.proc_grid[key][4][self.kp]
self.chunk_bounds[key][5] = self.proc_grid[key][5][self.kp]
self.chunk_size[key] = [None]*3
self.chunk_size[key][0] = self.chunk_bounds[key][1]-self.chunk_bounds[key][0]+1
self.chunk_size[key][1] = self.chunk_bounds[key][3]-self.chunk_bounds[key][2]+1
self.chunk_size[key][2] = self.chunk_bounds[key][5]-self.chunk_bounds[key][4]+1
# Verify that local grid size is not smaller than ghost cell size
assert (self.chunk_size[key][0]>=self.nghx and
self.chunk_size[key][1]>=self.nghy and
self.chunk_size[key][2]>=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'''
from .field import Field3d
import numpy as np
# Block execution of some processors if IO is limited
self._baton_wait(io_limit)
# Determine which chunks are to be loaded by the current processor
ip_beg_ext = self.ip*self.chunks_per_proc[0]
jp_beg_ext = self.jp*self.chunks_per_proc[1]
kp_beg_ext = self.kp*self.chunks_per_proc[2]
ip_end_ext = ip_beg_ext+self.chunks_per_proc[0]-1
jp_end_ext = jp_beg_ext+self.chunks_per_proc[1]-1
kp_end_ext = kp_beg_ext+self.chunks_per_proc[2]-1
# Get the total size of the field to be loaded
ib = self.proc_grid_ext[key][0][ip_beg_ext]
ie = self.proc_grid_ext[key][1][ip_end_ext]
jb = self.proc_grid_ext[key][2][jp_beg_ext]
je = self.proc_grid_ext[key][3][jp_end_ext]
kb = self.proc_grid_ext[key][4][kp_beg_ext]
ke = self.proc_grid_ext[key][5][kp_end_ext]
nxl = ie-ib+1
nyl = je-jb+1
nzl = ke-kb+1
# Allocate an array to hold the entire field
data = np.empty(
(nxl+2*self.nghx,
nyl+2*self.nghy,
nzl+2*self.nghz),dtype=self.precision)
# Compute origin of subfield
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][2]+(kb-1-self.nghz)*self.spacing[2])
# Create a Field3d
self.field[key] = Field3d(data,origin,self.spacing)
# Go through each chunk to be read and construct the field
for ip_ext in range(ip_beg_ext,ip_end_ext+1):
for jp_ext in range(jp_beg_ext,jp_end_ext+1):
for kp_ext in range(kp_beg_ext,kp_end_ext+1):
# Determine rank of the chunk to be read
rank_ext = self.rank_from_position(ip_ext,jp_ext,kp_ext,external=True)
# Compute bounds of this chunk
ib_ext = self.proc_grid_ext[key][0][ip_ext]
ie_ext = self.proc_grid_ext[key][1][ip_ext]
jb_ext = self.proc_grid_ext[key][2][jp_ext]
je_ext = self.proc_grid_ext[key][3][jp_ext]
kb_ext = self.proc_grid_ext[key][4][kp_ext]
ke_ext = self.proc_grid_ext[key][5][kp_ext]
nxl_ext = ie_ext-ib_ext+1
nyl_ext = je_ext-jb_ext+1
nzl_ext = ke_ext-kb_ext+1
# Read data and insert it
subfield = self.func_load(rank_ext,key)
self.field[key].insert_subfield(subfield)
# Continue execution of waiting processors if IO was limited
self._baton_pass(io_limit)
# Exchange ghost cells
self.exchange_ghost_cells(key)
# Initialize symmetries and impose BC
self.init_symmetries(key)
self.impose_boundary_conditions(key)
# Syncronize processes if requested
if barrier: self.comm.Barrier()
def differentiate(self,key,axis,key_out=None):
assert axis<3, "'axis' must be one of 0,1,2."
if key_out is None:
key_out = key+('x','y','z')[axis]
origin = list(self.origin)
shifting_state = self.shifting_state(key,axis=axis)
if shifting_state==-1:
padding = 'after'
origin[axis] += 0.5*self.spacing[axis]
elif shifting_state==0:
padding = 'before'
origin[axis] -= 0.5*self.spacing[axis]
elif shifting_state==1:
padding = 'before'
origin[axis] -= 0.5*self.spacing[axis]
else:
raise ValueError("Invalid shifting state.")
self.field[key_out] = self.field[key].derivative(axis,padding=padding)
self.origin[key_out] = tuple(origin)
self.spacing[key_out] = self.spacing[key].copy()
self.symmetries[key_out] = self.symmetries[key].copy()
# TBD: copy everything field specific
# TBD: make sure processor distribution is fine
if axis==0:
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]
elif axis==1:
self.symmetries[key_out][1,0,1] = -self.symmetries[key_out][1,0,1]
self.symmetries[key_out][1,2,1] = -self.symmetries[key_out][1,2,1]
elif axis==2:
self.symmetries[key_out][1,1,0] = -self.symmetries[key_out][1,1,0]
self.symmetries[key_out][1,1,2] = -self.symmetries[key_out][1,1,2]
# Exchange ghost cells and set boundary conditions
self.exchange_ghost_cells(key_out)
self.impose_boundary_conditions(key_out)
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.'''
import numpy as np
if key_out is None:
key_out = key
else:
self.origin[key_out] = self.origin[key].copy()
self.spacing[key_out] = self.spacing[key].copy()
# Compute radius of Gaussian filter
radius = self.field[key].gaussian_filter_radius(sigma,truncate=truncate)
if not iterate:
# Assert that we have sufficient amount of ghost cells
assert all([self.num_ghost[ii]>=radius[ii] for ii in range(3)]),\
"Too few ghost cells for stencil: {}, {}".format(self.num_ghost,radius)
niter = 1
else:
# Determine number of iterations required for current ghost cell setup
sigma = np.array(sigma)
niter = 1
while not all([self.num_ghost[ii]>=radius[ii] for ii in range(3)]):
sigma /= np.sqrt(2)
niter *= 2
radius = self.field[key].gaussian_filter_radius(sigma,truncate=truncate)
assert all([radius[ii]>0 if sigma[ii]>0.0 else True for ii in range(3)]),\
"Iterative procedure leads to invalid stencil radius: "\
"increase number of ghost cells. {}".format(radius)
print('Iterations: {}, stencil radius: {}'.format(niter,radius))
for iiter in range(niter):
# Filter field: if key_out is None, perform operation inplace
self.field[key_out] = self.field[key].gaussian_filter(sigma,
truncate=truncate,border='constant',const_val=0.0)
# Exchange ghost cells and set boundary conditions
self.exchange_ghost_cells(key_out)
self.impose_boundary_conditions(key_out)
# Iterate inplace from now on
key = key_out
def vtk_contour(self,key,val):
'''Compute isocontour for chunks.'''
if any([self.num_ghost[ii]>1 for ii in range(3)]):
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)
for ii in range(3))
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):
'''Extracts a plane from field.'''
if any([self.num_ghost[ii]>1 for ii in range(3)]):
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)
for ii in range(3))
return self.field[key].extract_subfield(
idx_origin,numpoints).vtk_slice(normal,origin)
else:
return self.field[key].vtk_slice(normal,origin)
return
def rank_from_position(self,ip,jp,kp,external=False):
if external:
nyp,nzp = self.nyp_ext,self.nzp_ext
else:
nyp,nzp = self.nyp,self.nzp
return ip*nyp*nzp+jp*nzp+kp
def position_from_rank(self,rank,external=False):
if external:
nyp,nzp = self.nyp_ext,self.nzp_ext
else:
nyp,nzp = self.nyp,self.nzp
ip = rank//(nyp*nzp)
jp = (rank//nzp)%nyp
kp = rank%nzp
return (ip,jp,kp)
def shifting_state(self,key,axis=None):
if axis is None:
return tuple(self.shifting_state(key,axis=ii) for axis in range(3))
assert axis<3, "'axis' must be one of 0,1,2."
return int(round((self.origin[key][axis]-self.bounds[2*axis])/(0.5*self.spacing[axis])))
def exchange_ghost_cells(self,key):
'''Communicates all ghost cells of specified field'''
# Trigger non-blocking communication:
# Communicate faces (6 faces)
self._communicate_ghost_cells(key,(-1,0,0)) # left
self._communicate_ghost_cells(key,(+1,0,0)) # right
self._communicate_ghost_cells(key,(0,-1,0)) # down
self._communicate_ghost_cells(key,(0,+1,0)) # up
self._communicate_ghost_cells(key,(0,0,-1)) # front
self._communicate_ghost_cells(key,(0,0,+1)) # back
# Communicate edges (12 edges)
self._communicate_ghost_cells(key,(-1,-1,0)) # left,down
self._communicate_ghost_cells(key,(-1,0,-1)) # left,front
self._communicate_ghost_cells(key,(-1,+1,0)) # left,up
self._communicate_ghost_cells(key,(-1,0,+1)) # left,back
self._communicate_ghost_cells(key,(+1,-1,0)) # right,down
self._communicate_ghost_cells(key,(+1,0,-1)) # right,front
self._communicate_ghost_cells(key,(+1,+1,0)) # right,up
self._communicate_ghost_cells(key,(+1,0,+1)) # right,back
self._communicate_ghost_cells(key,(0,-1,-1)) # down,front
self._communicate_ghost_cells(key,(0,-1,+1)) # down,back
self._communicate_ghost_cells(key,(0,+1,-1)) # up,front
self._communicate_ghost_cells(key,(0,+1,+1)) # up,back
# Communicate corners (8 corners)
self._communicate_ghost_cells(key,(-1,-1,-1)) # left,down,front
self._communicate_ghost_cells(key,(-1,-1,+1)) # left,down,back
self._communicate_ghost_cells(key,(-1,+1,-1)) # left,up,front
self._communicate_ghost_cells(key,(-1,+1,+1)) # left,up,back
self._communicate_ghost_cells(key,(+1,-1,-1)) # right,down,front
self._communicate_ghost_cells(key,(+1,-1,+1)) # right,down,back
self._communicate_ghost_cells(key,(+1,+1,-1)) # right,up,front
self._communicate_ghost_cells(key,(+1,+1,+1)) # right,up,back
def impose_boundary_conditions(self,key):
'''Imposes symmetry boundary conditions on each non-periodic wall'''
self._symmetrize_wall(key,(-1,0,0))
self._symmetrize_wall(key,(+1,0,0))
self._symmetrize_wall(key,(0,-1,0))
self._symmetrize_wall(key,(0,+1,0))
self._symmetrize_wall(key,(0,0,-1))
self._symmetrize_wall(key,(0,0,+1))
def _communicate_ghost_cells(self,key,positionDst):
'''Triggers communication of ghost cells'''
import numpy as np
assert np.max(positionDst)<=1 and np.min(positionDst)>=-1, "communicate_ghost_cells: "\
"invalid neighbor position {}".format(positionDst)
# [send/recv] get the rank of the neighbor where data is to be sent to
# The position is passed as values -1,0,+1, but need to be converted to array indices
ip_dst = positionDst[0]+1
jp_dst = positionDst[1]+1
kp_dst = positionDst[2]+1
ip_src = -positionDst[0]+1
jp_src = -positionDst[1]+1
kp_src = -positionDst[2]+1
rank_dst = self.nghbr[ip_dst,jp_dst,kp_dst]
rank_src = self.nghbr[ip_src,jp_src,kp_src]
# [send/recv] create a tag
tag = ip_dst*100+jp_dst*10+kp_dst
# [send/recv] get the indices of data to be sent/received
nxl,nyl,nzl = self.chunk_size[key]
if positionDst[0]==-1:
ii_src = slice(self.nghx,2*self.nghx)
ii_dst = slice(self.nghx+nxl,2*self.nghx+nxl)
elif positionDst[0]==0:
ii_src = slice(self.nghx,self.nghx+nxl)
ii_dst = slice(self.nghx,self.nghx+nxl)
elif positionDst[0]==1:
ii_src = slice(nxl,nxl+self.nghx)
ii_dst = slice(0,self.nghx)
else:
raise ValueError('Invalid direction for ghost cell exchange: {}'.format(positionDst[0]))
if positionDst[1]==-1:
jj_src = slice(self.nghy,2*self.nghy)
jj_dst = slice(self.nghy+nyl,2*self.nghy+nyl)
elif positionDst[1]==0:
jj_src = slice(self.nghy,self.nghy+nyl)
jj_dst = slice(self.nghy,self.nghy+nyl)
elif positionDst[1]==1:
jj_src = slice(nyl,nyl+self.nghy)
jj_dst = slice(0,self.nghy)
else:
raise ValueError('Invalid direction for ghost cell exchange: {}'.format(positionDst[1]))
if positionDst[2]==-1:
kk_src = slice(self.nghz,2*self.nghz)
kk_dst = slice(self.nghz+nzl,2*self.nghz+nzl)
elif positionDst[2]==0:
kk_src = slice(self.nghz,self.nghz+nzl)
kk_dst = slice(self.nghz,self.nghz+nzl)
elif positionDst[2]==1:
kk_src = slice(nzl,nzl+self.nghz)
kk_dst = slice(0,self.nghz)
else:
raise ValueError('Invalid direction for ghost cell exchange: {}'.format(positionDst[2]))
# [send/recv] Create a list for requests
reqsend = None
reqrecv = None
# [send] now send the data to the neighbor, but only if there is one!
# Communication must be done non-blocking, and we can use upper-case routines since this is a numpy array
if rank_dst>=0:
sendbuf = np.ascontiguousarray(self.field[key].data[ii_src,jj_src,kk_src])
reqsend = self.comm.Isend(sendbuf,dest=rank_dst,tag=tag)
# [recv] the corresponding receive: results are stored in a buffer which will be assigned to the parent array later
if rank_src>=0:
recvbuf = np.zeros(
(ii_dst.stop-ii_dst.start,
jj_dst.stop-jj_dst.start,
kk_dst.stop-kk_dst.start),dtype=self.precision)
reqrecv = self.comm.Irecv(recvbuf,source=rank_src,tag=tag)
# [recv] wait for data to be received
if reqrecv is not None:
reqrecv.wait()
self.field[key].data[ii_dst,jj_dst,kk_dst] = recvbuf
# [send] wait for data to be sent
if reqsend is not None:
reqsend.wait()
def _symmetrize_wall(self,key,positionBd):
import numpy as np
assert np.max(positionBd)<=1 and np.min(positionBd)>=-1, "symmetrize_wall: invalid neighbor "\
"position {}".format(positionBd)
assert np.count_nonzero(positionBd)==1, "symmetrize_wall: only face direction is accepted "\
"(no edges or corners)"
inghbr = positionBd[0]+1
jnghbr = positionBd[1]+1
knghbr = positionBd[2]+1
if self.nghbr[inghbr,jnghbr,knghbr]<0:
sl_dst = [slice(None),slice(None),slice(None)]
sl_src = [slice(None),slice(None),slice(None)]
for axis in range(3):
if positionBd[axis]==-1: # lower boundary
# index of first point within the domain
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
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
sl_dst[axis] = slice(0,idx,1)
sl_src[axis] = slice(2*idx-1,idx-1,-1)
else: # point on boundary
sl_dst[axis] = slice(0,idx,1)
sl_src[axis] = slice(2*idx,idx,-1)
break
elif positionBd[axis]==1: # upper boundary
# index of last point within the domain
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
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
sl_dst[axis] = slice(idx+1,self.field[key].numpoints[axis],1)
sl_src[axis] = slice(idx,2*idx+1-self.field[key].numpoints[axis],-1)
else: # point on boundary
sl_dst[axis] = slice(idx+1,self.field[key].numpoints[axis],1)
sl_src[axis] = slice(idx-1,2*idx-self.field[key].numpoints[axis],-1)
break
self.field[key].data[tuple(sl_dst)] = self.symmetries[key][inghbr,jnghbr,knghbr]*\
self.field[key].data[tuple(sl_src)]
def _baton_wait(self,batch_size,tag=420):
'''Block execution until an empty message from rank-batch_wait
is received (issued by _baton_pass)'''
from mpi4py import MPI
if batch_size is not None:
if self.rank>=batch_size:
source = self.rank-batch_size
self.comm.recv(source=source,tag=tag)
def _baton_pass(self,batch_size,tag=420):
'''Sends an empty message to rank+batch_wait to unblock its
execution (issued by _baton_wait)'''
from mpi4py import MPI
if batch_size is not None:
dest = self.rank+batch_size
if dest<self.comm.Get_size():
data = None
self.comm.send(data,dest=dest,tag=tag)
class GatherIterator:
'''Sends 'data' sequentially to 'root' which can iterate over it
without gathering all data at once. Every process which is not 'root'
receives None.'''
def __init__(self,data,comm=None,root=0,barrier=False):
from mpi4py import MPI
comm = MPI.COMM_WORLD if comm is None else comm
self.comm = comm
self.rank = comm.Get_rank()
self.size = comm.Get_size()
self.root = root
self.barrier = barrier
self.iter = 0
self.data = data
def __iter__(self):
return self
def __next__(self):
if self.iter>=self.size:
if self.barrier: self.comm.Barrier()
raise StopIteration
if self.rank==self.root:
if self.rank==self.iter:
r = self.data
else:
r = self.comm.recv(source=self.iter,tag=0)
else:
if self.rank==self.iter:
self.comm.send(self.data,dest=self.root,tag=0)
r = None
self.iter += 1
return r
def parprint(*args, **kwargs):
from mpi4py import MPI
rank = kwargs.pop('rank') if 'rank' in kwargs else 0
comm = kwargs.pop('comm') if 'comm' in kwargs else MPI.COMM_WORLD
if get_rank(comm=comm)==rank:
print(*args, **kwargs)
def get_rank(comm=None):
from mpi4py import MPI
comm = MPI.COMM_WORLD if comm is None else comm
return comm.Get_rank()
def is_root(comm=None,root=0):
from mpi4py import MPI
comm = MPI.COMM_WORLD if comm is None else comm
return comm.Get_rank()==root
def gather(data,comm=None):
from mpi4py import MPI
comm = MPI.COMM_WORLD if comm is None else comm
return comm.gather(data,root=0)

52
ucf.py
View File

@ -460,6 +460,7 @@ class UCF:
self.__currentSetParams = 0 self.__currentSetParams = 0
self.__currentSetNumElements = 0 self.__currentSetNumElements = 0
# TBD: refactor -> UCFCollection, UCFTarFile constructors etc
class UCFTarFile: class UCFTarFile:
'''Superclass for instances of ucf.tar files.''' '''Superclass for instances of ucf.tar files.'''
def __init__(self,file_tar,file_index=None): def __init__(self,file_tar,file_index=None):
@ -511,6 +512,33 @@ class UCFTarFile:
period = params.getfloat('geometry',key_bound[axis]) period = params.getfloat('geometry',key_bound[axis])
isPeriodic = params.getboolean('geometry',key_period[axis]) isPeriodic = params.getboolean('geometry',key_period[axis])
return period if isPeriodic else None return period if isPeriodic else None
def procgrid(self):
self._initialize_buffer_proc()
proc = {}
offset = 0
for key in ('u','v','w','p','s'):
proc[key] = tuple(self._buffer_proc[offset])
offset += 1
return proc
def origin(self):
self._initialize_buffer_grid()
origin = {}
offset = 0
for key in ('u','v','w','p','s'):
origin[key] = (
self._buffer_grid[offset][0][0],
self._buffer_grid[offset][1][0],
self._buffer_grid[offset][2][0])
offset +=1
return origin
def spacing(self):
self._initialize_buffer_grid()
spacing = (
self._buffer_grid[0][0][1]-self._buffer_grid[0][0][0],
self._buffer_grid[0][1][1]-self._buffer_grid[0][1][0],
self._buffer_grid[0][2][1]-self._buffer_grid[0][2][0],
)
return spacing
def gidx_from_key(self,key): def gidx_from_key(self,key):
if key[0]=='u': return 0 if key[0]=='u': return 0
elif key[0]=='v': return 1 elif key[0]=='v': return 1
@ -569,18 +597,10 @@ class UCFSnapshot(UCFTarFile):
dx,dy,dz = (grid[0][1]-grid[0][0], dx,dy,dz = (grid[0][1]-grid[0][0],
grid[1][1]-grid[1][0], grid[1][1]-grid[1][0],
grid[2][1]-grid[2][0]) grid[2][1]-grid[2][0])
xo,yo,zo = (grid[0][chunk['ibeg']]-chunk['ighost']*dx, xo,yo,zo = (grid[0][chunk['ibeg']-1]-(chunk['ighost'])*dx,
grid[1][chunk['jbeg']]-chunk['ighost']*dy, grid[1][chunk['jbeg']-1]-(chunk['ighost'])*dy,
grid[2][chunk['kbeg']]-chunk['ighost']*dz) grid[2][chunk['kbeg']-1]-(chunk['ighost'])*dz)
# Per convention periodicity requires ghost cells for Field3d return Field3d(chunk['data'],(xo,yo,zo),(dx,dy,dz))
if keep_ghost:
period = tuple(self.period(ii) if self.nproc(axis=ii)==1
else None for ii in range(0,3))
else:
period = (None,None,None)
return Field3d(chunk['data'],(xo,yo,zo),(dx,dy,dz),period)
def field_complete(self):
return
def particles(self,select_col=None): def particles(self,select_col=None):
'''Returns particle data as Particles object.''' '''Returns particle data as Particles object.'''
from .particle import Particles from .particle import Particles
@ -607,10 +627,10 @@ class UCFSnapshot(UCFTarFile):
return ijk[0]*self.nproc(axis=1)*self.nproc(axis=2)+ijk[1]*self.nproc(axis=2)+ijk[2] return ijk[0]*self.nproc(axis=1)*self.nproc(axis=2)+ijk[1]*self.nproc(axis=2)+ijk[2]
def dset_from_key(self,key): def dset_from_key(self,key):
assert self.has_key(key), "Snapshot does not contain requested key." assert self.has_key(key), "Snapshot does not contain requested key."
if key[0]=='u': return 0 if key[0]=='u': return 1
elif key[0]=='v': return 1 elif key[0]=='v': return 2
elif key[0]=='w': return 2 elif key[0]=='w': return 3
elif key[0]=='p': return 3 elif key[0]=='p': return 4
elif key[0]=='s': return int(key[1:]) elif key[0]=='s': return int(key[1:])
else: raise ValueError("Invalid value of 'key'.") else: raise ValueError("Invalid value of 'key'.")
def has_key(self,key): def has_key(self,key):