added integration routine: useful for statistics

This commit is contained in:
Michael Krayer 2021-05-28 14:26:22 +02:00
parent 91c5460b3a
commit 4a33b5f00a
2 changed files with 136 additions and 4 deletions

View File

@ -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.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 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."
endpoint = [origin[ii]+(numpoints[ii]-1)*spacing[ii] for ii in range(0,3)]

View File

@ -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: