feat(tidy3d): FXC-3693-triangle-mesh-support-for-adjoint

[marcorudolphflex]

• Implement TriangleMesh _compute_derivatives, wiring it into the existing permittivity-gradient machinery so mesh surfaces participate in adjoint workflows.
• Sample triangle faces adaptively within simulation bounds, evaluate boundary gradients via DerivativeInfo.evaluate_gradient_at_points, and accumulate sensitivities back onto per-vertex surface-mesh data arrays.
• Cache barycentric sampling patterns, reuse interpolators when available, and add early-outs for degenerate meshes and out-of-domain geometries to avoid unnecessary work.
• Expand autograd integration tests to include a TriangleMesh structure and ensure forward workflows still trace gradients for the full structure set.
• Add a dedicated test_autograd_triangle_mesh.py suite that checks analytic gradients, directional derivatives, permutation invariance, constant-field force balance, out-of-bounds zeroing, and non-watertight handling.
• Document TriangleMesh autograd support in the geometry API docs, including performance guidance for large meshes.
• Enhance TriangleMesh sampling to consider both face area and longest edge length, preventing undersampling of skinny, high-aspect triangles.
• Extend the watertight mesh fixtures with a slender tetrahedron case to exercise the new edge-aware subdivision logic inside the gradient tests.
• Add a regression that asserts the subdivision count always meets the minimum implied by the longest edge, guarding against future regressions in sampling density.

Greptile Overview

Greptile Summary

This PR implements autograd support for TriangleMesh geometries, enabling gradient-based optimization of arbitrary mesh surfaces in inverse design workflows.

Key changes:

• Implemented _compute_derivatives in TriangleMesh that samples triangle surfaces adaptively, evaluates boundary gradients via the existing permittivity-gradient infrastructure, and accumulates sensitivities onto per-vertex data arrays using barycentric interpolation
• Enhanced triangle subdivision logic to consider both face area and longest edge length, preventing undersampling of skinny high-aspect triangles
• Added specialized handling for collapsed simulation axes (2D) by computing plane-triangle intersections and sampling along intersection segments
• Extended DataArray.__init__ to convert numpy object arrays containing autograd boxes, enabling traced geometry parameters to flow through mesh construction
• Added comprehensive test coverage including analytic gradient validation, directional derivatives, permutation invariance, constant-field force balance, and comparison against PolySlab gradients for equivalent rectangular geometries
• Implemented to_stl export method for TriangleMesh with binary/ASCII format support

Issues found:

• Three floating-point equality comparisons should use tolerance-based checks (np.isclose or <=) instead of ==
• CHANGELOG entry missing backticks around code identifier

Confidence Score: 4/5

• This PR is safe to merge with minor syntax fixes needed for floating-point comparisons
• The implementation is mathematically sound with comprehensive test coverage validating gradients against analytic solutions and finite differences. The adaptive sampling strategy properly handles edge cases like collapsed axes, out-of-bounds geometries, and degenerate triangles. However, three floating-point equality checks need fixing to avoid potential numerical instability, and the CHANGELOG needs a formatting update.
• Pay attention to mesh.py:975 and test_autograd_triangle_mesh.py:149,190 for floating-point comparison fixes

Important Files Changed

File Analysis

Filename	Score	Overview
tidy3d/components/geometry/mesh.py	4/5	Core TriangleMesh autograd implementation with comprehensive surface sampling, gradient accumulation, and special handling for collapsed axes. One floating-point comparison issue.
tidy3d/components/data/data_array.py	5/5	Added autograd box conversion for numpy object arrays to properly handle traced geometry parameters during differentiation.
tests/test_components/autograd/test_autograd_triangle_mesh.py	4/5	Comprehensive test suite covering analytic gradients, directional derivatives, permutation invariance, and edge cases. Two floating-point comparison issues found.
tests/test_components/autograd/numerical/test_autograd_polyslab_trianglemesh_numerical.py	5/5	Validates TriangleMesh gradients against PolySlab and finite-difference for rectangular geometries across 2D/3D configurations.

Sequence Diagram

sequenceDiagram
    participant User
    participant Autograd
    participant TriangleMesh
    participant DerivativeInfo
    participant Sampling
    participant Interpolator
    participant GradAccum

    User->>Autograd: Define objective with TriangleMesh parameters
    Autograd->>TriangleMesh: Forward pass (construct geometry)
    TriangleMesh->>TriangleMesh: _triangles_to_trimesh (strip autograd boxes)
    Note over TriangleMesh: Convert ArrayBox to static arrays<br/>for trimesh compatibility
    
    Autograd->>TriangleMesh: Backward pass (_compute_derivatives)
    TriangleMesh->>TriangleMesh: Validate mesh_dataset exists
    TriangleMesh->>DerivativeInfo: Check simulation_bounds
    
    alt Mesh outside simulation
        TriangleMesh-->>Autograd: Return zero gradients
    else Mesh intersects simulation
        TriangleMesh->>Sampling: _collect_surface_samples(triangles, spacing, bounds)
        
        loop For each triangle face
            Sampling->>Sampling: Compute area and normal
            alt Collapsed axis (2D simulation)
                Sampling->>Sampling: _triangle_plane_segments (intersect with plane)
                Sampling->>Sampling: Sample along intersection segments
            else Full 3D
                Sampling->>Sampling: _subdivision_count (adaptive based on area + edges)
                Sampling->>Sampling: _get_barycentric_samples (cached patterns)
            end
            Sampling->>Sampling: Filter samples inside simulation bounds
        end
        
        Sampling-->>TriangleMesh: Return {points, normals, perps, weights, faces, barycentric}
        
        TriangleMesh->>DerivativeInfo: create_interpolators if needed
        DerivativeInfo-->>Interpolator: Build field interpolators
        
        TriangleMesh->>DerivativeInfo: evaluate_gradient_at_points(samples, interpolators)
        DerivativeInfo->>Interpolator: Interpolate fields at sample points
        Interpolator-->>DerivativeInfo: Return gradient values g
        DerivativeInfo-->>TriangleMesh: Boundary gradient g
        
        TriangleMesh->>GradAccum: Accumulate per-vertex contributions
        loop For each vertex (0, 1, 2)
            GradAccum->>GradAccum: scaled = (weights * g * normals) * barycentric[:, vertex]
            GradAccum->>GradAccum: np.add.at(triangle_grads[:, vertex], faces, scaled)
        end
        
        GradAccum-->>TriangleMesh: triangle_grads array
        TriangleMesh-->>Autograd: vjps[("mesh_dataset", "surface_mesh")]
    end
    
    Autograd-->>User: Final gradients w.r.t. mesh vertices

Loading

Context used:

• Rule from dashboard - Use a tolerance-based comparison (e.g., np.isclose) for floating-point numbers instead of direct equ... (source)
• Rule from dashboard - In changelogs, enclose code identifiers (class, function names) in backticks and use specific names ... (source)

[greptile-apps]

_{5 files reviewed, 4 comments}

_{Edit Code Review Agent Settings | Greptile}

[github-actions]

Diff Coverage

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

• tidy3d/components/data/data_array.py (66.7%): Missing lines 87
• tidy3d/components/geometry/mesh.py (90.0%): Missing lines 754,759,762,788-790,794,862,866,906,917,972,990,1011,1058,1090,1111,1121-1122,1128-1129,1132-1133,1137,1140,1145,1150,1168,1217,1238,1272,1278
• tidy3d/components/geometry/primitives.py (37.9%): Missing lines 274,279,289-290,292-293,295-296,299-300,302-305,307,313,316-318,320-326,328-330,332-335,339-340,348

Summary

• Total: 382 lines
• Missing: 69 lines
• Coverage: 81%

tidy3d/components/data/data_array.py

  83             data = TidyArrayBox.from_arraybox(data)
  84         # do the same for xr.Variable or xr.DataArray type
  85         elif isinstance(data, (xr.Variable, xr.DataArray)):
  86             if isbox(data.data) and not is_tidy_box(data.data):
! 87                 data.data = TidyArrayBox.from_arraybox(data.data)
  88         super().__init__(data, *args, **kwargs)
  89 
  90     @classmethod
  91     def __get_validators__(cls):

---

tidy3d/components/geometry/mesh.py

  750         """Compute adjoint derivatives for a ``TriangleMesh`` geometry."""
  751         vjps: AutogradFieldMap = {}
  752 
  753         if not self.mesh_dataset:
! 754             raise DataError("Can't compute derivatives without mesh data.")
  755 
  756         valid_paths = {("mesh_dataset", "surface_mesh")}
  757         for path in derivative_info.paths:
  758             if path not in valid_paths:
! 759                 raise ValueError(f"No derivative defined w.r.t. 'TriangleMesh' field '{path}'.")
  760 
  761         if ("mesh_dataset", "surface_mesh") not in derivative_info.paths:
! 762             return vjps
  763 
  764         triangles = np.asarray(self.triangles, dtype=config.adjoint.gradient_dtype_float)
  765 
  766         # early exit if geometry is completely outside simulation bounds

---

  784             sim_max=sim_max,
  785         )
  786 
  787         if samples["points"].shape[0] == 0:
! 788             zeros = np.zeros_like(triangles)
! 789             vjps[("mesh_dataset", "surface_mesh")] = zeros
! 790             return vjps
  791 
  792         interpolators = derivative_info.interpolators
  793         if interpolators is None:
! 794             interpolators = derivative_info.create_interpolators(
  795                 dtype=config.adjoint.gradient_dtype_float
  796             )
  797 
  798         g = derivative_info.evaluate_gradient_at_points(

---

  858         warning_msg = "Some triangles from the mesh lie outside the simulation bounds - this may lead to inaccurate gradients."
  859         for face_index, tri in enumerate(triangles_arr):
  860             area, normal = self._triangle_area_and_normal(tri)
  861             if area <= AREA_SIZE_THRESHOLD:
! 862                 continue
  863 
  864             perps = self._triangle_tangent_basis(tri, normal)
  865             if perps is None:
! 866                 continue
  867             perp1, perp2 = perps
  868 
  869             if collapsed_axis is not None and plane_value is not None:
  870                 samples, outside_bounds = self._collect_surface_samples_2d(

---

  902                 log.warning(warning_msg)
  903                 warned = True
  904 
  905             if samples is None:
! 906                 continue
  907 
  908             points_list.append(samples["points"])
  909             normals_list.append(samples["normals"])
  910             perps1_list.append(samples["perps1"])

---

  913             faces_list.append(samples["faces"])
  914             bary_list.append(samples["barycentric"])
  915 
  916         if not points_list:
! 917             return {
  918                 "points": np.zeros((0, 3), dtype=dtype),
  919                 "normals": np.zeros((0, 3), dtype=dtype),
  920                 "perps1": np.zeros((0, 3), dtype=dtype),
  921                 "perps2": np.zeros((0, 3), dtype=dtype),

---

  968         for start, end in segments:
  969             vec = end - start
  970             length = float(np.linalg.norm(vec))
  971             if length <= tol:
! 972                 continue
  973 
  974             subdivisions = max(1, int(np.ceil(length / spacing)))
  975             t_vals = (np.arange(subdivisions, dtype=dtype) + 0.5) / subdivisions
  976             sample_points = start[None, :] + t_vals[:, None] * vec[None, :]

---

  986                 )
  987 
  988             outside_bounds = outside_bounds or (not np.all(inside_mask))
  989             if not np.any(inside_mask):
! 990                 continue
  991 
  992             sample_points = sample_points[inside_mask]
  993             bary_inside = bary[inside_mask]
  994             n_inside = sample_points.shape[0]

---

  1007             faces.append(faces_tile)
  1008             barycentric.append(bary_inside)
  1009 
  1010         if not points:
! 1011             return None, outside_bounds
  1012 
  1013         samples = {
  1014             "points": np.concatenate(points, axis=0),
  1015             "normals": np.concatenate(normals, axis=0),

---

  1054             sample_points[:, valid_axes] >= (sim_min - tol)[valid_axes], axis=1
  1055         ) & np.all(sample_points[:, valid_axes] <= (sim_max + tol)[valid_axes], axis=1)
  1056         outside_bounds = not np.all(inside_mask)
  1057         if not np.any(inside_mask):
! 1058             return None, outside_bounds
  1059 
  1060         sample_points = sample_points[inside_mask]
  1061         bary_inside = barycentric[inside_mask]
  1062         n_samples_inside = sample_points.shape[0]

---

  1086         edge02 = triangle[2] - triangle[0]
  1087         cross = np.cross(edge01, edge02)
  1088         norm = np.linalg.norm(cross)
  1089         if norm <= 0.0:
! 1090             return 0.0, np.zeros(3, dtype=triangle.dtype)
  1091         normal = (cross / norm).astype(triangle.dtype, copy=False)
  1092         area = 0.5 * norm
  1093         return area, normal

---

  1107 
  1108         def add_point(pt: np.ndarray) -> None:
  1109             for existing in points:
  1110                 if np.linalg.norm(existing - pt) <= tol:
! 1111                     return
  1112             points.append(pt.copy())
  1113 
  1114         for i, j in edges:
  1115             di = distances[i]

---

  1117             vi = vertices[i]
  1118             vj = vertices[j]
  1119 
  1120             if abs(di) <= tol and abs(dj) <= tol:
! 1121                 segments.append((vi.copy(), vj.copy()))
! 1122                 continue
  1123 
  1124             if di * dj > 0.0:
  1125                 continue

---

  1124             if di * dj > 0.0:
  1125                 continue
  1126 
  1127             if abs(di) <= tol:
! 1128                 add_point(vi)
! 1129                 continue
  1130 
  1131             if abs(dj) <= tol:
! 1132                 add_point(vj)
! 1133                 continue
  1134 
  1135             denom = di - dj
  1136             if abs(denom) <= tol:
! 1137                 continue
  1138             t = di / denom
  1139             if t < 0.0 or t > 1.0:
! 1140                 continue
  1141             point = vi + t * (vj - vi)
  1142             add_point(point)
  1143 
  1144         if segments:
! 1145             return segments
  1146 
  1147         if len(points) >= 2:
  1148             return [(points[0], points[1])]
  1149 
! 1150         return []
  1151 
  1152     @staticmethod
  1153     def _barycentric_coordinates(triangle: NDArray, points: np.ndarray, tol: float) -> np.ndarray:
  1154         """Compute barycentric coordinates of ``points`` with respect to ``triangle``."""

---

  1164         d01 = float(np.dot(v0v1, v0v2))
  1165         d11 = float(np.dot(v0v2, v0v2))
  1166         denom = d00 * d11 - d01 * d01
  1167         if abs(denom) <= tol:
! 1168             return np.tile(
  1169                 np.array([1.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0], dtype=triangle.dtype), (pts.shape[0], 1)
  1170             )
  1171 
  1172         v0p = pts - v0

---

  1213     def _build_barycentric_samples(subdivisions: int) -> np.ndarray:
  1214         """Construct barycentric sampling points for a given subdivision level."""
  1215 
  1216         if subdivisions <= 1:
! 1217             return np.array([[1.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0]])
  1218 
  1219         bary = []
  1220         for i in range(subdivisions):
  1221             for j in range(subdivisions - i):

---

  1234 
  1235         def midpoint(i: int, j: int) -> int:
  1236             key = (i, j) if i < j else (j, i)
  1237             if key in midpoint_cache:
! 1238                 return midpoint_cache[key]
  1239             vm = 0.5 * (verts_list[i] + verts_list[j])
  1240             verts_list.append(vm)
  1241             idx = len(verts_list) - 1
  1242             midpoint_cache[key] = idx

---

  1268                 edge = (candidate / length).astype(triangle.dtype, copy=False)
  1269                 break
  1270 
  1271         if edge is None:
! 1272             return None
  1273 
  1274         perp1 = edge
  1275         perp2 = np.cross(normal, perp1)
  1276         perp2_norm = np.linalg.norm(perp2)

---

  1274         perp1 = edge
  1275         perp2 = np.cross(normal, perp1)
  1276         perp2_norm = np.linalg.norm(perp2)
  1277         if perp2_norm <= tol:
! 1278             return None
  1279         perp2 = (perp2 / perp2_norm).astype(triangle.dtype, copy=False)
  1280         return perp1, perp2

---

tidy3d/components/geometry/primitives.py

  270         subdivisions: Optional[int] = None,
  271     ) -> np.ndarray:
  272         """Return unit sphere triangles discretized via an icosphere."""
  273 
! 274         unit_tris = UNIT_SPHERE._unit_sphere_triangles(
  275             target_edge_length=target_edge_length,
  276             subdivisions=subdivisions,
  277             copy_result=True,
  278         )
! 279         return unit_tris
  280 
  281     def _unit_sphere_triangles(
  282         self,
  283         *,

---

  285         subdivisions: Optional[int] = None,
  286         copy_result: bool = True,
  287     ) -> np.ndarray:
  288         """Return cached unit-sphere triangles with optional copying."""
! 289         if target_edge_length is not None and subdivisions is not None:
! 290             raise ValueError("Specify either target_edge_length OR subdivisions, not both.")
  291 
! 292         if subdivisions is None:
! 293             subdivisions = self._subdivisions_for_edge(target_edge_length)
  294 
! 295         triangles, _ = self._icosphere_data(subdivisions)
! 296         return np.array(triangles, copy=copy_result)
  297 
  298     def _subdivisions_for_edge(self, target_edge_length: Optional[float]) -> int:
! 299         if target_edge_length is None or target_edge_length <= 0.0:
! 300             return 0
  301 
! 302         for subdiv in range(_MAX_ICOSPHERE_SUBDIVISIONS + 1):
! 303             _, max_edge = self._icosphere_data(subdiv)
! 304             if max_edge <= target_edge_length:
! 305                 return subdiv
  306 
! 307         log.warning(
  308             f"Requested sphere mesh edge length {target_edge_length:.3e} μm requires more than "
  309             f"{_MAX_ICOSPHERE_SUBDIVISIONS} subdivisions. "
  310             "Clipping to the finest available mesh.",
  311             log_once=True,

---

  309             f"{_MAX_ICOSPHERE_SUBDIVISIONS} subdivisions. "
  310             "Clipping to the finest available mesh.",
  311             log_once=True,
  312         )
! 313         return _MAX_ICOSPHERE_SUBDIVISIONS
  314 
  315     def _icosphere_data(self, subdivisions: int) -> tuple[np.ndarray, float]:
! 316         cache = self._icosphere_cache
! 317         if subdivisions in cache:
! 318             return cache[subdivisions]
  319 
! 320         vertices = np.asarray(_ICOSAHEDRON_VERTS, dtype=float)
! 321         faces = np.asarray(_ICOSAHEDRON_FACES, dtype=int)
! 322         if subdivisions > 0:
! 323             vertices = vertices.copy()
! 324             faces = faces.copy()
! 325             for _ in range(subdivisions):
! 326                 vertices, faces = TriangleMesh.subdivide_faces(vertices, faces)
  327 
! 328         norms = np.linalg.norm(vertices, axis=1, keepdims=True)
! 329         norms = np.where(norms == 0.0, 1.0, norms)
! 330         vertices = vertices / norms
  331 
! 332         triangles = vertices[faces]
! 333         max_edge = self._max_edge_length(triangles)
! 334         cache[subdivisions] = (triangles, max_edge)
! 335         return triangles, max_edge
  336 
  337     @staticmethod
  338     def _max_edge_length(triangles: np.ndarray) -> float:
! 339         v = triangles
! 340         edges = np.stack(
  341             [
  342                 v[:, 1] - v[:, 0],
  343                 v[:, 2] - v[:, 1],
  344                 v[:, 0] - v[:, 2],

---

  344                 v[:, 0] - v[:, 2],
  345             ],
  346             axis=1,
  347         )
! 348         return float(np.linalg.norm(edges, axis=2).max())
  349 
  350 
  351 UNIT_SPHERE = Sphere(center=(0.0, 0.0, 0.0), radius=1.0)

---

[greptile-apps]

_{10 files reviewed, 4 comments}

_{Edit Code Review Agent Settings | Greptile}

[yaugenst-flex]

Thanks @marcorudolphflex this is awesome!