adds autograd differentiation w.r.t. cylinder radius and center (non-…

…axis)

CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added value_and_grad function to the autograd plugin, importable via `from tidy3d.plugins.autograd import value_and_grad`. Supports differentiating functions with auxiliary data (`value_and_grad(f, has_aux=True)`).
 - `Simulation.num_computational_grid_points` property to examine the number of grid cells that compose the computational domain corresponding to the simulation. This can differ from `Simulation.num_cells` based on boundary conditions and symmetries.
 - Support for `dilation` argument in `JaxPolySlab`.
+- Support for autograd differentiation with respect to `Cylinder.radius` and `Cylinder.center` (for elements not along axis dimension).
+- `Cylinder.to_polyslab(num_pts_circumference, **kwargs)` to convert a cylinder into a discretized version represented by a `PolySlab`.
 
 ### Changed
 - `PolySlab` now raises error when differentiating and dilation causes damage to the polygon.

tests/test_components/test_autograd.py

@@ -392,6 +392,17 @@ def make_structures(params: anp.ndarray) -> dict[str, td.Structure]:
     custom_med_pole_res = td.CustomPoleResidue(eps_inf=eps_inf, poles=[(a1, c1), (a2, c2)])
     custom_pole_res = td.Structure(geometry=box, medium=custom_med_pole_res)
 
+    radius = 0.4 * (1 + anp.abs(vector @ params))
+    cyl_center_y = vector @ params
+    cyl_center_z = -vector @ params
+    cylinder_geo = td.Cylinder(
+        radius=radius,
+        center=(0, cyl_center_y, cyl_center_z),
+        axis=0,
+        length=LX / 2 if IS_3D else td.inf,
+    )
+    cylinder = td.Structure(geometry=cylinder_geo, medium=med)
+
     return dict(
         medium=medium,
         center_list=center_list,
@@ -402,6 +413,7 @@ def make_structures(params: anp.ndarray) -> dict[str, td.Structure]:
         geo_group=geo_group,
         pole_res=pole_res,
         custom_pole_res=custom_pole_res,
+        cylinder=cylinder,
     )
 
 
@@ -496,6 +508,7 @@ def plot_sim(sim: td.Simulation, plot_eps: bool = True) -> None:
     "geo_group",
     "pole_res",
     "custom_pole_res",
+    "cylinder",
 )
 monitor_keys_ = ("mode", "diff", "field_vol", "field_point")
 
@@ -586,7 +599,7 @@ def test_polyslab_axis_ops(axis):
 
 
 @pytest.mark.skipif(not RUN_NUMERICAL, reason="Numerical gradient tests runs through web API.")
-@pytest.mark.parametrize("structure_key, monitor_key", (("medium", "mode"),))
+@pytest.mark.parametrize("structure_key, monitor_key", (("cylinder", "mode"),))
 def test_autograd_numerical(structure_key, monitor_key):
     """Test an objective function through tidy3d autograd."""
 
@@ -610,7 +623,7 @@ def objective(*args):
     assert anp.all(grad != 0.0), "some gradients are 0"
 
     # numerical gradients
-    delta = 1e-1
+    delta = 1e-3
     sims_numerical = {}
 
     params_num = np.zeros((N_PARAMS, N_PARAMS))

tests/test_components/test_geometry.py

@@ -240,6 +240,10 @@ def test_slanted_cylinder_infinite_length_validate():
         )
 
 
+def test_cylinder_to_polyslab():
+    ps = CYLINDER.to_polyslab(num_pts_circumference=10, sidewall_angle=0.02)
+
+
 def test_box_from_bounds():
     b = td.Box.from_bounds(rmin=(-td.inf, 0, 0), rmax=(td.inf, 0, 0))
     assert b.center[0] == 0.0

tidy3d/components/geometry/primitives.py

@@ -5,23 +5,33 @@
 from math import isclose
 from typing import List
 
+import autograd.numpy as anp
 import numpy as np
 import pydantic.v1 as pydantic
 import shapely
 
-from ...constants import LARGE_NUMBER, MICROMETER
+from ...constants import C_0, LARGE_NUMBER, MICROMETER
 from ...exceptions import SetupError, ValidationError
 from ...packaging import verify_packages_import
+from ..autograd import AutogradFieldMap, TracedSize1D
+from ..autograd.derivative_utils import DerivativeInfo
 from ..base import cached_property, skip_if_fields_missing
 from ..types import Axis, Bound, Coordinate, MatrixReal4x4, Shapely, Tuple
 from . import base
+from .polyslab import PolySlab
 
 # for sampling conical frustum in visualization
 _N_SAMPLE_CURVE_SHAPELY = 40
 
 # for shapely circular shapes discretization in visualization
 _N_SHAPELY_QUAD_SEGS = 200
 
+# Default number of points to discretize polyslab in `Cylinder.to_polyslab()`
+_N_PTS_CYLINDER_POLYSLAB = 51
+
+# Default number of points per wvl in material for discretizing cylinder in autograd derivative
+_PTS_PER_WVL_MAT_CYLINDER_DISCRETIZE = 10
+
 
 class Sphere(base.Centered, base.Circular):
     """Spherical geometry.
@@ -185,7 +195,7 @@ class Cylinder(base.Centered, base.Circular, base.Planar):
     """
 
     # Provide more explanations on where radius is defined
-    radius: pydantic.NonNegativeFloat = pydantic.Field(
+    radius: TracedSize1D = pydantic.Field(
         ...,
         title="Radius",
         description="Radius of geometry at the ``reference_plane``.",
@@ -215,6 +225,123 @@ def _only_middle_for_infinite_length_slanted_cylinder(cls, val, values):
             )
         return val
 
+    def to_polyslab(
+        self, num_pts_circumference: int = _N_PTS_CYLINDER_POLYSLAB, **kwargs
+    ) -> PolySlab:
+        """Convert instance of ``Cylinder`` into a discretized version using ``PolySlab``.
+
+        Parameters
+        ----------
+        num_pts_circumference : int = 51
+            Number of points in the circumference of the discretized polyslab.
+        **kwargs:
+            Extra keyword arguments passed to ``PolySlab()``, such as ``sidewall_angle``.
+
+        Returns
+        -------
+        PolySlab
+            Extruded polygon representing a discretized version of the cylinder.
+        """
+
+        center_axis = self.center_axis
+        length_axis = self.length_axis
+        slab_bounds = (center_axis - length_axis / 2.0, center_axis + length_axis / 2.0)
+
+        if num_pts_circumference < 3:
+            raise ValueError("'PolySlab' from 'Cylinder' must have 3 or more radius points.")
+
+        _, (x0, y0) = self.pop_axis(self.center, axis=self.axis)
+
+        xs_, ys_ = self._points_unit_circle(num_pts_circumference=num_pts_circumference)
+
+        xs = x0 + self.radius * xs_
+        ys = y0 + self.radius * ys_
+
+        vertices = anp.stack((xs, ys), axis=-1)
+
+        return PolySlab(
+            vertices=vertices,
+            axis=self.axis,
+            slab_bounds=slab_bounds,
+            **kwargs,
+        )
+
+    def _points_unit_circle(
+        self, num_pts_circumference: int = _N_PTS_CYLINDER_POLYSLAB
+    ) -> np.ndarray:
+        """Set of x and y points for the unit circle when discretizing cylinder as a polyslab."""
+        angles = np.linspace(0, 1, num_pts_circumference + 1)[:-1]
+        xs = np.cos(angles)
+        ys = np.sin(angles)
+        return np.stack((xs, ys), axis=0)
+
+    def compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradFieldMap:
+        """Compute the adjoint derivatives for this object."""
+
+        # compute number of points in the circumference of the polyslab using resolution info
+        wvl0 = C_0 / derivative_info.frequency
+        wvl_mat = wvl0 / max(1.0, np.sqrt(abs(derivative_info.eps_in)))
+
+        circumference = 2 * np.pi * self.radius
+        wvls_in_circumference = circumference / wvl_mat
+
+        num_pts_circumference = int(
+            np.ceil(_PTS_PER_WVL_MAT_CYLINDER_DISCRETIZE * wvls_in_circumference)
+        )
+        num_pts_circumference = max(3, num_pts_circumference)
+
+        # construct equivalent polyslab and compute the derivatives
+        polyslab = self.to_polyslab(num_pts_circumference=num_pts_circumference)
+
+        derivative_info_polyslab = derivative_info.updated_copy(paths=[("vertices",)])
+        vjps_polyslab = polyslab.compute_derivatives(derivative_info_polyslab)
+
+        vjps_vertices_xs, vjps_vertices_ys = vjps_polyslab[("vertices",)].T
+
+        # transform polyslab vertices derivatives into Cylinder parameter derivatives
+        vjps = {}
+        for path in derivative_info.paths:
+            if path == ("radius",):
+                xs_, ys_ = self._points_unit_circle(num_pts_circumference=num_pts_circumference)
+
+                vjp_xs = np.sum(xs_ * vjps_vertices_xs)
+                vjp_ys = np.sum(ys_ * vjps_vertices_ys)
+
+                vjps[path] = vjp_xs + vjp_ys
+
+            elif "center" in path:
+                _, center_index = path
+                if center_index == self.axis:
+                    raise NotImplementedError(
+                        "Currently cannot differentiate Cylinder with respect to its 'center' along"
+                        " the axis. If you would like this feature added, please feel free to raise"
+                        " an issue on the tidy3d front end repository."
+                    )
+
+                _, (index_x, index_y) = self.pop_axis((0, 1, 2), axis=self.axis)
+                if center_index == index_x:
+                    vjps[path] = np.sum(vjp_xs)
+                elif center_index == index_y:
+                    vjps[path] = np.sum(vjp_ys)
+                else:
+                    raise ValueError(
+                        "Something unexpected happened. Was asked to differentiate "
+                        f"with respect to 'Cylinder.center[{center_index}]', but this was not "
+                        "detected as being one of the parallel axis with "
+                        f"'Cylinder.axis' of '{self.axis}'. If you received this error, please raise "
+                        "an issue on the tidy3d front end repository with details about how you "
+                        "defined your 'Cylinder' in the objective function."
+                    )
+
+            else:
+                raise NotImplementedError(
+                    f"Differentiation with respect to 'Cylinder' '{path}' field not supported. "
+                    "If you would like this feature added, please feel free to raise "
+                    "an issue on the tidy3d front end repository."
+                )
+
+        return vjps
+
     @property
     def center_axis(self):
         """Gets the position of the center of the geometry in the out of plane dimension."""

tidy3d/components/scene.py

@@ -240,6 +240,7 @@ def intersecting_media(
         List[:class:`.AbstractMedium`]
             Set of distinct mediums that intersect with the given planar object.
         """
+        structures = [s.to_static() for s in structures]
         if test_object.size.count(0.0) == 1:
             # get all merged structures on the test_object, which is already planar
             structures_merged = Scene._filter_structures_plane_medium(structures, test_object)