wip triangle mesh

Author: marcorudolphflex

tests/test_components/autograd/numerical/test_autograd_box_polyslab_numerical.py

@@ -11,6 +11,9 @@
 
 import tidy3d as td
 import tidy3d.web as web
+from tidy3d import config
+
+config.local_cache.enabled = True
 
 WL_UM = 0.65
 FREQ0 = td.C_0 / WL_UM
@@ -293,17 +296,17 @@ def squeeze_dimension(array: np.ndarray, is_3d: bool, infinite_dim: int | None)
     return np.delete(squeezed, infinite_dim)
 
 
-@pytest.mark.numerical
+# @pytest.mark.numerical
 @pytest.mark.parametrize(
     "is_3d, infinite_dim_2d",
     [
         (True, 2),
-        (False, 0),
-        (False, 1),
-        (False, 2),
+        # (False, 0),
+        # (False, 1),
+        # (False, 2),
     ],
 )
-@pytest.mark.parametrize("shift_box_center", (True, False))
+@pytest.mark.parametrize("shift_box_center", (True,))
 def test_box_and_polyslab_gradients_match(is_3d, infinite_dim_2d, shift_box_center, tmp_path):
     """Test that the box and polyslab gradients match for rectangular slab geometries. Allow
     comparison as well to finite difference values."""
@@ -428,16 +431,25 @@ def angled_overlap_deg(v1, v2):
     print(f"PolySlab Finite Difference vs. Adjoint: {polyslab_fd_adj_overlap_deg}")
 
     assert box_polyslab_overlap_deg < ANGLE_OVERLAP_THRESH_DEG, (
-        "Autograd gradients for Box and PolySlab disagree"
+        "Autograd gradients for Box and PolySlab disagree: "
+        f"angle = {box_polyslab_overlap_deg:.2f} deg, "
+        f"threshold = {ANGLE_OVERLAP_THRESH_DEG:.2f} deg"
     )
+
     assert fd_overlap_deg < ANGLE_OVERLAP_THRESH_DEG, (
-        "Finite-difference gradients for Box and PolySlab disagree"
+        "Finite-difference gradients for Box and PolySlab disagree: "
+        f"angle = {fd_overlap_deg:.2f} deg, "
+        f"threshold = {ANGLE_OVERLAP_THRESH_DEG:.2f} deg"
     )
 
     if COMPARE_TO_FINITE_DIFFERENCE:
         assert box_fd_adj_overlap_deg < ANGLE_OVERLAP_FD_ADJ_THRESH_DEG, (
-            "Autograd and finite-difference gradients for the Box geometry disagree"
+            "Autograd and finite-difference gradients for the Box geometry disagree: "
+            f"angle = {box_fd_adj_overlap_deg:.2f} deg, "
+            f"threshold = {ANGLE_OVERLAP_FD_ADJ_THRESH_DEG:.2f} deg"
         )
         assert polyslab_fd_adj_overlap_deg < ANGLE_OVERLAP_FD_ADJ_THRESH_DEG, (
-            "Autograd and finite-difference gradients for the PolySlab geometry disagree"
+            "Autograd and finite-difference gradients for the PolySlab geometry disagree: "
+            f"angle = {polyslab_fd_adj_overlap_deg:.2f} deg, "
+            f"threshold = {ANGLE_OVERLAP_FD_ADJ_THRESH_DEG:.2f} deg"
         )

tests/test_components/autograd/numerical/test_autograd_polyslab_trianglemesh_numerical.py

@@ -8,11 +8,13 @@
 import numpy as np
 import pytest
 from autograd import value_and_grad
+from autograd.builtins import dict as dict_ag
+from autograd.tracer import isbox
 
 import tidy3d as td
 import tidy3d.web as web
 from tidy3d import config
-from tidy3d.components.autograd import get_static
+from tidy3d.components.autograd import TidyArrayBox, get_static, is_tidy_box
 
 config.local_cache.enabled = True
 WL_UM = 0.65
@@ -205,15 +207,21 @@ def _triangles_from_params(params, box_center, xp):
     return vertices[xp.array(TRIANGLE_FACE_VERTEX_IDS)]
 
 
-def make_trianglemesh_geometry(params, box_center) -> td.TriangleMesh:
+def make_trianglemesh_geometry(params, box_center):
     params = anp.array(params)
     params_static = np.array(get_static(params))
     triangles_static = _triangles_from_params(params_static, box_center, np)
     mesh = td.TriangleMesh.from_triangles(triangles_static)
     triangles_traced = _triangles_from_params(params, box_center, anp)
-    surface_mesh = mesh.mesh_dataset.surface_mesh.copy(deep=False, data=triangles_traced)
-    mesh_dataset = mesh.mesh_dataset.copy(update={"surface_mesh": surface_mesh}, deep=True)
-    return mesh.updated_copy(validate=False, mesh_dataset=mesh_dataset)
+    traced_fields = None
+    if isbox(triangles_traced):
+        traced_data = triangles_traced
+        if not is_tidy_box(traced_data):
+            traced_data = TidyArrayBox.from_arraybox(traced_data)
+        traced_fields = dict_ag(
+            {("structures", 0, "geometry", "mesh_dataset", "surface_mesh"): traced_data}
+        )
+    return mesh, traced_fields
 
 
 def make_polyslab_geometry(params, box_center, axis: int) -> td.PolySlab:
@@ -247,13 +255,20 @@ def run_parameter_simulations(
     simulation_dict = {}
 
     for idx, param_values in enumerate(parameter_sets):
-        geometry = make_geometry(param_values, box_center)
+        geometry_data = make_geometry(param_values, box_center)
+        if isinstance(geometry_data, tuple):
+            geometry, traced_fields = geometry_data
+        else:
+            geometry = geometry_data
+            traced_fields = None
         structure = td.Structure(
             geometry=geometry,
             medium=td.Medium(permittivity=PERMITTIVITY),
         )
 
-        sim = base_sim.updated_copy(structures=[structure])
+        sim = base_sim.updated_copy(structures=[structure], validate=False)
+        if traced_fields is not None:
+            sim = sim._insert_traced_fields(traced_fields)
         simulation_dict[f"sim_{idx}"] = sim
 
     if len(simulation_dict) == 1:

tidy3d/components/geometry/mesh.py

@@ -8,7 +8,7 @@
 import numpy as np
 import pydantic.v1 as pydantic
 
-from tidy3d.components.autograd import AutogradFieldMap
+from tidy3d.components.autograd import AutogradFieldMap, get_static
 from tidy3d.components.autograd.derivative_utils import DerivativeInfo
 from tidy3d.components.base import cached_property
 from tidy3d.components.data.data_array import DATA_ARRAY_MAP, TriangleMeshDataArray
@@ -75,7 +75,9 @@ def _check_mesh(cls, val: TriangleMeshDataset) -> TriangleMeshDataset:
 
         import trimesh
 
-        mesh = cls._triangles_to_trimesh(val.surface_mesh)
+        surface_mesh = val.surface_mesh
+        triangles = get_static(surface_mesh.data)
+        mesh = cls._triangles_to_trimesh(triangles)
         if not all(np.array(mesh.area_faces) > AREA_SIZE_THRESHOLD):
             old_tol = trimesh.tol.merge
             trimesh.tol.merge = np.sqrt(2 * AREA_SIZE_THRESHOLD)