diff --git a/autotest/test_model_splitter.py b/autotest/test_model_splitter.py index 6e1b93d84..c08995c2c 100644 --- a/autotest/test_model_splitter.py +++ b/autotest/test_model_splitter.py @@ -437,7 +437,7 @@ def test_control_records(function_tmpdir): split_ws.mkdir() with set_dir(split_ws): mfsplit = flopy.mf6.utils.Mf6Splitter(sim) - new_sim = mfsplit.split_model(arr) + new_sim = mfsplit.split_model(arr, split_ws) ml1 = new_sim.get_model("model_1") diff --git a/flopy/mf6/utils/model_splitter.py b/flopy/mf6/utils/model_splitter.py index f55a6a738..d2b38de34 100644 --- a/flopy/mf6/utils/model_splitter.py +++ b/flopy/mf6/utils/model_splitter.py @@ -181,6 +181,7 @@ def __init__(self, sim, modelname=None): self._allow_splitting = True self._fdigits = 1 + self._keep_external = True # multi-model splitting attr self._multimodel_exchange_gwf_names = {} @@ -678,12 +679,13 @@ def optimize_splitting_mask(self, nparts, active_only=False, options=None, verbo else: cellids = package.packagedata.array.cellid if self._modelgrid.grid_type == "structured": - cellids = [(0, i[1], i[2]) for i in cellids] + # skip disconnected cells in SFR package + cellids = [(0, i[1], i[2]) for i in cellids if i != (-1, -1, -1)] nodes = self._modelgrid.get_node(cellids) elif self._modelgrid.grid_type == "vertex": - nodes = [i[1] for i in cellids] + nodes = [i[1] for i in cellids if i != (-1, -1)] else: - nodes = [i[0] for i in cellids] + nodes = [i[0] for i in cellids if i != (-1,)] if isinstance(package, (modflow.ModflowGwflak, modflow.ModflowGwtlkt, modflow.ModflowGwelke)): lakenos = package.connectiondata.array.ifno + 1 @@ -758,15 +760,24 @@ def optimize_splitting_mask(self, nparts, active_only=False, options=None, verbo cellids2 = recarray.cellid2 _, nodes1 = self._cellid_to_layer_node(cellids1) _, nodes2 = self._cellid_to_layer_node(cellids2) - mnums1 = membership[nodes1] - mnums2 = membership[nodes2] - ev = np.equal(mnums1, mnums2) - if np.all(ev): - continue - idx = np.asarray(~ev).nonzero()[0] - mnum_to = mnums1[idx] - adj_nodes = nodes2[idx] - membership[adj_nodes] = mnum_to + cnt = 0 + while cnt < len(nodes1): + mnums1 = membership[nodes1] + mnums2 = membership[nodes2] + ev = np.equal(mnums1, mnums2) + if np.all(ev): + break + idx = np.asarray(~ev).nonzero()[0] + mnum_to = mnums1[idx] + adj_nodes = np.array(nodes2)[idx] + membership[adj_nodes] = mnum_to + cnt += 1 + + if cnt == len(nodes1): + raise AssertionError( + "Cannot uniquely spilt around HFB boundaries, try another " + "value for nparts" + ) return membership.reshape(shape) @@ -1536,6 +1547,9 @@ def _remap_array(self, item, mfarray, mapped_data, **kwargs): # external array tmp = fnames[lay].split(".") filename = f"{'.'.join(tmp[:-1])}.{mkey :0{self._fdigits}d}.{tmp[-1]}" + folder_path = (self._new_sim.sim_path / filename).parent + if not folder_path.exists(): + folder_path.mkdir(parents=True) cr = { "filename": filename, @@ -1621,6 +1635,9 @@ def _remap_mflist( if how == 3 and new_recarray is not None: tmp = fname.split(".") filename = f"{'.'.join(tmp[:-1])}.{mkey :0{self._fdigits}d}.{tmp[-1]}" + folder_path = (self._new_sim.sim_path / filename).parent + if not folder_path.exists(): + folder_path.mkdir(parents=True) new_recarray = { "data": new_recarray, @@ -2040,6 +2057,10 @@ def _remap_sfr(self, package, mapped_data): cellids[messy_idx] = rcids layers, nodes = self._cellid_to_layer_node(cellids) + # adjust the messy_idx layer number + if layers is not None and messy_idx: + layers[messy_idx] = 0 + new_model, new_node = self._get_new_model_new_node(nodes) for mkey, model in self._model_dict.items(): @@ -3519,6 +3540,15 @@ def _remap_package(self, package, ismvr=False): if "stress_period_data" in data: if not data["stress_period_data"]: continue + + if self._keep_external: + shape = self._grid_info[mdl][0] + if len(shape) == 2: + max_cols = shape[1] + else: + max_cols = shape[0] + self._new_sim.simulation_data.max_columns_of_data = max_cols + paks[mdl] = pak_cls( self._model_dict[mdl], pname=package.name[0], **data ) @@ -3842,7 +3872,7 @@ def create_multi_model_exchanges(self, mname0, mname1): filename=filename, ) - def split_model(self, array): + def split_model(self, array, sim_ws=None): """ User method to split a model based on an array @@ -3852,6 +3882,10 @@ def split_model(self, array): integer array of new model numbers. Array must either be of dimension (NROW, NCOL), (NCPL), or (NNODES for unstructured grid models). + sim_ws : PathLike or str + optional directory path for writing the new simulation to. This parameter + is recommended when the model contains external files and the user would + like to preserve external linkages while splitting. Returns ------- @@ -3863,6 +3897,10 @@ def split_model(self, array): "is part of a split simulation" ) + if sim_ws is None: + self._keep_external = False + sim_ws = self._sim.sim_path + # set number formatting string for file paths array = np.array(array).astype(int) s = str(np.max(array)) @@ -3872,7 +3910,7 @@ def split_model(self, array): if self._new_sim is None: self._new_sim = modflow.MFSimulation( - version=self._sim.version, exe_name=self._sim.exe_name, sim_ws=self._sim.sim_path + version=self._sim.version, exe_name=self._sim.exe_name, sim_ws=sim_ws ) self._create_sln_tdis() @@ -3898,6 +3936,9 @@ def split_model(self, array): **nam_options[mkey], ) + if not self._keep_external: + self._model.set_all_data_internal(check_data=True) + for package in self._model.packagelist: paks = self._remap_package(package) @@ -3909,7 +3950,7 @@ def split_model(self, array): return self._new_sim - def split_multi_model(self, array): + def split_multi_model(self, array, sim_ws=None): """ Method to split integrated models such as GWF-GWT or GWF-GWE models. Note: this method will not work to split multiple connected GWF models @@ -3920,6 +3961,10 @@ def split_multi_model(self, array): integer array of new model numbers. Array must either be of dimension (NROW, NCOL), (NCPL), or (NNODES for unstructured grid models). + sim_ws : PathLike or str + optional directory path for writing the new simulation to. This parameter + is recommended when the model contains external files and the user would + like to preserve external linkages while splitting. Returns ------- @@ -3986,7 +4031,7 @@ def split_multi_model(self, array): new_sim = self.split_model(array) for mname in model_names[1:]: self.switch_models(modelname=mname, remap_nodes=False) - new_sim = self.split_model(array) + new_sim = self.split_model(array, sim_ws=sim_ws) for mbase in model_names[1:]: for label in model_labels: