added ConnectedRegions

field.py

@@ -1,9 +1,9 @@
-import numpy
+import numpy as np
 class Field3d:
     def __init__(self,data,origin,spacing):
         assert len(origin)==3,  "'origin' must be of length 3"
         assert len(spacing)==3, "'spacing' must be of length 3"
-        self.data      = numpy.array(data)
+        self.data      = np.array(data)
         self.origin    = tuple([float(x) for x in origin])
         self.spacing   = tuple([float(x) for x in spacing])
         self.eps_collapse = 1e-8
@@ -39,6 +39,13 @@ class Field3d:
         else:
             return Field3d(self.data*other,self.origin,self.spacing)
 
+    def __truediv__(self,other):
+        if isinstance(other,Field3d):
+            assert self.has_same_grid(other), "Grid mismatch."
+            return Field3d(self.data/other.data,self.origin,self.spacing)
+        else:
+            return Field3d(self.data/other,self.origin,self.spacing)
+
     def  __radd__(self,other):
         return Field3d(other+self.data,self.origin,self.spacing)
 
@@ -72,6 +79,14 @@ class Field3d:
             self.data *= other
         return self
 
+    def __itruediv__(self,other):
+        if isinstance(other,Field3d):
+            assert self.has_same_grid(other), "Grid mismatch."
+            self.data /= other.data
+        else:
+            self.data /= other
+        return self
+
 # TBD: this should return another Field3d object
 #    def __getitem__(self,val):
 #        assert isinstance(val,tuple) and len(val)==3, "Field3d must be indexed by [ii,jj,kk]."
@@ -113,24 +128,24 @@ class Field3d:
         return cls(data,origin,spacing)
 
     @classmethod
-    def allocate(cls,dim,origin,spacing,fill=None,dtype=numpy.float64,pseudo=False):
+    def allocate(cls,dim,origin,spacing,fill=None,dtype=np.float64,pseudo=False):
         '''Allocates an empty field, or a field filled with 'fill'.'''
-        assert isinstance(dim,(tuple,list,numpy.ndarray)) and len(dim)==3,\
+        assert isinstance(dim,(tuple,list,np.ndarray)) and len(dim)==3,\
             "'dim' must be a tuple/list of length 3."
-        assert isinstance(origin,(tuple,list,numpy.ndarray)) and len(origin)==3,\
+        assert isinstance(origin,(tuple,list,np.ndarray)) and len(origin)==3,\
             "'origin' must be a tuple/list of length 3."
-        assert isinstance(spacing,(tuple,list,numpy.ndarray)) and len(spacing)==3,\
+        assert isinstance(spacing,(tuple,list,np.ndarray)) and len(spacing)==3,\
             "'spacing' must be a tuple/list of length 3."
         if pseudo:
-            data = numpy.empty((0,0,0))
+            data = np.empty((0,0,0))
             r = cls(data,origin,spacing)
             r._dim = dim
             return r
         else:
             if fill is None:
-                data = numpy.empty(dim,dtype=dtype)
+                data = np.empty(dim,dtype=dtype)
             else:
-                data = numpy.full(dim,fill,dtype=dtype)
+                data = np.full(dim,fill,dtype=dtype)
             return cls(data,origin,spacing)
 
     def save(self,file,name='Field3d',truncate=False):
@@ -186,16 +201,16 @@ class Field3d:
             return tuple(self.nearest_gridpoint(coord[ii],axis=ii,lower=lower) for ii in range(3))
         assert axis<3, "'axis' must be one of 0,1,2."
         if lower:
-            return numpy.floor((coord+self.eps_collapse-self.origin[axis])/self.spacing[axis]).astype('int')
+            return np.floor((coord+self.eps_collapse-self.origin[axis])/self.spacing[axis]).astype('int')
         else:
-            return numpy.round((coord-self.origin[axis])/self.spacing[axis]).astype('int')
+            return np.round((coord-self.origin[axis])/self.spacing[axis]).astype('int')
 
     def distance_to_nearest_gridpoint(self,coord,axis=None,lower=False):
         if axis is None:
             assert len(coord)==3, "If 'axis' is None, 'coord' must be a tuple/list of length 3."
             return tuple(self.distance_to_nearest_gridpoint(coord[ii],axis=ii,lower=lower) for ii in range(3))
         assert axis<3, "'axis' must be one of 0,1,2."
-        val = numpy.remainder(coord+self.eps_collapse-self.origin[axis],self.spacing[axis])-self.eps_collapse
+        val = np.remainder(coord+self.eps_collapse-self.origin[axis],self.spacing[axis])-self.eps_collapse
         if not lower and val>0.5*self.spacing[axis]:
             val = self.spacing[axis]-val
         return val
@@ -273,8 +288,8 @@ class Field3d:
         else:
             corr_orig = -1
             origin[axis] += 0.5*self.spacing[axis]           
-        data = ndimage.correlate1d(self.data,numpy.array([-1.,1.])/self.spacing[axis],axis=axis,
+        data = ndimage.correlate1d(self.data,np.array([-1.,1.])/self.spacing[axis],axis=axis,
-                                    mode='constant',cval=numpy.nan,origin=corr_orig)
+                                    mode='constant',cval=np.nan,origin=corr_orig)
         if only_keep_interior:
             sl = 3*[slice(None)]
             if shift_origin=='before':
@@ -292,9 +307,9 @@ class Field3d:
     def laplacian(self,only_keep_interior=False):
         '''Computes the Laplacian of a field.'''
         from scipy import ndimage
-        data  = ndimage.correlate1d(self.data,numpy.array([1.,-2.,1.])/self.spacing[0]**2,axis=0,mode='constant',cval=numpy.nan,origin=0)
+        data  = ndimage.correlate1d(self.data,np.array([1.,-2.,1.])/self.spacing[0]**2,axis=0,mode='constant',cval=np.nan,origin=0)
-        data += ndimage.correlate1d(self.data,numpy.array([1.,-2.,1.])/self.spacing[1]**2,axis=1,mode='constant',cval=numpy.nan,origin=0)
+        data += ndimage.correlate1d(self.data,np.array([1.,-2.,1.])/self.spacing[1]**2,axis=1,mode='constant',cval=np.nan,origin=0)
-        data += ndimage.correlate1d(self.data,numpy.array([1.,-2.,1.])/self.spacing[2]**2,axis=2,mode='constant',cval=numpy.nan,origin=0)
+        data += ndimage.correlate1d(self.data,np.array([1.,-2.,1.])/self.spacing[2]**2,axis=2,mode='constant',cval=np.nan,origin=0)
         origin  = list(self.origin)
         if only_keep_interior:
             data = data[1:-1,1:-1,1:-1]
@@ -305,7 +320,7 @@ class Field3d:
     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,\
+        assert isinstance(integrate_axis,(list,tuple,np.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."
@@ -322,14 +337,14 @@ class Field3d:
         axes = tuple(axes)
         if ignore_nan:
             if average:
-                weight = numpy.sum(~numpy.isnan(self.data),axis=axes,keepdims=True)
+                weight = np.sum(~np.isnan(self.data),axis=axes,keepdims=True)
-            func_sum = numpy.nansum
+            func_sum = np.nansum
         else:
-            func_sum = numpy.sum
+            func_sum = np.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. "\
+            assert isinstance(ufunc,np.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)
@@ -341,11 +356,11 @@ class Field3d:
     def gaussian_filter(self,sigma,truncate=4.0,only_keep_interior=False):
         '''Applies a gaussian filter: sigma is standard deviation for Gaussian kernel for each axis.'''
         from scipy import ndimage
-        assert isinstance(sigma,(tuple,list,numpy.ndarray)) and len(sigma)==3,\
+        assert isinstance(sigma,(tuple,list,np.ndarray)) and len(sigma)==3,\
                     "'sigma' must be a tuple/list of length 3"
         # Convert sigma from simulation length scales to grid points as required by ndimage
         sigma_img = tuple(sigma[ii]/self.spacing[ii] for ii in range(3))
-        data = ndimage.gaussian_filter(self.data,sigma_img,truncate=truncate,mode='constant',cval=numpy.nan)
+        data = ndimage.gaussian_filter(self.data,sigma_img,truncate=truncate,mode='constant',cval=np.nan)
         origin  = list(self.origin)
         if only_keep_interior:
             r = self.gaussian_filter_radius(sigma,truncate=truncate)
@@ -356,7 +371,7 @@ class Field3d:
 
     def gaussian_filter_radius(self,sigma,truncate=4.0):
         '''Radius of Gaussian filter. Stencil width is 2*radius+1.'''
-        assert isinstance(sigma,(tuple,list,numpy.ndarray)) and len(sigma)==3,\
+        assert isinstance(sigma,(tuple,list,np.ndarray)) and len(sigma)==3,\
                     "'sigma' must be a tuple/list of length 3"
         # Convert sigma from simulation length scales to grid points as required by ndimage
         sigma_img = tuple(sigma[ii]/self.spacing[ii] for ii in range(3))
@@ -368,7 +383,7 @@ class Field3d:
     def shift_origin(self,rel_shift,only_keep_interior=False):
         '''Shifts the origin of a field by multiple of spacing.'''
         from scipy import ndimage
-        assert isinstance(rel_shift,(tuple,list,numpy.ndarray)) and len(rel_shift)==3,\
+        assert isinstance(rel_shift,(tuple,list,np.ndarray)) and len(rel_shift)==3,\
             "'shift' must be tuple/list with length 3."
         assert all([rel_shift[ii]>=-1.0-self.eps_collapse and rel_shift[ii]<=1.0+self.eps_collapse for ii in range(3)]),\
             "'shift' must be in (-1.0,1.0). {}".format(rel_shift)
@@ -381,14 +396,14 @@ class Field3d:
             elif rel_shift[axis]>0:
                 w = rel_shift[axis] if rel_shift[axis]<=1.0 else 1.0
                 weights = (1.0-w,w)
-                data = ndimage.correlate1d(data,weights,axis=axis,mode='constant',cval=numpy.nan,origin=-1)
+                data = ndimage.correlate1d(data,weights,axis=axis,mode='constant',cval=np.nan,origin=-1)
                 origin[axis] += w*self.spacing[axis]
                 if only_keep_interior:
                     sl[axis] = slice(0,-1)
             else:
                 w = rel_shift[axis] if rel_shift[axis]>=-1.0 else -1.0
                 weights = (-w,1.0+w)
-                data = ndimage.correlate1d(data,weights,axis=axis,mode='constant',cval=numpy.nan,origin=0)
+                data = ndimage.correlate1d(data,weights,axis=axis,mode='constant',cval=np.nan,origin=0)
                 origin[axis] += w*self.spacing[axis]
                 if only_keep_interior:
                     sl[axis] = slice(1,None)
@@ -413,9 +428,9 @@ class Field3d:
         assert all([endpoint[ii]<=self.endpoint(ii) for ii in range(0,3)]), "New end point is out of bounds."
         # Allocate (possibly padded array)
         origin_pad,dim_pad,sl_out = padding(origin,spacing,dim,padding,numpad)
-        data = numpy.zeros(dim_pad,dtype=self.data.dtype)
+        data = np.zeros(dim_pad,dtype=self.data.dtype)
         # Trilinear interpolation
-        if numpy.allclose(spacing,self.spacing):
+        if np.allclose(spacing,self.spacing):
             # spacing is the same: we can construct universal weights for the stencil
             i0,j0,k0 = self.nearest_gridpoint(origin,axis=None,lower=True)
             cx,cy,cz = [self.distance_to_nearest_gridpoint(origin[ii],axis=ii,lower=True)/self.spacing[ii] 
@@ -459,12 +474,12 @@ class Field3d:
     @staticmethod
     def padding(origin,spacing,dim,padding,numpad):
         if isinstance(numpad,int):
-            numpad = numpy.fill(3,numpad,dtype=int)
+            numpad = np.fill(3,numpad,dtype=int)
         else:
-            numpad = numpy.array(numpad,dtype=int)
+            numpad = np.array(numpad,dtype=int)
         assert len(numpad)==3, "'numpad' must be either an integer or tuple/list of length 3."
-        origin_pad    = numpy.array(origin)
+        origin_pad    = np.array(origin)
-        dim_pad = numpy.array(dim)
+        dim_pad = np.array(dim)
         sl_out = [slice(None),slice(None),slice(None)]
         if padding is not None:
             if padding=='before':
@@ -499,7 +514,7 @@ class Field3d:
         if cz>1.0 and cz<1.0+self.eps_collapse: cz=1.0
         assert cx>=0.0 and cy>=0.0 and cz>=0.0, "'rel_dist' must be >=0"
         assert cx<=1.0 and cy<=1.0 and cz<=1.0, "'rel_dist' must be <=1"
-        c = numpy.zeros((2,2,2))
+        c = np.zeros((2,2,2))
         c[0,0,0] = 1.0-(cx+cy+cz)+(cx*cy+cx*cz+cy*cz)-(cx*cy*cz)
         c[1,0,0] = cx-(cx*cy+cx*cz)+(cx*cy*cz)
         c[0,1,0] = cy-(cx*cy+cy*cz)+(cx*cy*cz)
@@ -552,6 +567,115 @@ def gaussian_filter_umean_channel(array,spacing,sigma,truncate=4.0):
     array = ndimage.gaussian_filter1d(array,sigma_img,axis=1,truncate=truncate,mode='mirror')
     return array
 
+class ConnectedRegions:
+    def __init__(self,boolarr,periodicity,connect_diagonals=False,bytes_label=32):
+        assert isinstance(boolarr,np.ndarray) and boolarr.dtype==np.dtype('bool'),\
+               "'boolarr' must be a numpy array of dtype('bool')."
+        assert all([isinstance(x,(bool,int)) for x in periodicity]),\
+                "'periodicity' requires bool values."
+        assert bytes_label in (8,16,32,64),\
+               "'bytes_label' must be one of {8,16,32,64}."
+        self._dim  = boolarr.shape
+        self._ndim = boolarr.ndim
+        assert self._ndim in (2,3),\
+               "'boolarr' must be either two or three dimensional."
+        assert len(periodicity)==self._ndim,\
+               "Length of 'periodicity' must match number of dimensions of data."
+        from scipy import ndimage
+        # Construct connectivity stencil
+        if self._ndim==2:
+            connectivity = np.ones((3,3),dtype='bool')
+            if not connect_diagonals:
+                connectivity[0,0] = False
+                connectivity[0,2] = False
+                connectivity[2,0] = False
+                connectivity[0,2] = False
+        else:
+            connectivity = np.ones((3,3,3),dtype='bool')
+            if not connect_diagonals:
+                connectivity[0,0,0] = False
+                connectivity[2,0,0] = False
+                connectivity[0,2,0] = False
+                connectivity[0,0,2] = False
+                connectivity[2,2,0] = False
+                connectivity[0,2,2] = False
+                connectivity[2,0,2] = False
+                connectivity[2,2,2] = False
+        # Compute labels:
+        # this does not take into account periodic wrapping
+        dtype_label = np.dtype('uint'+str(bytes_label))
+        self.label = np.empty(self._dim,dtype=dtype_label)
+        ndimage.label(boolarr,structure=connectivity,output=self.label)
+        self.count = np.max(self.label)
+        # Merge labels if there are periodic overlaps
+        map_tgt = np.array(range(0,self.count+1),dtype=dtype_label)
+        for axis in range(self._ndim):
+            if not periodicity[axis]:
+                continue
+            # Merge the first and last plane and compute connectivity
+            sl = self._ndim*[slice(None)]
+            sl[axis] = (-1,0)
+            boolarr_ = boolarr[tuple(sl)]
+            label_   = np.empty(boolarr_.shape,dtype=dtype_label)
+            ndimage.label(boolarr_,structure=connectivity,output=label_)
+            for val_ in np.unique(label_):
+                # Get all global labels which are associated to a region 
+                # connected over the boundary
+                global_labels = list(np.unique(self.label[tuple(sl)][label_==val_]))
+                # If there is only one label, nothing needs to be done
+                if len(global_labels)==1:
+                    continue
+                # Determine target label:
+                # this needs to be done recursively because the original
+                # target may already be reassigned
+                tgt = global_labels[0]
+                while tgt!=map_tgt[tgt]:
+                    tgt=map_tgt[tgt]
+                map_tgt[global_labels[1:]] = tgt
+        # Remove gaps from target mapping
+        map_tgt = np.unique(map_tgt,return_inverse=True)[1]
+        # Remap labels
+        self.label = map_tgt[self.label]
+        self.count = np.max(map_tgt)
+
+    def volume(self,label=None):
+        if label is None:
+            return np.bincount(self.label.ravel())
+        else:
+            return np.sum(self.label==label)
+
+    def volume_domain(self):
+        return np.prod(self._dim)
+
+    def discard_regions(self,selection):
+        selection = self._parse_selection(selection)
+        # Map tagged regions to zero in order to discard them
+        map_tgt = np.array(range(0,self.count+1),dtype=self.label.dtype)
+        map_tgt[selection] = 0
+        # Remove gaps from target mapping
+        map_tgt = np.unique(map_tgt,return_inverse=True)[1]
+        # Discard regions
+        self.label = map_tgt[self.label]
+        self.count = np.max(map_tgt)
+
+    def _parse_selection(self,selection):
+        dtype = self.label.dtype
+        if isinstance(selection,int):
+            return np.array(selection,dtype=dtype)
+        elif isinstance(selection,(list,tuple,np.ndarray)):
+            selection = np.array(selection)
+            if selection.dtype==np.dtype('bool'):
+                assert selection.ndim==1 and selection.shape[0]==self.count+1,\
+                        "Boolean indexing must provide count+1 values."
+            else:
+                selection = selection.astype(dtype)
+                assert np.max(selection)<=self.count and np.min(selection)>=0,\
+                        "Entry in selection is out-of-bounds."
+            return selection
+        else:
+            raise ValueError('Invalid input.')
+
+
 class ChunkIterator:
     '''Iterates through all chunks. 'snapshot' must be an instance
     of a class which returns a Field3d from the method call