feat get budget data by second package name paknam2

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,8 @@ 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"])
+            self.paknamlist_from.append(header["paknam"])
+            self.paknamlist_to.append(header["paknam2"])
             ipos = self.file.tell()
 
             if self.verbose:
@@ -1389,7 +1391,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 +1403,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 +1434,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 +1468,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 +1483,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 +1504,7 @@ def _unique_package_names(self):
             List of unique package names in the binary file.
 
         """
-        return self.paknamlist
+        return np.unique(self.paknamlist_to if to else self.paknamlist_from)
 
     def get_kstpkper(self):
         """
@@ -1576,6 +1579,7 @@ def get_data(
         totim=None,
         text=None,
         paknam=None,
+        paknam2=None,
         full3D=False,
     ):
         """
@@ -1593,6 +1597,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
@@ -1635,83 +1643,33 @@ 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
+        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)
-
+            selected = selected & (self.recordarray["kstp"] == kstp1)
+            selected = selected & (self.recordarray["kper"] == kper1)
+        if text16:
+            selected = selected & (self.recordarray["text"] == text16)
+        if paknam16:
+            selected = selected & (self.recordarray["paknam"] == paknam16)
+        if paknam16_2:
+            selected = selected & (self.recordarray["paknam2"] == paknam16_2)
+        if totim:
+            selected = selected & np.isclose(self.recordarray["totim"], totim)
+        if idx is not None:
+            selected[idx] = True
+
+        if np.any(selected):
+            return [self.get_record(idx, full3D=full3D) for idx, t in enumerate(selected) if t]
         else:
-            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
+            print('Return empty record.')
+            return None
 
     def get_ts(self, idx, text=None, times=None):
         """