From 6871d07bc31962de5f5d5126c06a17c9213baf21 Mon Sep 17 00:00:00 2001 From: "Michael Ou@SSPA" Date: Thu, 21 Dec 2023 18:35:46 -0500 Subject: [PATCH] feat get budget data by second package name `paknam2` --- flopy/utils/binaryfile.py | 153 +++++++++++++++++--------------------- 1 file changed, 69 insertions(+), 84 deletions(-) diff --git a/flopy/utils/binaryfile.py b/flopy/utils/binaryfile.py index 8ac1416ae4..1e3401c1f5 100644 --- a/flopy/utils/binaryfile.py +++ b/flopy/utils/binaryfile.py @@ -1012,7 +1012,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 @@ -1078,7 +1079,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"): @@ -1218,8 +1220,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: @@ -1389,7 +1393,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 @@ -1401,7 +1405,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 @@ -1432,11 +1436,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) @@ -1466,7 +1470,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 @@ -1481,17 +1485,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 @@ -1501,7 +1506,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): """ @@ -1576,6 +1581,7 @@ def get_data( totim=None, text=None, paknam=None, + paknam2=None, full3D=False, ): """ @@ -1593,6 +1599,10 @@ 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. 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 @@ -1622,8 +1632,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) @@ -1635,83 +1645,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): """