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

Author: marcorudolphflex

CHANGELOG.md

@@ -31,6 +31,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added multimode support to `WavePort` in the smatrix plugin, allowing multiple modes to be analyzed per port.
 - Added support for `.lydrc` files for design rule checking in the `klayout` plugin.
 - Added a Gaussian inverse design filter option with autograd gradients and complete padding mode coverage.
+- Added support of `TriangleMesh` for autograd.
 
 ### Breaking Changes
 - Edge singularity correction at PEC and lossy metal edges defaults to `True`.

docs/api/geometry.rst

@@ -295,5 +295,6 @@ Use the ``from_stl()`` class method to import from an external STL file, or ``fr
    + `Importing STL files <../notebooks/STLImport.html>`_
    + `Defining complex geometries using trimesh <../notebooks/CreatingGeometryUsingTrimesh.html>`_
 
-~~~~
+Shape gradients for ``TriangleMesh`` geometries are supported through the autograd workflow. When a mesh participates in an adjoint optimization, boundary sensitivities are evaluated on the triangle faces. The cost of the surface integral scales with the number of mesh faces; very fine meshes may require additional sampling to converge gradients, so consider simplifying or coarsening meshes when possible, or adjusting the autograd configuration to trade off accuracy and runtime.
 
+~~~~

tests/test_components/autograd/test_autograd.py

@@ -527,6 +527,31 @@ def make_structures(params: anp.ndarray) -> dict[str, td.Structure]:
     )
     cylinder = td.Structure(geometry=cylinder_geo, medium=polyslab.medium)
 
+    # triangle mesh geometry with param-dependent medium response
+    base_vertices = np.array(
+        [
+            (0.0, 0.0, 0.0),
+            (0.6, 0.0, 0.0),
+            (0.0, 0.6, 0.0),
+            (0.0, 0.0, 0.6),
+        ],
+    )
+    faces = np.array(
+        [
+            (0, 2, 1),
+            (0, 1, 3),
+            (0, 3, 2),
+            (1, 2, 3),
+        ],
+        dtype=int,
+    )
+    triangle_mesh_geo = td.TriangleMesh.from_vertices_faces(base_vertices, faces)
+    mesh_eps = 1.8 + 0.2 * anp.abs(vector @ params)
+    triangle_mesh = td.Structure(
+        geometry=triangle_mesh_geo,
+        medium=td.Medium(permittivity=mesh_eps),
+    )
+
     return {
         "medium": medium,
         "center_list": center_list,
@@ -541,6 +566,7 @@ def make_structures(params: anp.ndarray) -> dict[str, td.Structure]:
         "pole_res": pole_res,
         "custom_pole_res": custom_pole_res,
         "cylinder": cylinder,
+        "triangle_mesh": triangle_mesh,
     }
 
 
@@ -638,6 +664,7 @@ def plot_sim(sim: td.Simulation, plot_eps: bool = True) -> None:
     "pole_res",
     "custom_pole_res",
     "cylinder",
+    "triangle_mesh",
 )
 monitor_keys_ = ("mode", "diff", "field_vol", "field_point")
 

tests/test_components/autograd/test_autograd_triangle_mesh.py

@@ -0,0 +1,315 @@
+"""Tests for TriangleMesh autograd derivatives."""
+
+from __future__ import annotations
+
+import numpy as np
+import numpy.testing as npt
+import pytest
+
+import tidy3d as td
+
+VERTICES_TETRA = np.array(
+    [
+        (0.0, 0.0, 0.0),
+        (1.0, 0.0, 0.0),
+        (0.0, 1.0, 0.0),
+        (0.0, 0.0, 1.0),
+    ],
+    dtype=float,
+)
+
+FACES_TETRA = np.array(
+    [
+        (0, 2, 1),
+        (0, 1, 3),
+        (0, 3, 2),
+        (1, 2, 3),
+    ]
+)
+
+VERTICES_OCTA = np.array(
+    [
+        (1.0, 0.0, 0.0),
+        (-1.0, 0.0, 0.0),
+        (0.0, 1.0, 0.0),
+        (0.0, -1.0, 0.0),
+        (0.0, 0.0, 1.0),
+        (0.0, 0.0, -1.0),
+    ],
+    dtype=float,
+)
+
+FACES_OCTA = np.array(
+    [
+        (4, 0, 2),
+        (4, 2, 1),
+        (4, 1, 3),
+        (4, 3, 0),
+        (5, 2, 0),
+        (5, 1, 2),
+        (5, 3, 1),
+        (5, 0, 3),
+    ],
+    dtype=int,
+)
+
+VERTICES_SLENDER_TETRA = np.array(
+    [
+        (0.0, 0.0, 0.0),
+        (1.0, 0.0, 0.0),
+        (1.0e-4, 1.0e-4, 0.0),
+        (0.0, 0.0, 1.0),
+    ],
+    dtype=float,
+)
+
+FACES_SLENDER_TETRA = FACES_TETRA
+
+VERTICES_NON_WATERTIGHT = np.array(
+    [
+        (0.0, 0.0, 0.0),
+        (1.0, 0.0, 0.0),
+        (0.0, 1.0, 0.0),
+        (0.0, 0.0, 1.0),
+    ],
+    dtype=float,
+)
+
+FACES_NON_WATERTIGHT = np.array(
+    [
+        (0, 1, 2),
+        (0, 2, 3),
+    ],
+    dtype=int,
+)
+
+MESH_DEFINITIONS: dict[str, tuple[np.ndarray, np.ndarray]] = {
+    "tetrahedron": (VERTICES_TETRA, FACES_TETRA),
+    "octahedron": (VERTICES_OCTA, FACES_OCTA),
+    "slender_tetrahedron": (VERTICES_SLENDER_TETRA, FACES_SLENDER_TETRA),
+}
+
+
+@pytest.fixture(params=list(MESH_DEFINITIONS.keys()), ids=list(MESH_DEFINITIONS.keys()))
+def watertight_mesh(request) -> td.TriangleMesh:
+    """Parameterized fixture returning watertight meshes of varying complexity."""
+
+    vertices, faces = MESH_DEFINITIONS[request.param]
+    return td.TriangleMesh.from_vertices_faces(vertices, faces)
+
+
+@pytest.fixture
+def non_watertight_mesh() -> td.TriangleMesh:
+    """Simple non-watertight surface used to validate graceful handling."""
+
+    return td.TriangleMesh.from_vertices_faces(VERTICES_NON_WATERTIGHT, FACES_NON_WATERTIGHT)
+
+
+class DummyDerivativeInfo:
+    """Lightweight stand-in for ``DerivativeInfo`` used in unit tests."""
+
+    def __init__(
+        self,
+        grad_func,
+        spacing: float = 0.2,
+        simulation_bounds: tuple[tuple[float, float, float], tuple[float, float, float]]
+        | None = None,
+    ) -> None:
+        self.paths = [("mesh_dataset", "surface_mesh")]
+        self.frequencies = [200e12]
+        self.eps_in = 12.0
+        self.interpolators = {}
+        default_bounds = ((-10.0, -10.0, -10.0), (10.0, 10.0, 10.0))
+        self.simulation_bounds = simulation_bounds or default_bounds
+        self._grad_func = grad_func
+        self._spacing = spacing
+        self.bounds_intersect = self.simulation_bounds
+
+    def adaptive_vjp_spacing(self) -> float:
+        return self._spacing
+
+    def create_interpolators(self, dtype=None):
+        return {}
+
+    def evaluate_gradient_at_points(
+        self, spatial_coords, normals, perps1, perps2, interpolators=None
+    ):
+        return self._grad_func(spatial_coords)
+
+
+def area_and_normal(triangle: np.ndarray) -> tuple[float, np.ndarray]:
+    """Return signed area and unit normal for a triangle."""
+
+    edge01 = triangle[1] - triangle[0]
+    edge02 = triangle[2] - triangle[0]
+    cross = np.cross(edge01, edge02)
+    norm = np.linalg.norm(cross)
+    if np.isclose(norm, 0.0):
+        return 0.0, np.zeros(3, dtype=triangle.dtype)
+    return 0.5 * norm, cross / norm
+
+
+def linear_grad_func_factory(coeffs: np.ndarray, offset: float):
+    """Create a linear function g(x) = coeffs.x + offset."""
+
+    def grad_func(points: np.ndarray) -> np.ndarray:
+        return points @ coeffs + offset
+
+    return grad_func
+
+
+def test_triangle_mesh_gradient_linear_matches_analytic(watertight_mesh):
+    """Validate per-vertex gradients against analytic integrals for linear g."""
+
+    mesh = watertight_mesh
+    coeffs = np.array([0.6, -0.25, 0.4], dtype=float)
+    offset = -0.15
+    grad_func = linear_grad_func_factory(coeffs, offset)
+
+    spacing = 0.01 if mesh.triangles.shape[0] <= 4 else 0.005
+    derivative_info = DummyDerivativeInfo(grad_func, spacing=spacing)
+    grads = mesh._compute_derivatives(derivative_info)[("mesh_dataset", "surface_mesh")]
+
+    expected = np.zeros_like(grads)
+    for face_idx, tri in enumerate(mesh.triangles):
+        area, normal = area_and_normal(tri)
+        if np.isclose(area, 0.0):
+            continue
+
+        g_vals = grad_func(tri)
+        for local_idx in range(3):
+            others = [(local_idx + 1) % 3, (local_idx + 2) % 3]
+            gi = g_vals[local_idx]
+            gj = g_vals[others[0]]
+            gk = g_vals[others[1]]
+            integral = area / 12.0 * (2.0 * gi + gj + gk)
+            expected[face_idx, local_idx, :] = integral * normal
+
+    # surface integration uses adaptive sampling so allow a few-percent mismatch
+    npt.assert_allclose(grads, expected, rtol=8e-2, atol=1e-6)
+
+
+def test_triangle_mesh_gradient_directional_derivative_matches_quadrature(watertight_mesh):
+    """Directional derivative from gradients matches exact integral."""
+
+    mesh = watertight_mesh
+    coeffs = np.array([0.3, -0.45, 0.55], dtype=float)
+    offset = 0.2
+    grad_func = linear_grad_func_factory(coeffs, offset)
+
+    spacing = 0.01 if mesh.triangles.shape[0] <= 4 else 0.005
+    derivative_info = DummyDerivativeInfo(grad_func, spacing=spacing)
+    grads = mesh._compute_derivatives(derivative_info)[("mesh_dataset", "surface_mesh")]
+
+    rng = np.random.default_rng(1234)
+    delta = rng.normal(scale=5e-3, size=grads.shape)
+
+    total_pred = float(np.sum(grads * delta))
+
+    total_exact = 0.0
+    weight_matrix = np.full((3, 3), 1.0 / 12.0)
+    np.fill_diagonal(weight_matrix, 1.0 / 6.0)
+
+    for face_idx, tri in enumerate(mesh.triangles):
+        area, normal = area_and_normal(tri)
+        if np.isclose(area, 0.0):
+            continue
+
+        g_vals = grad_func(tri)
+        dot_vals = delta[face_idx] @ normal
+        total_exact += area * dot_vals @ weight_matrix @ g_vals
+
+    npt.assert_allclose(total_pred, total_exact, rtol=1e-3, atol=1e-6)
+
+
+def test_triangle_mesh_gradient_constant_field_integrates_to_zero(watertight_mesh):
+    """Constant surface gradient should integrate to zero net force on a watertight mesh."""
+
+    mesh = watertight_mesh
+
+    def constant_grad(points: np.ndarray) -> np.ndarray:
+        return np.ones(points.shape[0], dtype=float)
+
+    derivative_info = DummyDerivativeInfo(constant_grad, spacing=0.01)
+    grads = mesh._compute_derivatives(derivative_info)[("mesh_dataset", "surface_mesh")]
+
+    net_force = np.sum(grads, axis=(0, 1))
+    npt.assert_allclose(net_force, np.zeros(3), atol=5e-6, rtol=1e-3)
+
+
+def test_triangle_mesh_gradient_face_permutation_invariant(watertight_mesh):
+    """Reordering faces does not change the per-face gradients after reindexing."""
+
+    base_mesh = watertight_mesh
+    perm = np.random.default_rng(42).permutation(base_mesh.triangles.shape[0])
+    permuted_mesh = td.TriangleMesh.from_triangles(base_mesh.triangles[perm])
+
+    coeffs = np.array([0.2, 0.1, -0.3], dtype=float)
+    offset = 0.05
+    grad_func = linear_grad_func_factory(coeffs, offset)
+
+    derivative_info = DummyDerivativeInfo(grad_func, spacing=0.01)
+    grad_base = base_mesh._compute_derivatives(derivative_info)[("mesh_dataset", "surface_mesh")]
+    grad_perm = permuted_mesh._compute_derivatives(derivative_info)[
+        ("mesh_dataset", "surface_mesh")
+    ]
+
+    inv_perm = np.argsort(perm)
+    grad_perm_reordered = grad_perm[inv_perm]
+
+    npt.assert_allclose(grad_base, grad_perm_reordered, rtol=1e-3, atol=1e-6)
+
+
+def test_triangle_mesh_gradient_zero_when_outside_bounds(watertight_mesh):
+    """Gradients vanish when the mesh lies entirely outside the simulation bounds."""
+
+    mesh = watertight_mesh
+
+    def constant_grad(points: np.ndarray) -> np.ndarray:
+        return np.ones(points.shape[0], dtype=float)
+
+    far_bounds = ((100.0, 100.0, 100.0), (101.0, 101.0, 101.0))
+    derivative_info = DummyDerivativeInfo(
+        constant_grad,
+        spacing=0.05,
+        simulation_bounds=far_bounds,
+    )
+
+    grads = mesh._compute_derivatives(derivative_info)[("mesh_dataset", "surface_mesh")]
+    npt.assert_allclose(grads, np.zeros_like(grads))
+
+
+def test_triangle_mesh_non_watertight_warns_and_computes(non_watertight_mesh, caplog):
+    """Non-watertight meshes should warn but still return finite gradients."""
+
+    def constant_grad(points: np.ndarray) -> np.ndarray:
+        return np.ones(points.shape[0], dtype=float)
+
+    derivative_info = DummyDerivativeInfo(constant_grad, spacing=0.05)
+
+    with caplog.at_level("WARNING"):
+        grads = non_watertight_mesh._compute_derivatives(derivative_info)[
+            ("mesh_dataset", "surface_mesh")
+        ]
+
+    assert grads.shape == non_watertight_mesh.triangles.shape
+    assert np.all(np.isfinite(grads))
+    assert not non_watertight_mesh.trimesh.is_watertight
+
+
+def test_triangle_mesh_subdivision_respects_long_edges():
+    """Subdivision count should scale with the longest edge length."""
+
+    tri = VERTICES_SLENDER_TETRA[FACES_SLENDER_TETRA[0]]
+    area, _ = area_and_normal(tri)
+    edge_lengths = (
+        np.linalg.norm(tri[1] - tri[0]),
+        np.linalg.norm(tri[2] - tri[1]),
+        np.linalg.norm(tri[0] - tri[2]),
+    )
+
+    spacing = 0.05
+    subdivisions = td.TriangleMesh._subdivision_count(area, spacing, edge_lengths)
+    expected_min = int(np.ceil(max(edge_lengths) / spacing))
+
+    assert subdivisions >= expected_min