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

CHANGELOG.md

@@ -38,6 +38,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `interp_spec` in `ModeSpec` to allow downsampling and interpolation of waveguide modes in frequency.
 - Added warning if port mesh refinement is incompatible with the `GridSpec` in the `TerminalComponentModeler`.
 - Various types, e.g. different `Simulation` or `SimulationData` sub-classes, can be loaded from file directly with `Tidy3dBaseModel.from_file()`.
+- 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/numerical/test_autograd_box_polyslab_numerical.py

@@ -41,6 +41,22 @@
     sys.stdout = sys.stderr
 
 
+def angled_overlap_deg(v1, v2):
+    norm_v1 = np.linalg.norm(v1)
+    norm_v2 = np.linalg.norm(v2)
+
+    if np.isclose(norm_v1, 0.0) or np.isclose(norm_v2, 0.0):
+        if not (np.isclose(norm_v1, 0.0) and np.isclose(norm_v2, 0.0)):
+            return np.inf
+
+        return 0.0
+
+    dot = np.minimum(1.0, np.sum((v1 / np.linalg.norm(v1)) * (v2 / np.linalg.norm(v2))))
+    angle_deg = np.arccos(dot) * 180.0 / np.pi
+
+    return angle_deg
+
+
 def case_identifier(is_3d: bool, infinite_dim_2d: int | None, shift_box_center: bool) -> str:
     geometry_tag = "3d" if is_3d else f"2d_infinite_dim_{infinite_dim_2d}"
     shift_tag = "shifted" if shift_box_center else "centered"
@@ -422,21 +438,6 @@ def test_box_and_polyslab_gradients_match(
         )
         np.savez(npz_path, **test_data)
 
-    def angled_overlap_deg(v1, v2):
-        norm_v1 = np.linalg.norm(v1)
-        norm_v2 = np.linalg.norm(v2)
-
-        if np.isclose(norm_v1, 0.0) or np.isclose(norm_v2, 0.0):
-            if not (np.isclose(norm_v1, 0.0) and np.isclose(norm_v2, 0.0)):
-                return np.inf
-
-            return 0.0
-
-        dot = np.minimum(1.0, np.sum((v1 / np.linalg.norm(v1)) * (v2 / np.linalg.norm(v2))))
-        angle_deg = np.arccos(dot) * 180.0 / np.pi
-
-        return angle_deg
-
     box_polyslab_overlap_deg = angled_overlap_deg(box_grad_filtered, polyslab_grad_filtered)
     fd_overlap_deg = angled_overlap_deg(fd_box, fd_polyslab)
     box_fd_adj_overlap_deg = angled_overlap_deg(box_grad_filtered, fd_box)

tests/test_components/autograd/numerical/test_autograd_polyslab_trianglemesh_numerical.py

@@ -0,0 +1,303 @@
+# Test autograd and compare to numerically computed finite difference gradients for
+# PolySlab and TriangleMesh geometries representing the same rectangular slab.
+from __future__ import annotations
+
+import sys
+
+import autograd.numpy as anp
+import numpy as np
+import pytest
+from autograd import value_and_grad
+
+import tidy3d as td
+import tidy3d.web as web
+from tests.test_components.autograd.numerical.test_autograd_box_polyslab_numerical import (
+    angled_overlap_deg,
+    dimension_permutation,
+    finite_difference,
+    make_base_simulation,
+    run_parameter_simulations,
+    squeeze_dimension,
+)
+from tidy3d import config
+
+config.local_cache.enabled = True
+WL_UM = 0.65
+FREQ0 = td.C_0 / WL_UM
+PERIODS_UM = (3 * WL_UM, 4 * WL_UM)
+INFINITE_DIM_SIZE_UM = 0.1
+SRC_OFFSET = -2.5
+MONITOR_OFFSET = 2.5
+PERMITTIVITY = 2.5**2
+MESH_SPACING_UM = WL_UM / 40.0
+FINITE_DIFFERENCE_STEP = MESH_SPACING_UM
+LOCAL_GRADIENT = True
+VERBOSE = False
+SHOW_PRINT_STATEMENTS = True
+PLOT_FD_ADJ_COMPARISON = False
+SAVE_OUTPUT_DATA = True
+COMPARE_TO_FINITE_DIFFERENCE = True
+COMPARE_TO_POLYSLAB = True
+
+ANGLE_OVERLAP_THRESH_DEG = 10.0
+ANGLE_OVERLAP_FD_ADJ_THRESH_DEG = 10.0
+
+VERTEX_SIGNS = np.array(
+    [
+        (-1.0, -1.0, -1.0),
+        (-1.0, -1.0, 1.0),
+        (-1.0, 1.0, -1.0),
+        (-1.0, 1.0, 1.0),
+        (1.0, -1.0, -1.0),
+        (1.0, -1.0, 1.0),
+        (1.0, 1.0, -1.0),
+        (1.0, 1.0, 1.0),
+    ]
+)
+
+TRIANGLE_FACE_VERTEX_IDS = np.array(
+    [
+        (1, 3, 0),
+        (4, 1, 0),
+        (0, 3, 2),
+        (2, 4, 0),
+        (1, 7, 3),
+        (5, 1, 4),
+        (5, 7, 1),
+        (3, 7, 2),
+        (6, 4, 2),
+        (2, 7, 6),
+        (6, 5, 4),
+        (7, 5, 6),
+    ],
+    dtype=int,
+)
+
+if PLOT_FD_ADJ_COMPARISON:
+    pytestmark = pytest.mark.usefixtures("mpl_config_interactive")
+else:
+    pytestmark = pytest.mark.usefixtures("mpl_config_noninteractive")
+
+if SHOW_PRINT_STATEMENTS:
+    sys.stdout = sys.stderr
+
+
+def _triangles_from_params(params, box_center, xp):
+    params_arr = xp.array(params)
+    center_arr = xp.array(box_center)
+    half_size = 0.5 * params_arr
+    vertices = center_arr + xp.array(VERTEX_SIGNS) * half_size
+    return vertices[xp.array(TRIANGLE_FACE_VERTEX_IDS)]
+
+
+def make_trianglemesh_geometry(params, box_center):
+    triangles = _triangles_from_params(params, box_center, np)
+    mesh = td.TriangleMesh.from_triangles(triangles)
+    return mesh
+
+
+def make_polyslab_geometry(params, box_center, axis: int) -> td.PolySlab:
+    half_size = 0.5 * params
+    slab_bounds = (
+        box_center[axis] - half_size[axis],
+        box_center[axis] + half_size[axis],
+    )
+    plane_axes = [idx for idx in range(3) if idx != axis]
+
+    vertices = []
+    for sign_0, sign_1 in ((-1, -1), (-1, 1), (1, 1), (1, -1)):
+        coord_0 = box_center[plane_axes[0]] + sign_0 * half_size[plane_axes[0]]
+        coord_1 = box_center[plane_axes[1]] + sign_1 * half_size[plane_axes[1]]
+        vertices.append((coord_0, coord_1))
+
+    return td.PolySlab(vertices=tuple(vertices), slab_bounds=slab_bounds, axis=axis)
+
+
+def run_parameter_simulations2(
+    parameter_sets: list[anp.ndarray],
+    make_geometry,
+    box_center,
+    base_sim: td.Simulation,
+    fom,
+    tmp_path,
+    *,
+    local_gradient: bool,
+):
+    simulations = []
+
+    for param_values in parameter_sets:
+        geometry = make_geometry(param_values, box_center)
+        structure = td.Structure(
+            geometry=geometry,
+            medium=td.Medium(permittivity=PERMITTIVITY),
+        )
+        sim = base_sim.updated_copy(structures=[structure], validate=False)
+        simulations.append(sim)
+
+    sim_data = web.run(
+        simulations,
+        local_gradient=local_gradient,
+        path=tmp_path,
+        verbose=VERBOSE,
+    )
+    sim_fom = [fom(data) for data in sim_data]
+    if len(sim_fom) == 1:
+        sim_fom = sim_fom[0]
+    return sim_fom
+
+
+def make_objective(
+    make_geometry,
+    box_center,
+    tag: str,
+    base_sim: td.Simulation,
+    fom,
+    tmp_path,
+    *,
+    local_gradient: bool,
+):
+    def objective(parameters):
+        results = run_parameter_simulations(
+            parameters,
+            make_geometry,
+            box_center,
+            tag,
+            base_sim,
+            fom,
+            tmp_path,
+            local_gradient=local_gradient,
+        )
+
+        return results
+
+    return objective
+
+
+# @pytest.mark.numerical
+@pytest.mark.parametrize(
+    "is_3d, infinite_dim_2d",
+    [
+        (True, 2),
+        (False, 0),
+        (False, 1),
+        (False, 2),
+    ],
+)
+@pytest.mark.parametrize("shift_box_center", (True, False))
+def test_polyslab_and_trianglemesh_gradients_match(
+    is_3d, infinite_dim_2d, shift_box_center, tmp_path
+):
+    """Test that the triangle mesh and polyslab gradients match for rectangular slab geometries. Allow
+    comparison as well to finite difference values."""
+
+    base_sim, fom = make_base_simulation(is_3d, infinite_dim_2d if not is_3d else None)
+
+    if shift_box_center:
+        slab_init_size = [2.0 * WL_UM, 2.5 * WL_UM, 0.75 * WL_UM]
+    else:
+        slab_init_size = [1.0 * WL_UM, 1.25 * WL_UM, 0.75 * WL_UM]
+
+    initial_params = anp.array(slab_init_size)
+
+    polyslab_axis = 2 if is_3d else infinite_dim_2d
+
+    box_center = [0.0, 0.0, 0.0]
+    if shift_box_center:
+        # test what happens when part of the structure falls outside the simulation domain
+        # but don't shift along source axis
+        if is_3d:
+            box_center[0:2] = [0.5 * p for p in PERIODS_UM]
+        else:
+            _, final_dim_2d = dimension_permutation(infinite_dim_2d)
+            box_center[infinite_dim_2d] = 0.5 * INFINITE_DIM_SIZE_UM
+            box_center[final_dim_2d] = 0.5 * PERIODS_UM[0]
+
+    triangle_objective = make_objective(
+        make_trianglemesh_geometry,
+        box_center,
+        "trianglemesh",
+        base_sim,
+        fom,
+        tmp_path,
+        local_gradient=LOCAL_GRADIENT,
+    )
+
+    polyslab_objective = make_objective(
+        lambda p, box_center: make_polyslab_geometry(p, box_center, polyslab_axis),
+        box_center,
+        "polyslab",
+        base_sim,
+        fom,
+        tmp_path,
+        local_gradient=LOCAL_GRADIENT,
+    )
+
+    triangle_objective_fd = make_objective(
+        make_trianglemesh_geometry,
+        box_center,
+        "trianglemesh_fd",
+        base_sim,
+        fom,
+        tmp_path,
+        local_gradient=False,
+    )
+
+    _triangle_value, triangle_grad = value_and_grad(triangle_objective)([initial_params])
+    assert triangle_grad is not None
+    if is_3d or infinite_dim_2d not in [1, 2]:
+        grad_norm_triangle = np.linalg.norm(triangle_grad)
+        assert grad_norm_triangle > 1e-6, (
+            f"Assumed norm to be bigger than 1e-6, got {grad_norm_triangle}"
+        )
+    triangle_grad_filtered = squeeze_dimension(triangle_grad, is_3d, infinite_dim_2d)
+    polyslab_grad_filtered = None
+    if COMPARE_TO_POLYSLAB:
+        _polyslab_value, polyslab_grad = value_and_grad(polyslab_objective)([initial_params])
+        polyslab_grad_filtered = squeeze_dimension(polyslab_grad, is_3d, infinite_dim_2d)
+        print(
+            "polyslab_grad_filtered\t",
+            polyslab_grad_filtered.tolist()
+            if not isinstance(polyslab_grad_filtered, list)
+            else polyslab_grad_filtered,
+        )
+
+    fd_triangle = None
+    if COMPARE_TO_FINITE_DIFFERENCE:
+        fd_triangle = squeeze_dimension(
+            finite_difference(triangle_objective_fd, initial_params, is_3d, infinite_dim_2d),
+            is_3d,
+            infinite_dim_2d,
+        )
+
+    if SAVE_OUTPUT_DATA:
+        test_data = {
+            "fd trianglemesh": fd_triangle,
+            "grad trianglemesh": triangle_grad_filtered,
+            "grad polyslab": polyslab_grad_filtered,
+        }
+        np.savez(
+            f"test_diff_triangle_poly_{'3' if is_3d else '2'}d_infinite_dim_{infinite_dim_2d}.npz",
+            **test_data,
+        )
+
+    if COMPARE_TO_POLYSLAB:
+        triangle_polyslab_overlap_deg = angled_overlap_deg(
+            triangle_grad_filtered, polyslab_grad_filtered
+        )
+        print(f"TriangleMesh FD vs. polyslab overlap: {triangle_polyslab_overlap_deg:.3f}° ")
+        assert triangle_polyslab_overlap_deg < ANGLE_OVERLAP_THRESH_DEG, (
+            f"[TriangleMesh vs. PolySlab] Autograd gradients disagree: "
+            f"angle overlap = {triangle_polyslab_overlap_deg:.3f}° "
+            f"(threshold = {ANGLE_OVERLAP_THRESH_DEG:.3f}°, "
+            f"difference = {triangle_polyslab_overlap_deg - ANGLE_OVERLAP_THRESH_DEG:+.3f}°)"
+        )
+
+    if COMPARE_TO_FINITE_DIFFERENCE:
+        triangle_fd_adj_overlap_deg = angled_overlap_deg(triangle_grad_filtered, fd_triangle)
+        print(f"TriangleMesh FD vs. Adjoint angle overlap: {triangle_fd_adj_overlap_deg:.3f}° ")
+        assert triangle_fd_adj_overlap_deg < ANGLE_OVERLAP_FD_ADJ_THRESH_DEG, (
+            f"Autograd and finite-difference gradients disagree: "
+            f"angle overlap = {triangle_fd_adj_overlap_deg:.3f}° "
+            f"(threshold = {ANGLE_OVERLAP_FD_ADJ_THRESH_DEG:.3f}°, "
+            f"difference = {triangle_fd_adj_overlap_deg - ANGLE_OVERLAP_FD_ADJ_THRESH_DEG:+.3f}°)"
+        )