feat(adjoint): sidewall angle gradients

[yaugenst-flex]

Greptile Summary

Updated On: 2025-09-04 13:18:06 UTC

This PR implements sidewall angle gradients for automatic differentiation (adjoint optimization) in PolySlab and Cylinder geometries. The changes enable users to compute gradients with respect to sidewall angles, expanding Tidy3D's inverse design capabilities beyond existing vertex and bounds optimization.

The implementation consists of several key architectural changes:

Core Infrastructure: The sidewall_angle field in the Planar base class is upgraded from float to TracedFloat, enabling autograd tracking. A new reference_axis_pos property is added to standardize reference point calculations across different reference plane configurations.

Gradient Computation: A comprehensive _compute_derivative_sidewall_angle() method is implemented in PolySlab that calculates Vector-Jacobian Products (VJP) through surface integration over sidewall patches. The mathematical approach correctly handles the relationship between sidewall angle changes and surface normal velocities, with proper area element corrections for slanted surfaces.

Cylinder Support: The Cylinder class is enhanced to map sidewall angle gradients to its underlying PolySlab representation, with defensive programming to handle edge cases where derivative paths may not exist.

Web Platform Integration: New batch category types ('tidy3d_autograd', 'tidy3d_autograd_async', 'autograd_fwd', 'autograd_bwd') are added to properly categorize and route different phases of autograd computations.

Validation Compatibility: Existing validation functions are updated to handle autograd-traced sidewall angles using get_static() and getval() functions, ensuring validation logic remains functional while preserving autograd traceability.

The feature integrates seamlessly with Tidy3D's existing adjoint framework, following established patterns for parameter tracing and gradient computation. The implementation includes extensive mathematical handling of sidewall geometry transformations and proper numerical integration techniques.

Important Files Changed

Changed Files

Filename	Score	Overview
tidy3d/web/api/container.py	5/5	Adds new batch category types for autograd workflow classification
tests/test_components/autograd/test_autograd.py	5/5	Adds end-to-end test validating sidewall angle gradient propagation
CHANGELOG.md	5/5	Documents new autograd support for sidewall angles in Cylinder and PolySlab
tidy3d/components/geometry/primitives.py	4/5	Implements sidewall angle gradient mapping for Cylinder class with defensive programming
tidy3d/components/geometry/base.py	5/5	Upgrades sidewall_angle to TracedFloat and adds reference_axis_pos property
tidy3d/components/geometry/utils.py	5/5	Updates validation to handle autograd-traced sidewall angles using get_static
tidy3d/components/geometry/polyslab.py	4/5	Major implementation of sidewall angle gradient computation with VJP calculation
tests/test_components/autograd/numerical/test_autograd_polyslab_sidewall_numerical.py	4/5	New numerical validation test comparing adjoint vs finite difference gradients
tests/test_components/autograd/test_autograd_polyslab_sidewall.py	4/5	Comprehensive unit tests for sidewall angle gradient computation scenarios
tests/test_components/autograd/test_sidewall_edge_cases.py	4/5	Edge case tests covering geometric transformations and boundary conditions

Confidence score: 4/5

• This PR implements a complex mathematical feature with comprehensive testing, but the mathematical complexity and integration points create some risk
• Score reflects the thorough testing coverage and well-structured implementation, but lowered due to complex VJP calculations and potential for numerical precision issues
• Pay close attention to the PolySlab gradient computation method and numerical validation tests for mathematical correctness

Sequence Diagram

sequenceDiagram
    participant User
    participant JaxPolySlab
    participant ModeSolver as "Mode Solver"
    participant WebAPI as "Web API"
    participant AutogradEngine as "Autograd Engine"
    participant DerivativeComputer as "Derivative Computer"

    User->>JaxPolySlab: "Create PolySlab with traced sidewall_angle"
    JaxPolySlab->>JaxPolySlab: "Validate sidewall_angle bounds"
    User->>ModeSolver: "Create simulation with JaxPolySlab"
    User->>AutogradEngine: "Define objective function"
    AutogradEngine->>WebAPI: "Run forward simulation"
    WebAPI-->>AutogradEngine: "Return SimulationData"
    AutogradEngine->>DerivativeComputer: "Compute gradient w.r.t. sidewall_angle"
    DerivativeComputer->>DerivativeComputer: "_compute_derivative_sidewall_angle()"
    DerivativeComputer->>DerivativeComputer: "_collect_sidewall_patches()"
    DerivativeComputer->>DerivativeComputer: "_adaptive_edge_samples()"
    DerivativeComputer->>DerivativeComputer: "evaluate_gradient_at_points()"
    DerivativeComputer->>DerivativeComputer: "Compute shape-derivative factors"
    DerivativeComputer-->>AutogradEngine: "Return VJP gradients"
    AutogradEngine-->>User: "Return objective value and gradients"

Loading

[tylerflex]

OK!  A few comments and suggestions. Thanks @yaugenst-flex

[groberts-flex]

looks good to me, small note on testing!

[yaugenst-flex]

@groberts-flex @tylerflex ok polished it up, i think this is ready to go. waiting for dario's merge so i can run the numerical tests but other than that i think it's done

[greptile-apps]

_{10 files reviewed, 1 comment}

_{Edit Code Review Bot Settings | Greptile}

[github-actions]

Diff Coverage

Diff: origin/develop...HEAD, staged and unstaged changes

• tidy3d/components/geometry/base.py (100%)
• tidy3d/components/geometry/polyslab.py (96.3%): Missing lines 1523,1557,1630,1684,1696-1697,1721,1840
• tidy3d/components/geometry/primitives.py (42.9%): Missing lines 324,326,333-335,351-352,357-359,361-364,371,377,382-383,386,388
• tidy3d/components/geometry/utils.py (100%)

Summary

• Total: 264 lines
• Missing: 28 lines
• Coverage: 89%

tidy3d/components/geometry/polyslab.py

Lines 1519-1527

  1519 
  1520         z0 = max(self.slab_bounds[0], sim_min[self.axis])
  1521         z1 = min(self.slab_bounds[1], sim_max[self.axis])
  1522         if z1 <= z0:
! 1523             return np.array([], dtype=GRADIENT_DTYPE_FLOAT), 0.0, z0, z1
  1524 
  1525         n_z = max(1, int(np.ceil((z1 - z0) / dx)))
  1526         dz = (z1 - z0) / n_z
  1527         z_centers = np.linspace(z0 + dz / 2, z1 - dz / 2, n_z, dtype=GRADIENT_DTYPE_FLOAT)

---

Lines 1553-1561

  1553             if t_start >= t_end:
  1554                 return None
  1555 
  1556         if t_end <= t_start + EDGE_CLIP_TOLERANCE:
! 1557             return None
  1558 
  1559         return (t_start, t_end)
  1560 
  1561     @staticmethod

---

Lines 1626-1634

  1626         z_centers, dz, z0, z1 = self._z_slices(sim_min, sim_max, is_2d=is_2d, dx=dx * cos_th)
  1627 
  1628         # early exit: no slices
  1629         if (not is_2d) and len(z_centers) == 0:
! 1630             return {
  1631                 "centers": np.empty((0, 3), dtype=GRADIENT_DTYPE_FLOAT),
  1632                 "normals": np.empty((0, 3), dtype=GRADIENT_DTYPE_FLOAT),
  1633                 "perps1": np.empty((0, 3), dtype=GRADIENT_DTYPE_FLOAT),
  1634                 "perps2": np.empty((0, 3), dtype=GRADIENT_DTYPE_FLOAT),

---

Lines 1680-1688

  1680         for ei, (v0, v1) in enumerate(zip(vertices, next_v)):
  1681             edge_vec = v1 - v0
  1682             L = np.linalg.norm(edge_vec)
  1683             if np.isclose(L, 0.0):
! 1684                 continue
  1685 
  1686             # constant along edge: unit tangent in 3D (no axis component)
  1687             t_edge = basis["perp1"][ei]
  1688             # outward in-plane normal from canonical basis normal (axis-consistent)

---

Lines 1692-1701

  1692             if not np.isclose(nrm, 0.0):
  1693                 n2d = n2d / nrm
  1694             else:
  1695                 # fallback to right-handed construction if degenerate
! 1696                 tmp = np.cross(axis_vec, t_edge)
! 1697                 n2d = tmp / (np.linalg.norm(tmp) + 1e-20)
  1698 
  1699             for zc in z_centers:
  1700                 # offset at this slice along the in-plane outward normal
  1701                 d = -(zc - z_ref) * tan_th

---

Lines 1717-1725

  1717 
  1718                 # densify tangentially as |theta| grows: ds_phys scales with cos(theta)
  1719                 s_list, w_list = self._adaptive_edge_samples(L, denom_edge, t0, t1)
  1720                 if len(s_list) == 0:
! 1721                     continue
  1722 
  1723                 pts2d = v0 + np.outer(s_list, edge_vec)
  1724                 xyz = (
  1725                     self.unpop_axis_vect(

---

Lines 1836-1844

  1836             is_2d=False,
  1837             dx=dx,
  1838         )
  1839         if patch["centers"].shape[0] == 0:
! 1840             return 0.0
  1841 
  1842         # Shape-derivative factors:
  1843         # - Offset: d(z) = -(z - z_ref) * tan(theta)
  1844         # - Tangential rate: dd/dtheta = -(z - z_ref) * sec(theta)^2

---

tidy3d/components/geometry/primitives.py

Lines 320-330

  320         # build PolySlab derivative paths based on requested Cylinder paths
  321         ps_paths = set()
  322         for path in derivative_info.paths:
  323             if path == ("length",):
! 324                 ps_paths.update({("slab_bounds", 0), ("slab_bounds", 1)})
  325             elif path == ("radius",):
! 326                 ps_paths.add(("vertices",))
  327             elif "center" in path:
  328                 _, center_index = path
  329                 _, (index_x, index_y) = self.pop_axis((0, 1, 2), axis=self.axis)
  330                 if center_index in (index_x, index_y):

---

Lines 329-339

  329                 _, (index_x, index_y) = self.pop_axis((0, 1, 2), axis=self.axis)
  330                 if center_index in (index_x, index_y):
  331                     ps_paths.add(("vertices",))
  332                 else:
! 333                     ps_paths.update({("slab_bounds", 0), ("slab_bounds", 1)})
! 334             elif path == ("sidewall_angle",):
! 335                 ps_paths.add(("sidewall_angle",))
  336 
  337         # pass interpolators to PolySlab if available to avoid redundant conversions
  338         update_kwargs = {
  339             "paths": list(ps_paths),

---

Lines 347-368

  347 
  348         vjps = {}
  349         for path in derivative_info.paths:
  350             if path == ("length",):
! 351                 vjp_top = vjps_polyslab.get(("slab_bounds", 0), 0.0)
! 352                 vjp_bot = vjps_polyslab.get(("slab_bounds", 1), 0.0)
  353                 vjps[path] = vjp_top - vjp_bot
  354 
  355             elif path == ("radius",):
  356                 # transform polyslab vertices derivatives into radius derivative
! 357                 xs_, ys_ = self._points_unit_circle(num_pts_circumference=num_pts_circumference)
! 358                 if ("vertices",) not in vjps_polyslab:
! 359                     vjps[path] = 0.0
  360                 else:
! 361                     vjps_vertices_xs, vjps_vertices_ys = vjps_polyslab[("vertices",)].T
! 362                     vjp_xs = np.sum(xs_ * vjps_vertices_xs)
! 363                     vjp_ys = np.sum(ys_ * vjps_vertices_ys)
! 364                     vjps[path] = vjp_xs + vjp_ys
  365 
  366             elif "center" in path:
  367                 _, center_index = path
  368                 _, (index_x, index_y) = self.pop_axis((0, 1, 2), axis=self.axis)

---

Lines 367-375

  367                 _, center_index = path
  368                 _, (index_x, index_y) = self.pop_axis((0, 1, 2), axis=self.axis)
  369                 if center_index == index_x:
  370                     if ("vertices",) not in vjps_polyslab:
! 371                         vjps[path] = 0.0
  372                     else:
  373                         vjps_vertices_xs = vjps_polyslab[("vertices",)][:, 0]
  374                         vjps[path] = np.sum(vjps_vertices_xs)
  375                 elif center_index == index_y:

---

Lines 373-392

  373                         vjps_vertices_xs = vjps_polyslab[("vertices",)][:, 0]
  374                         vjps[path] = np.sum(vjps_vertices_xs)
  375                 elif center_index == index_y:
  376                     if ("vertices",) not in vjps_polyslab:
! 377                         vjps[path] = 0.0
  378                     else:
  379                         vjps_vertices_ys = vjps_polyslab[("vertices",)][:, 1]
  380                         vjps[path] = np.sum(vjps_vertices_ys)
  381                 else:
! 382                     vjp_top = vjps_polyslab.get(("slab_bounds", 0), 0.0)
! 383                     vjp_bot = vjps_polyslab.get(("slab_bounds", 1), 0.0)
  384                     vjps[path] = vjp_top + vjp_bot
  385 
! 386             elif path == ("sidewall_angle",):
  387                 # direct mapping: cylinder angle equals polyslab angle
! 388                 vjps[path] = vjps_polyslab.get(("sidewall_angle",), 0.0)
  389 
  390             else:
  391                 raise NotImplementedError(
  392                     f"Differentiation with respect to 'Cylinder' '{path}' field not supported. "

---

[groberts-flex]

thanks for adding this! it looks good to me if numerical tests are looking good!