From 87fda17de956bb54196975ab26c8510681b62768 Mon Sep 17 00:00:00 2001 From: jlarsen-usgs Date: Fri, 4 Sep 2026 11:42:59 -0700 Subject: [PATCH 1/7] fix(ModflowSfr2): support single reach models * single reach models crashed due to np.diff(reach_data).max() > 0 comparison, fix for this issue breaks up this comparison by taking the diff first and evaluating if the diff returns an empty result prior to taking the max() --- flopy/modflow/mfsfr2.py | 52 +++++++++++++++++++++++------------------ 1 file changed, 29 insertions(+), 23 deletions(-) diff --git a/flopy/modflow/mfsfr2.py b/flopy/modflow/mfsfr2.py index ecb735107..ea0b9b6ed 100644 --- a/flopy/modflow/mfsfr2.py +++ b/flopy/modflow/mfsfr2.py @@ -458,10 +458,13 @@ def __init__( self.reach_data[n] = reach_data[n] # assign node numbers if there are none (structured grid) - if np.diff(self.reach_data.node).max() == 0 and self.parent.has_package("DIS"): - # first make kij list - lrc = np.array(self.reach_data)[["k", "i", "j"]].tolist() - self.reach_data["node"] = self.parent.dis.get_node(lrc) + diff = np.diff(self.reach_data.node) + if len(diff) > 0: + if diff.max() == 0 and self.parent.has_package("DIS"): + # first make kij list + lrc = np.array(self.reach_data)[["k", "i", "j"]].tolist() + self.reach_data["node"] = self.parent.dis.get_node(lrc) + # assign unique ID and outreach columns to each reach self.reach_data.sort(order=["iseg", "ireach"]) new_cols = { @@ -493,27 +496,30 @@ def __init__( for n in segment_data[i].dtype.names: self.segment_data[i][n] = segment_data[i][n] # compute outreaches if nseg and outseg columns have non-default values - if ( - np.diff(self.reach_data.iseg).max() != 0 - and np.max(list(set(self.graph.keys()))) != 0 - and np.max(list(set(self.graph.values()))) != 0 - ): - if len(self.graph) == 1: - self.segment_data[0]["nseg"] = 1 - self.reach_data["iseg"] = 1 - - consistent_seg_numbers = ( - len(set(self.reach_data.iseg).difference(set(self.graph.keys()))) == 0 - ) - if not consistent_seg_numbers: - warnings.warn( - "Inconsistent segment numbers of reach_data and segment_data" + diff = np.diff(self.reach_data.iseg) + if len(diff) > 0: + if ( + np.diff(diff).max() != 0 + and np.max(list(set(self.graph.keys()))) != 0 + and np.max(list(set(self.graph.values()))) != 0 + ): + if len(self.graph) == 1: + self.segment_data[0]["nseg"] = 1 + self.reach_data["iseg"] = 1 + + consistent_seg_numbers = ( + len(set(self.reach_data.iseg).difference(set(self.graph.keys()))) + == 0 ) + if not consistent_seg_numbers: + warnings.warn( + "Inconsistent segment numbers of reach_data and segment_data" + ) - # first convert any not_a_segment_values to 0 - for v in self.not_a_segment_values: - self.segment_data[0].outseg[self.segment_data[0].outseg == v] = 0 - self.set_outreaches() + # first convert any not_a_segment_values to 0 + for v in self.not_a_segment_values: + self.segment_data[0].outseg[self.segment_data[0].outseg == v] = 0 + self.set_outreaches() self.channel_geometry_data = channel_geometry_data self.channel_flow_data = channel_flow_data From 1b2228e74c81f3cd4bab8315074989d761cdc4f0 Mon Sep 17 00:00:00 2001 From: jlarsen-usgs Date: Fri, 4 Sep 2026 11:56:20 -0700 Subject: [PATCH 2/7] update np.diff(diff).max() to diff.max() --- flopy/modflow/mfsfr2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flopy/modflow/mfsfr2.py b/flopy/modflow/mfsfr2.py index ea0b9b6ed..628b0b38a 100644 --- a/flopy/modflow/mfsfr2.py +++ b/flopy/modflow/mfsfr2.py @@ -499,7 +499,7 @@ def __init__( diff = np.diff(self.reach_data.iseg) if len(diff) > 0: if ( - np.diff(diff).max() != 0 + diff.max() != 0 and np.max(list(set(self.graph.keys()))) != 0 and np.max(list(set(self.graph.values()))) != 0 ): From cce415afb026c59c01592ac164317bebbef241ca Mon Sep 17 00:00:00 2001 From: jlarsen-usgs Date: Fri, 4 Sep 2026 12:04:33 -0700 Subject: [PATCH 3/7] Add single reach test case to test_sfr.py --- autotest/test_sfr.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/autotest/test_sfr.py b/autotest/test_sfr.py index 999342f61..77022567d 100644 --- a/autotest/test_sfr.py +++ b/autotest/test_sfr.py @@ -935,3 +935,26 @@ def test_mf2005(function_tmpdir, namfile): np.array_equal(str2.segment_data[0][name], m.str.segment_data[0][name]) is True ) + + +def test_single_reach_sfr_pkg(function_tmpdir): + """Test for a single segment/reach SFR package""" + sfrfiletxt = ( + "1 1 0 0 86400.00000000 0.00010000 0 0\n" + "1 5 1 1 1 100.0\n" + "1 0 0\n" + "1 1 0 0 100.0 0 0 0.3 0.025\n" + "1.0 1.0 87.5 10.0\n" + "1.0 1.0 86.5 10.0\n" + ) + sfrfile = io.StringIO(sfrfiletxt) + m = Modflow("junk", model_ws=function_tmpdir) + sfr = ModflowSfr2.load(sfrfile, model=m) + assert len(sfr.segment_data[0]) == 1 + assert len(sfr.reach_data) == 1 + + sfrfile2 = function_tmpdir / "junk.sfr" + sfr.write_file() + sfr = ModflowSfr2.load(sfrfile2, model=m) + assert len(sfr.segment_data[0]) == 1 + assert len(sfr.reach_data) == 1 From 4983bb4d0795c09d2875843ad8d46def7b04fbb4 Mon Sep 17 00:00:00 2001 From: jlarsen-usgs Date: Fri, 4 Sep 2026 14:55:45 -0700 Subject: [PATCH 4/7] Fix(_remap_obs): fix support for GWF observation types with 2 cellids --- flopy/mf6/utils/model_splitter.py | 31 +++++++++++-------------------- 1 file changed, 11 insertions(+), 20 deletions(-) diff --git a/flopy/mf6/utils/model_splitter.py b/flopy/mf6/utils/model_splitter.py index d2676a712..07ddd7512 100644 --- a/flopy/mf6/utils/model_splitter.py +++ b/flopy/mf6/utils/model_splitter.py @@ -2929,10 +2929,10 @@ def _remap_obs(self, package, mapped_data, remapper, pkg_type=None): "structured", "vertex", ): - layers2 = [ + layers2 = np.array([ cid[0] if cid is not None else None for cid in cellid2 - ] + ], dtype=object) if self._modelgrid.grid_type == "structured": cellid2 = [ ( @@ -2949,7 +2949,7 @@ def _remap_obs(self, package, mapped_data, remapper, pkg_type=None): ] node2 = self._modelgrid.get_node( - list(cellid2[conv_idx]) + [cellid2[cv_ix] for cv_ix in conv_idx] ) new_node2 = np.full( (len(recarray),), None, dtype=object @@ -2964,10 +2964,10 @@ def _remap_obs(self, package, mapped_data, remapper, pkg_type=None): if ix in conv_idx: continue else: - new_node2.append(None) - new_model2.append(new_model1[ix]) + new_node2[ix] = None + new_model2[ix] = int(new_model1[ix]) - if not np.allclose(new_model1, new_model2): + if not np.allclose(new_model1, new_model2.astype(int)): raise AssertionError( "One or more observation records cross model boundaries" ) @@ -2978,20 +2978,11 @@ def _remap_obs(self, package, mapped_data, remapper, pkg_type=None): for mkey, model in self._model_dict.items(): idx = np.asarray(new_model2 == mkey).nonzero() tmp_node = new_node2[idx] - cidx = np.asarray(tmp_node != None).nonzero() # noqa: E711 - tmp_cellid = model.modelgrid.get_lrc( - tmp_node[cidx].to_list() + tmp_layers = layers2[idx] + cidx = np.asarray(tmp_node != None).nonzero() # noqa: E711 + tmp_cellid = self._new_node_to_cellid( + model, tmp_node, tmp_layers, cidx ) - if self._modelgrid.grid_type in ( - "structured", - "vertex", - ): - tmp_layers = layers2[cidx] - tmp_cellid = [ - (tmp_layers[ix],) + cid[1:] - for ix, cid in enumerate(tmp_cellid) - ] - tmp_node[cidx] = tmp_cellid new_cellid2[idx] = tmp_node else: @@ -3169,7 +3160,7 @@ def _new_node_to_cellid(self, model, new_node, layers, idx): new_node = new_node[idx].astype(int) if self._modelgrid.grid_type == "structured": - new_node += layers[idx] * model.modelgrid.ncpl + new_node += layers[idx].astype(int) * model.modelgrid.ncpl new_cellids = model.modelgrid.get_lrc(new_node.astype(int)) elif self._modelgrid.grid_type == "vertex": new_cellids = [tuple(cid) for cid in zip(layers[idx], new_node)] From 323b5fe8b25ac9a04fd744f84b463ebaff78cf0b Mon Sep 17 00:00:00 2001 From: jlarsen-usgs Date: Fri, 4 Sep 2026 16:04:46 -0700 Subject: [PATCH 5/7] add testing for vertex and structured scenarios --- autotest/test_model_splitter.py | 131 ++++++++++++++++++++++++++++++++ 1 file changed, 131 insertions(+) diff --git a/autotest/test_model_splitter.py b/autotest/test_model_splitter.py index 4246d0513..10df3ff02 100644 --- a/autotest/test_model_splitter.py +++ b/autotest/test_model_splitter.py @@ -1981,3 +1981,134 @@ def test_reconstruct_recarray(function_tmpdir): vrec = vrecarray[ix] for name in recarray.dtype.names: assert rec[name] == vrec[name], "Recarray reconstruction failed" + + +@requires_exe("mf6") +def test_fjf_obs_structured(function_tmpdir): + sim_ws = function_tmpdir + split_ws = function_tmpdir / "split_model" + + nlay = 1 + nrow = 10 + ncol = 10 + delc = np.full((nrow,), 10) + delr = np.full((ncol,), 10) + top = np.full((nrow, ncol), 50) + botm = np.full((nlay, nrow, ncol), 0) + idomain = np.ones(botm.shape, dtype=int) + + sim = flopy.mf6.MFSimulation(sim_ws=sim_ws) + ims = flopy.mf6.ModflowIms(sim, complexity="SIMPLE") + tdis = flopy.mf6.ModflowTdis(sim) + + gwf = flopy.mf6.ModflowGwf(sim) + dis = flopy.mf6.ModflowGwfdis( + gwf, + nlay=nlay, + nrow=nrow, + ncol=ncol, + delr=delr, + delc=delc, + top=top, + botm=botm, + idomain=idomain, + ) + npf = flopy.mf6.ModflowGwfnpf(gwf, k=5, k33=1) + ic = flopy.mf6.ModflowGwfic(gwf, strt=top - 5) + + chd_rec = [(0, i, 0, 45) for i in range(nrow)] + ghb_rec = [(0, i, 9, 35, 25) for i in range(nrow)] + chd = flopy.mf6.ModflowGwfchd( + gwf, stress_period_data={0: chd_rec}, pname="chd_left" + ) + ghb = flopy.mf6.ModflowGwfghb( + gwf, stress_period_data={0: ghb_rec}, pname="ghb_right" + ) + + obs = flopy.mf6.ModflowUtlobs( + gwf, + continuous={ + "obs.csv": [ + ("faceflow1", "flow-ja-face", (0, 3, 3), (0, 3, 4)), + ("h1", "head", (0, 3, 4)), + ("faceflow2", "flow-ja-face", (0, 6, 6), (0, 6, 7)), + ("h2", "head", (0, 6, 7)), + ] + }, + ) + + sim.write_simulation() + sim.run_simulation() + + array = np.zeros(top.shape, dtype=int) + array[:, (ncol // 2) :] = 1 + mfs = flopy.mf6.utils.Mf6Splitter(sim) + new_sim = mfs.split_model(array, sim_ws=split_ws) + new_sim.write_simulation() + success, _ = new_sim.run_simulation() + assert success, "split model run failed, obs remapping may have failed" + + for mdl in range(2): + obs_ra = new_sim.get_model(f"{gwf.name}_{mdl}").obs.continuous.data[ + f"obs_{mdl}.csv" + ] + assert len(obs_ra) == 2, "number remapped of observation records incorrect" + id2 = obs_ra[obs_ra.obsname == f"faceflow{mdl + 1}"].id2[0] + assert isinstance(id2, tuple), "obs id2 datatype is incorrect" + + if mdl == 0: + v_ra = obs.continuous.data["obs.csv"] + v_id = v_ra[v_ra.obsname == f"faceflow{mdl + 1}"].id2[0] + + assert id2 == v_id, "id2 cellid not properly set" + + +@requires_exe("mf6") +def test_fjf_obs_vertex(function_tmpdir): + sim_ws = get_example_data_path() / "mf6" / "test003_gwfs_disv" + split_ws = function_tmpdir / "split_model" + + sim = flopy.mf6.MFSimulation.load(sim_ws=sim_ws) + sim.set_sim_path(function_tmpdir) + + gwf = sim.get_model() + modelgrid = gwf.modelgrid + ncpl = modelgrid.ncpl + + obs = flopy.mf6.ModflowUtlobs( + gwf, + continuous={ + "obs.csv": [ + ("faceflow1", "flow-ja-face", (0, 24), (0, 34)), + ("h1", "head", (0, 34)), + ("faceflow2", "flow-ja-face", (0, 64), (0, 74)), + ("h2", "head", (0, 74)), + ] + }, + ) + + sim.write_simulation() + sim.run_simulation() + + array = np.zeros((ncpl,), dtype=int) + array[(ncpl // 2) :] = 1 + + mfs = flopy.mf6.utils.Mf6Splitter(sim) + new_sim = mfs.split_model(array, sim_ws=split_ws) + new_sim.write_simulation() + success, _ = new_sim.run_simulation() + assert success, "split model run failed, obs remapping may have failed" + + for mdl in range(2): + obs_ra = new_sim.get_model(f"{gwf.name}_{mdl}").obs.continuous.data[ + f"obs_{mdl}.csv" + ] + assert len(obs_ra) == 2, "number remapped of observation records incorrect" + id2 = obs_ra[obs_ra.obsname == f"faceflow{mdl + 1}"].id2[0] + assert isinstance(id2, tuple), "obs id2 datatype is incorrect" + + if mdl == 0: + v_ra = obs.continuous.data["obs.csv"] + v_id = v_ra[v_ra.obsname == f"faceflow{mdl + 1}"].id2[0] + + assert id2 == v_id, "id2 cellid not properly set" From e57fa916eb6e477705406a881fac924828c2ba98 Mon Sep 17 00:00:00 2001 From: jlarsen-usgs Date: Fri, 18 Sep 2026 11:58:14 -0700 Subject: [PATCH 6/7] linting --- autotest/test_model_splitter.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/autotest/test_model_splitter.py b/autotest/test_model_splitter.py index 844fad134..8c0fbe262 100644 --- a/autotest/test_model_splitter.py +++ b/autotest/test_model_splitter.py @@ -2002,7 +2002,7 @@ def test_fjf_obs_structured(function_tmpdir): tdis = flopy.mf6.ModflowTdis(sim) gwf = flopy.mf6.ModflowGwf(sim) - + dis = flopy.mf6.ModflowGwfdis( gwf, nlay=nlay, @@ -2090,7 +2090,7 @@ def test_fjf_obs_vertex(function_tmpdir): sim.write_simulation() sim.run_simulation() - + array = np.zeros((ncpl,), dtype=int) array[(ncpl // 2) :] = 1 @@ -2113,7 +2113,7 @@ def test_fjf_obs_vertex(function_tmpdir): v_id = v_ra[v_ra.obsname == f"faceflow{mdl + 1}"].id2[0] assert id2 == v_id, "id2 cellid not properly set" - + def test_sfr_none_cells(function_tmpdir): sim_ws = function_tmpdir / "sfr_none_test" From a0ba8691c12460a0b336f6f26311ebdf61523536 Mon Sep 17 00:00:00 2001 From: jlarsen-usgs Date: Wed, 23 Sep 2026 11:41:46 -0700 Subject: [PATCH 7/7] update(optimize_splitting_mask): add edge weighting to graph for LAK packages * refactor weighting code for HFB edge weights * add testing --- autotest/test_model_splitter.py | 25 ++++++++++ flopy/mf6/utils/model_splitter.py | 79 ++++++++++++++++++++++++------- 2 files changed, 86 insertions(+), 18 deletions(-) diff --git a/autotest/test_model_splitter.py b/autotest/test_model_splitter.py index 8c0fbe262..041ecc430 100644 --- a/autotest/test_model_splitter.py +++ b/autotest/test_model_splitter.py @@ -291,6 +291,14 @@ def test_metis_splitting_with_lak_sfr(function_tmpdir): "optimize_splitting_mask is not correcting for lakes properly" ) + counts = [np.count_nonzero(array == i) for i in np.unique(array)] + ratio = np.max(counts) / np.mean(counts) + if ratio > 1.05: + raise AssertionError( + "load balancing ratio is higher than expected, LAK_EDGE_WEIGHTS " + "may be too low or not properly applied" + ) + new_sim = mfsplit.split_model(array) new_sim.set_sim_path(function_tmpdir / "split_model") new_sim.write_simulation() @@ -2272,3 +2280,20 @@ def test_sfr_none_cells(function_tmpdir): if cid == (-1, -1, -1): none_cnt += 1 assert none_cnt == none_cells, "Splitter not correctly assigning SFR None cells" + + +@requires_pkg("pymetis") +def test_optimal_mask_force_lak_rebalancing(): + sim_path = get_example_data_path() / "mf6" / "test045_lake2tr" + + sim = MFSimulation.load(sim_ws=sim_path) + mfsplit = Mf6Splitter(sim) + array = mfsplit.optimize_splitting_mask(nparts=15) + + counts = [np.count_nonzero(array == i) for i in np.unique(array)] + ratio = np.max(counts) / np.mean(counts) + if ratio > 1.75: + raise AssertionError( + "load balancing ratio is higher than expected, lak remapping " + "rebalancing adjustments should be checked" + ) diff --git a/flopy/mf6/utils/model_splitter.py b/flopy/mf6/utils/model_splitter.py index c6f1f24ad..19df3656e 100644 --- a/flopy/mf6/utils/model_splitter.py +++ b/flopy/mf6/utils/model_splitter.py @@ -123,9 +123,10 @@ } -# cost of cutting a cell face that a horizontal flow barrier crosses, -# relative to the cost of one for every other face +# cost of cutting a graph edge for special conditions, +# relative to the cost of one for every other graph edge HFB_EDGE_WEIGHT = 1000 +LAK_EDGE_WEIGHT = 1000 class Mf6Splitter: @@ -611,10 +612,6 @@ def optimize_splitting_mask(self, nparts, active_only=False, options=None, verbo ------- np.ndarray """ - if active_only: - import_optional_dependency("sklearn") - from sklearn.neighbors import NearestNeighbors - pymetis = import_optional_dependency( "pymetis", "please install pymetis using: " @@ -691,17 +688,21 @@ def optimize_splitting_mask(self, nparts, active_only=False, options=None, verbo # for k in inactive: [neighbors.pop(k) for k in inactive] node_map = {i: ix for ix, i in enumerate(np.where(iact > 0)[0])} + # todo: this might need be better as node_map[k]: [existing logic] for k, v in neighbors.items() neighbors = { - k : [node_map[i] for i in v if i not in inactive] for k, v in neighbors.items() + node_map[k] : [node_map[i] for i in v if i not in inactive] for k, v in neighbors.items() } + neighbors = dict(sorted(neighbors.items())) + else: + neighbors = dict(sorted(neighbors.items())) + node_map = {i: ix for ix, i in enumerate(neighbors.keys())} if verbose: print("Creating graph and weights") - neighbors = dict(sorted(neighbors.items())) - weights = [np.count_nonzero(idomain[:, nn]) + adv_pkg_weights[nn] for nn in neighbors.keys()] + weights = [np.count_nonzero(idomain[:, nn]) + adv_pkg_weights[nn] for nn in node_map.keys()] graph = [np.array(neigh, dtype=int) for neigh in neighbors.values()] - eweights = None + hfb_faces = [] if hfbs: if verbose: print("Weighting horizontal flow barrier faces") @@ -709,24 +710,55 @@ def optimize_splitting_mask(self, nparts, active_only=False, options=None, verbo # the faces it crosses expensive to cut steers the partition around # it. Barriers that are still cut are moved after the partition is # built. - gnode = {node: ix for ix, node in enumerate(neighbors.keys())} - hfb_faces = set() for hfb in hfbs: for recarray in hfb.stress_period_data.data.values(): _, nodes1 = self._cellid_to_layer_node(recarray.cellid1) _, nodes2 = self._cellid_to_layer_node(recarray.cellid2) - for node1, node2 in zip(nodes1, nodes2): - if node1 not in gnode or node2 not in gnode: - continue + nodes1 = [node_map[i] if i in node_map else -1 for i in nodes1] + nodes2 = [node_map[i] if i in node_map else -1 for i in nodes2] + hfb_nodes = np.array([nodes1, nodes2]) + + # filter out barriers that touch inactive cells + vidx = np.where(np.min(hfb_nodes, axis=0) > -1)[0] + hfb_nodes = hfb_nodes[:, vidx] + + # filter out vertical flow barriers + not_vfb = np.where((hfb_nodes[0] - hfb_nodes[1]) != 0)[0] + hfb_nodes = hfb_nodes[:, not_vfb] - hfb_faces.add((gnode[node1], gnode[node2])) - hfb_faces.add((gnode[node2], gnode[node1])) + hfb_faces.extend(zip(hfb_nodes[0], hfb_nodes[1])) + hfb_faces.extend(zip(hfb_nodes[1], hfb_nodes[0])) + hfb_faces = set(hfb_faces) + + lak_faces = [] + if laks: + if verbose: + print("Weighting lak faces") + for lakeno in laks: + lak_nodes = np.where(lak_array == lakeno)[0] + lak_nodes = [node_map[i] for i in lak_nodes] + nodes1, nodes2 = [], [] + for lak_node in lak_nodes: + neighs = neighbors[lak_node] + for nn in neighs: + if nn in lak_nodes: + nodes1.append(lak_node) + nodes2.append(nn) + + lak_faces.extend(zip(nodes1, nodes2)) + + lak_faces = set(lak_faces) + + eweights = None + if hfbs or laks: eweights = [] for node, conns in enumerate(graph): for conn in conns: if (node, int(conn)) in hfb_faces: eweights.append(HFB_EDGE_WEIGHT) + elif (node, int(conn)) in lak_faces: + eweights.append(LAK_EDGE_WEIGHT) else: eweights.append(1) @@ -753,7 +785,18 @@ def optimize_splitting_mask(self, nparts, active_only=False, options=None, verbo if laks: for lak in laks: idx = np.asarray(lak_array == lak).nonzero()[0] - mnum = np.unique(membership[idx])[0] + mnums = np.unique(membership[idx]) #[0] + mnum = mnums[0] + + if len(mnums) > 1: + # if the lake is in multiple parts of the membership array, + # reset it so it's in the part of the array that has the + # greatest coverage of the lake. + lak_membership = membership[idx] + counts = [np.count_nonzero(lak_membership == i) for i in mnums] + ix = np.argmax(counts) + mnum = mnums[ix] + membership[idx] = mnum if hfbs: