feat(plot_bc): add subset argument (#2694)

* address review comments and support plotAll with subset
* Add assertions for the number of subset bcs

Author: jdhughes-dev

autotest/test_plot_cross_section.py

@@ -26,6 +26,26 @@ def test_cross_section_bc_gwfs_disv(example_data_path):
         )
 
 
+@pytest.mark.mf6
+@pytest.mark.xfail(reason="sometimes get LineCollections instead of PatchCollections")
+def test_cross_section_bc_gwfs_disv_subset(example_data_path):
+    mpath = example_data_path / "mf6" / "test003_gwfs_disv"
+    sim = MFSimulation.load(sim_ws=mpath)
+    ml6 = sim.get_model("gwf_1")
+    xc = flopy.plot.PlotCrossSection(ml6, line={"line": ([0, 5.5], [10, 5.5])})
+    xc.plot_bc("CHD", subset=[(0, 49)])
+    ax = xc.ax
+
+    assert len(ax.collections) != 0, "Boundary condition was not drawn"
+
+    for col in ax.collections:
+        assert isinstance(col, PatchCollection), (
+            f"Unexpected collection type: {type(col)}"
+        )
+        count = col.get_array().count()
+        assert count == 1, f"More than one CHD plotted ({count})"
+
+
 @pytest.mark.mf6
 @pytest.mark.xfail(reason="sometimes get LineCollections instead of PatchCollections")
 def test_cross_section_bc_lake2tr(example_data_path):

autotest/test_plot_map_view.py

@@ -101,6 +101,96 @@ def test_map_view_bc_gwfs_disv(example_data_path):
         )
 
 
+@pytest.mark.mf2005
+@pytest.mark.xfail(reason="sometimes get wrong collection type")
+def test_map_view_bc_freyberg_subset(example_data_path):
+    mpath = example_data_path / "freyberg"
+    name_file = "freyberg.nam"
+    ml = Modflow.load(name_file, model_ws=mpath, verbose=True)
+    mapview = flopy.plot.PlotMapView(model=ml)
+    mapview.plot_bc(
+        "RIV",
+        subset=[
+            (0, 34, 14),
+            (0, 35, 14),
+            (0, 36, 14),
+            (0, 37, 14),
+            (0, 38, 14),
+            (0, 39, 14),
+        ],
+    )
+    ax = mapview.ax
+
+    if len(ax.collections) == 0:
+        raise AssertionError("Boundary condition was not drawn")
+
+    for col in ax.collections:
+        assert isinstance(col, (QuadMesh, PathCollection)), (
+            f"Unexpected collection type: {type(col)}"
+        )
+        if isinstance(col, QuadMesh):
+            count = col.get_array().count()
+            assert count == 6, f"More than six river cells plotted ({count})"
+
+    # plt.show(block=True)
+
+
+@pytest.mark.mf2005
+@pytest.mark.xfail(reason="sometimes get wrong collection type")
+def test_map_view_bc_freyberg_ml_subset_plotAll(example_data_path):
+    mpath = example_data_path / "freyberg_multilayer_transient"
+    name_file = "freyberg.nam"
+    ml = Modflow.load(name_file, model_ws=mpath, verbose=True)
+    mapview = flopy.plot.PlotMapView(model=ml, layer=2)
+    mapview.plot_bc(
+        "WEL",
+        plotAll=True,
+        subset=[
+            (0, 8, 15),
+            (0, 28, 5),
+        ],
+    )
+
+    ax = mapview.ax
+
+    if len(ax.collections) == 0:
+        raise AssertionError("Boundary condition was not drawn")
+
+    for col in ax.collections:
+        assert isinstance(col, (QuadMesh, PathCollection, LineCollection)), (
+            f"Unexpected collection type: {type(col)}"
+        )
+        if isinstance(col, QuadMesh):
+            count = col.get_array().count()
+            assert count == 2, f"More than two wells plotted ({count})"
+
+    # plt.show(block=True)
+
+
+@pytest.mark.mf6
+@pytest.mark.xfail(reason="sometimes get wrong collection type")
+def test_map_view_bc_gwfs_disv_subset(example_data_path):
+    mpath = example_data_path / "mf6" / "test003_gwfs_disv"
+    sim = MFSimulation.load(sim_ws=mpath)
+    ml6 = sim.get_model("gwf_1")
+    ml6.modelgrid.set_coord_info(angrot=-14)
+    mapview = flopy.plot.PlotMapView(model=ml6)
+    mapview.plot_bc(
+        "CHD", subset=[(0, 10), (0, 30), (0, 50), (0, 70), (0, 90), (0, 49)]
+    )
+    ax = mapview.ax
+
+    if len(ax.collections) == 0:
+        raise AssertionError("Boundary condition was not drawn")
+
+    for col in ax.collections:
+        assert isinstance(col, (QuadMesh, PathCollection)), (
+            f"Unexpected collection type: {type(col)}"
+        )
+
+    # plt.show(block=True)
+
+
 @pytest.mark.mf6
 @pytest.mark.xfail(reason="sometimes get wrong collection type")
 def test_map_view_bc_lake2tr(example_data_path):

flopy/plot/crosssection.py

@@ -881,7 +881,16 @@ def _plot_vertical_hfb_lines(self, color=None, **kwargs):
 
         return lc
 
-    def plot_bc(self, name=None, package=None, kper=0, color=None, head=None, **kwargs):
+    def plot_bc(
+        self,
+        name=None,
+        package=None,
+        kper=0,
+        color=None,
+        head=None,
+        subset=None,
+        **kwargs,
+    ):
         """
         Plot boundary conditions locations for a specific boundary
         type from a flopy model
@@ -902,6 +911,15 @@ def plot_bc(self, name=None, package=None, kper=0, color=None, head=None, **kwar
             to set top of patches to the minimum of the top of a\
             layer or the head value. Used to create
             patches that conform to water-level elevations.
+        subset : int, tuple of ints, or list of such
+            Subset of valid cellids. Acceptable values depend on grid type:
+
+            - Structured grids (DIS): (layer, row, column) or list of such
+            - Vertex grids (DISV): (layer, cellid) or list of such
+            - Unstructured grids (DISU): node number or list of such
+
+            All indices must be zero-based.
+
         **kwargs : dictionary
             keyword arguments passed to matplotlib.collections.PatchCollection
 
@@ -1013,6 +1031,20 @@ def plot_bc(self, name=None, package=None, kper=0, color=None, head=None, **kwar
             idx = idx.flatten()
             plotarray[idx] = 1
 
+        if subset is not None:
+            if isinstance(subset, (int, tuple)):
+                subset = [subset]
+            subset = tuple(np.array(subset).T)
+            if len(subset) != len(plotarray.shape):
+                msg = (
+                    f"The subset dimensions ({len(subset)}) is not equal to the "
+                    + f"grid dimensions ({len(plotarray.shape)})"
+                )
+                raise IndexError(msg)
+            mask = np.zeros(plotarray.shape, dtype=plotarray.dtype)
+            mask[subset] = 1
+            plotarray *= mask
+
         plotarray = np.ma.masked_equal(plotarray, 0)
         if color is None:
             key = name[:3].upper()

flopy/plot/map.py

@@ -566,6 +566,7 @@ def plot_bc(
         color=None,
         plotAll=False,
         boundname=None,
+        subset=None,
         **kwargs,
     ):
         """
@@ -586,6 +587,16 @@ def plot_bc(
             Boolean used to specify that boundary condition locations for all
             layers will be plotted on the current ModelMap layer.
             (Default is False)
+        boundname : string
+            select boundary conditions with specific boundname
+        subset : int, tuple of ints, or list of such
+            Subset of valid cellids. Acceptable values depend on grid type:
+
+            - Structured grids (DIS): (layer, row, column) or list of such
+            - Vertex grids (DISV): (layer, cellid) or list of such
+            - Unstructured grids (DISU): node number or list of such
+
+            All indices must be zero-based.
         **kwargs : dictionary
             keyword arguments passed to matplotlib.collections.PatchCollection
 
@@ -690,6 +701,25 @@ def plot_bc(
         else:
             plotarray[idx] = 1
 
+        if subset is not None:
+            if isinstance(subset, (int, tuple)):
+                subset = [subset]
+            subset = tuple(np.array(subset).T)
+            if len(subset) != len(plotarray.shape):
+                msg = (
+                    f"The subset dimensions ({len(subset)}) is not equal to the "
+                    + f"grid dimensions ({len(plotarray.shape)})"
+                )
+                raise IndexError(msg)
+            mask = np.zeros(plotarray.shape, dtype=plotarray.dtype)
+            mask[subset] = 1
+            if plotAll and len(self.mg.shape) > 1:
+                arr_sum = np.sum(mask, axis=0)
+                arr_sum[arr_sum > 0] = 1
+                for k in range(nlay):
+                    mask[k] = arr_sum.copy()
+            plotarray *= mask
+
         # mask the plot array
         plotarray = np.ma.masked_equal(plotarray, 0)