added striding to parallel routines in order to gather a complete field at rank 0

This commit is contained in:
Michael Krayer 2021-06-02 09:35:16 +02:00
parent 3366816adf
commit 0e48c9b382
2 changed files with 93 additions and 63 deletions

View File

@ -236,7 +236,6 @@ class Field3d:
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'''
@ -342,42 +341,6 @@ class Field3d:
radius.append(int(truncate*sigma_img[ii]+0.5))
return tuple(radius)
# def shift_origin(self,rel_shift):
# raise NotImplementedError("Routine has not been verified yet.")
# #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

View File

@ -491,11 +491,6 @@ class PPP:
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)
@ -516,33 +511,105 @@ class PPP:
with open(fout,"wb") as f:
pickle.dump(payload,f)
def to_vtk(self,key,no_ghost=False):
def to_vtk(self,key,stride=(1,1,1),merge_at_root=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()
as a vtk mesh. If other ghost cells are required, use one of the _subfield
methods and apply .to_vtk() to the result.'''
from .field import Field3d
if merge_at_root:
if self.rank==0:
stride = np.array(stride)
# Allocate the full output field
origin = np.array(self.origin[key])
spacing = stride*np.array(self.spacing)
dim = (np.array(self.dim(key,axis=None))/stride+1).astype(int)
# the following seems necessary and seems to work, but i didn't think much about it
for axis in range(3):
if self.dim(key,axis=axis)%stride[axis]!=0:
dim[axis]+=1
output = Field3d.allocate(dim,origin,spacing,dtype=self.field[key].dtype())
# Recieve subfields and insert them
output.insert_subfield(self._subfield_for_vtk_merge(key,stride=stride))
for rank_src in range(1,self.nproc):
output.insert_subfield(self.comm.recv(source=rank_src))
return output.to_vtk()
else:
self.comm.send(self._subfield_for_vtk_merge(key,stride=stride),dest=0)
return None
else:
return self._subfield_for_vtk(key,stride=stride).to_vtk()
def vtk_contour(self,key,val):
def vtk_contour(self,key,val,stride=(1,1,1)):
'''Compute isocontour for chunks.'''
return self._subfield_for_vtk(key,False).vtk_contour(val)
return self._subfield_for_vtk(key,stride=stride).vtk_contour(val)
def vtk_slice(self,key,normal,origin):
def vtk_slice(self,key,normal,origin,stride=(1,1,1)):
'''Extracts a plane from field.'''
return self._subfield_for_vtk(key,False).vtk_slice(normal,origin)
return self._subfield_for_vtk(key,stride=stride).vtk_slice(normal,origin)
def _subfield_for_vtk(self,key,no_ghost):
num_ghost = np.array(self.num_ghost)
idx_origin = num_ghost
dim = np.array(self.field[key].dim(axis=None))
dim -= 2*num_ghost
if not no_ghost:
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 _subfield_for_vtk(self,key,stride=(1,1,1)):
'''Extracts a subfield from chunk which has one leading and no trailing
ghost cell except when there is no trailing neighbor. This guarantees
non-overlapping contours.'''
no_upper_ghost = [True,True,True]
if self.ip==self.nxp-1: no_upper_ghost[0]=False
if self.jp==self.nyp-1: no_upper_ghost[1]=False
if self.kp==self.nzp-1: no_upper_ghost[2]=False
return self._subfield(key,stride,(1,1,1),no_upper_ghost=no_upper_ghost)
def _subfield_for_vtk_merge(self,key,stride=(1,1,1)):
'''Extracts a subfield from chunk which has one leading and no trailing
ghost cell except when there is no trailing neighbor. This guarantees
non-overlapping contours.'''
no_lower_ghost = [True,True,True]
no_upper_ghost = [True,True,True]
if self.ip==self.nxp-1: no_upper_ghost[0]=False
if self.jp==self.nyp-1: no_upper_ghost[1]=False
if self.kp==self.nzp-1: no_upper_ghost[2]=False
return self._subfield(key,stride,(1,1,1),no_lower_ghost=no_lower_ghost,
no_upper_ghost=no_upper_ghost)
def _subfield_interior(self,key,stride=(1,1,1)):
return self._subfield(key,stride,(0,0,0))
def _subfield(self,key,stride,num_ghost,no_lower_ghost=(False,False,False),
no_upper_ghost=(False,False,False)):
'''Returns the field with a stride applied.'''
stride = np.array(stride,dtype=int)
num_ghost = np.array(num_ghost,dtype=int)
no_lower_ghost = np.array(no_lower_ghost,dtype=bool)
no_upper_ghost = np.array(no_upper_ghost,dtype=bool)
assert len(stride)==3, "'stride' needs to be of length 3."
assert len(num_ghost)==3, "'num_ghost' needs to be of length 3."
assert len(no_lower_ghost)==3, "'no_lower_ghost' needs to be of length 3."
assert len(no_upper_ghost)==3, "'no_upper_ghost' needs to be of length 3."
assert all(np.array(self.num_ghost)>=stride*num_ghost), "At least stride*num_ghost "\
"ghost cells are required in the original field. Have {}, need {}.".format(
tuple(self.num_ghost),tuple(stride*num_ghost))
ib = self.proc_grid[key][0][self.ip]-1
ie = self.proc_grid[key][1][self.ip]-1
jb = self.proc_grid[key][2][self.jp]-1
je = self.proc_grid[key][3][self.jp]-1
kb = self.proc_grid[key][4][self.kp]-1
ke = self.proc_grid[key][5][self.kp]-1
ijkb = np.array((ib,jb,kb))
ijke = np.array((ie,je,ke))
# Determine the first point to consider in this chunk for the case of output
# without ghost cells
idx_origin = (np.ceil(ijkb/stride)*stride).astype(int)-ijkb+self.num_ghost
# Now the last point to consider under the same circumstances
idx_endpoint = (np.floor(ijke/stride)*stride).astype(int)-ijkb+self.num_ghost
# If we want the subfield to have ghost cells, shift the origin and endpoint
idx_origin -= (~no_lower_ghost).astype(int)*stride*num_ghost
idx_endpoint += (~no_upper_ghost).astype(int)*stride*num_ghost
# Compute the number of points
num_points = ((idx_endpoint-idx_origin)/stride+1).astype(int)
# Some last assertions which should never fail
# print(self.rank,idx_origin,idx_endpoint,num_points,self.field[key].dim(),self.field[key].dim())
assert all(idx_origin>=0)
assert all(idx_endpoint<np.array(self.field[key].dim()))
# Return the subfield
return self.field[key].extract_subfield(idx_origin,num_points,stride=stride)
def rank_from_position(self,ip,jp,kp,external=False):
if external: