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) self.data = self.data.astype(dtype,copy=False)
return return
def derivative(self,axis,only_keep_interior=False,shift_origin='before'): def derivative(self,axis,only_keep_interior=False,shift_origin='before'):
'''Computes derivative wrt to direction 'axis' with 2nd order finite differences '''Computes derivative wrt to direction 'axis' with 2nd order finite differences
centered between the origin grid points''' centered between the origin grid points'''
@ -342,42 +341,6 @@ 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):
# 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): def shift_origin(self,rel_shift,only_keep_interior=False):
'''Shifts the origin of a field by multiple of spacing.''' '''Shifts the origin of a field by multiple of spacing.'''
from scipy import ndimage from scipy import ndimage

View File

@ -491,11 +491,6 @@ class PPP:
gc.collect() gc.collect()
return return
def merge(self,key,stride=(1,1,1)):
'''Returns the complete (possibly downsampled) field at rank 0.'''
return
def save_state(self,filename): def save_state(self,filename):
import pickle import pickle
fout = filename+'.{:05d}.pickle'.format(self.rank) fout = filename+'.{:05d}.pickle'.format(self.rank)
@ -516,33 +511,105 @@ class PPP:
with open(fout,"wb") as f: with open(fout,"wb") as f:
pickle.dump(payload,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) '''Returns the field (only its interior + some ghost cells for plotting)
as a vtk mesh.''' as a vtk mesh. If other ghost cells are required, use one of the _subfield
return self._subfield_for_vtk(key,no_ghost).to_vtk() 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.''' '''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.''' '''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): def _subfield_for_vtk(self,key,stride=(1,1,1)):
num_ghost = np.array(self.num_ghost) '''Extracts a subfield from chunk which has one leading and no trailing
idx_origin = num_ghost ghost cell except when there is no trailing neighbor. This guarantees
dim = np.array(self.field[key].dim(axis=None)) non-overlapping contours.'''
dim -= 2*num_ghost no_upper_ghost = [True,True,True]
if not no_ghost: if self.ip==self.nxp-1: no_upper_ghost[0]=False
idx_origin -= 1 if self.jp==self.nyp-1: no_upper_ghost[1]=False
dim += 1 if self.kp==self.nzp-1: no_upper_ghost[2]=False
if self.ip==self.nxp-1: dim[0]+=1 return self._subfield(key,stride,(1,1,1),no_upper_ghost=no_upper_ghost)
if self.jp==self.nyp-1: dim[1]+=1
if self.kp==self.nzp-1: dim[2]+=1 def _subfield_for_vtk_merge(self,key,stride=(1,1,1)):
idx_origin = tuple(idx_origin) '''Extracts a subfield from chunk which has one leading and no trailing
dim = tuple(dim) ghost cell except when there is no trailing neighbor. This guarantees
return self.field[key].extract_subfield(idx_origin,dim) 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): def rank_from_position(self,ip,jp,kp,external=False):
if external: if external: