diff --git a/autotest/test_grid.py b/autotest/test_grid.py index c390f9b15..23e60c536 100644 --- a/autotest/test_grid.py +++ b/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 diff --git a/flopy/discretization/structuredgrid.py b/flopy/discretization/structuredgrid.py index 0d2df58fd..95b75cddb 100644 --- a/flopy/discretization/structuredgrid.py +++ b/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 diff --git a/flopy/discretization/unstructuredgrid.py b/flopy/discretization/unstructuredgrid.py index f34764dd6..12e061acb 100644 --- a/flopy/discretization/unstructuredgrid.py +++ b/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 diff --git a/flopy/discretization/vertexgrid.py b/flopy/discretization/vertexgrid.py index dc5be8710..26541335e 100644 --- a/flopy/discretization/vertexgrid.py +++ b/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