added striding to parallel routines in order to gather a complete field at rank 0
This commit is contained in:
parent
3366816adf
commit
0e48c9b382
37
field.py
37
field.py
|
|
@ -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
|
||||
|
|
|
|||
119
parallel.py
119
parallel.py
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Reference in New Issue