added integration routine: useful for statistics

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)]

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) 
+            numpoints = self.chunk_size(axis=None)
-                            for ii in range(3))
             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: