diff --git a/autotest/test_model_splitter.py b/autotest/test_model_splitter.py index 11785c59e..c34c7b699 100644 --- a/autotest/test_model_splitter.py +++ b/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() diff --git a/flopy/discretization/grid.py b/flopy/discretization/grid.py index 1af13b345..b8a387a6e 100644 --- a/flopy/discretization/grid.py +++ b/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: diff --git a/flopy/mf6/utils/model_splitter.py b/flopy/mf6/utils/model_splitter.py index 3603dc02d..97d5afdab 100644 --- a/flopy/mf6/utils/model_splitter.py +++ b/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: