fix(model_splitter): multiple bug fixes (#2728)

* set data internal for model splitting

* fix(split_model): update external file handling for model splitter

Maintain external file paths when user passes optional sim_ws parameter

* updates(model_splitter.py): multiple bug fixes

* filter (-1,) SFR cellids from splitting mask creation in optimize_splitting_mask
* adjust (-1,) SFR cellid layer number to 0 in _remap_sfr to avoid index errors
* handle external model files by allowing user to specify new simulation workspace `sim_ws` to `split_model` and `split_multi_model`
* dynamically adjust `max_columns_of_data` when external model files are maintained in the split models

* update external ascii test

* Add "angrot" to the offsets dictionary

Author: jlarsen-usgs

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")
 

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)
 
@@ -1012,6 +1023,7 @@ def _remap_nodes(self, array):
                 self._offsets[m] = {
                     "xorigin": self._modelgrid.xvertices[rmax + 1, cmin],
                     "yorigin": self._modelgrid.yvertices[rmax + 1, cmin],
+                    "angrot": self._modelgrid.angrot
                 }
                 # get new nrow and ncol information
                 nrow = (rmax - rmin) + 1
@@ -1035,6 +1047,7 @@ def _remap_nodes(self, array):
                 self._offsets[m] = {
                     "xorigin": self._modelgrid.xoffset,
                     "yorigin": self._modelgrid.yoffset,
+                    "angrot": self._modelgrid.angrot
                 }
 
         new_ncpl = {}
@@ -1536,6 +1549,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 +1637,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 +2059,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 +3542,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 +3874,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 +3884,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 +3899,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 +3912,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 +3938,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 +3952,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 +3963,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 +4033,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: