Merge pull request #560 from xylar/add-grow-mask-to-flood-fill

Add a grow mask to constrain flood fill mask

Author: xylar

conda_package/docs/mesh_conversion.rst

@@ -413,10 +413,14 @@ Computing a Flood-fill Mask
 
 The function :py:func:`mpas_tools.mesh.mask.compute_mpas_flood_fill_mask()`
 and the command-line tool ``compute_mpas_flood_fill_mask``
-fill in a mask, starting with the ocean points closest to the seed points
+fill in a mask, starting with the cell centers closest to the seed points
 given in :py:class:`geometric_features.FeatureCollection` ``fcSeed``.  This
 algorithm runs in serial, and will be more efficient the more seed points
-are provided and the more widely scattered over the ocean they are.
+are provided and the more widely scattered over the mesh they are.
+
+An optional ``daGrow`` argument to the function (not currently available from
+the command-line tool) provides a mask into which the flood fill is allowed to
+grow.  The default is all ones.
 
 The resulting dataset contains a single variable:
 

conda_package/mpas_tools/mesh/mask.py

@@ -387,7 +387,8 @@ def entry_point_compute_mpas_transect_masks():
                  engine=args.engine)
 
 
-def compute_mpas_flood_fill_mask(dsMesh, fcSeed, logger=None, workers=-1):
+def compute_mpas_flood_fill_mask(dsMesh, fcSeed, daGrow=None, logger=None,
+                                 workers=-1):
     """
     Flood fill from the given set of seed points to create a contiguous mask.
     The flood fill operates using cellsOnCell, starting from the cells
@@ -401,6 +402,11 @@ def compute_mpas_flood_fill_mask(dsMesh, fcSeed, logger=None, workers=-1):
     fcSeed : geometric_features.FeatureCollection
         A feature collection containing points at which to start the flood fill
 
+    daGrow : xarray.DataArray, optional
+        A data array of size ``nCells`` with a mask that is 1 anywhere the
+        flood fill is allowed to grow.  The default is that the mask is all
+        ones.
+
     logger : logging.Logger, optional
         A logger for the output if not stdout
 
@@ -426,17 +432,22 @@ def compute_mpas_flood_fill_mask(dsMesh, fcSeed, logger=None, workers=-1):
     if logger is not None:
         logger.info('  Computing flood fill mask on cells:')
 
-    mask = _compute_seed_mask(fcSeed, lon, lat, workers)
+    seedMask = _compute_seed_mask(fcSeed, lon, lat, workers)
 
     cellsOnCell = dsMesh.cellsOnCell.values - 1
 
-    mask = _flood_fill_mask(mask, cellsOnCell)
+    if daGrow is not None:
+        growMask = daGrow.values
+    else:
+        growMask = numpy.ones(dsMesh.sizes['nCells'])
+
+    seedMask = _flood_fill_mask(seedMask, growMask, cellsOnCell)
 
     if logger is not None:
         logger.info('  Adding masks to dataset...')
     # create a new data array for the mask
     masksVarName = 'cellSeedMask'
-    dsMasks[masksVarName] = (('nCells',), numpy.array(mask, dtype=int))
+    dsMasks[masksVarName] = (('nCells',), numpy.array(seedMask, dtype=int))
 
     if logger is not None:
         logger.info('  Done.')
@@ -1183,30 +1194,31 @@ def _compute_seed_mask(fcSeed, lon, lat, workers):
     return mask
 
 
-def _flood_fill_mask(mask, cellsOnCell):
+def _flood_fill_mask(seedMask, growMask, cellsOnCell):
     """
     Flood fill starting with a mask of seed points
     """
 
     maxNeighbors = cellsOnCell.shape[1]
 
     while True:
-        neighbors = cellsOnCell[mask == 1, :]
+        neighbors = cellsOnCell[seedMask == 1, :]
         maskCount = 0
         for iNeighbor in range(maxNeighbors):
             indices = neighbors[:, iNeighbor]
-            # we only want to mask valid neighbors and locations that aren't
-            # already masked
+            # we only want to mask valid neighbors, locations that aren't
+            # already masked, and locations that we're allowed to flood
             indices = indices[indices >= 0]
-            localMask = mask[indices] == 0
+            localMask = numpy.logical_and(seedMask[indices] == 0,
+                                          growMask[indices] == 1)
             maskCount += numpy.count_nonzero(localMask)
             indices = indices[localMask]
-            mask[indices] = 1
+            seedMask[indices] = 1
 
         if maskCount == 0:
             break
 
-    return mask
+    return seedMask
 
 
 def _compute_edge_sign(dsMesh, edgeMask, shape):