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

Updated On: 2025-11-05 16:16:52 UTC

Greptile Summary

Implemented autograd support for TriangleMesh geometries by adding _compute_derivatives method that evaluates boundary sensitivities using surface integral formulation. Reused the permittivity-gradient infrastructure from PolySlab, enabling per-vertex gradient accumulation via barycentric weighting.

Key Changes

• Added _compute_derivatives method in tidy3d/components/geometry/mesh.py:705 that computes adjoint derivatives for triangle mesh vertices
• Implemented adaptive surface sampling with barycentric subdivision based on triangle area and spacing requirements
• Added helper methods: _collect_surface_samples, _triangle_area_and_normal, _subdivision_count, _get_barycentric_samples, _build_barycentric_samples, and _triangle_tangent_basis
• Introduced class-level _barycentric_cache to avoid recomputing sampling patterns for repeated subdivision levels
• Integrated TriangleMesh into existing autograd test suite with comprehensive validation tests
• Added dedicated test file with 8 tests validating gradient accuracy against analytic solutions for watertight meshes
• Updated documentation with performance notes about mesh complexity impact on gradient computation

Issues Found

• Multiple instances of floating-point equality comparison (==) instead of tolerance-based comparison (<=) violate custom rule for floating-point precision handling

Confidence Score: 4/5

• Safe to merge with minor syntax improvements needed for floating-point comparisons
• Implementation is well-structured and thoroughly tested with comprehensive validation against analytic solutions. The surface integral approach correctly reuses existing infrastructure. Minor issues with floating-point equality comparisons should be fixed but don't affect correctness significantly. Documentation is clear and helpful.
• Pay attention to floating-point comparison issues in tidy3d/components/geometry/mesh.py:806 and tests/test_components/autograd/test_autograd_triangle_mesh.py:134,163,202

Important Files Changed

File Analysis

Filename	Score	Overview
tidy3d/components/geometry/mesh.py	4/5	Implemented _compute_derivatives method for mesh-based adjoint optimization using surface integral evaluation with barycentric sampling
tests/test_components/autograd/test_autograd_triangle_mesh.py	5/5	Comprehensive test suite validating gradient accuracy against analytic solutions for watertight meshes

Sequence Diagram

sequenceDiagram
    participant User
    participant TriangleMesh
    participant DerivativeInfo
    participant SurfaceSampler
    participant Interpolator
    participant GradientAccumulator

    User->>TriangleMesh: _compute_derivatives(derivative_info)
    TriangleMesh->>TriangleMesh: validate mesh_dataset exists
    TriangleMesh->>TriangleMesh: check paths for "surface_mesh"
    
    TriangleMesh->>TriangleMesh: convert triangles to gradient dtype
    TriangleMesh->>TriangleMesh: check if mesh outside sim bounds
    
    alt mesh outside bounds
        TriangleMesh-->>User: return zero gradients
    end
    
    TriangleMesh->>DerivativeInfo: adaptive_vjp_spacing()
    DerivativeInfo-->>TriangleMesh: spacing (dx)
    
    TriangleMesh->>SurfaceSampler: _collect_surface_samples(triangles, spacing, bounds)
    
    loop for each triangle face
        SurfaceSampler->>SurfaceSampler: compute area and normal
        SurfaceSampler->>SurfaceSampler: compute tangent basis (perp1, perp2)
        SurfaceSampler->>SurfaceSampler: determine subdivisions from area/spacing
        SurfaceSampler->>SurfaceSampler: get barycentric samples (cached)
        SurfaceSampler->>SurfaceSampler: compute sample points = barycentric @ triangle
        SurfaceSampler->>SurfaceSampler: filter points inside sim bounds
        SurfaceSampler->>SurfaceSampler: accumulate samples with weights
    end
    
    SurfaceSampler-->>TriangleMesh: samples (points, normals, perps, weights, faces, barycentric)
    
    alt no samples collected
        TriangleMesh-->>User: return zero gradients
    end
    
    TriangleMesh->>DerivativeInfo: create_interpolators(dtype)
    DerivativeInfo-->>TriangleMesh: interpolators
    
    TriangleMesh->>DerivativeInfo: evaluate_gradient_at_points(samples, interpolators)
    DerivativeInfo->>Interpolator: interpolate E_der_map at points
    DerivativeInfo->>Interpolator: interpolate D_der_map at points
    DerivativeInfo-->>TriangleMesh: gradient values (g)
    
    TriangleMesh->>GradientAccumulator: initialize triangle_grads (zeros)
    TriangleMesh->>GradientAccumulator: compute contrib_vec = weights * g * normals
    
    loop for each vertex (0, 1, 2)
        GradientAccumulator->>GradientAccumulator: scale by barycentric weight
        GradientAccumulator->>GradientAccumulator: np.add.at per-vertex gradients
    end
    
    GradientAccumulator-->>TriangleMesh: triangle_grads
    TriangleMesh-->>User: vjps[("mesh_dataset", "surface_mesh")]

Loading

[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/geometry/mesh.py (90.0%): Missing lines 711,716,719,745-747,751,807,811,829,850,879,920,948,954

Summary

• Total: 150 lines
• Missing: 15 lines
• Coverage: 90%

tidy3d/components/geometry/mesh.py

  707 
  708         vjps: AutogradFieldMap = {}
  709 
  710         if not self.mesh_dataset:
! 711             raise DataError("Can't compute derivatives without mesh data.")
  712 
  713         valid_paths = {("mesh_dataset", "surface_mesh")}
  714         for path in derivative_info.paths:
  715             if path not in valid_paths:
! 716                 raise ValueError(f"No derivative defined w.r.t. 'TriangleMesh' field '{path}'.")
  717 
  718         if ("mesh_dataset", "surface_mesh") not in derivative_info.paths:
! 719             return vjps
  720 
  721         triangles = np.asarray(self.triangles, dtype=config.adjoint.gradient_dtype_float)
  722 
  723         # early exit if geometry is completely outside simulation bounds

---

  741             sim_max=sim_max,
  742         )
  743 
  744         if samples["points"].shape[0] == 0:
! 745             zeros = np.zeros_like(triangles, dtype=config.adjoint.gradient_dtype_float)
! 746             vjps[("mesh_dataset", "surface_mesh")] = zeros
! 747             return vjps
  748 
  749         interpolators = derivative_info.interpolators
  750         if interpolators is None:
! 751             interpolators = derivative_info.create_interpolators(
  752                 dtype=config.adjoint.gradient_dtype_float
  753             )
  754 
  755         g = derivative_info.evaluate_gradient_at_points(

---

  803 
  804         for face_index, tri in enumerate(np.asarray(triangles, dtype=dtype)):
  805             area, normal = self._triangle_area_and_normal(tri)
  806             if area <= AREA_SIZE_THRESHOLD:
! 807                 continue
  808 
  809             perps = self._triangle_tangent_basis(tri, normal)
  810             if perps is None:
! 811                 continue
  812             perp1, perp2 = perps
  813 
  814             edge_lengths = (
  815                 np.linalg.norm(tri[1] - tri[0]),

---

  825             inside_mask = np.all(sample_points >= (sim_min - tol), axis=1) & np.all(
  826                 sample_points <= (sim_max + tol), axis=1
  827             )
  828             if not np.any(inside_mask):
! 829                 continue
  830 
  831             sample_points = sample_points[inside_mask]
  832             bary_inside = barycentric[inside_mask]
  833             n_samples_inside = sample_points.shape[0]

---

  846             faces_list.append(faces_tile)
  847             bary_list.append(bary_inside)
  848 
  849         if not points_list:
! 850             return {
  851                 "points": np.zeros((0, 3), dtype=dtype),
  852                 "normals": np.zeros((0, 3), dtype=dtype),
  853                 "perps1": np.zeros((0, 3), dtype=dtype),
  854                 "perps2": np.zeros((0, 3), dtype=dtype),

---

  875         edge02 = triangle[2] - triangle[0]
  876         cross = np.cross(edge01, edge02)
  877         norm = np.linalg.norm(cross)
  878         if norm <= 0.0:
! 879             return 0.0, np.zeros(3, dtype=triangle.dtype)
  880         normal = (cross / norm).astype(triangle.dtype, copy=False)
  881         area = 0.5 * norm
  882         return area, normal

---

  916     def _build_barycentric_samples(subdivisions: int) -> np.ndarray:
  917         """Construct barycentric sampling points for a given subdivision level."""
  918 
  919         if subdivisions <= 1:
! 920             return np.array([[1.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0]])
  921 
  922         bary = []
  923         for i in range(subdivisions):
  924             for j in range(subdivisions - i):

---

  944                 edge = (candidate / length).astype(triangle.dtype, copy=False)
  945                 break
  946 
  947         if edge is None:
! 948             return None
  949 
  950         perp1 = edge
  951         perp2 = np.cross(normal, perp1)
  952         perp2_norm = np.linalg.norm(perp2)

---

  950         perp1 = edge
  951         perp2 = np.cross(normal, perp1)
  952         perp2_norm = np.linalg.norm(perp2)
  953         if perp2_norm <= tol:
! 954             return None
  955         perp2 = (perp2 / perp2_norm).astype(triangle.dtype, copy=False)
  956         return perp1, perp2

---