From 5ec612a1a5b21c094bd910c9ed39f3ea5fef5084 Mon Sep 17 00:00:00 2001 From: "Michael Ou@SSPA" Date: Thu, 14 Mar 2024 08:29:34 -0400 Subject: [PATCH] feat(binaryfile): get budget by second package name `paknam2` (#2050) * feat get budget data by second package name `paknam2` * Add a test and update docstring --------- Co-authored-by: Langevin, Christian D --- autotest/test_binaryfile.py | 19 +++++ flopy/utils/binaryfile.py | 157 +++++++++++++++++------------------- 2 files changed, 92 insertions(+), 84 deletions(-) diff --git a/autotest/test_binaryfile.py b/autotest/test_binaryfile.py index 1089e22ad0..b320a5c20e 100644 --- a/autotest/test_binaryfile.py +++ b/autotest/test_binaryfile.py @@ -559,3 +559,22 @@ def test_read_mf2005_freyberg(example_data_path, function_tmpdir, compact): assert len(cbb_data) == len(cbb_data_kstpkper) for i in range(len(cbb_data)): assert np.array_equal(cbb_data[i], cbb_data_kstpkper[i]) + + +def test_read_mf6_budgetfile(example_data_path): + cbb_file = ( + example_data_path + / "mf6" + / "test005_advgw_tidal" + / "expected_output" + / "AdvGW_tidal.cbc" + ) + cbb = CellBudgetFile(cbb_file) + rch_zone_1 = cbb.get_data(paknam2="rch-zone_1".upper()) + rch_zone_2 = cbb.get_data(paknam2="rch-zone_2".upper()) + rch_zone_3 = cbb.get_data(paknam2="rch-zone_3".upper()) + + # ensure there is a record for each time step + assert len(rch_zone_1) == 120 * 3 + 1 + assert len(rch_zone_2) == 120 * 3 + 1 + assert len(rch_zone_3) == 120 * 3 + 1 diff --git a/flopy/utils/binaryfile.py b/flopy/utils/binaryfile.py index 414a387987..899d6ffc44 100644 --- a/flopy/utils/binaryfile.py +++ b/flopy/utils/binaryfile.py @@ -1013,7 +1013,8 @@ def __init__( self.iposarray = [] self.textlist = [] self.imethlist = [] - self.paknamlist = [] + self.paknamlist_from = [] + self.paknamlist_to = [] self.nrecords = 0 self.compact = True # compact budget file flag @@ -1079,7 +1080,8 @@ def __reset(self): self.iposarray = [] self.textlist = [] self.imethlist = [] - self.paknamlist = [] + self.paknamlist_from = [] + self.paknamlist_to = [] self.nrecords = 0 def _set_precision(self, precision="single"): @@ -1219,8 +1221,10 @@ def _build_index(self): raise BudgetIndexError("Improper precision") self.textlist.append(header["text"]) self.imethlist.append(header["imeth"]) - if header["paknam"] not in self.paknamlist: - self.paknamlist.append(header["paknam"]) + if header["paknam"] not in self.paknamlist_from: + self.paknamlist_from.append(header["paknam"]) + if header["paknam2"] not in self.paknamlist_to: + self.paknamlist_to.append(header["paknam2"]) ipos = self.file.tell() if self.verbose: @@ -1390,7 +1394,7 @@ def _find_text(self, text): raise Exception(errmsg) return text16 - def _find_paknam(self, paknam): + def _find_paknam(self, paknam, to=False): """ Determine if selected record name is in budget file @@ -1402,7 +1406,7 @@ def _find_paknam(self, paknam): tpaknam = paknam.decode() else: tpaknam = paknam - for t in self._unique_package_names(): + for t in self._unique_package_names(to): if tpaknam.upper() in t.decode(): paknam16 = t break @@ -1433,11 +1437,11 @@ def list_unique_records(self): rec = rec.decode() print(f"{rec.strip():16} {imeth:5d}") - def list_unique_packages(self): + def list_unique_packages(self, to=False): """ Print a list of unique package names """ - for rec in self._unique_package_names(): + for rec in self._unique_package_names(to): if isinstance(rec, bytes): rec = rec.decode() print(rec) @@ -1467,7 +1471,7 @@ def get_unique_record_names(self, decode=False): names = self.textlist return names - def get_unique_package_names(self, decode=False): + def get_unique_package_names(self, decode=False, to=False): """ Get a list of unique package names in the file @@ -1482,17 +1486,18 @@ def get_unique_package_names(self, decode=False): List of unique package names in the binary file. """ + if decode: names = [] - for text in self.paknamlist: + for text in self._unique_package_names(to): if isinstance(text, bytes): text = text.decode() names.append(text) else: - names = self.paknamlist + names = self._unique_package_names(to) return names - def _unique_package_names(self): + def _unique_package_names(self, to=False): """ Get a list of unique package names in the file @@ -1502,7 +1507,7 @@ def _unique_package_names(self): List of unique package names in the binary file. """ - return self.paknamlist + return self.paknamlist_to if to else self.paknamlist_from def get_kstpkper(self): """ @@ -1577,6 +1582,7 @@ def get_data( totim=None, text=None, paknam=None, + paknam2=None, full3D=False, ) -> Union[List, np.ndarray]: """ @@ -1594,6 +1600,14 @@ def get_data( text : str The text identifier for the record. Examples include 'RIVER LEAKAGE', 'STORAGE', 'FLOW RIGHT FACE', etc. + paknam : str + The `from` package name for the record. + paknam2 : str + The `to` package name for the record. This argument can be + useful for MODFLOW 6 budget files if multiple packages of + the same type are specified. The paknam2 argument can be + specified as the package name (not the package type) in + order to retrieve budget data for a specific named package. full3D : boolean If true, then return the record as a three dimensional numpy array, even for those list-style records written as part of a @@ -1623,8 +1637,8 @@ def get_data( if totim is not None: if len(self.times) == 0: errmsg = """This is an older style budget file that - does not have times in it. Use the MODFLOW - compact budget format if you want to work with + does not have times in it. Use the MODFLOW + compact budget format if you want to work with times. Or you may access this file using the kstp and kper arguments or the idx argument.""" raise Exception(errmsg) @@ -1636,83 +1650,58 @@ def get_data( paknam16 = None if paknam is not None: paknam16 = self._find_paknam(paknam) - + paknam16_2 = None + if paknam2 is not None: + paknam16_2 = self._find_paknam(paknam2, to=True) + + # build the selection mask + select_indices = np.array([True] * len(self.recordarray)) + selected = False + if idx is not None: + select_indices[idx] = False + select_indices = ~select_indices + selected = True if kstpkper is not None: kstp1 = kstpkper[0] + 1 kper1 = kstpkper[1] + 1 - if text is None and paknam is None: - select_indices = np.where( - (self.recordarray["kstp"] == kstp1) - & (self.recordarray["kper"] == kper1) - ) - else: - if paknam is None and text is not None: - select_indices = np.where( - (self.recordarray["kstp"] == kstp1) - & (self.recordarray["kper"] == kper1) - & (self.recordarray["text"] == text16) - ) - elif text is None and paknam is not None: - select_indices = np.where( - (self.recordarray["kstp"] == kstp1) - & (self.recordarray["kper"] == kper1) - & (self.recordarray["paknam"] == paknam16) - ) - else: - select_indices = np.where( - (self.recordarray["kstp"] == kstp1) - & (self.recordarray["kper"] == kper1) - & (self.recordarray["text"] == text16) - & (self.recordarray["paknam"] == paknam16) - ) - - elif totim is not None: - if text is None and paknam is None: - select_indices = np.where(self.recordarray["totim"] == totim) - else: - if paknam is None and text is not None: - select_indices = np.where( - (self.recordarray["totim"] == totim) - & (self.recordarray["text"] == text16) - ) - elif text is None and paknam is not None: - select_indices = np.where( - (self.recordarray["totim"] == totim) - & (self.recordarray["paknam"] == paknam16) - ) - else: - select_indices = np.where( - (self.recordarray["totim"] == totim) - & (self.recordarray["text"] == text16) - & (self.recordarray["paknam"] == paknam16) - ) - - # allow for idx to be a list or a scalar - elif idx is not None: - if isinstance(idx, list): - select_indices = idx - else: - select_indices = [idx] - - # case where only text is entered - elif text is not None: - select_indices = np.where(self.recordarray["text"] == text16) + select_indices = select_indices & ( + self.recordarray["kstp"] == kstp1 + ) + select_indices = select_indices & ( + self.recordarray["kper"] == kper1 + ) + selected = True + if text16 is not None: + select_indices = select_indices & ( + self.recordarray["text"] == text16 + ) + selected = True + if paknam16 is not None: + select_indices = select_indices & ( + self.recordarray["paknam"] == paknam16 + ) + selected = True + if paknam16_2 is not None: + select_indices = select_indices & ( + self.recordarray["paknam2"] == paknam16_2 + ) + selected = True + if totim is not None: + select_indices = select_indices & np.isclose( + self.recordarray["totim"], totim + ) + selected = True - else: + if not selected: raise TypeError( "get_data() missing 1 required argument: 'kstpkper', 'totim', " "'idx', or 'text'" ) - - # build and return the record list - if isinstance(select_indices, tuple): - select_indices = select_indices[0] - recordlist = [] - for idx in select_indices: - rec = self.get_record(idx, full3D=full3D) - recordlist.append(rec) - - return recordlist + return [ + self.get_record(idx, full3D=full3D) + for idx, t in enumerate(select_indices) + if t + ] def get_ts(self, idx, text=None, times=None): """