From 43d5fa8b4f2e1d48a2d35de02ebc2a7adf591168 Mon Sep 17 00:00:00 2001
From: Jan Petykiewicz <jan@mpxd.net>
Date: Mon, 20 Apr 2026 10:21:22 -0700
Subject: [PATCH] [draw] fix handling of Nx3 vertex arrays

---
 gridlock/draw.py           | 11 +++++------
 gridlock/test/test_grid.py | 35 +++++++++++++++++++++++++++++++++--
 2 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/gridlock/draw.py b/gridlock/draw.py
index 9ba4623..864468f 100644
--- a/gridlock/draw.py
+++ b/gridlock/draw.py
@@ -61,16 +61,18 @@ class GridDrawMixin(GridPosMixin):
         for ii in range(len(poly_list)):
             polygon = poly_list[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):
                 raise GridError(malformed + 'must be a Nx2 or Nx3 ndarray')
             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
 
             if not polygon.shape[0] > 2:
                 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
         foregrounds: Sequence[foreground_callable_t] | Sequence[float]
@@ -296,8 +298,6 @@ class GridDrawMixin(GridPosMixin):
         if isinstance(z, dict):
             z = Extent(**z)
 
-        center = numpy.asarray([x.center, y.center, z.center])
-
         p = numpy.array([[x.min, y.max],
                          [x.max, y.max],
                          [x.max, y.min],
@@ -398,4 +398,3 @@ class GridDrawMixin(GridPosMixin):
 
         slab = Slab(axis=direction, center=center[direction], span=thickness)
         self.draw_polygon(cell_data, slab=slab, polygon=poly, foreground=foreground_func, offset2d=center[surface])
-
diff --git a/gridlock/test/test_grid.py b/gridlock/test/test_grid.py
index 8d9ca92..6cb9edc 100644
--- a/gridlock/test/test_grid.py
+++ b/gridlock/test/test_grid.py
@@ -1,8 +1,8 @@
-# import pytest
+import pytest
 import numpy
 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:
@@ -116,3 +116,34 @@ def test_draw_2shift_4x4() -> None:
                            [0, 0.125, 0.125, 0]])[None, :, :, None]
 
     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)
+