fix(get_transmissivities): support DISU (#2685)

Follow up on #2611 to support DISU. Closes #2170.

Plus forward-looking refactoring in case we add more grid types before 4.x comes out. Rather than conditionals everywhere to access/modify grid properties differently based on the grid type, add an adapter exposing what's needed consistently. This is necessary because grid properties may have different shapes depending on the grid type. In 4.x this can be avoided by design I think. But for now this makes the essence of the transmissivities calculation clearer.

Author: wpbonelli

autotest/test_postprocessing.py

@@ -584,6 +584,55 @@ def test_get_transmissivities_mf6_vertex(function_tmpdir):
         assert "r, c parameters only valid for structured grids" in str(e.value)
 
 
+def test_get_transmissivities_mf6_unstructured(function_tmpdir):
+    nl = 1
+    nr = 8
+    nc = 8
+    botm = np.ones((nl, nr, nc), dtype=float)
+    top = np.ones((nr, nc), dtype=float) * 2.1
+    hk = np.ones((nl, nr, nc), dtype=float) * 2.0
+    for i in range(nl):
+        botm[nl - i - 1, :, :] = i
+
+    gridprops = get_disv_kwargs(
+        nlay=nl, nrow=nr, ncol=nc, tp=top, botm=botm, delr=1.0, delc=1.0
+    )
+    from flopy.utils.gridutil import get_disu_kwargs
+
+    gridprops_disu = get_disu_kwargs(
+        nlay=nl,
+        nrow=nr,
+        ncol=nc,
+        tp=top,
+        botm=botm,
+        delr=1.0,
+        delc=1.0,
+        return_vertices=True,
+    )
+    nodes = gridprops_disu["nodes"]
+
+    ws = function_tmpdir
+    name = "test_mf6_unstructured_transmissivity"
+    sim = flopy.mf6.MFSimulation(sim_name=name, sim_ws=ws, exe_name="mf6")
+    gwf = flopy.mf6.ModflowGwf(sim, modelname=name, save_flows=True)
+    disu = flopy.mf6.ModflowGwfdisu(gwf, **gridprops_disu)
+    npf = flopy.mf6.ModflowGwfnpf(gwf, k=hk, save_specific_discharge=True)
+
+    # test specific point
+    T = get_transmissivities(disu.bot.array, gwf, x=[0.5], y=[0.5])
+    assert T.shape == (nl, 1)
+
+    # test all cell centers
+    T = get_transmissivities(
+        disu.bot.array, gwf, x=gwf.modelgrid.xcellcenters, y=gwf.modelgrid.ycellcenters
+    )
+    assert T.shape == (nl, nodes)
+
+    with pytest.raises(ValueError) as e:
+        get_transmissivities(disu.bot.array, gwf, r=[0], c=[0])
+        assert "r, c parameters only valid for structured grids" in str(e.value)
+
+
 def test_get_water_table():
     hdry = -1e30
     hds = np.ones((3, 3, 3), dtype=float) * hdry

flopy/utils/postprocessing.py

@@ -1,6 +1,185 @@
 import numpy as np
 
 
+class _GridAdapter:
+    """
+    Temporary workaround to normalize grid-specific array access patterns.
+
+    Different grid types (DIS, DISV, DISU) store arrays in different shapes.
+    Adapters provide a uniformly shaped interface, (nlay, ncells_per_layer).
+
+    Currently this is used only in get_transmissivities(). It should not be
+    necessary anymore in 4.x.
+    """
+
+    def __init__(self, model, dis_package, nlay):
+        self.model = model
+        self.dis = dis_package
+        self.nlay = nlay
+
+    def reshape_if_needed(self, array, target_shape):
+        """Reshape 1D array to 2D if needed."""
+        if array.ndim == 1 and len(target_shape) == 2:
+            return array.reshape(target_shape)
+        return array
+
+    def get_k_array(self, paklist):
+        """
+        Return hydraulic conductivity array in shape (nlay, ncells_per_layer).
+
+        Parameters
+        ----------
+        paklist : list
+            List of package names in the model
+        """
+        # Check for flow packages in order of preference
+        if "LPF" in paklist:
+            return self.model.lpf.hk.array
+        elif "UPW" in paklist:
+            return self.model.upw.hk.array
+        elif "NPF" in paklist:
+            return self._get_npf_k_array()
+        else:
+            raise ValueError("No LPF, UPW, or NPF package.")
+
+    def _get_npf_k_array(self):
+        """Get NPF k array. Subclasses may override for grid-specific handling."""
+        return self.model.npf.k.array
+
+    def get_bottom_array(self):
+        """Return bottom elevation array in shape (nlay, ncells_per_layer)."""
+        raise NotImplementedError
+
+    def get_top_for_slice(self, indices, grid_type):
+        """Return model top elevation for the given cell indices."""
+        raise NotImplementedError
+
+    def get_layer_tops(self, botm_sliced, indices, grid_type):
+        """Return top elevation for each layer at the given indices."""
+        raise NotImplementedError
+
+    def normalize_heads_array(self, heads):
+        """
+        Normalize heads array to (nlay, ncells_per_layer) shape if needed.
+
+        Parameters
+        ----------
+        heads : ndarray
+            Heads array in any valid format for this grid type
+
+        Returns
+        -------
+        ndarray
+            Heads array in (nlay, ncells_per_layer) shape
+        """
+        # Default: no normalization needed
+        return heads
+
+
+class _DisAdapter(_GridAdapter):
+    """Adapter for structured (DIS) grids."""
+
+    def get_bottom_array(self):
+        return self.dis.botm.array
+
+    def get_top_for_slice(self, indices, grid_type):
+        return self.dis.top.array[indices]
+
+    def get_layer_tops(self, botm_sliced, indices, grid_type):
+        tops = np.empty_like(botm_sliced, dtype=float)
+        tops[0, :] = self.dis.top.array[indices]
+        tops[1:, :] = botm_sliced[:-1]
+        return tops
+
+
+class _DisvAdapter(_GridAdapter):
+    """Adapter for vertex (DISV) grids."""
+
+    def get_bottom_array(self):
+        return self.dis.botm.array
+
+    def get_top_for_slice(self, indices, grid_type):
+        # DISV top is (ncpl,), indices is a tuple with one element
+        return self.dis.top.array[indices[0]]
+
+    def get_layer_tops(self, botm_sliced, indices, grid_type):
+        tops = np.empty_like(botm_sliced, dtype=float)
+        tops[0, :] = self.dis.top.array[indices[0]]
+        tops[1:, :] = botm_sliced[:-1]
+        return tops
+
+
+class _DisuAdapter(_GridAdapter):
+    """Adapter for unstructured (DISU) grids."""
+
+    def _get_npf_k_array(self):
+        """Get NPF k array, reshaping from (nodes,) to (nlay, ncpl) if needed."""
+        k_array = self.model.npf.k.array
+        if k_array.ndim == 1:
+            return k_array.reshape((self.nlay, -1))
+        return k_array
+
+    def get_bottom_array(self):
+        bot_array = self.dis.bot.array
+        if bot_array.ndim == 1:
+            return bot_array.reshape((self.nlay, -1))
+        return bot_array
+
+    def get_top_for_slice(self, indices, grid_type):
+        # DISU top is per-node (nodes,), reshape to (nlay, ncpl) and get layer 0
+        top_array = self.dis.top.array
+        if top_array.ndim == 1:
+            top_array = top_array.reshape((self.nlay, -1))
+        return top_array[0, indices[0]]
+
+    def get_layer_tops(self, botm_sliced, indices, grid_type):
+        # DISU has per-node tops, so each layer has different top values
+        tops = np.empty_like(botm_sliced, dtype=float)
+        top_array = self.dis.top.array
+        if top_array.ndim == 1:
+            top_array = top_array.reshape((self.nlay, -1))
+        tops[0, :] = top_array[0, indices[0]]
+        tops[1:, :] = botm_sliced[:-1]
+        return tops
+
+    def normalize_heads_array(self, heads):
+        """Normalize heads from flat (nnodes,) to (nlay, ncpl) if needed."""
+        if heads.ndim == 1:
+            return heads.reshape((self.nlay, -1))
+        return heads
+
+
+def _get_grid_adapter(model, nlay):
+    """
+    Get the appropriate grid adapter for the model.
+
+    Parameters
+    ----------
+    model : flopy.modflow.Modflow or flopy.mf6.ModflowGwf object
+        Model object
+    nlay : int
+        Number of layers
+
+    Returns
+    -------
+    adapter : _GridAdapter subclass instance
+        Grid-specific adapter
+    dis_package : discretization package
+        The DIS, DISV, or DISU package
+    """
+    paklist = model.get_package_list()
+
+    if "DISU" in paklist:
+        return _DisuAdapter(model, model.disu, nlay)
+    elif "DISV" in paklist:
+        return _DisvAdapter(model, model.disv, nlay)
+    elif "DIS" in paklist:
+        return _DisAdapter(model, model.dis, nlay)
+    else:
+        # For older MODFLOW versions, default to DIS adapter
+        return _DisAdapter(model, model.dis, nlay)
+
+
 def get_transmissivities(
     heads,
     m,
@@ -24,9 +203,9 @@ def get_transmissivities(
     heads : 2D array OR 3D array
         numpy array of shape nlay by n locations (2D) OR complete heads array
         with the correct shape for structured grids (nlay, nrow, ncol) or for
-        vertex grids (nlay, ncpl).
+        vertex grids (nlay, ncpl) or unstructured grids (nnodes).
     m : flopy.modflow.Modflow or flopy.mf6.ModflowGwf object
-        Must have dis and lpf, upw, or npf packages.
+        Must have dis, disv, or disu and lpf, upw, or npf packages.
     r : 1D array-like of ints, of length n locations
         row indices (optional; alternately specify x, y).
         Only valid for structured grids.
@@ -55,7 +234,7 @@ def get_transmissivities(
     - r, c (row, column indices)
     - x, y (real world coordinates)
 
-    For vertex grids only x, y coordinates are supported.
+    For vertex and unstructured grids, only x, y coordinates are supported.
 
     Examples
     --------
@@ -78,6 +257,9 @@ def get_transmissivities(
         nrow = m.nrow
         ncol = m.ncol
 
+    # get grid adapter
+    adapter = _get_grid_adapter(m, nlay)
+
     # get slicing indices
     if r is not None and c is not None:
         if grid_type != "structured":
@@ -93,21 +275,18 @@ def get_transmissivities(
     else:
         raise ValueError("Must specify r, c indices or x, y locations.")
 
-    # slice k
+    # get k array using adapter (handles all flow packages)
     paklist = m.get_package_list()
-    if "LPF" in paklist:
-        hk = m.lpf.hk.array[(slice(None),) + indices]
-    elif "UPW" in paklist:
-        hk = m.upw.hk.array[(slice(None),) + indices]
-    elif "NPF" in paklist:
-        hk = m.npf.k.array[(slice(None),) + indices]
-    else:
-        raise ValueError("No LPF, UPW, or NPF package.")
+    k_array = adapter.get_k_array(paklist)
+    hk = k_array[(slice(None),) + indices]
+
+    # get and slice bottom array
+    botm_array = adapter.get_bottom_array()
+    botm = botm_array[(slice(None),) + indices]
 
-    # slice botm
-    botm = m.dis.botm.array[(slice(None),) + indices]
+    # normalize and slice heads using adapter
+    heads = adapter.normalize_heads_array(heads)
 
-    # slice heads
     if grid_type == "structured" and heads.shape == (nlay, nrow, ncol):
         heads = heads[(slice(None),) + indices]
     elif grid_type != "structured" and heads.shape == (nlay, ncpl):
@@ -117,14 +296,11 @@ def get_transmissivities(
 
     # open interval tops/bottoms default to model top/bottom
     if sctop is None:
-        sctop = m.dis.top.array[indices]
+        sctop = adapter.get_top_for_slice(indices, grid_type)
     if scbot is None:
-        scbot = m.dis.botm.array[(-1,) + indices]
+        scbot = botm[-1, :]
 
-    # make an array of layer tops
-    tops = np.empty_like(botm, dtype=float)
-    tops[0, :] = m.dis.top.array[indices]
-    tops[1:, :] = botm[:-1]
+    tops = adapter.get_layer_tops(botm, indices, grid_type)
 
     # expand top and bottom arrays to be same shape as botm, thickness, etc.
     # (so we have an open interval value for each layer)