From 24d5837034c42eade2c8f4b8a0a8f40b3a10a78f Mon Sep 17 00:00:00 2001
From: Scott Paulinski <spaulinski@usgs.gov>
Date: Fri, 17 Dec 2021 12:57:32 -0800
Subject: [PATCH 1/8] feat(inspect cells): new feature that returns model data
 associated with specified model cells (#1140)

---
 flopy/mf6/data/mfdatautil.py |  37 ++++-
 flopy/mf6/mfmodel.py         | 113 +++++++++++++++
 flopy/mf6/mfpackage.py       | 269 ++++++++++++++++++++++++++++++++++-
 3 files changed, 416 insertions(+), 3 deletions(-)

diff --git a/flopy/mf6/data/mfdatautil.py b/flopy/mf6/data/mfdatautil.py
index c041ca75e4..8afd80aff7 100644
--- a/flopy/mf6/data/mfdatautil.py
+++ b/flopy/mf6/data/mfdatautil.py
@@ -8,8 +8,15 @@
 import struct
 
 
-def iterable(obj):
-    return isinstance(obj, Iterable)
+def iterable(obj, any_iterator=False):
+    if any_iterator:
+        try:
+            my_iter = iter(obj)
+        except TypeError as te:
+            return False
+        return True
+    else:
+        return isinstance(obj, Iterable)
 
 
 def get_first_val(arr):
@@ -18,6 +25,15 @@ def get_first_val(arr):
     return arr
 
 
+def cellids_equal(cellid_1, cellid_2):
+    if len(cellid_1) != len(cellid_2):
+        return False
+    for id1, id2 in zip(cellid_1, cellid_2):
+        if id1 != id2:
+            return False
+    return True
+
+
 # convert_data(data, type) : type
 #    converts data "data" to type "type" and returns the converted data
 def convert_data(data, data_dimensions, data_type, data_item=None, sub_amt=1):
@@ -224,6 +240,23 @@ def to_string(
         return str(val)
 
 
+class DataSearchOutput:
+    def __init__(self, path_to_data=None, data_header=None):
+        self.path_to_data = path_to_data
+        self.data_header = data_header
+        self.data_entry_ids = []
+        self.data_entry_cellids = []
+        self.data_entry_stress_period = []
+        self.data_entries = []
+
+    @property
+    def transient(self):
+        if len(self.data_entry_stress_period) > 0:
+            if self.data_entry_stress_period[0] != -1:
+                return True
+        return False
+
+
 class MFComment:
     """
     Represents a variable in a MF6 input file
diff --git a/flopy/mf6/mfmodel.py b/flopy/mf6/mfmodel.py
index b99444eba0..5a295a5b30 100644
--- a/flopy/mf6/mfmodel.py
+++ b/flopy/mf6/mfmodel.py
@@ -20,6 +20,7 @@
 from ..mbase import ModelInterface
 from .utils.mfenums import DiscretizationType
 from .data import mfstructure
+from .data.mfdatautil import iterable
 from .utils.output_util import MF6Output
 from ..utils.check import mf6check
 
@@ -825,6 +826,118 @@ def load_base(
 
         return instance
 
+    def inspect_cells(
+        self, cell_list, stress_period=None, output_file_path=None
+    ):
+        """
+        Inspect model cells.  Returns model data associated with cells.
+
+        Parameters
+        ----------
+        cell_list : list of tuples
+            List of model cells.  Each model cell is a tuple of integers.
+            ex: [(1,1,1), (2,4,3)]
+        stress_period : int
+            For transient data qnly return data from this stress period.  If
+            not specified or None, all stress period data will be returned.
+        output_file_path: str
+            path to output file that will contain the inspection results
+
+        Returns
+        -------
+        output : dict
+            Dictionary containing inspection results
+
+        Examples
+        --------
+
+        >>> import flopy
+        >>> sim = flopy.mf6.MFSimulation.load("name", "mf6", "mf6.exe", ".")
+        >>> model = sim.get_model()
+        >>> inspect_list = [(2, 3, 2), (0, 4, 2), (0, 2, 4)]
+        >>> out_file = os.path.join("temp", "inspect_AdvGW_tidal.csv")
+        >>> model.inspect_cells(inspect_list, output_file_path=out_file)
+        """
+        output_by_package = {}
+        # loop through all packages
+        for pp in self.packagelist:
+            # call the package's "inspect_cells" method
+            package_output = pp.inspect_cells(cell_list, stress_period)
+            if len(package_output) > 0:
+                output_by_package[pp.package_name] = package_output
+        if len(output_by_package) > 0 and output_file_path is not None:
+            with open(output_file_path, "w") as fd:
+                # write document header
+                fd.write(f"Inspect cell results for model {self.name}\n")
+                output = []
+                for cell in cell_list:
+                    output.append(" ".join([str(i) for i in cell]))
+                output = ",".join(output)
+                fd.write(f"Model cells inspected,{output}\n\n")
+
+                for package_name, matches in output_by_package.items():
+                    fd.write(f"Results from {package_name} package\n")
+                    for search_output in matches:
+                        # write header line with data name
+                        fd.write(
+                            f",Results from "
+                            f"{search_output.path_to_data[-1]}\n"
+                        )
+                        # write data header
+                        if search_output.transient:
+                            fd.write(",stress_period/key")
+                        if search_output.data_header is not None:
+                            if len(search_output.data_entry_cellids) > 0:
+                                fd.write(",cellid")
+                            h_columns = ",".join(search_output.data_header)
+                            fd.write(f",{h_columns}\n")
+                        else:
+                            fd.write(f",cellid,data\n")
+                        # write data found
+                        for index, data_entry in enumerate(
+                            search_output.data_entries
+                        ):
+                            if search_output.transient:
+                                sp = search_output.data_entry_stress_period[
+                                    index
+                                ]
+                                fd.write(f",{sp}")
+                            if search_output.data_header is not None:
+                                if len(search_output.data_entry_cellids) > 0:
+                                    cells = search_output.data_entry_cellids[
+                                        index
+                                    ]
+                                    output = " ".join([str(i) for i in cells])
+                                    fd.write(f",{output}")
+                                fd.write(self._format_data_entry(data_entry))
+                            else:
+                                output = " ".join(
+                                    [
+                                        str(i)
+                                        for i in search_output.data_entry_ids[
+                                            index
+                                        ]
+                                    ]
+                                )
+                                fd.write(f",{output}")
+                                fd.write(self._format_data_entry(data_entry))
+                    fd.write(f"\n")
+        return output_by_package
+
+    @staticmethod
+    def _format_data_entry(data_entry):
+        output = ""
+        if iterable(data_entry, True):
+            for item in data_entry:
+                if isinstance(item, tuple):
+                    formatted = " ".join([str(i) for i in item])
+                    output = f"{output},{formatted}"
+                else:
+                    output = f"{output},{item}"
+            return f"{output}\n"
+        else:
+            return f",{data_entry}\n"
+
     def write(self, ext_file_action=ExtFileAction.copy_relative_paths):
         """
         Writes out model's package files.
diff --git a/flopy/mf6/mfpackage.py b/flopy/mf6/mfpackage.py
index 17c9f78021..5fd399eaf2 100644
--- a/flopy/mf6/mfpackage.py
+++ b/flopy/mf6/mfpackage.py
@@ -3,6 +3,7 @@
 import errno
 import inspect
 import datetime
+import copy
 import numpy as np
 
 from .mfbase import PackageContainer, ExtFileAction, PackageContainerType
@@ -18,9 +19,10 @@
 from .data import mfstructure, mfdata
 from ..utils import datautil
 from .data import mfdataarray, mfdatalist, mfdatascalar
+from .data.mfstructure import MFDataItemStructure
 from .coordinates import modeldimensions
 from ..pakbase import PackageInterface
-from .data.mfdatautil import MFComment
+from .data.mfdatautil import MFComment, cellids_equal, DataSearchOutput
 from ..utils.check import mf6check
 from .utils.output_util import MF6Output
 from ..mbase import ModelInterface
@@ -1983,6 +1985,271 @@ def _update_size_defs(self):
                                     )
                                 )
 
+    def inspect_cells(self, cell_list, stress_period=None):
+        """
+        Inspect model cells.  Returns package data associated with cells.
+
+        Parameters
+        ----------
+        cell_list : list of tuples
+            List of model cells.  Each model cell is a tuple of integers.
+            ex: [(1,1,1), (2,4,3)]
+        stress_period : int
+            For transient data, only return data from this stress period.  If
+            not specified or None, all stress period data will be returned.
+
+        Returns
+        -------
+        output : array
+            Array containing inspection results
+
+        """
+        data_found = []
+
+        # loop through blocks
+        local_index_names = []
+        local_index_blocks = []
+        local_index_values = []
+        local_index_cellids = []
+        # loop through blocks in package
+        for block in self.blocks.values():
+            # loop through data in block
+            for dataset in block.datasets.values():
+                if isinstance(dataset, mfdatalist.MFList):
+                    # handle list data
+                    cellid_column = None
+                    local_index_name = None
+                    # loop through list data column definitions
+                    for index, data_item in enumerate(
+                        dataset.structure.data_item_structures
+                    ):
+                        if index == 0 and data_item.type == DatumType.integer:
+                            local_index_name = data_item.name
+                        # look for cellid column in list data row
+                        if isinstance(data_item, MFDataItemStructure) and (
+                            data_item.is_cellid or data_item.possible_cellid
+                        ):
+                            cellid_column = index
+                            break
+                    if cellid_column is not None:
+                        data_output = DataSearchOutput(dataset.path)
+                        local_index_vals = []
+                        local_index_cells = []
+                        # get data
+                        if isinstance(dataset, mfdatalist.MFTransientList):
+                            # data may be in multiple transient blocks, get
+                            # data from appropriate blocks
+                            main_data = dataset.get_data(stress_period)
+                            if stress_period is not None:
+                                main_data = {stress_period: main_data}
+                        else:
+                            # data is all in one block, get data
+                            main_data = {-1: dataset.get_data()}
+
+                        # loop through each dataset
+                        for key, value in main_data.items():
+                            if data_output.data_header is None:
+                                data_output.data_header = value.dtype.names
+                            # loop through list data rows
+                            for line in value:
+                                # loop through list of cells we are searching
+                                # for
+                                for cell in cell_list:
+                                    if isinstance(
+                                        line[cellid_column], tuple
+                                    ) and cellids_equal(
+                                        line[cellid_column], cell
+                                    ):
+                                        # save data found
+                                        data_output.data_entries.append(line)
+                                        data_output.data_entry_ids.append(cell)
+                                        data_output.data_entry_stress_period.append(
+                                            key
+                                        )
+                                        if datautil.DatumUtil.is_int(line[0]):
+                                            # save index data for further
+                                            # processing.  assuming index is
+                                            # always first entry
+                                            local_index_vals.append(line[0])
+                                            local_index_cells.append(cell)
+
+                        if (
+                            local_index_name is not None
+                            and len(local_index_vals) > 0
+                        ):
+                            # capture index lookups for scanning related data
+                            local_index_names.append(local_index_name)
+                            local_index_blocks.append(block.path[-1])
+                            local_index_values.append(local_index_vals)
+                            local_index_cellids.append(local_index_cells)
+                        if len(data_output.data_entries) > 0:
+                            data_found.append(data_output)
+                elif isinstance(dataset, mfdataarray.MFArray):
+                    # handle array data
+                    data_shape = copy.deepcopy(
+                        dataset.structure.data_item_structures[0].shape
+                    )
+                    if dataset.path[-1] == "top":
+                        # top is a special case where the two datasets
+                        # need to be combined to get the correct layer top
+                        model_grid = self.model_or_sim.modelgrid
+                        main_data = {-1: model_grid.top_botm}
+                        data_shape.append("nlay")
+                    else:
+                        if isinstance(dataset, mfdataarray.MFTransientArray):
+                            # data may be in multiple blocks, get data from
+                            # appropriate blocks
+                            main_data = dataset.get_data(stress_period)
+                            if stress_period is not None:
+                                main_data = {stress_period: main_data}
+                        else:
+                            # data is all in one block, get a process data
+                            main_data = {-1: dataset.get_data()}
+                    if main_data is None:
+                        continue
+                    data_output = DataSearchOutput(dataset.path)
+                    # loop through datasets
+                    for key, array_data in main_data.items():
+                        if array_data is None:
+                            continue
+                        # loop through list of cells we are searching for
+                        for cell in cell_list:
+                            if (
+                                len(data_shape) == 3
+                                or data_shape[0] == "nodes"
+                            ):
+                                # data is by cell
+                                if array_data.ndim == 3 and len(cell) == 3:
+                                    data_output.data_entries.append(
+                                        array_data[cell[0], cell[1], cell[2]]
+                                    )
+                                    data_output.data_entry_ids.append(cell)
+                                    data_output.data_entry_stress_period.append(
+                                        key
+                                    )
+                                elif array_data.ndim == 2 and len(cell) == 2:
+                                    data_output.data_entries.append(
+                                        array_data[cell[0], cell[1]]
+                                    )
+                                    data_output.data_entry_ids.append(cell)
+                                    data_output.data_entry_stress_period.append(
+                                        key
+                                    )
+                                elif array_data.ndim == 1 and len(cell) == 1:
+                                    data_output.data_entries.append(
+                                        array_data[cell[0]]
+                                    )
+                                    data_output.data_entry_ids.append(cell)
+                                    data_output.data_entry_stress_period.append(
+                                        key
+                                    )
+                                else:
+                                    if (
+                                        self.simulation_data.verbosity_level.value
+                                        >= VerbosityLevel.normal.value
+                                    ):
+                                        warning_str = (
+                                            'WARNING: CellID "{}" not same '
+                                            "number of dimensions as data "
+                                            "{}.".format(cell, dataset.path)
+                                        )
+                                        print(warning_str)
+                            elif len(data_shape) == 2:
+                                # get data based on row/col
+                                if array_data.ndim == 2 and len(cell) == 3:
+                                    data_output.data_entries.append(
+                                        array_data[cell[1], cell[2]]
+                                    )
+                                    data_output.data_entry_ids.append(cell)
+                                    data_output.data_entry_stress_period.append(
+                                        key
+                                    )
+                                elif array_data.ndim == 1 and len(cell) == 2:
+                                    data_output.data_entries.append(
+                                        array_data[cell[1]]
+                                    )
+                                    data_output.data_entry_ids.append(cell)
+                                    data_output.data_entry_stress_period.append(
+                                        key
+                                    )
+                    if len(data_output.data_entries) > 0:
+                        data_found.append(data_output)
+
+        if len(local_index_names) > 0:
+            # look for data that shares the index value with data found
+            # for example a shared well or reach number
+            for block in self.blocks.values():
+                # loop through data
+                for dataset in block.datasets.values():
+                    if isinstance(dataset, mfdatalist.MFList):
+                        data_item = dataset.structure.data_item_structures[0]
+                        data_output = DataSearchOutput(dataset.path)
+                        # loop through previous data found
+                        for (
+                            local_index_name,
+                            local_index_vals,
+                            cell_ids,
+                            local_block_name,
+                        ) in zip(
+                            local_index_names,
+                            local_index_values,
+                            local_index_cellids,
+                            local_index_blocks,
+                        ):
+                            if local_block_name == block.path[-1]:
+                                continue
+                            if (
+                                isinstance(data_item, MFDataItemStructure)
+                                and data_item.name == local_index_name
+                                and data_item.type == DatumType.integer
+                            ):
+                                # matching data index type found, get data
+                                if isinstance(
+                                    dataset, mfdatalist.MFTransientList
+                                ):
+                                    # data may be in multiple blocks, get data
+                                    # from appropriate blocks
+                                    main_data = dataset.get_data(stress_period)
+                                    if stress_period is not None:
+                                        main_data = {stress_period: main_data}
+                                else:
+                                    # data is all in one block
+                                    main_data = {-1: dataset.get_data()}
+                                # loop through the data
+                                for key, value in main_data.items():
+                                    if value is None:
+                                        continue
+                                    if data_output.data_header is None:
+                                        data_output.data_header = (
+                                            value.dtype.names
+                                        )
+                                    # loop through each row of data
+                                    for line in value:
+                                        # loop through the index values we are
+                                        # looking for
+                                        for index_val, cell_id in zip(
+                                            local_index_vals, cell_ids
+                                        ):
+                                            # try to match index values we are
+                                            # looking for to the data
+                                            if index_val == line[0]:
+                                                # save data found
+                                                data_output.data_entries.append(
+                                                    line
+                                                )
+                                                data_output.data_entry_ids.append(
+                                                    index_val
+                                                )
+                                                data_output.data_entry_cellids.append(
+                                                    cell_id
+                                                )
+                                                data_output.data_entry_stress_period.append(
+                                                    key
+                                                )
+                        if len(data_output.data_entries) > 0:
+                            data_found.append(data_output)
+        return data_found
+
     def remove(self):
         """Removes this package from the simulation/model it is currently a
         part of.

From 442d9a7a6bab8328fc5df6248968c6918ad3d0b4 Mon Sep 17 00:00:00 2001
From: Scott Paulinski <spaulinski@usgs.gov>
Date: Fri, 17 Dec 2021 12:58:09 -0800
Subject: [PATCH 2/8] feat(inspect cells): added tests for new feature (#1140)

---
 autotest/t505_test.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/autotest/t505_test.py b/autotest/t505_test.py
index fe7b1c3ad5..720b7edc9a 100644
--- a/autotest/t505_test.py
+++ b/autotest/t505_test.py
@@ -908,6 +908,10 @@ def test_np001():
         ],
     )
 
+    cell_list = [(0, 0, 4), (0, 0, 9)]
+    out_file = os.path.join("temp", "inspect_test_np001.csv")
+    model.inspect_cells(cell_list, output_file_path=out_file)
+
     # make folder to save simulation
     sim.set_sim_path(run_folder)
 
@@ -2184,6 +2188,10 @@ def test005_advgw_tidal():
     ts_path = os.path.join(run_folder, "well-rates", "well-rates.ts")
     assert os.path.exists(ts_path)
 
+    cell_list = [(2, 3, 2), (0, 4, 2), (0, 2, 4), (0, 5, 5), (0, 9, 9)]
+    out_file = os.path.join("temp", "inspect_AdvGW_tidal.csv")
+    model.inspect_cells(cell_list, output_file_path=out_file)
+
     # run simulation
     sim.run_simulation()
 
@@ -3521,6 +3529,10 @@ def test028_sfr():
     assert sfr_package.connectiondata.get_data()[2][1] == 1.0
     assert sfr_package.packagedata.get_data()[1][1].lower() == "none"
 
+    cell_list = [(0, 2, 3), (0, 3, 4), (0, 4, 5)]
+    out_file = os.path.join("temp", "inspect_test028_sfr.csv")
+    model.inspect_cells(cell_list, output_file_path=out_file)
+
     sim.set_sim_path(run_folder)
     sim.write_simulation()
     sim.load(
@@ -3851,16 +3863,16 @@ def test_transport():
 
 
 if __name__ == "__main__":
+    test_np001()
+    test028_sfr()
     test_array()
     test_multi_model()
-    test_np001()
     test_np002()
     test004_bcfss()
     test005_advgw_tidal()
     test006_2models_gnc()
     test006_gwf3_disv()
     test021_twri()
-    test028_sfr()
     test035_fhb()
     test050_circle_island()
     test_transport()

From 6e386a032a98e3d5f0130e109b501ced193df222 Mon Sep 17 00:00:00 2001
From: Scott Paulinski <spaulinski@usgs.gov>
Date: Fri, 17 Dec 2021 13:42:35 -0800
Subject: [PATCH 3/8] feat(inspect cells)

---
 flopy/mf6/mfpackage.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/flopy/mf6/mfpackage.py b/flopy/mf6/mfpackage.py
index 5fd399eaf2..bd7c01a62a 100644
--- a/flopy/mf6/mfpackage.py
+++ b/flopy/mf6/mfpackage.py
@@ -2048,6 +2048,8 @@ def inspect_cells(self, cell_list, stress_period=None):
 
                         # loop through each dataset
                         for key, value in main_data.items():
+                            if value is None:
+                                continue
                             if data_output.data_header is None:
                                 data_output.data_header = value.dtype.names
                             # loop through list data rows

From acbae8e64e7d31aa014f407c817ae1a0a7a00320 Mon Sep 17 00:00:00 2001
From: Scott Paulinski <spaulinski@usgs.gov>
Date: Tue, 21 Dec 2021 08:12:39 -0800
Subject: [PATCH 4/8] feat(inspect cells): dependent variable and budget
 inspection

---
 flopy/mf6/data/mfdatautil.py |   1 +
 flopy/mf6/mfmodel.py         | 176 +++++++++++++++++++++++++++++++++--
 flopy/mf6/mfpackage.py       |  63 +------------
 3 files changed, 172 insertions(+), 68 deletions(-)

diff --git a/flopy/mf6/data/mfdatautil.py b/flopy/mf6/data/mfdatautil.py
index 8afd80aff7..7a7528377f 100644
--- a/flopy/mf6/data/mfdatautil.py
+++ b/flopy/mf6/data/mfdatautil.py
@@ -248,6 +248,7 @@ def __init__(self, path_to_data=None, data_header=None):
         self.data_entry_cellids = []
         self.data_entry_stress_period = []
         self.data_entries = []
+        self.output = False
 
     @property
     def transient(self):
diff --git a/flopy/mf6/mfmodel.py b/flopy/mf6/mfmodel.py
index 5a295a5b30..cbbef59b2a 100644
--- a/flopy/mf6/mfmodel.py
+++ b/flopy/mf6/mfmodel.py
@@ -20,7 +20,7 @@
 from ..mbase import ModelInterface
 from .utils.mfenums import DiscretizationType
 from .data import mfstructure
-from .data.mfdatautil import iterable
+from .data.mfdatautil import iterable, DataSearchOutput
 from .utils.output_util import MF6Output
 from ..utils.check import mf6check
 
@@ -827,7 +827,12 @@ def load_base(
         return instance
 
     def inspect_cells(
-        self, cell_list, stress_period=None, output_file_path=None
+        self,
+        cell_list,
+        stress_period=None,
+        output_file_path=None,
+        inspect_budget=True,
+        inspect_dependent_var=True,
     ):
         """
         Inspect model cells.  Returns model data associated with cells.
@@ -841,8 +846,11 @@ def inspect_cells(
             For transient data qnly return data from this stress period.  If
             not specified or None, all stress period data will be returned.
         output_file_path: str
-            path to output file that will contain the inspection results
-
+            Path to output file that will contain the inspection results
+        inspect_budget: bool
+            Inpect budget file
+        inspect_dependent_var: bool
+            Inpect head file
         Returns
         -------
         output : dict
@@ -858,13 +866,106 @@ def inspect_cells(
         >>> out_file = os.path.join("temp", "inspect_AdvGW_tidal.csv")
         >>> model.inspect_cells(inspect_list, output_file_path=out_file)
         """
+        # handle no cell case
+        if cell_list is None or len(cell_list) == 0:
+            return None
+
         output_by_package = {}
         # loop through all packages
         for pp in self.packagelist:
             # call the package's "inspect_cells" method
             package_output = pp.inspect_cells(cell_list, stress_period)
             if len(package_output) > 0:
-                output_by_package[pp.package_name] = package_output
+                output_by_package[
+                    f"{pp.package_name} package"
+                ] = package_output
+        # get dependent variable
+        if inspect_dependent_var:
+            try:
+                if self.model_type == "gwf6":
+                    heads = self.output.head()
+                    name = "heads"
+                elif self.model_type == "gwt6":
+                    heads = self.output.concentration()
+                    name = "concentration"
+                else:
+                    inspect_dependent_var = False
+            except Exception:
+                inspect_dependent_var = False
+        if inspect_dependent_var and heads is not None:
+            kstp_kper_lst = heads.get_kstpkper()
+            data_output = DataSearchOutput((name,))
+            data_output.output = True
+            for kstp_kper in kstp_kper_lst:
+                head_array = np.array(heads.get_data(kstpkper=kstp_kper))
+                # flatten output data in disv and disu cases
+                if len(cell_list[0]) == 2:
+                    head_array = head_array[0, :, :]
+                elif len(cell_list[0]) == 1:
+                    head_array = head_array[0, 0, :]
+                # find data matches
+                self.match_array_cells(
+                    cell_list,
+                    head_array.shape,
+                    head_array,
+                    kstp_kper,
+                    data_output,
+                )
+            if len(data_output.data_entries) > 0:
+                output_by_package[f"{name} output"] = [data_output]
+
+        # get model dimensions
+        model_shape = self.modelgrid.shape
+
+        # get budgets
+        if inspect_budget:
+            try:
+                bud = self.output.budget()
+            except Exception:
+                inspect_budget = False
+        if inspect_budget and bud is not None:
+            kstp_kper_lst = bud.get_kstpkper()
+            rec_names = bud.get_unique_record_names()
+            budget_matches = []
+            for rec_name in rec_names:
+                # clean up binary string name
+                string_name = str(rec_name)[3:-1].strip()
+                data_output = DataSearchOutput((string_name,))
+                data_output.output = True
+                for kstp_kper in kstp_kper_lst:
+                    budget_array = np.array(
+                        bud.get_data(
+                            kstpkper=kstp_kper,
+                            text=rec_name,
+                            full3D=True,
+                        )
+                    )
+                    if len(budget_array.shape) == 4:
+                        # get rid of 4th "time" dimension
+                        budget_array = budget_array[0, :, :, :]
+                    # flatten output data in disv and disu cases
+                    if len(cell_list[0]) == 2:
+                        budget_array = budget_array[0, :, :]
+                    elif len(cell_list[0]) == 1:
+                        budget_array = budget_array[0, :]
+                    # find data matches
+                    if budget_array.shape != model_shape:
+                        # no support yet for different shaped budgets like
+                        # flow_ja_face
+                        continue
+
+                    self.match_array_cells(
+                        cell_list,
+                        budget_array.shape,
+                        budget_array,
+                        kstp_kper,
+                        data_output,
+                    )
+                if len(data_output.data_entries) > 0:
+                    budget_matches.append(data_output)
+            if len(budget_matches) > 0:
+                output_by_package["budget output"] = budget_matches
+
         if len(output_by_package) > 0 and output_file_path is not None:
             with open(output_file_path, "w") as fd:
                 # write document header
@@ -876,7 +977,7 @@ def inspect_cells(
                 fd.write(f"Model cells inspected,{output}\n\n")
 
                 for package_name, matches in output_by_package.items():
-                    fd.write(f"Results from {package_name} package\n")
+                    fd.write(f"Results from {package_name}\n")
                     for search_output in matches:
                         # write header line with data name
                         fd.write(
@@ -885,7 +986,10 @@ def inspect_cells(
                         )
                         # write data header
                         if search_output.transient:
-                            fd.write(",stress_period/key")
+                            if search_output.output:
+                                fd.write(",stress_period,time_step")
+                            else:
+                                fd.write(",stress_period/key")
                         if search_output.data_header is not None:
                             if len(search_output.data_entry_cellids) > 0:
                                 fd.write(",cellid")
@@ -901,7 +1005,10 @@ def inspect_cells(
                                 sp = search_output.data_entry_stress_period[
                                     index
                                 ]
-                                fd.write(f",{sp}")
+                                if search_output.output:
+                                    fd.write(f",{sp[1]},{sp[0]}")
+                                else:
+                                    fd.write(f",{sp}")
                             if search_output.data_header is not None:
                                 if len(search_output.data_entry_cellids) > 0:
                                     cells = search_output.data_entry_cellids[
@@ -924,6 +1031,59 @@ def inspect_cells(
                     fd.write(f"\n")
         return output_by_package
 
+    def match_array_cells(
+        self, cell_list, data_shape, array_data, key, data_output
+    ):
+        # loop through list of cells we are searching for
+        for cell in cell_list:
+            if len(data_shape) == 3 or data_shape[0] == "nodes":
+                # data is by cell
+                if array_data.ndim == 3 and len(cell) == 3:
+                    data_output.data_entries.append(
+                        array_data[cell[0], cell[1], cell[2]]
+                    )
+                    data_output.data_entry_ids.append(cell)
+                    data_output.data_entry_stress_period.append(key)
+                elif array_data.ndim == 2 and len(cell) == 2:
+                    data_output.data_entries.append(
+                        array_data[cell[0], cell[1]]
+                    )
+                    data_output.data_entry_ids.append(cell)
+                    data_output.data_entry_stress_period.append(key)
+                elif array_data.ndim == 1 and len(cell) == 1:
+                    data_output.data_entries.append(array_data[cell[0]])
+                    data_output.data_entry_ids.append(cell)
+                    data_output.data_entry_stress_period.append(key)
+                else:
+                    if (
+                        self.simulation_data.verbosity_level.value
+                        >= VerbosityLevel.normal.value
+                    ):
+                        warning_str = (
+                            'WARNING: CellID "{}" not same '
+                            "number of dimensions as data "
+                            "{}.".format(cell, data_output.path_to_data)
+                        )
+                        print(warning_str)
+            elif len(data_shape) == 2:
+                # get data based on ncpl/lay
+                if array_data.ndim == 2 and len(cell) == 2:
+                    data_output.data_entries.append(
+                        array_data[cell[0], cell[1]]
+                    )
+                    data_output.data_entry_ids.append(cell)
+                    data_output.data_entry_stress_period.append(key)
+                elif array_data.ndim == 1 and len(cell) == 1:
+                    data_output.data_entries.append(array_data[cell[0]])
+                    data_output.data_entry_ids.append(cell)
+                    data_output.data_entry_stress_period.append(key)
+            elif len(data_shape) == 1:
+                # get data based on nodes
+                if len(cell) == 1 and array_data.ndim == 1:
+                    data_output.data_entries.append(array_data[cell[0]])
+                    data_output.data_entry_ids.append(cell)
+                    data_output.data_entry_stress_period.append(key)
+
     @staticmethod
     def _format_data_entry(data_entry):
         output = ""
diff --git a/flopy/mf6/mfpackage.py b/flopy/mf6/mfpackage.py
index bd7c01a62a..10c93cd00e 100644
--- a/flopy/mf6/mfpackage.py
+++ b/flopy/mf6/mfpackage.py
@@ -2114,66 +2114,9 @@ def inspect_cells(self, cell_list, stress_period=None):
                     for key, array_data in main_data.items():
                         if array_data is None:
                             continue
-                        # loop through list of cells we are searching for
-                        for cell in cell_list:
-                            if (
-                                len(data_shape) == 3
-                                or data_shape[0] == "nodes"
-                            ):
-                                # data is by cell
-                                if array_data.ndim == 3 and len(cell) == 3:
-                                    data_output.data_entries.append(
-                                        array_data[cell[0], cell[1], cell[2]]
-                                    )
-                                    data_output.data_entry_ids.append(cell)
-                                    data_output.data_entry_stress_period.append(
-                                        key
-                                    )
-                                elif array_data.ndim == 2 and len(cell) == 2:
-                                    data_output.data_entries.append(
-                                        array_data[cell[0], cell[1]]
-                                    )
-                                    data_output.data_entry_ids.append(cell)
-                                    data_output.data_entry_stress_period.append(
-                                        key
-                                    )
-                                elif array_data.ndim == 1 and len(cell) == 1:
-                                    data_output.data_entries.append(
-                                        array_data[cell[0]]
-                                    )
-                                    data_output.data_entry_ids.append(cell)
-                                    data_output.data_entry_stress_period.append(
-                                        key
-                                    )
-                                else:
-                                    if (
-                                        self.simulation_data.verbosity_level.value
-                                        >= VerbosityLevel.normal.value
-                                    ):
-                                        warning_str = (
-                                            'WARNING: CellID "{}" not same '
-                                            "number of dimensions as data "
-                                            "{}.".format(cell, dataset.path)
-                                        )
-                                        print(warning_str)
-                            elif len(data_shape) == 2:
-                                # get data based on row/col
-                                if array_data.ndim == 2 and len(cell) == 3:
-                                    data_output.data_entries.append(
-                                        array_data[cell[1], cell[2]]
-                                    )
-                                    data_output.data_entry_ids.append(cell)
-                                    data_output.data_entry_stress_period.append(
-                                        key
-                                    )
-                                elif array_data.ndim == 1 and len(cell) == 2:
-                                    data_output.data_entries.append(
-                                        array_data[cell[1]]
-                                    )
-                                    data_output.data_entry_ids.append(cell)
-                                    data_output.data_entry_stress_period.append(
-                                        key
-                                    )
+                        self.model_or_sim.match_array_cells(
+                            cell_list, data_shape, array_data, key, data_output
+                        )
                     if len(data_output.data_entries) > 0:
                         data_found.append(data_output)
 

From 6d39d543a9ce3a8e026cd03b34d7086ca3e04eef Mon Sep 17 00:00:00 2001
From: Scott Paulinski <spaulinski@usgs.gov>
Date: Tue, 21 Dec 2021 08:14:05 -0800
Subject: [PATCH 5/8] feat(inspect cells): tests

---
 autotest/t504_test.py |  5 +++++
 autotest/t505_test.py | 45 ++++++++++++++++++++++++++++++-------------
 2 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/autotest/t504_test.py b/autotest/t504_test.py
index 339f51b21f..96f90caf47 100644
--- a/autotest/t504_test.py
+++ b/autotest/t504_test.py
@@ -386,6 +386,11 @@ def test006_gwf3():
         success, buff = sim.run_simulation()
         assert success, f"simulation {sim.name} rerun did not run"
 
+        # inspect cells
+        cell_list = [(0,), (7,), (14,)]
+        out_file = os.path.join("temp", "inspect_test006_gwf3.csv")
+        model.inspect_cells(cell_list, output_file_path=out_file)
+
         budget_obj = bf.CellBudgetFile(expected_cbc_file_a, precision="double")
         budget_fjf_valid = np.array(
             budget_obj.get_data(text="    FLOW JA FACE", full3D=True)
diff --git a/autotest/t505_test.py b/autotest/t505_test.py
index 720b7edc9a..84b38df460 100644
--- a/autotest/t505_test.py
+++ b/autotest/t505_test.py
@@ -908,10 +908,6 @@ def test_np001():
         ],
     )
 
-    cell_list = [(0, 0, 4), (0, 0, 9)]
-    out_file = os.path.join("temp", "inspect_test_np001.csv")
-    model.inspect_cells(cell_list, output_file_path=out_file)
-
     # make folder to save simulation
     sim.set_sim_path(run_folder)
 
@@ -932,6 +928,13 @@ def test_np001():
     if run:
         sim.run_simulation()
 
+        # inspect cells
+        cell_list = [(0, 0, 4), (0, 0, 9)]
+        out_file = os.path.join("temp", "inspect_test_np001.csv")
+        model.inspect_cells(
+            cell_list, output_file_path=out_file, stress_period=1
+        )
+
         # get expected results
         budget_obj = bf.CellBudgetFile(expected_cbc_file, precision="double")
         budget_frf_valid = np.array(
@@ -2188,13 +2191,14 @@ def test005_advgw_tidal():
     ts_path = os.path.join(run_folder, "well-rates", "well-rates.ts")
     assert os.path.exists(ts_path)
 
+    # run simulation
+    sim.run_simulation()
+
+    # inspect cells
     cell_list = [(2, 3, 2), (0, 4, 2), (0, 2, 4), (0, 5, 5), (0, 9, 9)]
     out_file = os.path.join("temp", "inspect_AdvGW_tidal.csv")
     model.inspect_cells(cell_list, output_file_path=out_file)
 
-    # run simulation
-    sim.run_simulation()
-
     # compare output to expected results
     head_new = os.path.join(run_folder, "AdvGW_tidal.hds")
     outfile = os.path.join(run_folder, "head_compare.dat")
@@ -2853,6 +2857,11 @@ def test006_gwf3_disv():
     if run:
         sim.run_simulation()
 
+        # inspect cells
+        cell_list = [(0, 0), (0, 7), (0, 17)]
+        out_file = os.path.join("temp", "inspect_test_gwf3_disv.csv")
+        model.inspect_cells(cell_list, output_file_path=out_file)
+
         # compare output to expected results
         head_new = os.path.join(run_folder, "flow.hds")
         outfile = os.path.join(run_folder, "head_compare.dat")
@@ -3529,10 +3538,6 @@ def test028_sfr():
     assert sfr_package.connectiondata.get_data()[2][1] == 1.0
     assert sfr_package.packagedata.get_data()[1][1].lower() == "none"
 
-    cell_list = [(0, 2, 3), (0, 3, 4), (0, 4, 5)]
-    out_file = os.path.join("temp", "inspect_test028_sfr.csv")
-    model.inspect_cells(cell_list, output_file_path=out_file)
-
     sim.set_sim_path(run_folder)
     sim.write_simulation()
     sim.load(
@@ -3602,6 +3607,11 @@ def test028_sfr():
     if run:
         sim.run_simulation()
 
+        # inspect cells
+        cell_list = [(0, 2, 3), (0, 3, 4), (0, 4, 5)]
+        out_file = os.path.join("temp", "inspect_test028_sfr.csv")
+        model.inspect_cells(cell_list, output_file_path=out_file)
+
         # compare output to expected results
         head_new = os.path.join(run_folder, "test1tr.hds")
         outfile = os.path.join(run_folder, "head_compare.dat")
@@ -3839,6 +3849,15 @@ def test_transport():
     if run:
         sim.run_simulation()
 
+        # inspect cells
+        cell_list = [
+            (0, 0, 0),
+        ]
+        out_file = os.path.join("temp", "inspect_transport_gwf.csv")
+        gwf.inspect_cells(cell_list, output_file_path=out_file)
+        out_file = os.path.join("temp", "inspect_transport_gwt.csv")
+        gwt.inspect_cells(cell_list, output_file_path=out_file)
+
         # compare output to expected results
         head_new = os.path.join(run_folder, "gwf_mst03.hds")
         outfile = os.path.join(run_folder, "head_compare.dat")
@@ -3863,16 +3882,16 @@ def test_transport():
 
 
 if __name__ == "__main__":
-    test_np001()
-    test028_sfr()
     test_array()
     test_multi_model()
+    test_np001()
     test_np002()
     test004_bcfss()
     test005_advgw_tidal()
     test006_2models_gnc()
     test006_gwf3_disv()
     test021_twri()
+    test028_sfr()
     test035_fhb()
     test050_circle_island()
     test_transport()

From 37755d9019416b7e5860b28d843e4e966495efe8 Mon Sep 17 00:00:00 2001
From: Scott Paulinski <spaulinski@usgs.gov>
Date: Tue, 21 Dec 2021 09:12:56 -0800
Subject: [PATCH 6/8] feat(inspect cells): doc update

---
 flopy/mf6/mfmodel.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/flopy/mf6/mfmodel.py b/flopy/mf6/mfmodel.py
index cbbef59b2a..d3800f3f67 100644
--- a/flopy/mf6/mfmodel.py
+++ b/flopy/mf6/mfmodel.py
@@ -848,9 +848,9 @@ def inspect_cells(
         output_file_path: str
             Path to output file that will contain the inspection results
         inspect_budget: bool
-            Inpect budget file
+            Inspect budget file
         inspect_dependent_var: bool
-            Inpect head file
+            Inspect head file
         Returns
         -------
         output : dict

From 058617ea0de283fc0c0e7ba39a4ef06e97db958a Mon Sep 17 00:00:00 2001
From: Scott Paulinski <spaulinski@usgs.gov>
Date: Tue, 11 Jan 2022 09:18:50 -0800
Subject: [PATCH 7/8] feat(inspect_cells)

---
 flopy/mf6/mfmodel.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/flopy/mf6/mfmodel.py b/flopy/mf6/mfmodel.py
index d3800f3f67..748034942e 100644
--- a/flopy/mf6/mfmodel.py
+++ b/flopy/mf6/mfmodel.py
@@ -897,6 +897,8 @@ def inspect_cells(
             data_output = DataSearchOutput((name,))
             data_output.output = True
             for kstp_kper in kstp_kper_lst:
+                if stress_period is not None and stress_period != kstp_kper[1]:
+                    continue
                 head_array = np.array(heads.get_data(kstpkper=kstp_kper))
                 # flatten output data in disv and disu cases
                 if len(cell_list[0]) == 2:
@@ -933,20 +935,27 @@ def inspect_cells(
                 data_output = DataSearchOutput((string_name,))
                 data_output.output = True
                 for kstp_kper in kstp_kper_lst:
+                    if (
+                        stress_period is not None
+                        and stress_period != kstp_kper[1]
+                    ):
+                        continue
                     budget_array = np.array(
                         bud.get_data(
                             kstpkper=kstp_kper,
                             text=rec_name,
                             full3D=True,
-                        )
+                        )[0]
                     )
                     if len(budget_array.shape) == 4:
                         # get rid of 4th "time" dimension
                         budget_array = budget_array[0, :, :, :]
                     # flatten output data in disv and disu cases
-                    if len(cell_list[0]) == 2:
+                    if len(cell_list[0]) == 2 and len(budget_array.shape) >= 3:
                         budget_array = budget_array[0, :, :]
-                    elif len(cell_list[0]) == 1:
+                    elif (
+                        len(cell_list[0]) == 1 and len(budget_array.shape) >= 2
+                    ):
                         budget_array = budget_array[0, :]
                     # find data matches
                     if budget_array.shape != model_shape:

From c75b88f3f9fba4f2e61d6ddd8983142533b9b04d Mon Sep 17 00:00:00 2001
From: Scott Paulinski <spaulinski@usgs.gov>
Date: Tue, 11 Jan 2022 09:19:37 -0800
Subject: [PATCH 8/8] feat(inspect_cells)

---
 autotest/t504_test.py | 14 ++++++++++++++
 autotest/t505_test.py |  8 ++++++--
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/autotest/t504_test.py b/autotest/t504_test.py
index 96f90caf47..cc8f5ac17f 100644
--- a/autotest/t504_test.py
+++ b/autotest/t504_test.py
@@ -747,6 +747,10 @@ def test006_2models_mvr():
         success, buff = sim.run_simulation()
         assert success, f"simulation {sim.name} rerun did not run"
 
+        cell_list = [(0, 3, 1)]
+        out_file = os.path.join("temp", "inspect_test006_2models_mvr.csv")
+        models[0].inspect_cells(cell_list, output_file_path=out_file)
+
         # compare output to expected results
         head_new = os.path.join(save_folder, "model1.hds")
         assert pymake.compare_heads(
@@ -855,6 +859,11 @@ def test001e_uzf_3lay():
         success, buff = sim.run_simulation()
         assert success, f"simulation {sim.name} rerun did not run"
 
+        # inspect cells
+        cell_list = [(0, 0, 1), (0, 0, 2), (2, 0, 8)]
+        out_file = os.path.join("temp", "inspect_test001e_uzf_3lay.csv")
+        model.inspect_cells(cell_list, output_file_path=out_file)
+
     # test load_only
     model_package_check = ["chd", "ic", "npf", "oc", "sto", "uzf"]
     load_only_lists = [
@@ -946,6 +955,11 @@ def test045_lake2tr():
         success, buff = sim.run_simulation()
         assert success, f"simulation {sim.name} rerun did not run"
 
+        # inspect cells
+        cell_list = [(0, 6, 5), (0, 8, 5), (1, 18, 6)]
+        out_file = os.path.join("temp", "inspect_test045_lake2tr.csv")
+        model.inspect_cells(cell_list, output_file_path=out_file)
+
         # compare output to expected results
         head_new = os.path.join(save_folder, "lakeex2a.hds")
         assert pymake.compare_heads(
diff --git a/autotest/t505_test.py b/autotest/t505_test.py
index 84b38df460..7bbcec32f4 100644
--- a/autotest/t505_test.py
+++ b/autotest/t505_test.py
@@ -929,10 +929,10 @@ def test_np001():
         sim.run_simulation()
 
         # inspect cells
-        cell_list = [(0, 0, 4), (0, 0, 9)]
+        cell_list = [(0, 0, 0), (0, 0, 4), (0, 0, 9)]
         out_file = os.path.join("temp", "inspect_test_np001.csv")
         model.inspect_cells(
-            cell_list, output_file_path=out_file, stress_period=1
+            cell_list, output_file_path=out_file, stress_period=0
         )
 
         # get expected results
@@ -1344,6 +1344,10 @@ def test_np002():
         # run simulation
         sim.run_simulation()
 
+        cell_list = [(0, 0, 0), (0, 0, 3), (0, 0, 4), (0, 0, 9)]
+        out_file = os.path.join("temp", "inspect_test_np002.csv")
+        model.inspect_cells(cell_list, output_file_path=out_file)
+
         sim2 = MFSimulation.load(sim_ws=run_folder)
         model_ = sim2.get_model(model_name)
         npf_package = model_.get_package("npf")