jan/gridlock
2
0
Fork 0
Code Issues Pull requests Releases Wiki Activity

[draw] fix handling of Nx3 vertex arrays

Browse source
This commit is contained in:
Jan Petykiewicz 2026-04-20 10:21:22 -07:00
commit 43d5fa8b4f
2 changed files with 38 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

11
gridlock/draw.py
View file

@ -61,16 +61,18 @@ class GridDrawMixin(GridPosMixin):
for ii in range(len(poly_list)): for ii in range(len(poly_list)):
polygon = poly_list[ii] polygon = poly_list[ii]
malformed = f'Malformed polygon: ({ii})' malformed = f'Malformed polygon: ({ii})'
if polygon.ndim != 2:
raise GridError(malformed + 'must be a 2-dimensional ndarray')
if polygon.shape[1] not in (2, 3): if polygon.shape[1] not in (2, 3):
raise GridError(malformed + 'must be a Nx2 or Nx3 ndarray') raise GridError(malformed + 'must be a Nx2 or Nx3 ndarray')
if polygon.shape[1] == 3: if polygon.shape[1] == 3:
polygon = polygon[surface, :] if numpy.unique(polygon[:, slab.axis]).size != 1:
raise GridError(malformed + 'must be in plane with surface normal ' + 'xyz'[slab.axis])
polygon = polygon[:, surface]
poly_list[ii] = polygon poly_list[ii] = polygon
if not polygon.shape[0] > 2: if not polygon.shape[0] > 2:
raise GridError(malformed + 'must consist of more than 2 points') raise GridError(malformed + 'must consist of more than 2 points')
if polygon.ndim > 2 and not numpy.unique(polygon[:, slab.axis]).size == 1:
raise GridError(malformed + 'must be in plane with surface normal ' + 'xyz'[slab.axis])
# Broadcast foreground where necessary # Broadcast foreground where necessary
foregrounds: Sequence[foreground_callable_t] | Sequence[float] foregrounds: Sequence[foreground_callable_t] | Sequence[float]
@ -296,8 +298,6 @@ class GridDrawMixin(GridPosMixin):
if isinstance(z, dict): if isinstance(z, dict):
z = Extent(**z) z = Extent(**z)
center = numpy.asarray([x.center, y.center, z.center])
p = numpy.array([[x.min, y.max], p = numpy.array([[x.min, y.max],
[x.max, y.max], [x.max, y.max],
[x.max, y.min], [x.max, y.min],
@ -398,4 +398,3 @@ class GridDrawMixin(GridPosMixin):
slab = Slab(axis=direction, center=center[direction], span=thickness) slab = Slab(axis=direction, center=center[direction], span=thickness)
self.draw_polygon(cell_data, slab=slab, polygon=poly, foreground=foreground_func, offset2d=center[surface]) self.draw_polygon(cell_data, slab=slab, polygon=poly, foreground=foreground_func, offset2d=center[surface])

35
gridlock/test/test_grid.py
View file

@ -1,8 +1,8 @@
# import pytest import pytest
import numpy import numpy
from numpy.testing import assert_allclose #, assert_array_equal from numpy.testing import assert_allclose #, assert_array_equal
from .. import Grid, Extent #, Slab, Plane from .. import Grid, Extent, GridError, Plane
def test_draw_oncenter_2x2() -> None: def test_draw_oncenter_2x2() -> None:
@ -116,3 +116,34 @@ def test_draw_2shift_4x4() -> None:
[0, 0.125, 0.125, 0]])[None, :, :, None] [0, 0.125, 0.125, 0]])[None, :, :, None]
assert_allclose(arr, correct) assert_allclose(arr, correct)
def test_draw_polygon_accepts_coplanar_nx3_vertices() -> None:
grid = Grid([[0, 1, 2], [0, 1, 2], [0, 1]], shifts=[[0, 0, 0]])
arr_2d = grid.allocate(0)
arr_3d = grid.allocate(0)
slab = dict(axis='z', center=0.5, span=1.0)
polygon_2d = numpy.array([[0, 0], [1, 0], [1, 1], [0, 1]], dtype=float)
polygon_3d = numpy.array([[0, 0, 0.5],
[1, 0, 0.5],
[1, 1, 0.5],
[0, 1, 0.5]], dtype=float)
grid.draw_polygon(arr_2d, slab=slab, polygon=polygon_2d, foreground=1)
grid.draw_polygon(arr_3d, slab=slab, polygon=polygon_3d, foreground=1)
assert_allclose(arr_3d, arr_2d)
def test_draw_polygon_rejects_noncoplanar_nx3_vertices() -> None:
grid = Grid([[0, 1, 2], [0, 1, 2], [0, 1]], shifts=[[0, 0, 0]])
arr = grid.allocate(0)
polygon = numpy.array([[0, 0, 0.5],
[1, 0, 0.5],
[1, 1, 0.75],
[0, 1, 0.5]], dtype=float)
with pytest.raises(GridError):
grid.draw_polygon(arr, slab=dict(axis='z', center=0.5, span=1.0), polygon=polygon, foreground=1)