Create test for extracting a MALI subdomain from a larger domain

compass/landice/__init__.py

@@ -12,6 +12,7 @@
 from compass.landice.tests.ismip6_run import Ismip6Run
 from compass.landice.tests.kangerlussuaq import Kangerlussuaq
 from compass.landice.tests.koge_bugt_s import KogeBugtS
+from compass.landice.tests.mesh_modifications import MeshModifications
 from compass.landice.tests.mismipplus import MISMIPplus
 from compass.landice.tests.thwaites import Thwaites
 from compass.mpas_core import MpasCore
@@ -42,5 +43,6 @@ def __init__(self):
         self.add_test_group(Ismip6Run(mpas_core=self))
         self.add_test_group(Kangerlussuaq(mpas_core=self))
         self.add_test_group(KogeBugtS(mpas_core=self))
+        self.add_test_group(MeshModifications(mpas_core=self))
         self.add_test_group(MISMIPplus(mpas_core=self))
         self.add_test_group(Thwaites(mpas_core=self))

compass/landice/mesh.py

@@ -1,3 +1,4 @@
+import sys
 import time
 
 import jigsawpy
@@ -12,6 +13,61 @@
 from netCDF4 import Dataset
 
 
+def mpas_flood_fill(seed_mask, grow_mask, cellsOnCell, nEdgesOnCell,
+                    grow_iters=sys.maxsize):
+    """
+    Flood-fill for mpas meshes using mpas cells.
+
+    Parameters
+    ----------
+    seed_mask : numpy.ndarray
+        Integer array of locations from which to flood fill
+        0 = invalid, 1 = valid
+
+    grow_mask : numpy.ndarray
+        Integer array of locations valid for growing into
+        0 = invalid, 1 = valid
+
+    cellsOnCell : numpy.ndarray
+        cellsOnCell array from the mpas mesh
+
+    nEdgesOnCell : numpy.ndarray
+        nEdgesOnCell array from the mpas mesh
+
+    grow_iters : integer
+        optional argument limiting the number of iterations
+        over which to extend the mask
+
+    Returns
+    -------
+    keep_mask : numpy.ndarray
+        mask calculated by the flood fill routine,
+        where cells connected to seed_mask
+        are 1 and everything else is 0.
+    """
+
+    iter = 0
+    keep_mask = seed_mask.copy()
+    n_mask_cells = keep_mask.sum()
+    for iter in range(grow_iters):
+        mask_ind = np.nonzero(keep_mask == 1)[0]
+        print(f'iter={iter}, keep_mask size={keep_mask.sum()}')
+        new_keep_mask = keep_mask.copy()
+        for iCell in mask_ind:
+            neighs = cellsOnCell[iCell, :nEdgesOnCell[iCell]] - 1
+            neighs = neighs[neighs >= 0]  # drop garbage cell
+            for jCell in neighs:
+                if grow_mask[jCell] == 1:
+                    new_keep_mask[jCell] = 1
+        keep_mask = new_keep_mask.copy()
+        n_mask_cells_new = keep_mask.sum()
+        if n_mask_cells_new == n_mask_cells:
+            break
+        n_mask_cells = n_mask_cells_new
+        iter += 1
+    return keep_mask
+
+
 def gridded_flood_fill(field, iStart=None, jStart=None):
     """
     Generic flood-fill routine to create mask of connected elements

compass/landice/tests/mesh_modifications/__init__.py

@@ -0,0 +1,19 @@
+from compass.landice.tests.mesh_modifications.subdomain_extractor import (
+    SubdomainExtractor,
+)
+from compass.testgroup import TestGroup
+
+
+class MeshModifications(TestGroup):
+    """
+    A test group for automating modifications to existing meshes
+    """
+    def __init__(self, mpas_core):
+        """
+        mpas_core : compass.landice.Landice
+            the MPAS core that this test group belongs to
+        """
+        super().__init__(mpas_core=mpas_core,
+                         name='mesh_modifications')
+
+        self.add_test_case(SubdomainExtractor(test_group=self))

compass/landice/tests/mesh_modifications/subdomain_extractor/__init__.py

@@ -0,0 +1,36 @@
+import sys
+from importlib import resources
+
+import numpy as np
+
+from compass.landice.tests.mesh_modifications.subdomain_extractor.extract_region import (  # noqa
+    ExtractRegion,
+)
+from compass.testcase import TestCase
+
+
+class SubdomainExtractor(TestCase):
+    """
+    A class for a test case that extracts a subdomain from a larger domain
+    """
+
+    def __init__(self, test_group):
+        """
+        Create the test case
+
+        Parameters
+        ----------
+        test_group : compass.landice.tests.mesh_modifications.MeshModifications
+            The test group that this test case belongs to
+
+        """
+        name = 'subdomain_extractor'
+        super().__init__(test_group=test_group, name=name)
+
+        self.add_step(ExtractRegion(test_case=self))
+
+    # no configure() method is needed
+
+    # no run() method is needed
+
+    # no validate() method is needed