Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
91 changes: 91 additions & 0 deletions autotest/test_grid.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
35 changes: 30 additions & 5 deletions flopy/discretization/structuredgrid.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Expand All @@ -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()
Expand Down
57 changes: 57 additions & 0 deletions flopy/discretization/unstructuredgrid.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
50 changes: 50 additions & 0 deletions flopy/discretization/vertexgrid.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
11 changes: 5 additions & 6 deletions flopy/export/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
21 changes: 4 additions & 17 deletions flopy/plot/crosssection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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
Expand Down
14 changes: 5 additions & 9 deletions flopy/utils/binaryfile/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down
Loading