added integration routine: useful for statistics
This commit is contained in:
parent
91c5460b3a
commit
4a33b5f00a
71
field.py
71
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.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)]
|
||||
|
|
|
|||
69
parallel.py
69
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:
|
||||
|
|
|
|||
Loading…
Reference in New Issue