feat(grid): add get_node() method to all grid types (#2680)

Add a get_node() method to VertexGrid and UnstructuredGrid, consistent with StructuredGrid, for cell ID -> node number conversions. Rename the lrc_list parameter on StructuredGrid.get_node() to cellids and deprecate lrc_list. Add an optional node2d parameter to determine whether to get the 2D or 3D node number, useful for structured and vertex grids, ignored for unstructured.

The cell ID -> node number conversion was done ad hoc in a few places in the codebase, this allows some simplification. Maybe in 4.x we can standardize the signature and pull an abstract method into the base Grid class.

Author: wpbonelli

autotest/test_grid.py

@@ -1761,3 +1761,94 @@ def test_structured_grid_intersect_edge_case(simple_structured_grid):
     assert lay is None, f"Expected lay=None, got {lay}"
     assert np.isnan(row), f"Expected row=nan, got {row}"
     assert np.isnan(col), f"Expected col=nan, got {col}"
+
+
+def test_structured_grid_get_node():
+    sg = StructuredGrid(nlay=3, nrow=4, ncol=5, delr=np.ones(5), delc=np.ones(4))
+
+    # 3D node numbers
+    assert sg.get_node((0, 0, 0)) == [0]
+    assert sg.get_node((1, 0, 0)) == [20]
+    assert sg.get_node((0, 1, 2)) == [7]
+    assert sg.get_node([(0, 0, 0), (1, 0, 0)]) == [0, 20]
+    assert sg.get_node([(0, 0, 0), (2, 3, 4)]) == [0, 59]
+
+    # 2D node numbers (layer ignored)
+    assert sg.get_node((0, 1, 2), node2d=True) == [7]
+    assert sg.get_node((1, 1, 2), node2d=True) == [7]
+    assert sg.get_node((2, 1, 2), node2d=True) == [7]
+    assert sg.get_node([(0, 0, 0), (1, 0, 0), (2, 1, 2)], node2d=True) == [0, 0, 7]
+
+
+def test_vertex_grid_get_node():
+    nlay = 3
+    ncpl = 100
+    vertices = [[i, float(i % 10), float(i // 10)] for i in range(121)]
+    cell2d = []
+    for i in range(ncpl):
+        # Simple quad cells
+        iv = [i, i + 1, i + 11, i + 10]
+        cell2d.append([i, 0.0, 0.0, 4] + iv)
+    vg = VertexGrid(
+        vertices=vertices,
+        cell2d=cell2d,
+        nlay=nlay,
+        ncpl=ncpl,
+    )
+
+    # 3D node number
+    assert vg.get_node((0, 5)) == [5]
+    assert vg.get_node((1, 5)) == [105]
+    assert vg.get_node((2, 99)) == [299]
+    assert vg.get_node([(0, 5), (1, 5)]) == [5, 105]
+
+    # 2D node number
+    assert vg.get_node((0, 5), node2d=True) == [5]
+    assert vg.get_node((1, 5), node2d=True) == [5]
+    assert vg.get_node((2, 5), node2d=True) == [5]
+    assert vg.get_node([(0, 10), (1, 20), (2, 30)], node2d=True) == [10, 20, 30]
+
+    with pytest.raises(ValueError, match="VertexGrid cellid must be"):
+        vg.get_node((0, 1, 2))
+
+    with pytest.raises(IndexError, match=r"Layer .* out of range"):
+        vg.get_node((5, 10))
+
+    with pytest.raises(IndexError, match=r"Cell2d .* out of range"):
+        vg.get_node((0, 200))
+
+
+def test_unstructured_grid_get_node():
+    ncpl = [100]
+    vertices = [[i, float(i % 10), float(i // 10)] for i in range(121)]
+    iverts = [[i, i + 1, i + 11, i + 10] for i in range(100)]
+    ug = UnstructuredGrid(
+        vertices=vertices,
+        iverts=iverts,
+        ncpl=ncpl,
+    )
+
+    # Test plain integers
+    assert ug.get_node(5) == [5]
+    assert ug.get_node(10) == [10]
+
+    # Test tuples
+    assert ug.get_node((5,)) == [5]
+    assert ug.get_node((10,)) == [10]
+
+    # Test list of integers
+    assert ug.get_node([5, 10, 99]) == [5, 10, 99]
+
+    # Test list of tuples
+    assert ug.get_node([(5,), (10,), (99,)]) == [5, 10, 99]
+
+    # node2d parameter should have no effect
+    assert ug.get_node(5, node2d=True) == [5]
+    assert ug.get_node((5,), node2d=True) == [5]
+    assert ug.get_node([5, 10], node2d=True) == [5, 10]
+
+    with pytest.raises(ValueError, match="UnstructuredGrid cellid"):
+        ug.get_node((0, 1))
+
+    with pytest.raises(IndexError, match=r"Node .* out of range"):
+        ug.get_node(200)

flopy/discretization/structuredgrid.py

@@ -1272,7 +1272,7 @@ def get_lrc(self, nodes):
             shape = tuple(dim or 1 for dim in shape)
         return list(zip(*np.unravel_index(nodes, shape)))
 
-    def get_node(self, lrc_list):
+    def get_node(self, lrc_list, cellids=None, node2d=False):
         """
         Get node number from a list of zero-based MODFLOW
         layer, row, column tuples.
@@ -1282,6 +1282,15 @@ def get_node(self, lrc_list):
         lrc_list : tuple of int or list of tuple of int
             Zero-based layer, row, column tuples
 
+            .. deprecated:: 3.10
+                This parameter is deprecated and will be
+                removed in FloPy 3.12. Use cellids instead.
+        cellids : tuple of int or list of tuple of int
+            Zero-based layer, row, column tuples
+        node2d : bool, optional
+            If True, return 2D node numbers (ignore layer).
+            If False (default), return 3D node numbers.
+
         Returns
         -------
         list
@@ -1297,10 +1306,26 @@ def get_node(self, lrc_list):
         >>> sg.get_node([(0, 2, 20), (0, 25, 0), (8, 10, 0)])
         [100, 1000, 10000]
         """
-        if not isinstance(lrc_list, list):
-            lrc_list = [lrc_list]
-        multi_index = tuple(np.array(lrc_list).T)
-        shape = self.shape
+
+        if lrc_list is not None:
+            if cellids is not None:
+                raise TypeError("lrc_list and cellids are mutually exclusive")
+            cellids = lrc_list
+
+        if cellids is None:
+            raise ValueError("Expected a value for cellids")
+
+        if not isinstance(cellids, list):
+            cellids = [cellids]
+
+        if node2d:
+            rc_list = [(row, col) for lay, row, col in cellids]
+            multi_index = tuple(np.array(rc_list).T)
+            shape = (self.nrow, self.ncol)
+        else:
+            multi_index = tuple(np.array(cellids).T)
+            shape = self.shape
+
         if shape[0] is None:
             shape = tuple(dim or 1 for dim in shape)
         return np.ravel_multi_index(multi_index, shape).tolist()

flopy/discretization/unstructuredgrid.py

@@ -898,6 +898,63 @@ def get_cell_vertices(self, cellid=None, node=None):
         self._copy_cache = True
         return cell_vert
 
+    def get_node(self, cellids, node2d=False):
+        """
+        Get node number from cellids.
+
+        For DISU grids, cellid IS the node number. The node2d
+        parameter is accepted for API consistency but has no effect.
+
+        Parameters
+        ----------
+        cellid_list : int, tuple of int, or list of int/tuple
+            DISU cellid(s). Can be a plain integer, a tuple (node,),
+            or a list of integers or tuples.
+        node2d : bool, optional
+            Accepted for API consistency. Has no effect for
+            unstructured grids (no layer concept).
+
+        Returns
+        -------
+        list
+            list of MODFLOW node numbers
+
+        Examples
+        --------
+        >>> import flopy
+        >>> ug = flopy.discretization.UnstructuredGrid(ncpl=[100], ...)
+        >>> ug.get_node(5)
+        [5]
+        >>> ug.get_node((5,))
+        [5]
+        >>> ug.get_node([5, 10])
+        [5, 10]
+        >>> ug.get_node([(5,), (10,)])
+        [5, 10]
+        """
+        if not isinstance(cellids, list):
+            cellids = [cellids]
+
+        nodes = []
+        for cellid in cellids:
+            # Accept both plain integers and tuples
+            if isinstance(cellid, (int, np.integer)):
+                node = int(cellid)
+            elif isinstance(cellid, (tuple, list)):
+                if len(cellid) != 1:
+                    raise ValueError(
+                        "UnstructuredGrid cellid tuple must have 1 element"
+                    )
+                node = cellid[0]
+            else:
+                raise TypeError(f"Expected int or tuple, got {type(cellid).__name__}")
+
+            if node < 0 or node >= self.nnodes:
+                raise IndexError(f"Node {node} out of range [0, {self.nnodes})")
+            nodes.append(node)
+
+        return nodes
+
     def plot(self, **kwargs):
         """
         Plot the grid lines.

flopy/discretization/vertexgrid.py

@@ -544,6 +544,56 @@ def get_cell_vertices(self, cellid=None, node=None):
         self._copy_cache = True
         return cell_verts
 
+    def get_node(self, cellids, node2d=False):
+        """
+        Get node number from a list of zero-based MODFLOW
+        (layer, cell2d) tuples.
+
+        Parameters
+        ----------
+        cellid_list : tuple of int or list of tuple of int
+            Zero-based (layer, cell2d) tuples
+        node2d : bool, optional
+            If True, return 2D node numbers (cell2d values).
+            If False (default), return 3D node numbers.
+
+        Returns
+        -------
+        list
+            list of MODFLOW nodes for each (layer, cell2d) tuple
+            in the input list
+
+        Examples
+        --------
+        >>> import flopy
+        >>> vg = flopy.discretization.VertexGrid(nlay=3, ncpl=100, ...)
+        >>> vg.get_node((0, 5))
+        [5]
+        >>> vg.get_node((1, 5))
+        [105]
+        >>> vg.get_node([(0, 5), (1, 5)], node2d=True)
+        [5, 5]
+        """
+        if not isinstance(cellids, list):
+            cellids = [cellids]
+
+        # Validate
+        for cellid in cellids:
+            if len(cellid) != 2:
+                raise ValueError("VertexGrid cellid must be (layer, cell2d) tuple")
+
+        if node2d:
+            return [cell2d for lay, cell2d in cellids]
+        else:
+            nodes = []
+            for lay, cell2d in cellids:
+                if lay < 0 or lay >= self.nlay:
+                    raise IndexError(f"Layer {lay} out of range [0, {self.nlay})")
+                if cell2d < 0 or cell2d >= self.ncpl:
+                    raise IndexError(f"Cell2d {cell2d} out of range [0, {self.ncpl})")
+                nodes.append(lay * self.ncpl + cell2d)
+            return nodes
+
     def plot(self, **kwargs):
         """
         Plot the grid lines.

flopy/export/utils.py

@@ -924,13 +924,12 @@ def mflist_export(f: Union[str, PathLike, NetCdf], mfl, **kwargs):
                     cell_index_name = (
                         "cellid_cell" if "cellid_cell" in df.columns else "cell"
                     )
+                    cellids = list(
+                        zip(df[layer_index_name].values, df[cell_index_name].values)
+                    )
+                    nodes = modelgrid.get_node(cellids)
                     verts = np.array(
-                        [
-                            modelgrid.get_cell_vertices(layer * modelgrid.ncpl + cell)
-                            for layer, cell in zip(
-                                df[layer_index_name].values, df[cell_index_name].values
-                            )
-                        ]
+                        [modelgrid.get_cell_vertices(node) for node in nodes]
                     )
                 else:
                     row_index_name = "i" if "i" in df.columns else "cellid_row"

flopy/plot/crosssection.py

@@ -795,17 +795,7 @@ def _cellid_to_node(self, cellid):
         int
             Node number
         """
-        if len(cellid) == 3:
-            # Structured grid: (layer, row, col)
-            layer, row, col = cellid
-            return layer * self.mg.nrow * self.mg.ncol + row * self.mg.ncol + col
-        elif len(cellid) == 2:
-            # Vertex grid: (layer, cell2d)
-            layer, cell2d = cellid
-            return layer * self._ncpl + cell2d
-        else:
-            # Unstructured grid: (node,)
-            return cellid[0]
+        return self.mg.get_node([cellid])[0]
 
     def _plot_vertical_hfb_lines(self, color=None, **kwargs):
         """
@@ -835,12 +825,9 @@ def _plot_vertical_hfb_lines(self, color=None, **kwargs):
 
         for cellid1, cellid2 in self._vertical_hfbs_to_plot:
             # Get the 2D cell identifier (row, col for DIS or cell2d for DISV)
-            if len(cellid1) == 3:
-                # Structured grid
-                node_2d = cellid1[1] * self.mg.ncol + cellid1[2]
-            elif len(cellid1) == 2:
-                # Vertex grid
-                node_2d = cellid1[1]
+            if len(cellid1) == 3 or len(cellid1) == 2:
+                # Structured or vertex grid
+                node_2d = self.mg.get_node([cellid1], node2d=True)[0]
             else:
                 # Unstructured - skip for now
                 continue

flopy/utils/binaryfile/__init__.py

@@ -1878,15 +1878,11 @@ def _cellids(self, idx):
         return cellid
 
     def _cellid_to_node(self, cellids) -> list[int]:
-        if self.modelgrid.grid_type == "structured":
-            return [
-                k * (self.nrow * self.ncol) + i * self.ncol + j + 1
-                for k, i, j in cellids
-            ]
-        elif self.modelgrid.grid_type == "vertex":
-            return [k * self.modelgrid.ncpl + cell + 1 for k, cell in cellids]
-        else:
-            return [node + 1 for node in cellids]
+        """Convert 0-based cellids to 1-based MODFLOW node numbers."""
+        # Get 0-based nodes from grid, then convert to 1-based (vectorized)
+        # UnstructuredGrid.get_node() accepts both plain ints and tuples
+        nodes_0based = self.modelgrid.get_node(cellids)
+        return (np.array(nodes_0based) + 1).tolist()
 
     def get_record(self, idx, full3D=False):
         """