Skip to content

Commit 2a769cd

Browse files
committed
refactor(plot): better shared face finding for HFB plotting
1 parent 905d20b commit 2a769cd

5 files changed

Lines changed: 475 additions & 203 deletions

File tree

autotest/test_faceutil.py

Lines changed: 208 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,208 @@
1+
import numpy as np
2+
import pytest
3+
4+
from flopy.utils.faceutil import get_shared_face_3d
5+
6+
7+
def test_get_shared_face_3d_structured_horizontal():
8+
"""Test get_shared_face_3d for horizontal barrier in structured grid.
9+
10+
Tests same layer, adjacent cells.
11+
"""
12+
from flopy.discretization import StructuredGrid
13+
14+
# Create a simple 2x2x1 structured grid
15+
nlay, nrow, ncol = 1, 2, 2
16+
delr = np.array([1.0, 1.0])
17+
delc = np.array([1.0, 1.0])
18+
top = np.array([[10.0, 10.0], [10.0, 10.0]])
19+
botm = np.array([[[0.0, 0.0], [0.0, 0.0]]])
20+
21+
grid = StructuredGrid(delc=delc, delr=delr, top=top, botm=botm)
22+
23+
# Test horizontal barrier between cells (0, 0, 0) and (0, 0, 1)
24+
# These are horizontally adjacent
25+
cellid1 = (0, 0, 0)
26+
cellid2 = (0, 0, 1)
27+
28+
face = get_shared_face_3d(grid, cellid1, cellid2)
29+
30+
assert face is not None
31+
assert len(face) == 4 # Vertical face should have 4 vertices
32+
33+
# Check that face is vertical (two z-coordinates: top and bottom)
34+
z_coords = sorted({v[2] for v in face})
35+
assert len(z_coords) == 2
36+
assert z_coords[0] == 0.0 # bottom
37+
assert z_coords[1] == 10.0 # top
38+
39+
# Check x-coordinate is at the shared boundary (x=1.0)
40+
x_coords = {v[0] for v in face}
41+
assert 1.0 in x_coords
42+
43+
44+
def test_get_shared_face_3d_structured_vertical():
45+
"""Test get_shared_face_3d for vertical barrier in structured grid.
46+
47+
Tests different layers, same horizontal position.
48+
"""
49+
from flopy.discretization import StructuredGrid
50+
51+
# Create a simple 2x2x2 structured grid (2 layers)
52+
nlay, nrow, ncol = 2, 2, 2
53+
delr = np.array([1.0, 1.0])
54+
delc = np.array([1.0, 1.0])
55+
top = np.array([[10.0, 10.0], [10.0, 10.0]])
56+
botm = np.array([[[5.0, 5.0], [5.0, 5.0]], [[0.0, 0.0], [0.0, 0.0]]])
57+
58+
grid = StructuredGrid(delc=delc, delr=delr, top=top, botm=botm)
59+
60+
# Test vertical barrier between cells (0, 0, 0) and (1, 0, 0)
61+
# These are vertically adjacent
62+
cellid1 = (0, 0, 0)
63+
cellid2 = (1, 0, 0)
64+
65+
face = get_shared_face_3d(grid, cellid1, cellid2)
66+
67+
assert face is not None
68+
assert len(face) == 2 # Horizontal face should have 2 vertices (the edge)
69+
70+
# Check that face is horizontal (all z-coordinates the same at interface)
71+
z_coords = {v[2] for v in face}
72+
assert len(z_coords) == 1
73+
assert next(iter(z_coords)) == 5.0 # Interface between layers
74+
75+
76+
def test_get_shared_face_3d_vertex():
77+
"""Test get_shared_face_3d for vertex grid."""
78+
from flopy.discretization import VertexGrid
79+
80+
# Create a simple 2-cell vertex grid
81+
nlay = 2
82+
ncpl = 2
83+
84+
# Two square cells side by side
85+
vertices = [
86+
[0, 0.0, 0.0],
87+
[1, 1.0, 0.0],
88+
[2, 2.0, 0.0],
89+
[3, 0.0, 1.0],
90+
[4, 1.0, 1.0],
91+
[5, 2.0, 1.0],
92+
]
93+
94+
cell2d = [
95+
[0, 0.5, 0.5, 4, 0, 1, 4, 3],
96+
[1, 1.5, 0.5, 4, 1, 2, 5, 4],
97+
]
98+
99+
top = np.array([10.0, 10.0])
100+
botm = np.array([[5.0, 5.0], [0.0, 0.0]])
101+
102+
grid = VertexGrid(vertices=vertices, cell2d=cell2d, top=top, botm=botm, nlay=nlay)
103+
104+
# Test horizontal barrier between cells (0, 0) and (0, 1) in same layer
105+
cellid1 = (0, 0)
106+
cellid2 = (0, 1)
107+
108+
face = get_shared_face_3d(grid, cellid1, cellid2)
109+
110+
assert face is not None
111+
assert len(face) == 4 # Vertical face
112+
113+
# Test vertical barrier between cells (0, 0) and (1, 0) in different layers
114+
cellid1 = (0, 0)
115+
cellid2 = (1, 0)
116+
117+
face = get_shared_face_3d(grid, cellid1, cellid2)
118+
119+
assert face is not None
120+
assert len(face) == 2 # Horizontal face
121+
122+
# Check all z-coordinates are the same (interface elevation)
123+
z_coords = {v[2] for v in face}
124+
assert len(z_coords) == 1
125+
assert next(iter(z_coords)) == 5.0
126+
127+
128+
def test_get_shared_face_3d_unstructured():
129+
"""Test get_shared_face_3d for unstructured grid with explicit 3D vertices."""
130+
from flopy.discretization import UnstructuredGrid
131+
132+
# Create a simple 2-cell unstructured grid with 3D vertices
133+
# Two hexahedral cells stacked vertically
134+
vertices = [
135+
# Bottom layer vertices (z=0)
136+
[0, 0.0, 0.0, 0.0],
137+
[1, 1.0, 0.0, 0.0],
138+
[2, 1.0, 1.0, 0.0],
139+
[3, 0.0, 1.0, 0.0],
140+
# Middle layer vertices (z=5)
141+
[4, 0.0, 0.0, 5.0],
142+
[5, 1.0, 0.0, 5.0],
143+
[6, 1.0, 1.0, 5.0],
144+
[7, 0.0, 1.0, 5.0],
145+
# Top layer vertices (z=10)
146+
[8, 0.0, 0.0, 10.0],
147+
[9, 1.0, 0.0, 10.0],
148+
[10, 1.0, 1.0, 10.0],
149+
[11, 0.0, 1.0, 10.0],
150+
]
151+
152+
iverts = [
153+
[0, 1, 2, 3, 4, 5, 6, 7], # Bottom cell
154+
[4, 5, 6, 7, 8, 9, 10, 11], # Top cell
155+
]
156+
157+
xc = np.array([0.5, 0.5])
158+
yc = np.array([0.5, 0.5])
159+
top = np.array([10.0, 10.0])
160+
botm = np.array([0.0, 5.0])
161+
162+
grid = UnstructuredGrid(
163+
vertices=vertices,
164+
iverts=iverts,
165+
xcenters=xc,
166+
ycenters=yc,
167+
top=top,
168+
botm=botm,
169+
ncpl=[2],
170+
)
171+
172+
# Test shared face between two vertically stacked cells
173+
cellid1 = (0,)
174+
cellid2 = (1,)
175+
176+
face = get_shared_face_3d(grid, cellid1, cellid2)
177+
178+
assert face is not None
179+
assert len(face) == 4 # Should have 4 shared vertices at the interface
180+
181+
# Check all shared vertices are at z=5.0 (the interface)
182+
z_coords = {v[2] for v in face}
183+
assert len(z_coords) == 1
184+
assert next(iter(z_coords)) == 5.0
185+
186+
187+
def test_get_shared_face_3d_error_cases():
188+
"""Test error handling in get_shared_face_3d."""
189+
from flopy.discretization import StructuredGrid
190+
191+
nlay, nrow, ncol = 1, 2, 2
192+
delr = np.array([1.0, 1.0])
193+
delc = np.array([1.0, 1.0])
194+
top = np.array([[10.0, 10.0], [10.0, 10.0]])
195+
botm = np.array([[[0.0, 0.0], [0.0, 0.0]]])
196+
197+
grid = StructuredGrid(delc=delc, delr=delr, top=top, botm=botm)
198+
199+
# Test with same cellid
200+
with pytest.raises(ValueError, match="cellid1 and cellid2 must be different"):
201+
get_shared_face_3d(grid, (0, 0, 0), (0, 0, 0))
202+
203+
# Test with non-adjacent cells (should return None)
204+
cellid1 = (0, 0, 0)
205+
cellid2 = (0, 1, 1) # Diagonally opposite, not adjacent
206+
207+
face = get_shared_face_3d(grid, cellid1, cellid2)
208+
assert face is None

flopy/plot/crosssection.py

Lines changed: 6 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
from . import plotutil
1414
from .plotutil import (
1515
get_shared_face_3d,
16-
is_vertical_barrier,
16+
is_vertical,
1717
to_mp7_endpoints,
1818
to_mp7_pathlines,
1919
)
@@ -781,22 +781,6 @@ def plot_grid(self, **kwargs):
781781
ax.add_collection(col)
782782
return col
783783

784-
def _cellid_to_node(self, cellid):
785-
"""
786-
Convert a cellid tuple to a node number.
787-
788-
Parameters
789-
----------
790-
cellid : tuple
791-
Cell identifier
792-
793-
Returns
794-
-------
795-
int
796-
Node number
797-
"""
798-
return self.mg.get_node([cellid])[0]
799-
800784
def _plot_vertical_hfb_lines(self, color=None, **kwargs):
801785
"""
802786
Plot vertical HFBs as lines at layer interfaces.
@@ -848,8 +832,8 @@ def _plot_vertical_hfb_lines(self, color=None, **kwargs):
848832

849833
# Get the x-coordinates along the cross section for this cell
850834
# Need to find this cell in projpts - convert both cellids to nodes
851-
node1 = self._cellid_to_node(cellid1)
852-
node2 = self._cellid_to_node(cellid2)
835+
node1 = self.mg.get_node([cellid1])[0]
836+
node2 = self.mg.get_node([cellid2])[0]
853837

854838
# Get x-coordinates from either node's projection
855839
xcoords = []
@@ -958,11 +942,11 @@ def plot_bc(self, name=None, package=None, kper=0, color=None, head=None, **kwar
958942
cellid1 = tuple(entry["cellid1"])
959943
cellid2 = tuple(entry["cellid2"])
960944

961-
if is_vertical_barrier(self.mg, cellid1, cellid2):
962-
# Store vertical HFBs for line plotting
945+
if not is_vertical(self.mg, cellid1, cellid2):
946+
# horizontal face, vertical barrier
963947
vertical_hfbs.append((cellid1, cellid2))
964948
else:
965-
# Horizontal barriers - plot both cells as patches
949+
# vertical face, horizontal barrier
966950
cellids.append(list(entry["cellid1"]))
967951
cellids.append(list(entry["cellid2"]))
968952

flopy/plot/map.py

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
from .plotutil import (
1414
get_shared_face,
1515
get_shared_face_3d,
16-
is_vertical_barrier,
16+
is_vertical,
1717
to_mp7_pathlines,
1818
)
1919

@@ -469,20 +469,18 @@ def _plot_barrier_bc(self, barrier_data, color=None, name=None, **kwargs):
469469
"""
470470
ax = kwargs.pop("ax", self.ax)
471471

472-
# Separate horizontal and vertical barriers
473472
horizontal_line_segments = []
474473
vertical_cell_indices = []
475474

476475
for cellid1, cellid2 in barrier_data:
477476
# Only plot barriers on the current layer (for layered grids)
478477
# For DISU (len==1), plot all barriers since there's no layer filtering
479-
if len(cellid1) >= 2: # DIS or DISV (has layer info)
478+
if len(cellid1) >= 2:
480479
if cellid1[0] != self.layer and cellid2[0] != self.layer:
481480
continue
482481

483-
# Check if this is a vertical barrier
484-
if is_vertical_barrier(self.mg, cellid1, cellid2):
485-
# For vertical barriers, plot the cells on the current layer
482+
# horizontal face, vertical barrier
483+
if not is_vertical(self.mg, cellid1, cellid2):
486484
if len(cellid1) >= 2:
487485
if cellid1[0] == self.layer:
488486
if len(cellid1) == 3:
@@ -499,7 +497,7 @@ def _plot_barrier_bc(self, barrier_data, color=None, name=None, **kwargs):
499497
else:
500498
vertical_cell_indices.append([self.layer, cellid2[1]])
501499
else:
502-
# Horizontal barrier - plot as line on shared face
500+
# vertical face, horizontal barrier
503501
shared_face = get_shared_face(self.mg, cellid1, cellid2)
504502
if shared_face is not None:
505503
horizontal_line_segments.append(shared_face)

0 commit comments

Comments
 (0)