Differentiation support w.r.t. slab_bounds, dilation, and sidewall_angle for td.PolySlab

tidy3d/components/geometry/base.py

@@ -23,7 +23,7 @@
 )
 from ...log import log
 from ...packaging import check_import, verify_packages_import
-from ..autograd import AutogradFieldMap, TracedCoordinate, TracedSize, get_static
+from ..autograd import AutogradFieldMap, TracedCoordinate, TracedFloat, TracedSize, get_static
 from ..autograd.derivative_utils import DerivativeInfo, integrate_within_bounds
 from ..base import Tidy3dBaseModel, cached_property
 from ..transformation import RotationAroundAxis
@@ -1603,7 +1603,7 @@ class Planar(SimplePlaneIntersection, Geometry, ABC):
         2, title="Axis", description="Specifies dimension of the planar axis (0,1,2) -> (x,y,z)."
     )
 
-    sidewall_angle: float = pydantic.Field(
+    sidewall_angle: TracedFloat = pydantic.Field(
         0.0,
         title="Sidewall angle",
         description="Angle of the sidewall. "

tidy3d/components/geometry/polyslab.py

@@ -16,7 +16,7 @@
 from ...exceptions import SetupError, ValidationError
 from ...log import log
 from ...packaging import verify_packages_import
-from ..autograd import AutogradFieldMap, TracedVertices, get_static
+from ..autograd import AutogradFieldMap, TracedFloat, TracedVertices, get_static
 from ..autograd.derivative_utils import DerivativeInfo
 from ..base import cached_property, skip_if_fields_missing
 from ..types import (
@@ -53,14 +53,14 @@ class PolySlab(base.Planar):
     >>> p = PolySlab(vertices=vertices, axis=2, slab_bounds=(-1, 1))
     """
 
-    slab_bounds: Tuple[float, float] = pydantic.Field(
+    slab_bounds: tuple[TracedFloat, TracedFloat] = pydantic.Field(
         ...,
         title="Slab Bounds",
         description="Minimum and maximum positions of the slab along axis dimension.",
         units=MICROMETER,
     )
 
-    dilation: float = pydantic.Field(
+    dilation: TracedFloat = pydantic.Field(
         0.0,
         title="Dilation",
         description="Dilation of the supplied polygon by shifting each edge along its "
@@ -87,6 +87,7 @@ def make_shapely_polygon(vertices: ArrayLike) -> shapely.Polygon:
     @pydantic.validator("slab_bounds", always=True)
     def slab_bounds_order(cls, val):
         """Maximum position of the slab should be no smaller than its minimal position."""
+        lb, ub = map(get_static, val)
         if val[1] < val[0]:
             raise SetupError(
                 "Polyslab.slab_bounds must be specified in the order of "
@@ -133,12 +134,12 @@ def no_complex_self_intersecting_polygon_at_reference_plane(cls, val, values):
         For 1), we issue an error since it is yet to be supported;
         For 2), we heal the polygon, and warn that the polygon has been cleaned up.
         """
+        dist = get_static(values["dilation"])
         # no need to validate anything here
-        if isclose(values["dilation"], 0):
+        if isclose(dist, 0):
             return val
 
         val_np = PolySlab._proper_vertices(val)
-        dist = values["dilation"]
 
         # 0) fully eroded
         if dist < 0 and dist < -PolySlab._maximal_erosion(val_np):
@@ -187,20 +188,23 @@ def no_self_intersecting_polygon_during_extrusion(cls, val, values):
         To detect this, we sample _N_SAMPLE_POLYGON_INTERSECT cross sections to see if any creation
         of polygons/holes, and changes in vertices number.
         """
+        sidewall_angle = get_static(values["sidewall_angle"])
+        dilation = get_static(values["dilation"])
+        slab_bounds = tuple(map(get_static, values["slab_bounds"]))
 
         # no need to validate anything here
-        if isclose(values["sidewall_angle"], 0):
+        if isclose(sidewall_angle, 0):
             return val
 
         # apply dilation
         poly_ref = PolySlab._proper_vertices(val)
-        if not isclose(values["dilation"], 0):
-            poly_ref = PolySlab._shift_vertices(poly_ref, values["dilation"])[0]
+        if not isclose(dilation, 0):
+            poly_ref = PolySlab._shift_vertices(poly_ref, dilation)[0]
             poly_ref = PolySlab._heal_polygon(poly_ref)
 
         # Fist, check vertex-vertex crossing at any point during extrusion
-        length = values["slab_bounds"][1] - values["slab_bounds"][0]
-        dist = [-length * np.tan(values["sidewall_angle"])]
+        length = slab_bounds[1] - slab_bounds[0]
+        dist = [-length * np.tan(sidewall_angle)]
         # reverse the dilation value if it's defined on the top
         if values["reference_plane"] == "top":
             dist = [-dist[0]]
@@ -247,7 +251,7 @@ def from_gds(
         gds_dtype: int = None,
         gds_scale: pydantic.PositiveFloat = 1.0,
         dilation: float = 0.0,
-        sidewall_angle: float = 0,
+        sidewall_angle: float = 0.0,
         reference_plane: PlanePosition = "middle",
     ) -> List[PolySlab]:
         """Import :class:`PolySlab` from a ``gdstk.Cell`` or a ``gdspy.Cell``.