From b774d8a61e2f3de779c8d17ca35ab6ef2c5d1797 Mon Sep 17 00:00:00 2001 From: Mike Taves Date: Sat, 30 Sep 2023 22:38:07 +1300 Subject: [PATCH] Accommodate different reach index names for future flopy versions (#83) --- .github/workflows/tests.yml | 3 +- swn/modflow/_base.py | 67 +++++----- swn/modflow/_swnmf6.py | 242 ++++++++++++++++++++---------------- tests/test_modflow6.py | 106 +++++++++------- 4 files changed, 229 insertions(+), 189 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 5b47283..08d9524 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -48,8 +48,9 @@ jobs: run: | pytest -v swn --doctest-modules - - name: Run tests with older packages + - name: Run tests with develop flopy and other older packages run: | + pip install https://github.com/modflowpy/flopy/archive/refs/heads/develop.zip pip install "shapely<2.0" "pandas<2.0" pytest -v -n2 --cov --cov-append diff --git a/swn/modflow/_base.py b/swn/modflow/_base.py index 065a7c3..c04a8fc 100644 --- a/swn/modflow/_base.py +++ b/swn/modflow/_base.py @@ -44,19 +44,34 @@ def __init__(self, logger=None): logger : logging.Logger, optional Logger to show messages. """ - from importlib.util import find_spec - - if not find_spec("flopy"): - raise ImportError(f"{self.__class__.__name__} requires flopy") from ..logger import get_logger, logging + class_name = self.__class__.__name__ + try: + import flopy + except ImportError: + raise ImportError(f"{class_name} requires flopy") + if logger is None: - self.logger = get_logger(self.__class__.__name__) + self.logger = get_logger(class_name) elif isinstance(logger, logging.Logger): self.logger = logger else: raise ValueError(f"expected 'logger' to be Logger; found {type(logger)!r}") - self.logger.info("creating new %s object", self.__class__.__name__) + self.logger.info("creating new %s object", class_name) + if class_name == "SwnModflow": + self.domain_label = "ibound" + self.reach_index_name = "reachID" + elif class_name == "SwnMf6": + self.domain_label = "idomain" + for block in flopy.mf6.ModflowGwfsfr.dfn: + if block[0] == "block packagedata" and block[1] != "name packagedata": + self.reach_index_name = block[1][5:] + break + else: + self.logger.error("cannot determine reach index name for GWF-SFR") + else: + self.logger.error("unsupported subclass %r", class_name) def __iter__(self): """Return object datasets with an iterator.""" @@ -213,7 +228,7 @@ def reaches(self): Attributes ---------- - reachID (SwnModflow) or rno (SwnMf6) : int, index + reachID (SwnModflow), rno or ifno (SwnMf6) : int, index Reach index number, starting from 1. geometry : geometry LineString segments, one per model cell. @@ -312,11 +327,9 @@ def model(self, model): modelgrid = model.modelgrid modeltime = model.modeltime if this_class == "SwnModflow": - domain_label = "ibound" domain = model.bas6.ibound[0].array.copy() perlen = pd.Series(model.dis.perlen.array) elif this_class == "SwnMf6": - domain_label = "idomain" domain = dis.idomain.array[0].copy() nper = sim.tdis.nper.data perlen = pd.Series(sim.tdis.perioddata.array.perlen) @@ -392,7 +405,7 @@ def model(self, model): cols, rows = np.meshgrid(np.arange(ncol), np.arange(nrow)) grid_df = pd.DataFrame({"i": rows.flatten(), "j": cols.flatten()}) grid_df.set_index(["i", "j"], inplace=True) - grid_df[domain_label] = domain.flatten() + grid_df[self.domain_label] = domain.flatten() # Note: modelgrid.get_cell_vertices(i, j) is slow! xv = modelgrid.xvertices yv = modelgrid.yvertices @@ -439,13 +452,9 @@ def from_swn_flopy( """ this_class = cls.__name__ if this_class == "SwnModflow": - domain_label = "ibound" - reach_index_name = "reachID" reach_length_name = "rchlen" uses_segments = True elif this_class == "SwnMf6": - domain_label = "idomain" - reach_index_name = "rno" reach_length_name = "rlen" uses_segments = False else: @@ -466,17 +475,15 @@ def from_swn_flopy( dis = model.dis grid_cells = obj.grid_cells.copy() if domain_action == "freeze": - sel = grid_cells[domain_label] != 0 + sel = grid_cells[obj.domain_label] != 0 if sel.any(): # Remove any inactive grid cells from analysis grid_cells = grid_cells.loc[sel] _ = grid_cells.sindex # create spatial index num_domain_modified = 0 if this_class == "SwnModflow": - domain_label = "ibound" domain = model.bas6.ibound[0].array.copy() elif this_class == "SwnMf6": - domain_label = "idomain" domain = dis.idomain.array[0].copy() else: raise TypeError(f"unsupported subclass {cls!r}") @@ -814,27 +821,27 @@ def do_linemerge(ij, df, drop_reach_ids): if domain_action == "modify" and domain[i, j] == 0: num_domain_modified += 1 domain[i, j] = 1 - obj.grid_cells[domain_label].at[i, j] = 1 + obj.grid_cells[obj.domain_label].at[i, j] = 1 if domain_action == "modify": if num_domain_modified: obj.logger.debug( "updating %d cells from %s array for top layer", num_domain_modified, - domain_label.upper(), + obj.domain_label.upper(), ) - if domain_label == "ibound": + if obj.domain_label == "ibound": obj.model.bas6.ibound[0] = domain - elif domain_label == "idomain": + elif obj.domain_label == "idomain": obj.model.dis.idomain.set_data(domain, layer=0) reaches = reaches.merge( - grid_cells[[domain_label]], left_on=["i", "j"], right_index=True + grid_cells[[obj.domain_label]], left_on=["i", "j"], right_index=True ) reaches.rename( - columns={domain_label: f"prev_{domain_label}"}, inplace=True + columns={obj.domain_label: f"prev_{obj.domain_label}"}, inplace=True ) else: - reaches[f"prev_{domain_label}"] = 1 + reaches[f"prev_{obj.domain_label}"] = 1 # Mark segments that are not used segments["in_model"] = True @@ -1012,7 +1019,7 @@ def do_linemerge(ij, df, drop_reach_ids): reaches.loc[sel, reach_length_name] = 1.0 reaches.reset_index(inplace=True, drop=True) - reaches.index.name = reach_index_name + reaches.index.name = obj.reach_index_name reaches.index += 1 # flopy series starts at one if not hasattr(reaches.geometry, "geom_type"): @@ -1352,7 +1359,7 @@ def get_location_frame_reach_info( Returns ------- pandas.DataFrame - - reachID (SwnModflow) or rno (SwnMf6) + - reachID (SwnModflow), rno or ifno (SwnMf6) - zero-based cell index: k, i, j, - one-based reach index: iseg, ireach - dist_to_reach @@ -1434,7 +1441,7 @@ def get_location_frame_reach_info( has_orig_geom = True reach_index_name = self.reaches.index.name - if reach_index_name in loc_df.columns: + if self.reach_index_name in loc_df.columns: self.logger.info("resetting %s from location frame", reach_index_name) del loc_df[reach_index_name] loc_df[reach_index_name] = -1 @@ -1526,11 +1533,7 @@ def plot(self, column=None, cmap="viridis_r", colorbar=False, ax=None): grid_cells = getattr(self, "grid_cells", None) if grid_cells is not None: - domain_label = { - "SwnModflow": "ibound", - "SwnMf6": "idomain", - }[self.__class__.__name__] - sel = self.grid_cells[domain_label] != 0 + sel = self.grid_cells[self.domain_label] != 0 if sel.any(): self.grid_cells.loc[sel].plot( ax=ax, color="whitesmoke", edgecolor="gainsboro" diff --git a/swn/modflow/_swnmf6.py b/swn/modflow/_swnmf6.py index 7df4eca..0612e9b 100644 --- a/swn/modflow/_swnmf6.py +++ b/swn/modflow/_swnmf6.py @@ -34,7 +34,7 @@ class SwnMf6(SwnModflowBase): segments : geopandas.GeoDataFrame Copied from swn.segments, but with additional columns added reaches : geopandas.GeoDataFrame - Similar to structure in model.sfr.reaches with index "rno", + Similar to structure in model.sfr.reaches with index "rno" or "ifno", ordered and starting from 1. Contains geometry and other columns not used by flopy. tsvar : dict @@ -112,48 +112,54 @@ def from_swn_flopy( else: reaches_segnum_s = set(obj.reaches.segnum) - def find_next_rno(segnum): + def find_next_ridx(segnum): if segnum in to_segnums_d: to_segnum = to_segnums_d[segnum] if to_segnum in reaches_segnum_s: sel = obj.reaches["segnum"] == to_segnum return obj.reaches[sel].index[0] else: # recurse downstream - return find_next_rno(to_segnum) + return find_next_ridx(to_segnum) else: return 0 - def get_to_rno(): + def get_to_ridx(): if segnum == next_segnum: - return next_rno + return next_ridx else: - return find_next_rno(segnum) + return find_next_ridx(segnum) - obj.reaches["to_rno"] = -1 + ridxname = obj.reach_index_name + to_ridxname = f"to_{ridxname}" + from_ridxsname = f"from_{ridxname}s" + obj.reaches[to_ridxname] = -1 if has_diversions: + from_ridxname = f"from_{ridxname}" + div_to_ridxsname = f"div_to_{ridxname}s" + div_from_ridxname = f"div_from_{ridxname}" segnum_iter = obj.reaches.loc[~obj.reaches.diversion, "segnum"].items() else: segnum_iter = obj.reaches["segnum"].items() - rno, segnum = next(segnum_iter) - for next_rno, next_segnum in segnum_iter: - obj.reaches.at[rno, "to_rno"] = get_to_rno() - rno, segnum = next_rno, next_segnum + ridx, segnum = next(segnum_iter) + for next_ridx, next_segnum in segnum_iter: + obj.reaches.at[ridx, to_ridxname] = get_to_ridx() + ridx, segnum = next_ridx, next_segnum next_segnum = swn.END_SEGNUM - obj.reaches.at[rno, "to_rno"] = get_to_rno() + obj.reaches.at[ridx, to_ridxname] = get_to_ridx() if has_diversions: - obj.reaches.loc[obj.reaches["diversion"], "to_rno"] = 0 - assert obj.reaches.to_rno.min() >= 0 + obj.reaches.loc[obj.reaches["diversion"], to_ridxname] = 0 + assert obj.reaches[to_ridxname].min() >= 0 - # Populate from_rnos set - obj.reaches["from_rnos"] = [set() for _ in range(len(obj.reaches))] - to_rnos = obj.reaches.loc[obj.reaches["to_rno"] > 0, "to_rno"] - for k, v in to_rnos.items(): - obj.reaches.at[v, "from_rnos"].add(k) + # Populate from_ set + obj.reaches[from_ridxsname] = [set() for _ in range(len(obj.reaches))] + to_ridxs = obj.reaches.loc[obj.reaches[to_ridxname] > 0, to_ridxname] + for k, v in to_ridxs.items(): + obj.reaches.at[v, from_ridxsname].add(k) # Refresh diversions if set if has_diversions: div_sel = obj.diversions["in_model"] - # populate (rno, idv) from their match to non-diversion reaches + # populate (ridx, idv) from their match to non-diversion reaches diversions_in_model = obj.diversions[div_sel] r_df = obj.get_location_frame_reach_info( diversions_in_model.rename(columns={"from_segnum": "segnum"})[ @@ -162,41 +168,43 @@ def get_to_rno(): downstream_bias=diversion_downstream_bias, geom_loc_df=getattr(diversions_in_model, "geometry", None), ) - obj.diversions["rno"] = 0 # valid from 1 - obj.diversions.loc[r_df.index, "rno"] = r_df.rno + obj.diversions[ridxname] = 0 # valid from 1 + obj.diversions.loc[r_df.index, ridxname] = r_df[ridxname] # evaluate idv, which is betwen 1 and ndv obj.diversions["idv"] = 0 obj.diversions.loc[div_sel, "idv"] = 1 - rno_counts = obj.diversions[div_sel].groupby("rno").count()["in_model"] - for rno, count in rno_counts[rno_counts > 1].items(): + ridx_counts = obj.diversions[div_sel].groupby(ridxname).count()["in_model"] + for ridx, count in ridx_counts[ridx_counts > 1].items(): obj.diversions.loc[ - obj.diversions["rno"] == rno, "idv" - ] = obj.diversions.idv[obj.diversions["rno"] == rno].cumsum() - # cross-reference iconr to rno used as a reach + obj.diversions[ridxname] == ridx, "idv" + ] = obj.diversions.idv[obj.diversions[ridxname] == ridx].cumsum() + # cross-reference iconr to ridx used as a reach diversion_reaches = ( obj.reaches.loc[obj.reaches.diversion].reset_index().set_index("divid") ) - obj.diversions["iconr"] = diversion_reaches["rno"] + obj.diversions["iconr"] = diversion_reaches[ridxname] # Also put data into reaches frame - obj.reaches["div_from_rno"] = 0 + obj.reaches[div_from_ridxname] = 0 rdiv = ( - obj.diversions.loc[div_sel, ["rno", "iconr"]] - .rename(columns={"rno": "from_rno", "iconr": "rno"}) + obj.diversions.loc[div_sel, [ridxname, "iconr"]] + .rename(columns={ridxname: from_ridxname, "iconr": ridxname}) .reset_index() - .set_index("rno") + .set_index(ridxname) ) - obj.reaches.loc[rdiv.index, "div_from_rno"] = rdiv["from_rno"] - obj.reaches["div_to_rnos"] = [set() for _ in range(len(obj.reaches))] - to_rnos = obj.reaches.loc[obj.reaches["div_from_rno"] > 0, "div_from_rno"] - for k, v in to_rnos.items(): - obj.reaches.at[v, "div_to_rnos"].add(k) + obj.reaches.loc[rdiv.index, div_from_ridxname] = rdiv[from_ridxname] + obj.reaches[div_to_ridxsname] = [set() for _ in range(len(obj.reaches))] + to_ridxs = obj.reaches.loc[ + obj.reaches[div_from_ridxname] > 0, div_from_ridxname + ] + for k, v in to_ridxs.items(): + obj.reaches.at[v, div_to_ridxsname].add(k) # Workaround potential MODFLOW6 bug where diversions cannot attach # to outlet reach, so add another reach sel = ( - (obj.reaches["to_rno"] == 0) + (obj.reaches[to_ridxname] == 0) & (~obj.reaches["diversion"]) - & (obj.reaches["div_to_rnos"].apply(len) > 0) + & (obj.reaches[div_to_ridxsname].apply(len) > 0) ) if sel.any() and "mask" not in obj.reaches.columns: obj.reaches["mask"] = False @@ -204,29 +212,29 @@ def get_to_rno(): empty_geom = wkt.loads("linestring z empty") else: empty_geom = wkt.loads("linestring empty") - for rno in sel[sel].index: - new_rno = len(obj.reaches) + 1 + for ridx in sel[sel].index: + new_ridx = len(obj.reaches) + 1 # Correction to this reach - obj.reaches.loc[rno, "to_rno"] = new_rno + obj.reaches.loc[ridx, to_ridxname] = new_ridx # Use as template for new reach - reach_d = obj.reaches.loc[rno].to_dict() + reach_d = obj.reaches.loc[ridx].to_dict() reach_d.update( { "geometry": empty_geom, "mask": True, "ireach": reach_d["ireach"] + 1, - "to_rno": 0, - "from_rnos": {rno}, - "div_to_rnos": set(), + to_ridxname: 0, + from_ridxsname: {ridx}, + div_to_ridxsname: set(), } ) with ignore_shapely_warnings_for_object_array(): - obj.reaches.loc[new_rno] = reach_d + obj.reaches.loc[new_ridx] = reach_d # Set 1.0 for most, 0.0 for head and diversion nodes obj.reaches["ustrf"] = 1.0 - zero_from_rnos = obj.reaches["from_rnos"].apply(len) == 0 - obj.reaches.loc[zero_from_rnos, "ustrf"] = 0.0 + zero_from_ridxs = obj.reaches[from_ridxsname].apply(len) == 0 + obj.reaches.loc[zero_from_ridxs, "ustrf"] = 0.0 return obj @@ -327,7 +335,7 @@ def _init_package_df(self, *, style, defcols_names, auxiliary): dat["cellid"] = dat[kij_l].to_records(index=False).tolist() if cellid_none.any(): dat.loc[cellid_none, "cellid"] = "NONE" - # Convert rno from one-based to zero-based notation + # Convert ridx from one-based to zero-based notation dat.index -= 1 else: raise ValueError("'style' must be either 'native' or 'flopy'") @@ -370,7 +378,7 @@ def packagedata_frame(self, style: str, auxiliary: list = [], boundname=None): style : str If "native", all indicies (including kij) use one-based notation. Also use k,i,j columns (as str) rather than cellid. - If "flopy", all indices (including rno) use zero-based notation. + If "flopy", all indices (including rno/ifno) use zero-based notation. Also use cellid as a tuple. auxiliary : str, list, optional String or list of auxiliary variable names. Default []. @@ -434,21 +442,27 @@ def packagedata_frame(self, style: str, auxiliary: list = [], boundname=None): from flopy.mf6 import ModflowGwfsfr as Mf6Sfr defcols_names = [dt[0] for dt in Mf6Sfr.packagedata.dtype(self.model)] - defcols_names.remove("rno") # this is the index + ridxname = defcols_names.pop(0) # this is the index, either "rno" or "ifno" + to_ridxname = f"to_{ridxname}" + from_ridxsname = f"from_{ridxname}s" dat = self._init_package_df( style=style, defcols_names=defcols_names, auxiliary=auxiliary ) if "rlen" not in dat.columns: dat.loc[:, "rlen"] = dat.geometry.length - dat["ncon"] = dat["from_rnos"].apply(len) + (dat["to_rno"] > 0).astype(int) + dat["ncon"] = dat[from_ridxsname].apply(len) + (dat[to_ridxname] > 0).astype( + int + ) dat["ndv"] = 0 if self.diversions is not None: - dat["ncon"] += dat["div_to_rnos"].apply(len) + ( - dat["div_from_rno"] > 0 + div_to_ridxsname = f"div_to_{ridxname}s" + div_from_ridxname = f"div_from_{ridxname}" + dat["ncon"] += dat[div_to_ridxsname].apply(len) + ( + dat[div_from_ridxname] > 0 ).astype(int) ndv = ( self.diversions[self.diversions["in_model"]] - .groupby("rno") + .groupby(ridxname) .count() .in_model ) @@ -514,7 +528,7 @@ def connectiondata_series(self, style: str): ---------- style : str If "native", all indicies (including kij) use one-based notation. - If "flopy", all indices (including rno) use zero-based notation. + If "flopy", all indices (including rno/ifno) use zero-based notation. """ def nonzerolst(x, neg=False): @@ -523,25 +537,30 @@ def nonzerolst(x, neg=False): else: return [x] if x > 0 else [] + ridxname = self.reach_index_name + from_ridxsname = f"from_{ridxname}s" + to_ridxname = f"to_{ridxname}" if self.diversions is not None: + div_from_ridxname = f"div_from_{ridxname}" + div_to_ridxsname = f"div_to_{ridxname}s" res = ( - self.reaches["from_rnos"].apply(sorted) - + self.reaches["div_from_rno"].apply(nonzerolst, neg=False) - + self.reaches["to_rno"].apply(nonzerolst, neg=True) - + self.reaches["div_to_rnos"] + self.reaches[from_ridxsname].apply(sorted) + + self.reaches[div_from_ridxname].apply(nonzerolst, neg=False) + + self.reaches[to_ridxname].apply(nonzerolst, neg=True) + + self.reaches[div_to_ridxsname] .apply(sorted) .apply(lambda x: [-i for i in x]) ) else: - res = self.reaches["from_rnos"].apply(sorted) + self.reaches[ - "to_rno" + res = self.reaches[from_ridxsname].apply(sorted) + self.reaches[ + to_ridxname ].apply(nonzerolst, neg=True) if style == "native": # keep one-based notation, but convert list to str return res elif style == "flopy": - # Convert rno from one-based to zero-based notation + # Convert ridx from one-based to zero-based notation res.index -= 1 return res.apply(lambda x: [v - 1 if v > 0 else v + 1 for v in x]) else: @@ -564,12 +583,12 @@ def write_connectiondata(self, fname: str): cn = self.connectiondata_series("native") icn = [f"ic{n+1}" for n in range(cn.apply(len).max())] rowfmt = f"{{:>{len(str(cn.index.max()))}}} {{}}\n" - rnolen = 1 + len(str(len(self.reaches))) - cn = cn.apply(lambda x: " ".join(str(v).rjust(rnolen) for v in x)) + ridxlen = 1 + len(str(len(self.reaches))) + cn = cn.apply(lambda x: " ".join(str(v).rjust(ridxlen) for v in x)) with open(fname, "w") as f: - f.write(f"# rno {' '.join(icn)}\n") - for rno, ic in cn.items(): - f.write(rowfmt.format(rno, ic)) + f.write(f"# {self.reach_index_name} {' '.join(icn)}\n") + for ridx, ic in cn.items(): + f.write(rowfmt.format(ridx, ic)) def flopy_connectiondata(self): """Return list of lists for flopy. @@ -665,8 +684,8 @@ def diversions_frame(self, style: str): if style == "native": pass elif style == "flopy": - # Convert rno from one-based to zero-based notation - dat[["rno", "idv", "iconr"]] -= 1 + # Convert ridx from one-based to zero-based notation + dat[[self.reach_index_name, "idv", "iconr"]] -= 1 else: raise ValueError("'style' must be either 'native' or 'flopy'") return dat[defcols_names] @@ -721,7 +740,7 @@ def package_period_frame( style : str If "native", all indicies (including kij) use one-based notation. Also use k,i,j columns (as str) rather than cellid. - If "flopy", all indices (including rno) use zero-based notation. + If "flopy", all indices (including rno/ifno) use zero-based notation. Also use cellid as a tuple. auxiliary : str, list, optional String or list of auxiliary variable names. Default []. @@ -801,7 +820,7 @@ def package_period_frame( dat["per"] = 1 if style == "flopy": dat["per"] -= 1 - return dat.reset_index().set_index(["per", "rno"]) + return dat.reset_index().set_index(["per", self.reach_index_name]) def write_package_period( self, package: str, fname, auxiliary: list = [], boundname=None @@ -1483,6 +1502,9 @@ def _check_reach_v_laybot(r, botms, buffer=1.0, rbed_elev=None): botms[0, r.i, r.j] = new_elev return botms + ridxname = self.reach_index_name + from_ridxsname = f"from_{ridxname}s" + to_ridxname = f"to_{ridxname}" if direction == "upstream": ustrm = True else: @@ -1499,7 +1521,7 @@ def _check_reach_v_laybot(r, botms, buffer=1.0, rbed_elev=None): self.logger.info("%s reaches above model top", reachsel.sum()) # copy of layer 1 bottom (for updating to fit in stream reaches) layerbots = self.model.dis.botm.array.copy() - sel = self.reaches["from_rnos"] == set() + sel = self.reaches[from_ridxsname] == set() if self.diversions is not None: sel &= ~self.reaches["diversion"] headreaches = self.reaches.loc[sel] @@ -1519,10 +1541,10 @@ def _check_reach_v_laybot(r, botms, buffer=1.0, rbed_elev=None): reaches = self.reaches.loc[self.reaches.segnum.isin(segs)].sort_index() # get outflow reach for this profile # maybe can't rely on it being the last one - # the sort_index() should order (assuming rno increases downstream) + # the sort_index() should order (assuming ridx increases downstream) # so last should be to_rno == 0 assert ( - reaches.iloc[-1].to_rno == 0 + reaches.iloc[-1][to_ridxname] == 0 ), "reach numbers possibly not increasing downstream" outflow = reaches.iloc[-1] # check if outflow above model top @@ -1723,8 +1745,8 @@ def _to_rno_elevs( 0. get a list of cells with to_rtp > rtp-minslope*(delr+delc)/2 1. drop rtp of downstream reach (to_rno) when higher than rtp - (of rno) - 2. grab all offensive reaches downstream of rno + (of ridx) + 2. grab all offensive reaches downstream of ridx 3. check and fix all layer bottoms to be above minthick+buffer 4. adjust top, up only, to accomodate minincision or rtp above cell top @@ -1751,6 +1773,7 @@ def _to_rno_elevs( None """ + to_ridxname = f"to_{self.reach_index_name}" # copy some data top = self.model.dis.top.array.copy() @@ -1780,9 +1803,9 @@ def _to_rno_elevs( rdf.loc[idx, "rbth"] = np.max([r["rbth"], minthick]) rdf.loc[idx, "rtp"] = top[r["ij"]] - minincise for idx, r in rdf.iterrows(): - trno = int(r["to_rno"]) - if trno != 0: - rdf.loc[idx, "to_rtp"] = rdf.loc[trno, "rtp"] + to_ridx = int(r[to_ridxname]) + if to_ridx != 0: + rdf.loc[idx, "to_rtp"] = rdf.loc[to_ridx, "rtp"] # start loop loop = 0 cont = True @@ -1795,26 +1818,26 @@ def _to_rno_elevs( loop += 1 chg = 0 for br in bad_reaches: - rno = br - trno = int(rdf.loc[br, "to_rno"]) + ridx = br + to_ridx = int(rdf.loc[br, to_ridxname]) chglist = [] - if trno != 0: + if to_ridx != 0: # count how many downstream reaches offend # keep track of changes in elev - dz = rdf.loc[rno, "mindz"] - check = rdf.loc[trno, "rtp"] > rdf.loc[rno, "rtp"] - dz - while trno != 0 and check: + dz = rdf.loc[ridx, "mindz"] + check = rdf.loc[to_ridx, "rtp"] > rdf.loc[ridx, "rtp"] - dz + while to_ridx != 0 and check: # keep list of dz in case another inflowing stream is # even lower - chglist.append(trno) - nelev = rdf.loc[rno, "rtp"] - dz + chglist.append(to_ridx) + nelev = rdf.loc[ridx, "rtp"] - dz # set to_rtp and rtp - rdf.loc[rno, "to_rtp"] = nelev - rdf.loc[trno, "rtp"] = nelev + rdf.loc[ridx, "to_rtp"] = nelev + rdf.loc[to_ridx, "rtp"] = nelev # move downstream - rno = trno - trno = rdf.loc[rno, "to_rno"] - dz = rdf.loc[rno, "mindz"] + ridx = to_ridx + to_ridx = rdf.loc[ridx, to_ridxname] + dz = rdf.loc[ridx, "mindz"] # now adjust layering if necessary if len(chglist) > 0 and fix_dis: @@ -1927,7 +1950,7 @@ def route_reaches(self, start, end, *, allow_indirect=False): Parameters ---------- start, end : any - Start and end reach numbers (rno). + Start and end reach numbers (rno or ifno). allow_indirect : bool, default False If True, allow the route to go downstream from start to a confluence, then route upstream to end. Defalut False allows @@ -1984,17 +2007,18 @@ def route_reaches(self, start, end, *, allow_indirect=False): 7 2 1 10.000000 """ # noqa if start not in self.reaches.index: - raise IndexError(f"invalid start rno {start}") + raise IndexError(f"invalid start {self.reach_index_name} {start}") if end not in self.reaches.index: - raise IndexError(f"invalid end rno {end}") + raise IndexError(f"invalid end {self.reach_index_name} {end}") if start == end: return [start] - to_rnos = dict(self.reaches.loc[self.reaches["to_rno"] != 0, "to_rno"]) + to_ridxname = f"to_{self.reach_index_name}" + to_ridxs = dict(self.reaches.loc[self.reaches[to_ridxname] != 0, to_ridxname]) - def go_downstream(rno): - yield rno - if rno in to_rnos: - yield from go_downstream(to_rnos[rno]) + def go_downstream(ridx): + yield ridx + if ridx in to_ridxs: + yield from go_downstream(to_ridxs[ridx]) con1 = list(go_downstream(start)) try: @@ -2018,16 +2042,16 @@ def go_downstream(rno): if not common: msg += " -- reach networks are disjoint" raise ConnectionError(msg) - # find the most upstream common rno or "confluence" - rno = common.pop() - idx1 = con1.index(rno) - idx2 = con2.index(rno) + # find the most upstream common ridx or "confluence" + ridx = common.pop() + idx1 = con1.index(ridx) + idx2 = con2.index(ridx) while common: - rno = common.pop() - tmp1 = con1.index(rno) + ridx = common.pop() + tmp1 = con1.index(ridx) if tmp1 < idx1: idx1 = tmp1 - idx2 = con2.index(rno) + idx2 = con2.index(ridx) return con1[:idx1] + list(reversed(con2[:idx2])) diff --git a/tests/test_modflow6.py b/tests/test_modflow6.py index 57b07c0..5c9cf3e 100644 --- a/tests/test_modflow6.py +++ b/tests/test_modflow6.py @@ -18,6 +18,18 @@ try: import flopy + + for block in flopy.mf6.ModflowGwfsfr.dfn: + if block[0] == "block packagedata" and block[1] != "name packagedata": + ridxname = block[1][5:] + break + else: + raise ValueError("cannot determine reach index name for GWF-SFR") + to_ridxname = f"to_{ridxname}" + from_ridxsname = f"from_{ridxname}s" + div_to_ridxsname = f"div_to_{ridxname}s" + div_from_ridxname = f"div_from_{ridxname}" + except ImportError: pytest.skip("skipping tests that require flopy", allow_module_level=True) @@ -202,10 +214,10 @@ def test_n3d_defaults(tmp_path, has_diversions): "iseg": [1, 1, 1, 2, 2, 3, 3], "ireach": [1, 2, 3, 1, 2, 1, 2], "rlen": [18.027756, 6.009252, 12.018504, 21.081851, 10.5409255, 10.0, 10.0], - "to_rno": [2, 3, 6, 5, 6, 7, 0], - "from_rnos": [set(), {1}, {2}, set(), {4}, {3, 5}, {6}], + to_ridxname: [2, 3, 6, 5, 6, 7, 0], + from_ridxsname: [set(), {1}, {2}, set(), {4}, {3, 5}, {6}], }, - index=pd.RangeIndex(1, 8, name="rno"), + index=pd.RangeIndex(1, 8, name=ridxname), ) if not has_diversions: pd.testing.assert_frame_equal( @@ -224,17 +236,17 @@ def test_n3d_defaults(tmp_path, has_diversions): "iseg": 0, "ireach": 0, "rlen": 1.0, - "to_rno": 0, - "from_rnos": [set(), set(), set(), set()], + to_ridxname: 0, + from_ridxsname: [set(), set(), set(), set()], }, - index=pd.RangeIndex(8, 12, name="rno"), + index=pd.RangeIndex(8, 12, name=ridxname), ), ] ) div_expected["diversion"] = [False] * 7 + [True] * 4 div_expected["divid"] = [0] * 7 + [0, 1, 2, 3] - div_expected["div_from_rno"] = [0] * 7 + [3, 5, 6, 6] - div_expected["div_to_rnos"] = [ + div_expected[div_from_ridxname] = [0] * 7 + [3, 5, 6, 6] + div_expected[div_to_ridxsname] = [ set(), set(), {8}, @@ -273,7 +285,7 @@ def test_n3d_defaults(tmp_path, has_diversions): "ncon": [1, 2, 2, 1, 2, 3, 1], "ndv": 0, }, - index=pd.RangeIndex(1, 8, name="rno"), + index=pd.RangeIndex(1, 8, name=ridxname), ) if not has_diversions: pd.testing.assert_frame_equal( @@ -294,7 +306,7 @@ def test_n3d_defaults(tmp_path, has_diversions): "ncon": 1, "ndv": 0, }, - index=pd.RangeIndex(8, 12, name="rno"), + index=pd.RangeIndex(8, 12, name=ridxname), ), ] ) @@ -341,7 +353,7 @@ def test_n3d_defaults(tmp_path, has_diversions): nm.diversions_frame("native"), pd.DataFrame( { - "rno": [3, 5, 6, 6], + ridxname: [3, 5, 6, 6], "idv": [1, 1, 1, 2], "iconr": [8, 9, 10, 11], "cprior": "upto", @@ -357,16 +369,16 @@ def test_n3d_defaults(tmp_path, has_diversions): assert nm.flopy_diversions() == expected if not has_diversions: assert repr(nm) == dedent( - """\ + f"""\ """ ) else: assert repr(nm) == dedent( - """\ + f"""\ """ ) @@ -605,7 +617,7 @@ def test_n2d_defaults(tmp_path): assert list(r.i) == [0, 0, 1, 0, 1, 1, 2] assert list(r.j) == [0, 1, 1, 1, 1, 1, 1] assert list(r.segnum) == [1, 1, 1, 2, 2, 0, 0] - assert list(r.to_rno) == [2, 3, 6, 5, 6, 7, 0] + assert list(r[to_ridxname]) == [2, 3, 6, 5, 6, 7, 0] nm.default_packagedata(hyd_cond1=2.0, thickness1=2.0) nm.set_reach_data_from_array("rtp", m.dis.top.array - 1.0) np.testing.assert_array_almost_equal( @@ -660,7 +672,7 @@ def test_packagedata(tmp_path): ] assert ( list(m.sfr.packagedata.array.dtype.names) - == ["rno", "cellid"] + partial_expected_cols + == [ridxname, "cellid"] + partial_expected_cols ) # Check pandas frames @@ -693,7 +705,7 @@ def test_packagedata(tmp_path): nm.reaches["var1"] = np.arange(len(nm.reaches), dtype=float) * 12.0 nm.set_sfr_obj(auxiliary="var1") assert list(m.sfr.packagedata.array.dtype.names) == [ - "rno", + ridxname, "cellid", ] + partial_expected_cols + ["var1"] np.testing.assert_array_almost_equal( @@ -707,7 +719,7 @@ def test_packagedata(tmp_path): nm.reaches["var2"] = np.arange(len(nm.reaches)) * 11 nm.set_sfr_obj(auxiliary=["var1", "var2"]) assert list(m.sfr.packagedata.array.dtype.names) == [ - "rno", + ridxname, "cellid", ] + partial_expected_cols + ["var1", "var2"] np.testing.assert_array_almost_equal( @@ -725,7 +737,7 @@ def test_packagedata(tmp_path): nm.reaches["boundname"] = nm.reaches["segnum"] nm.set_sfr_obj() assert list(m.sfr.packagedata.array.dtype.names) == [ - "rno", + ridxname, "cellid", ] + partial_expected_cols + ["boundname"] assert list(m.sfr.packagedata.array["boundname"]) == ( @@ -741,7 +753,7 @@ def test_packagedata(tmp_path): nm.reaches["var1"] = np.arange(len(nm.reaches), dtype=float) * 12.0 nm.set_sfr_obj(auxiliary=["var1"]) assert list(m.sfr.packagedata.array.dtype.names) == [ - "rno", + ridxname, "cellid", ] + partial_expected_cols + ["var1", "boundname"] np.testing.assert_array_almost_equal( @@ -826,9 +838,9 @@ def test_coastal(tmp_path, coastal_lines_gdf, coastal_flow_m): # Check other packages check_number_sum_hex(m.dis.idomain.array, 509, "c4135a084b2593e0b69c148136a3ad6d") assert repr(nm) == dedent( - """\ + f"""\ """ ) if matplotlib: @@ -924,9 +936,9 @@ def test_coastal_reduced(coastal_lines_gdf, coastal_flow_m, tmp_path): np.testing.assert_almost_equal(reach_geom.length, 237.72893664132727) assert len(nm.reaches) == 154 assert repr(nm) == dedent( - """\ + f"""\ """ ) if matplotlib: @@ -980,9 +992,9 @@ def test_coastal_idomain_modify(coastal_swn, coastal_flow_m, tmp_path): # Check other packages check_number_sum_hex(m.dis.idomain.array, 572, "d353560128577b37f730562d2f89c025") assert repr(nm) == dedent( - """\ + f"""\ """ ) if matplotlib: @@ -1041,7 +1053,7 @@ def test_include_downstream_reach_outside_model(tmp_path): ) assert m.sfr.nreaches.data == 5 dat = m.sfr.packagedata.array - np.testing.assert_array_equal(dat.rno, [0, 1, 2, 3, 4]) + np.testing.assert_array_equal(dat[ridxname], [0, 1, 2, 3, 4]) cellid = np.array(dat.cellid, dtype=[("k", int), ("i", int), ("j", int)]) np.testing.assert_array_equal(cellid["k"], [0, 0, 0, 0, 0]) np.testing.assert_array_equal(cellid["i"], [0, 1, 0, 1, 1]) @@ -1111,8 +1123,8 @@ def test_n3d_defaults_with_div_on_outlet(tmp_path): 1.0, 10.0, ], - "to_rno": [2, 3, 6, 5, 6, 7, 12, 0, 0, 0, 0, 0], - "from_rnos": [ + to_ridxname: [2, 3, 6, 5, 6, 7, 12, 0, 0, 0, 0, 0], + from_ridxsname: [ set(), {1}, {2}, @@ -1126,8 +1138,8 @@ def test_n3d_defaults_with_div_on_outlet(tmp_path): set(), {7}, ], - "div_from_rno": [0, 0, 0, 0, 0, 0, 0, 3, 5, 7, 7, 0], - "div_to_rnos": [ + div_from_ridxname: [0, 0, 0, 0, 0, 0, 0, 3, 5, 7, 7, 0], + div_to_ridxsname: [ set(), set(), {8}, @@ -1143,7 +1155,7 @@ def test_n3d_defaults_with_div_on_outlet(tmp_path): ], "mask": [False] * 11 + [True], }, - index=pd.RangeIndex(1, 13, name="rno"), + index=pd.RangeIndex(1, 13, name=ridxname), ) pd.testing.assert_frame_equal(nm.reaches[expected.columns], expected) nm.default_packagedata(hyd_cond1=1e-4) @@ -1160,7 +1172,7 @@ def test_n3d_defaults_with_div_on_outlet(tmp_path): "ncon": [1, 2, 3, 1, 3, 3, 4, 1, 1, 1, 1, 1], "ndv": [0, 0, 1, 0, 1, 0, 2, 0, 0, 0, 0, 0], }, - index=pd.RangeIndex(1, 13, name="rno"), + index=pd.RangeIndex(1, 13, name=ridxname), ) pd.testing.assert_frame_equal( nm.packagedata_frame("native")[expected.columns], expected @@ -1193,7 +1205,7 @@ def test_n3d_defaults_with_div_on_outlet(tmp_path): nm.diversions_frame("native"), pd.DataFrame( { - "rno": [3, 5, 7, 7], + ridxname: [3, 5, 7, 7], "idv": [1, 1, 1, 2], "iconr": [8, 9, 10, 11], "cprior": "upto", @@ -1208,9 +1220,9 @@ def test_n3d_defaults_with_div_on_outlet(tmp_path): ] assert nm.flopy_diversions() == expected assert repr(nm) == dedent( - """\ + f"""\ """ ) @@ -1325,13 +1337,13 @@ def test_diversions(tmp_path): nm = swn.SwnMf6.from_swn_flopy(n, m) nm.reaches["boundname"] = nm.reaches["divid"] - assert list(nm.diversions.rno) == [3, 5, 6, 6] - assert list(nm.diversions.idv) == [1, 1, 1, 2] + assert list(nm.diversions[ridxname]) == [3, 5, 6, 6] + assert list(nm.diversions["idv"]) == [1, 1, 1, 2] # Check optional parameter nm2 = swn.SwnMf6.from_swn_flopy(n, m, diversion_downstream_bias=0.3) - assert list(nm2.diversions.rno) == [3, 5, 7, 7] - assert list(nm2.diversions.idv) == [1, 1, 1, 2] + assert list(nm2.diversions[ridxname]) == [3, 5, 7, 7] + assert list(nm2.diversions["idv"]) == [1, 1, 1, 2] # Assemble MODFLOW SFR data nm.default_packagedata(hyd_cond1=0.0) @@ -1345,7 +1357,7 @@ def test_diversions(tmp_path): df[col] = df[col].astype(np.int64) pd.testing.assert_frame_equal( df, - swn.file.read_formatted_frame(tmp_path / "packagedata.dat").set_index("rno"), + swn.file.read_formatted_frame(tmp_path / "packagedata.dat").set_index(ridxname), ) pd.testing.assert_frame_equal( nm.diversions_frame("native").reset_index(drop=True), @@ -1453,9 +1465,9 @@ def test_route_reaches(): assert nm.route_reaches(7, 1) == [7, 3, 2, 1] assert nm.route_reaches(2, 4, allow_indirect=True) == [2, 3, 5, 4] # errors - with pytest.raises(IndexError, match="invalid start rno 0"): + with pytest.raises(IndexError, match=f"invalid start {ridxname} 0"): nm.route_reaches(0, 1) - with pytest.raises(IndexError, match="invalid end rno 0"): + with pytest.raises(IndexError, match=f"invalid end {ridxname} 0"): nm.route_reaches(1, 0) with pytest.raises(ConnectionError, match="1 does not connect to 4"): nm.route_reaches(1, 4) @@ -1499,7 +1511,7 @@ def test_package_period_frame(): "cond": [180.28, 60.09, 120.19, 210.82, 105.41, 100.0, 100.0], }, index=pd.MultiIndex.from_tuples( - [(1, rno + 1) for rno in range(7)], names=["per", "rno"] + [(1, ridx + 1) for ridx in range(7)], names=["per", ridxname] ), ) exp_flopy = pd.DataFrame( @@ -1517,7 +1529,7 @@ def test_package_period_frame(): "cond": [180.28, 60.09, 120.19, 210.82, 105.41, 100.0, 100.0], }, index=pd.MultiIndex.from_tuples( - [(0, rno) for rno in range(7)], names=["per", "rno"] + [(0, ridx) for ridx in range(7)], names=["per", ridxname] ), ) pd.testing.assert_frame_equal(nm.package_period_frame("drn", "native"), exp_native) @@ -1670,7 +1682,7 @@ def df2dic(df): "cond": [180.28, 60.09, 120.19, 210.82, 105.41, 100.0, 100.0], }, index=pd.MultiIndex.from_tuples( - [(0, rno) for rno in range(7)], names=["per", "rno"] + [(0, ridx) for ridx in range(7)], names=["per", ridxname] ), ) assert nm.flopy_package_period("drn") == df2dic(exp_flopy)