refactor(grid): consistent API for get_cell_vertices

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,63 @@ 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:
+                        # (layer, row, col) - ignore layer
+                        _, i, j = cellid
+                    else:
+                        raise ValueError(
+                            f"cellid tuple must have 2 or 3 elements, got {len(cellid)}"
+                        )
+                else:
+                    # cellid is a node number
+                    _, 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,60 @@ 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
         """
+        # 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")
+
+        # Use cellid if provided, otherwise use node
+        if node is not None:
+            node_idx = node
+        else:
+            node_idx = cellid
+
+        # Handle tuple form
+        if isinstance(node_idx, (tuple, list)):
+            if len(node_idx) == 1:
+                node_idx = node_idx[0]
+            else:
+                raise ValueError(
+                    f"cellid tuple must have 1 element for "
+                    f"unstructured grids, got {len(node_idx)}"
+                )
+
         self._copy_cache = False
-        cell_vert = list(zip(self.xvertices[cellid], self.yvertices[cellid]))
+        cell_vert = list(zip(self.xvertices[node_idx], self.yvertices[node_idx]))
         self._copy_cache = True
         return cell_vert
 

flopy/discretization/vertexgrid.py

@@ -474,23 +474,71 @@ 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")
+
+        # Use cellid if provided, otherwise use node
+        if node is not None:
+            cell2d = node
+        else:
+            cell2d = cellid
+
+        # Handle tuple forms
+        if isinstance(cell2d, (tuple, list)):
+            if len(cell2d) == 1:
+                # (cell2d,) or (node,)
+                cell2d = cell2d[0]
+            elif len(cell2d) == 2:
+                # (layer, cell2d) - ignore layer since vertices are 2D
+                _, cell2d = cell2d
+            else:
+                raise ValueError(
+                    f"cellid tuple must have 1 or 2 elements, got {len(cell2d)}"
+                )
 
-            cellid -= self.ncpl
+        # Convert node to cell2d if necessary
+        while cell2d >= self.ncpl:
+            if cell2d > self.nnodes:
+                raise IndexError(f"cellid {cell2d} out of index for size {self.nnodes}")
+            cell2d -= self.ncpl
 
         self._copy_cache = False
-        cell_verts = list(zip(self.xvertices[cellid], self.yvertices[cellid]))
+        cell_verts = list(zip(self.xvertices[cell2d], self.yvertices[cell2d]))
         self._copy_cache = True
         return cell_verts