polyslab accepts numpy arrays for vertices

Author: tylerflex

tests/test_components.py

@@ -127,6 +127,8 @@ def test_geometry():
     s = Sphere(radius=1, center=(0, 0, 0))
     s = Cylinder(radius=1, center=(0, 0, 0), axis=1, length=1)
     s = PolySlab(vertices=((1, 2), (3, 4), (5, 4)), slab_bounds=(-1, 1), axis=1)
+    vertices_np = np.array(s.vertices)
+    s_np = PolySlab(vertices=vertices_np, slab_bounds=(-1, 1), axis=1)
 
     # make sure wrong axis arguments error
     with pytest.raises(pydantic.ValidationError) as e_info:

tidy3d/components/geometry.py

@@ -11,10 +11,10 @@
 from descartes import PolygonPatch
 
 from .base import Tidy3dBaseModel
-from .types import Literal, Bound, Size, Coordinate, Axis, Coordinate2D
+from .types import Literal, Bound, Size, Coordinate, Axis, Coordinate2D, Array
 from .types import Vertices, Ax, Shapely
 from .viz import add_ax_if_none
-from ..log import Tidy3dKeyError
+from ..log import Tidy3dKeyError, SetupError
 
 # add this around extents of plots
 PLOT_BUFFER = 0.3
@@ -780,8 +780,9 @@ class PolySlab(Planar):
 
     Parameters
     ----------
-    vertices : List[Tuple[float, float]]
+    vertices : List[Tuple[float, float]] or np.ndarray
         List of vertices defining the polygon face along dimensions parallel to slab normal axis.
+        If numpy.ndarray, must have shape (N, 2) for N vertices.
     axis : int
         Integer index into the polygon's slab axis. (0,1,2) -> (x,y,z)
     slab_bounds: Tuple[float, float]
@@ -793,7 +794,7 @@ class PolySlab(Planar):
     """
 
     slab_bounds: Tuple[float, float]
-    vertices: Vertices
+    vertices: Union[Vertices, Array[float]]
     type: Literal["PolySlab"] = "PolySlab"
 
     @pydantic.validator("slab_bounds", always=True)
@@ -803,6 +804,18 @@ def set_length(cls, val, values):
         values["length"] = zmax - zmin
         return val
 
+    @pydantic.validator("vertices", always=True)
+    def correct_shape(cls, val):
+        """makes sure vertices is correct shape if numpy array"""
+        if isinstance(val, np.ndarray):
+            shape = val.shape
+            if len(shape) != 2 or shape[1] != 2:
+                raise SetupError(
+                    "PolySlab.vertices must be a 2 dimensional array shaped (N, 2).  "
+                    f"Given array with shape of {shape}."
+                )
+        return val
+
     @pydantic.validator("vertices", always=True)
     def set_center(cls, val, values):
         """sets the .center field using zmin, zmax, and polygon vertices"""