diff --git a/field.py b/field.py index 3819f65..38ddeb1 100644 --- a/field.py +++ b/field.py @@ -168,6 +168,42 @@ class Field3d: origin = self.origin data = numpy.gradient(self.data,self.spacing[axis],axis=axis,edge_order=2) return Field3d(data,origin,self.spacing) + + def integral(self,integrate_axis,average=False,ignore_nan=False,return_weights=False,ufunc=None): + '''Computes the integral or average along a given axis applying the + function 'ufunc' to each node.''' + assert isinstance(integrate_axis,(list,tuple,numpy.ndarray)) and len(integrate_axis)==3,\ + "'integrate_axis' must be a tuple/list of length 3." + assert all([isinstance(integrate_axis[ii],(bool,int)) for ii in range(3)]),\ + "'integrate_axis' requires bool values." + assert any(integrate_axis), "'integrate_axis' must contain at least one True." + axes = [] + weight = 1.0 + for axis in range(3): + if integrate_axis[axis]: + axes.append(axis) + if average: + weight *= self.numpoints[axis] + else: + weight *= self.spacing[axis] + axes = tuple(axes) + if ignore_nan: + if average: + weight = numpy.sum(~numpy.isnan(self.data),axis=axes,keepdims=True) + func_sum = numpy.nansum + else: + func_sum = numpy.sum + if ufunc is None: + out = func_sum(self.data,axis=axes,keepdims=True) + else: + assert isinstance(ufunc,numpy.ufunc), "'ufunc' needs to be a numpy ufunc. "\ + "Check out https://numpy.org/doc/stable/reference/ufuncs.html for reference." + assert ufunc.nin==1, "Only ufunc with single input argument are supported for now." + out = func_sum(ufunc(self.data),axis=axes,keepdims=True) + if return_weights: + return (out,weight) + else: + return out/weight def gradient(self,axis,preserve_origin=False,padding=None): return [self.derivative(axis,preserve_origin=preserve_origin,padding=padding) for axis in range(0,3)] @@ -196,6 +232,41 @@ class Field3d: radius.append(int(truncate*sigma_img[ii]+0.5)) return tuple(radius) + def shift_origin(self,rel_shift): + #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.numpoints,numpy.nan) + data = self.data.copy() + origin = list(self.origin) + for axis in range(3): + if abs(rel_shift[axis])=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)] diff --git a/parallel.py b/parallel.py index f33a336..5d9757a 100644 --- a/parallel.py +++ b/parallel.py @@ -358,7 +358,7 @@ class PPP: for axis in range(3): if arg.shape[axis]==1: continue - elif arg.shape[axis]==self.proc_grid[key][2*axis+1][-1]: + elif arg.shape[axis]==self.numpoints(key,axis=axis): pos = self.position_from_rank(self.rank,external=False)[axis] sl_arg[axis] = slice( self.proc_grid[key][2*axis][pos]-1, @@ -366,7 +366,7 @@ class PPP: else: raise ValueError("'arg' must either be singular or match global "\ "grid dimension. (axis={}: got {:d}, expected {:d}".format( - axis,arg.shape[axis],self.proc_grid[key][2*axis+1][-1])) + axis,arg.shape[axis],self.numpoints(key,axis=axis))) # 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_arg = tuple(sl_arg) @@ -378,12 +378,65 @@ class PPP: op(self.field[key].data,arg) return + def integrate(self,key,integrate_axis,ufunc=None,ignore_nan=False,average=False): + '''Computes a global integral or average along a given axis applying the + function 'ufunc' to each node. The result is returned by rank 0, all other + ranks return None.''' + from mpi4py import MPI + import numpy as np + if integrate_axis is None: + integrate_axis = tuple(self.periodicity) + # Compute local integral, but get weights separately + idx_origin = tuple(self.num_ghost[ii]-1 for ii in range(3)) + (integ_local,weights_local) = self.field[key].extract_subfield( + idx_origin,self.chunk_size(key,axis=None)).integral( + integrate_axis,ufunc=ufunc,ignore_nan=ignore_nan, + average=average,return_weights=True) + # Make sure weights_local is a numpy array + weights_local = np.array(weights_local) + # Simple implementation: all processor communicate directly to root + if self.rank==0: + dim_final = tuple(1 if integrate_axis[axis] else self.numpoints(key,axis=axis) for axis in range(3)) + # Allocate a nice spot of memory + integ = np.zeros(dim_final,dtype=integ_local.dtype) + if average: + weights = np.zeros(dim_final,dtype=weights_local.dtype) + # Receive from peers + for rank_src in range(0,self.nproc): + # Determine target array position + pos = self.position_from_rank(rank_src,external=False) + sl = 3*[slice(None)] + for axis in range(3): + if integrate_axis[axis]: + sl[axis] = slice(0,1) + else: + sl[axis] = slice( + self.proc_grid[key][2*axis][pos[axis]]-1, + self.proc_grid[key][2*axis+1][pos[axis]]) + # Receive data and put it in a good spot + if rank_src==0: + integ[sl] += integ_local + if average: + weights[sl] += weights_local + else: + integ[sl] += self.comm.recv(source=rank_src,tag=1) + if average: + weights[sl] += self.comm.recv(source=rank_src,tag=2) + if average: + return integ/weights + else: + return integ/weights_local + else: + self.comm.send(integ_local,dest=0,tag=1) + if average: + self.comm.send(weights_local,dest=0,tag=2) + return None + 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)) + numpoints = self.chunk_size(axis=None) return self.field[key].extract_subfield( idx_origin,numpoints).vtk_contour(val) else: @@ -432,6 +485,14 @@ class PPP: assert axis<3, "'axis' must be one of 0,1,2." return self.field[key].numpoints[axis]-2*self.num_ghost[axis] + def numpoints(self,key,axis=None): + '''Returns the total number of gridpoints across all processors + without ghost cells.''' + if axis is None: + return tuple(self.numpoints(key,axis=ii) for ii in range(3)) + assert axis<3, "'axis' must be one of 0,1,2." + return self.proc_grid[key][2*axis+1][-1] + def exchange_ghost_cells(self,key): '''Communicates all ghost cells of specified field''' # Trigger non-blocking communication: