refactor(grid): consistent API for get_cell_vertices (#2678)

Previously the Grid.get_cell_vertices() method looked/worked different depending on the grid type. Make it consistent across all three grid types by supporting both cellid and node keyword arguments and being lenient about what's passed to cellid i.e. accept a node number, a fully-sized tuple, or an "abbreviated" tuple lacking layer index. Cell ID and node number are conceptually distinct but til now VertexGrid and UnstructuredGrid methods had a single parameter cellid supporting (for vertex) or expecting (for unstructured) a node number, which we don't want to break. And continue also supporting i and j for StructuredGrid.

Author: wpbonelli

autotest/test_grid.py

@@ -25,6 +25,7 @@
     centroid_of_polygon,
     to_cvfd,
 )
+from flopy.utils.gridutil import get_disu_kwargs, get_disv_kwargs
 from flopy.utils.lgrutil import Lgr
 from flopy.utils.triangle import Triangle
 from flopy.utils.voronoi import VoronoiGrid
@@ -190,6 +191,92 @@ def test_get_cell_vertices():
         mg.get_cell_vertices(nn=0)
 
 
+def test_structured_grid_get_cell_vertices():
+    """Test StructuredGrid.get_cell_vertices() with various input forms"""
+    delr, delc = np.array([10.0] * 3), np.array([10.0] * 4)
+    sg = StructuredGrid(delr=delr, delc=delc)
+
+    # Test node kwarg
+    v1 = sg.get_cell_vertices(node=0)
+    expected = [
+        (np.float64(0.0), np.float64(40.0)),
+        (np.float64(10.0), np.float64(40.0)),
+        (np.float64(10.0), np.float64(30.0)),
+        (np.float64(0.0), np.float64(30.0)),
+    ]
+    assert v1 == expected
+
+    # Test positional args (i, j)
+    v2 = sg.get_cell_vertices(3, 0)
+
+    # Test cellid as tuple (i, j)
+    v3 = sg.get_cell_vertices((3, 0))
+    assert v2 == v3, "Positional and tuple forms should match"
+
+    # Test cellid as 3-element tuple (layer, i, j) - layer ignored
+    v4 = sg.get_cell_vertices(cellid=(0, 3, 0))
+    assert v2 == v4, "2-element and 3-element forms should match"
+
+    # Test named i, j kwargs (backward compatibility)
+    v5 = sg.get_cell_vertices(i=3, j=0)
+    assert v2 == v5, "Named i,j should match"
+
+
+def test_vertex_grid_get_cell_vertices():
+    """Test VertexGrid.get_cell_vertices() with various input forms"""
+    disv_props = get_disv_kwargs(2, 10, 10, 10.0, 10.0, 100.0, [50.0, 0.0])
+    # Remove nvert which is not needed for VertexGrid constructor
+    disv_props.pop("nvert", None)
+    vg = VertexGrid(**disv_props)
+
+    # Test cell2d index as positional arg
+    v1 = vg.get_cell_vertices(5)
+
+    # Test (layer, cell2d) tuple - layer ignored for 2D vertices
+    v2 = vg.get_cell_vertices((0, 5))
+    assert v1 == v2, "cell2d and (layer, cell2d) should match"
+
+    # Test named cellid kwarg
+    v3 = vg.get_cell_vertices(cellid=5)
+    assert v1 == v3, "Positional and kwarg should match"
+
+    # Test node number (>= ncpl, should be converted to cell2d)
+    # Node 105 = layer 1, cell2d 5 (ncpl=100)
+    v4 = vg.get_cell_vertices(node=105)
+
+    # Verify it's the same as (1, 5)
+    v5 = vg.get_cell_vertices((1, 5))
+    assert v4 == v5, "Node and (layer, cell2d) should match"
+
+
+def test_unstructured_grid_get_cell_vertices():
+    """Test UnstructuredGrid.get_cell_vertices() with various input forms"""
+    disu_props = get_disu_kwargs(
+        1, 10, 10, 10.0, 10.0, 100.0, [0.0], return_vertices=True
+    )
+    # Extract only the parameters needed for UnstructuredGrid
+    ug = UnstructuredGrid(
+        vertices=disu_props["vertices"],
+        cell2d=disu_props["cell2d"],
+        top=disu_props["top"],
+    )
+
+    # Test node as positional arg
+    v1 = ug.get_cell_vertices(5)
+
+    # Test (node,) single-element tuple
+    v2 = ug.get_cell_vertices((5,))
+    assert v1 == v2, "Int and tuple forms should match"
+
+    # Test node kwarg
+    v3 = ug.get_cell_vertices(node=5)
+    assert v1 == v3, "cellid and node should match"
+
+    # Test cellid kwarg
+    v4 = ug.get_cell_vertices(cellid=5)
+    assert v1 == v4, "Positional and kwarg should match"
+
+
 def test_get_lrc_get_node():
     nlay, nrow, ncol = 3, 4, 5
     nnodes = nlay * nrow * ncol

flopy/discretization/structuredgrid.py

@@ -1148,15 +1148,20 @@ def get_cell_vertices(self, *args, **kwargs):
 
         Parameters
         ----------
+        cellid : int or tuple, optional
+            Cell identifier. Can be:
+            - node number (int)
+            - (row, col) tuple
+            - (layer, row, col) tuple (layer is ignored, vertices are 2D)
         node : int, optional
-            Node index, mutually exclusive with i and j
+            Node index, mutually exclusive with cellid, i, and j
         i, j : int, optional
-            Row and column index, mutually exclusive with node
+            Row and column index, mutually exclusive with cellid and node
 
         Returns
         -------
         list
-            list of tuples with x,y coordinates to cell vertices
+            list of (x, y) cell vertex coordinates
 
         Examples
         --------
@@ -1168,26 +1173,61 @@ def get_cell_vertices(self, *args, **kwargs):
         [(0.0, 40.0), (10.0, 40.0), (10.0, 30.0), (0.0, 30.0)]
         >>> sg.get_cell_vertices(3, 0)
         [(0.0, 10.0), (10.0, 10.0), (10.0, 0.0), (0.0, 0.0)]
+        >>> sg.get_cell_vertices((3, 0))
+        [(0.0, 10.0), (10.0, 10.0), (10.0, 0.0), (0.0, 0.0)]
+        >>> sg.get_cell_vertices(cellid=(0, 3, 0))
+        [(0.0, 10.0), (10.0, 10.0), (10.0, 0.0), (0.0, 0.0)]
         """
         if kwargs:
             if args:
                 raise TypeError("mixed positional and keyword arguments not supported")
+            elif "cellid" in kwargs:
+                cellid = kwargs.pop("cellid")
+                # Handle cellid as int, tuple of 2, or tuple of 3
+                if isinstance(cellid, (tuple, list)):
+                    if len(cellid) == 2:
+                        i, j = cellid
+                    elif len(cellid) == 3:
+                        _, i, j = cellid  # ignore layer
+                    else:
+                        raise ValueError(
+                            f"cellid tuple must have 2 or 3 elements, got {len(cellid)}"
+                        )
+                else:
+                    _, i, j = self.get_lrc(cellid)[0]
             elif "node" in kwargs:
                 _, i, j = self.get_lrc(kwargs.pop("node"))[0]
             elif "i" in kwargs and "j" in kwargs:
                 i = kwargs.pop("i")
                 j = kwargs.pop("j")
+            else:
+                raise TypeError(
+                    "expected cellid, node, or i and j as keyword arguments"
+                )
             if kwargs:
                 unused = ", ".join(kwargs.keys())
                 raise TypeError(f"unused keyword arguments: {unused}")
         elif len(args) == 0:
             raise TypeError("expected one or more arguments")
-
-        if len(args) == 1:
-            _, i, j = self.get_lrc(args[0])[0]
+        elif len(args) == 1:
+            # Single arg could be node number or (row, col) or (layer, row, col) tuple
+            arg = args[0]
+            if isinstance(arg, (tuple, list)):
+                if len(arg) == 2:
+                    i, j = arg
+                elif len(arg) == 3:
+                    # (layer, row, col) - ignore layer
+                    _, i, j = arg
+                else:
+                    raise ValueError(
+                        f"cellid tuple must have 2 or 3 elements, got {len(arg)}"
+                    )
+            else:
+                # Node number
+                _, i, j = self.get_lrc(arg)[0]
         elif len(args) == 2:
             i, j = args
-        elif len(args) > 2:
+        else:
             raise TypeError("too many arguments")
 
         self._copy_cache = False

flopy/discretization/unstructuredgrid.py

@@ -847,16 +847,54 @@ def top_botm(self):
         new_botm = np.expand_dims(self._botm, 0)
         return np.concatenate((new_top, new_botm), axis=0)
 
-    def get_cell_vertices(self, cellid):
+    def get_cell_vertices(self, cellid=None, node=None):
         """
-        Method to get a set of cell vertices for a single cell
-            used in the Shapefile export utilities
-        :param cellid: (int) cellid number
+        Get a set of cell vertices for a single cell.
+
+        Parameters
+        ----------
+        cellid : int or tuple, optional
+            Cell identifier. Can be:
+            - node number (int)
+            - (node,) single-element tuple
+        node : int, optional
+            Node number, mutually exclusive with cellid
+
         Returns
-        ------- list of x,y cell vertices
+        -------
+        list
+            list of (x, y) cell vertex coordinates
+
+        Examples
+        --------
+        >>> import flopy
+        >>> from flopy.utils.gridutil import get_disu_kwargs
+        >>> disu_props = get_disu_kwargs(1, 10, 10, 1.0, 1.0, 1.0, [0.0])
+        >>> ug = flopy.discretization.UnstructuredGrid(**disu_props)
+        >>> ug.get_cell_vertices(5)  # node number
+        >>> ug.get_cell_vertices((5,))  # (node,) tuple
+        >>> ug.get_cell_vertices(node=5)  # explicit node kwarg
+        >>> ug.get_cell_vertices(cellid=5)  # explicit cellid kwarg
         """
+
+        if cellid is not None and node is not None:
+            raise ValueError("cellid and node are mutually exclusive")
+
+        if cellid is None and node is None:
+            raise TypeError("expected cellid or node argument")
+
+        idx = node if node is not None else cellid
+        if isinstance(idx, (tuple, list)):
+            if len(idx) == 1:
+                idx = idx[0]
+            else:
+                raise ValueError(
+                    f"cellid tuple must have 1 element for "
+                    f"unstructured grids, got {len(idx)}"
+                )
+
         self._copy_cache = False
-        cell_vert = list(zip(self.xvertices[cellid], self.yvertices[cellid]))
+        cell_vert = list(zip(self.xvertices[idx], self.yvertices[idx]))
         self._copy_cache = True
         return cell_vert
 

flopy/discretization/vertexgrid.py

@@ -474,23 +474,73 @@ def intersect(self, x, y, z=None, local=False, forgive=False):
                     results.astype(int) if np.all(valid_mask) else results,
                 )
 
-    def get_cell_vertices(self, cellid):
+    def get_cell_vertices(self, cellid=None, node=None):
         """
-        Method to get a set of cell vertices for a single cell
-            used in the Shapefile export utilities
-        :param cellid: (int) cellid number
+        Get a set of cell vertices for a single cell.
+
+        Parameters
+        ----------
+        cellid : int or tuple, optional
+            Cell identifier. Can be:
+            - cell2d index (int, 0 to ncpl-1)
+            - node number (int, >= ncpl) - will be converted to cell2d
+            - (cell2d,) single-element tuple
+            - (layer, cell2d) tuple (layer is ignored, vertices are 2D)
+        node : int, optional
+            Node number, mutually exclusive with cellid
+
         Returns
-        ------- list of x,y cell vertices
+        -------
+        list
+            list of (x, y) cell vertex coordinates
+
+        Examples
+        --------
+        >>> import flopy
+        >>> from flopy.utils.gridutil import get_disv_kwargs
+        >>> disv_props = get_disv_kwargs(1, 10, 10, 1.0, 1.0, 1.0, [0.0])
+        >>> vg = flopy.discretization.VertexGrid(**disv_props)
+        >>> vg.get_cell_vertices(5)  # cell2d index
+        >>> vg.get_cell_vertices((0, 5))  # (layer, cell2d) tuple
+        >>> vg.get_cell_vertices(node=105)  # node number
+        >>> vg.get_cell_vertices(cellid=(1, 5))  # explicit cellid kwarg
         """
-        while cellid >= self.ncpl:
-            if cellid > self.nnodes:
-                err = f"cellid {cellid} out of index for size {self.nnodes}"
-                raise IndexError(err)
+        # Handle arguments
+        if cellid is not None and node is not None:
+            raise ValueError("cellid and node are mutually exclusive")
+
+        if cellid is None and node is None:
+            raise TypeError("expected cellid or node argument")
 
-            cellid -= self.ncpl
+        # Use cellid if provided, otherwise use node
+        if node is not None:
+            idx = node
+        else:
+            idx = cellid
+
+        # Handle tuple forms
+        if isinstance(idx, (tuple, list)):
+            if len(idx) == 1:
+                # (cell2d,) or (node,)
+                idx = idx[0]
+            elif len(idx) == 2:
+                # (layer, cell2d) - ignore layer since vertices are 2D
+                _, idx = idx
+            else:
+                raise ValueError(
+                    f"cellid tuple must have 1 or 2 elements, got {len(idx)}"
+                )
+
+        # Convert node to cell2d if necessary
+        while idx >= self.ncpl:
+            if idx > self.nnodes:
+                raise IndexError(
+                    f"node number {idx} exceeds grid node count {self.nnodes}"
+                )
+            idx -= self.ncpl
 
         self._copy_cache = False
-        cell_verts = list(zip(self.xvertices[cellid], self.yvertices[cellid]))
+        cell_verts = list(zip(self.xvertices[idx], self.yvertices[idx]))
         self._copy_cache = True
         return cell_verts