update(optimize_splitting_mask): add active_only flag

autotest/test_model_splitter.py

@@ -221,7 +221,10 @@ def test_metis_splitting_with_lak_sfr(function_tmpdir):
 @requires_exe("mf6")
 @requires_pkg("pymetis")
 @requires_pkg("h5py")
+@requires_pkg("scipy")
 def test_save_load_node_mapping_structured(function_tmpdir):
+    import pymetis
+
     sim_path = get_example_data_path() / "mf6-freyberg"
     new_sim_path = function_tmpdir / "mf6-freyberg/split_model"
     hdf_file = new_sim_path / "node_map.hdf5"
@@ -235,7 +238,9 @@ def test_save_load_node_mapping_structured(function_tmpdir):
     original_heads = sim.get_model().output.head().get_alldata()[-1]
 
     mfsplit = Mf6Splitter(sim)
-    array = mfsplit.optimize_splitting_mask(nparts=nparts)
+    array = mfsplit.optimize_splitting_mask(
+        nparts=nparts, active_only=True, options=pymetis.Options(seed=42, contig=1)
+    )
     new_sim = mfsplit.split_model(array)
     new_sim.set_sim_path(new_sim_path)
     new_sim.write_simulation()

flopy/discretization/grid.py

@@ -647,7 +647,7 @@ def _set_neighbors(self, reset=False, method="rook"):
         ----------
         reset : bool
             flag to recalculate neighbors
-        method: str
+        method : str
             "rook" for shared edges and "queen" for shared vertex
 
         Returns
@@ -721,6 +721,7 @@ def neighbors(self, node=None, **kwargs):
         """
         method = kwargs.pop("method", None)
         reset = kwargs.pop("reset", False)
+
         if method is None:
             self._set_neighbors(reset=reset)
         else:

flopy/mf6/utils/model_splitter.py

@@ -602,7 +602,7 @@ def construct_modelgrid(f, name, grid_type):
         mfs._allow_splitting = False
         return mfs
 
-    def optimize_splitting_mask(self, nparts):
+    def optimize_splitting_mask(self, nparts, active_only=False, options=None):
         """
         Method to create a splitting array with a balanced number of active
         cells per model. This method uses the program METIS and pymetis to
@@ -611,11 +611,22 @@ def optimize_splitting_mask(self, nparts):
         Parameters
         ----------
         nparts: int
+            number of parts to split the model in to
+        active_only : bool
+            only consider active cells when building adjacency graph. Default is False
+        options : None or pymetis.Options
+            optional pymetis.Options class that gets passed through to the
+            pymetis.part_graph() function. Example
+            `options=pymetis.Options(seed=42, contig=1)`
 
         Returns
         -------
             np.ndarray
         """
+        if active_only:
+            import_optional_dependency("scipy")
+            from scipy.interpolate import griddata
+
         pymetis = import_optional_dependency(
             "pymetis",
             "please install pymetis using: "
@@ -629,6 +640,7 @@ def optimize_splitting_mask(self, nparts):
             ncpl = self._modelgrid.ncpl
             shape = self._modelgrid.shape[1:]
         else:
+            nlay = 1
             ncpl = self._modelgrid.nnodes
             shape = self._modelgrid.shape
         idomain = self._modelgrid.idomain
@@ -679,16 +691,53 @@ def optimize_splitting_mask(self, nparts):
                 else:
                     adv_pkg_weights[nodes] += 1
 
-        for nn, neighbors in self._modelgrid.neighbors().items():
+        # filter active cells here to avoid changes in the mapping algorithm
+        neighbors = self._modelgrid.neighbors(reset=True)
+
+        if active_only:
+            node_map = {}
+            iact = np.where(np.sum(idomain, axis=0), 1, 0)
+            inactive = np.where(iact == 0)[0]
+            for k in inactive:
+                neighbors.pop(k)
+            neighbors = {
+                k : [i for i in v if i not in inactive] for k, v in neighbors.items()
+            }
+
+            cnt = 0
+            for ix, isact in enumerate(iact):
+                if isact > 0:
+                    node_map[ix] = cnt
+                    cnt += 1
+
+        for nn, neigh in sorted(neighbors.items()):
             weight = np.count_nonzero(idomain[:, nn])
             adv_weight = adv_pkg_weights[nn]
             weights.append(weight + adv_weight)
-            graph.append(np.array(neighbors, dtype=int))
+            if active_only:
+                neigh = [node_map[n] for n in neigh]
+            graph.append(np.array(neigh, dtype=int))
 
         n_cuts, membership = pymetis.part_graph(
-            nparts, adjacency=graph, vweights=weights
+            nparts, adjacency=graph, vweights=weights, options=options
         )
         membership = np.array(membership, dtype=int)
+
+        if active_only:
+            if len(inactive) > 0:
+                xc = self._modelgrid.xcellcenters.ravel()
+                yc = self._modelgrid.ycellcenters.ravel()
+                axc = xc[list(node_map.keys())]
+                ayc = yc[list(node_map.keys())]
+                ixc = xc[inactive]
+                iyc = yc[inactive]
+                data = griddata((axc, ayc), membership, (ixc, iyc), method="nearest")
+
+                member_array = np.full((ncpl,), -1)
+                member_array[list(node_map.keys())] = membership[list(node_map.values())]
+                member_array[inactive] = data
+                membership = member_array
+
         if laks:
             for lak in laks:
                 idx = np.asarray(lak_array == lak).nonzero()[0]
@@ -892,7 +941,10 @@ def _get_nlay_shape_models(self, model, array):
             tuple : (nlay, grid_shape)
         """
         if array.size == model.modelgrid.size:
-            nlay = model.modelgrid.nlay
+            if self._modelgrid.grid_type in ("structured", "vertex"):
+                nlay = model.modelgrid.nlay
+            else:
+                nlay = 1
             shape = self._shape
 
         elif array.size == model.modelgrid.ncpl: