diff --git a/CHANGELOG.md b/CHANGELOG.md index e9e7cd2363..efe5d4635c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,6 +37,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Introduced `BroadbandPulse` for exciting simulations across a wide frequency spectrum. - 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`. +- Added support of `TriangleMesh` for autograd. ### Breaking Changes - Edge singularity correction at PEC and lossy metal edges defaults to `True`. diff --git a/docs/api/geometry.rst b/docs/api/geometry.rst index 2032661f50..99465c7052 100644 --- a/docs/api/geometry.rst +++ b/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. +~~~~ diff --git a/tests/test_components/autograd/numerical/test_autograd_box_polyslab_numerical.py b/tests/test_components/autograd/numerical/test_autograd_box_polyslab_numerical.py index fedb04d589..79db98708d 100644 --- a/tests/test_components/autograd/numerical/test_autograd_box_polyslab_numerical.py +++ b/tests/test_components/autograd/numerical/test_autograd_box_polyslab_numerical.py @@ -12,6 +12,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 @@ -41,6 +44,29 @@ sys.stdout = sys.stderr +def angled_overlap_deg(v1, v2): + # zero-out tiny vectors elementwise + thresh = 1e-6 + v1 = np.where(np.abs(v1) < thresh, 0.0, v1) + v2 = np.where(np.abs(v2) < thresh, 0.0, v2) + + norm_v1 = np.linalg.norm(v1) + norm_v2 = np.linalg.norm(v2) + + # handle degenerate vectors + if norm_v1 < thresh or norm_v2 < thresh: + # one is effectively zero, the other is not → undefined angle + if not (norm_v1 < thresh and norm_v2 < thresh): + 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" @@ -213,12 +239,12 @@ def run_parameter_simulations( medium=td.Medium(permittivity=PERMITTIVITY), ) - sim = base_sim.updated_copy(structures=[structure]) + sim = base_sim.updated_copy(structures=[structure], validate=False) simulation_dict[f"sim_{idx}"] = sim if len(simulation_dict) == 1: key, sim = next(iter(simulation_dict.items())) - result_path = output_dir / f"{key}.hdf5" + result_path = output_dir / f"{sim._hash_self()}.hdf5" sim_data = web.run( sim, task_name=key, @@ -422,21 +448,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) @@ -448,17 +459,14 @@ def angled_overlap_deg(v1, v2): print(f"Box Finite Difference vs. Adjoint: {box_fd_adj_overlap_deg}") 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" - ) - assert fd_overlap_deg < ANGLE_OVERLAP_THRESH_DEG, ( - "Finite-difference gradients for Box and PolySlab disagree" - ) - 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" ) diff --git a/tests/test_components/autograd/numerical/test_autograd_polyslab_trianglemesh_numerical.py b/tests/test_components/autograd/numerical/test_autograd_polyslab_trianglemesh_numerical.py new file mode 100644 index 0000000000..c4ad511fa0 --- /dev/null +++ b/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}°)" + ) diff --git a/tests/test_components/autograd/numerical/test_autograd_sphere_trianglemesh_numerical.py b/tests/test_components/autograd/numerical/test_autograd_sphere_trianglemesh_numerical.py new file mode 100644 index 0000000000..71ce3db3f7 --- /dev/null +++ b/tests/test_components/autograd/numerical/test_autograd_sphere_trianglemesh_numerical.py @@ -0,0 +1,271 @@ +# Test autograd gradients for stretched spheres represented as both native Sphere and TriangleMesh +# geometries, comparing to finite differences for validation. +from __future__ import annotations + +import sys +from collections.abc import Sequence +from typing import Callable + +import autograd.numpy as anp +import numpy as np +import pytest +from autograd import value_and_grad + +import tidy3d as td +from tests.test_components.autograd.numerical.test_autograd_box_polyslab_numerical import ( + ANGLE_OVERLAP_FD_ADJ_THRESH_DEG, + COMPARE_TO_FINITE_DIFFERENCE, + LOCAL_GRADIENT, + MESH_SPACING_UM, + PERIODS_UM, + SAVE_OUTPUT_DATA, + SHOW_PRINT_STATEMENTS, + WL_UM, + angled_overlap_deg, + make_base_simulation, + run_parameter_simulations, +) +from tidy3d import config + +config.local_cache.enabled = True + +if SHOW_PRINT_STATEMENTS: + sys.stdout = sys.stderr + +SPHERE_RADIUS_UM = 0.5 * WL_UM +TRIANGLE_EDGE_TARGET = WL_UM / 25.0 +# For now run single stretch factor to keep runtime manageable. +# STRETCH_FACTORS = (0.1, 1.0, 10.0) +STRETCH_FACTORS = (1.0,) +TRANSLATION_AXES = (0, 1) + + +def _unit_sphere_triangles(target_edge_unit: float) -> np.ndarray: + """Generate a triangulated unit sphere once at import time.""" + + def base_icosahedron() -> tuple[np.ndarray, np.ndarray]: + phi = (1.0 + np.sqrt(5.0)) / 2.0 + vertices = np.array( + [ + (-1, phi, 0), + (1, phi, 0), + (-1, -phi, 0), + (1, -phi, 0), + (0, -1, phi), + (0, 1, phi), + (0, -1, -phi), + (0, 1, -phi), + (phi, 0, -1), + (phi, 0, 1), + (-phi, 0, -1), + (-phi, 0, 1), + ], + dtype=float, + ) + vertices /= np.linalg.norm(vertices, axis=1)[:, None] + faces = np.array( + [ + (0, 11, 5), + (0, 5, 1), + (0, 1, 7), + (0, 7, 10), + (0, 10, 11), + (1, 5, 9), + (5, 11, 4), + (11, 10, 2), + (10, 7, 6), + (7, 1, 8), + (3, 9, 4), + (3, 4, 2), + (3, 2, 6), + (3, 6, 8), + (3, 8, 9), + (4, 9, 5), + (2, 4, 11), + (6, 2, 10), + (8, 6, 7), + (9, 8, 1), + ], + dtype=np.int32, + ) + return vertices, faces + + def max_edge_length(vertices: np.ndarray, faces: np.ndarray) -> float: + edges = set() + for tri in faces: + edges.add(tuple(sorted((tri[0], tri[1])))) + edges.add(tuple(sorted((tri[1], tri[2])))) + edges.add(tuple(sorted((tri[2], tri[0])))) + return max(np.linalg.norm(vertices[i] - vertices[j]) for (i, j) in edges) + + def subdivide(vertices: np.ndarray, faces: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + midpoint_cache: dict[tuple[int, int], int] = {} + verts_list = [v.copy() for v in vertices] + + def midpoint(i: int, j: int) -> int: + key = (i, j) if i < j else (j, i) + if key in midpoint_cache: + return midpoint_cache[key] + v = (vertices[i] + vertices[j]) / 2.0 + v /= np.linalg.norm(v) + verts_list.append(v) + idx = len(verts_list) - 1 + midpoint_cache[key] = idx + return idx + + new_faces: list[tuple[int, int, int]] = [] + for tri in faces: + a = midpoint(tri[0], tri[1]) + b = midpoint(tri[1], tri[2]) + c = midpoint(tri[2], tri[0]) + new_faces.extend( + ( + (tri[0], a, c), + (tri[1], b, a), + (tri[2], c, b), + (a, b, c), + ) + ) + return np.array(verts_list, dtype=float), np.array(new_faces, dtype=np.int32) + + vertices, faces = base_icosahedron() + max_subdivisions = 5 + subdiv_count = 0 + while max_edge_length(vertices, faces) > target_edge_unit and subdiv_count < max_subdivisions: + vertices, faces = subdivide(vertices, faces) + subdiv_count += 1 + + return vertices[faces] + + +UNIT_SPHERE_TRIANGLES = _unit_sphere_triangles( + target_edge_unit=TRIANGLE_EDGE_TARGET / SPHERE_RADIUS_UM +) + + +def make_sphere_triangle_geometry( + params: anp.ndarray, center: Sequence[float], stretch_factor: float, translation_axis: int +) -> td.Geometry: + radius = params[0] + translation = params[1] + + triangles = anp.array(UNIT_SPHERE_TRIANGLES) + triangles = triangles * radius + stretch_vec = anp.array([stretch_factor, 1.0, 1.0]) + triangles = triangles * stretch_vec + center_arr = anp.array(center, dtype=float) + center_arr = center_arr + translation * np.eye(3)[translation_axis] + triangles = triangles + center_arr + return td.TriangleMesh.from_triangles(triangles) + + +def finite_difference_params(objective, params: anp.ndarray) -> np.ndarray: + step = anp.array([MESH_SPACING_UM, MESH_SPACING_UM]) + perturbations = [] + valid_indices = [] + + for idx in range(params.size): + params_up = anp.array(params) + params_down = anp.array(params) + params_up = params_up.copy() + params_down = params_down.copy() + params_up[idx] += step[idx] + params_down[idx] -= step[idx] + perturbations.extend([params_up, params_down]) + valid_indices.append(idx) + + objectives = objective(anp.stack(perturbations)) + objectives = np.asarray(objectives, dtype=float) + fd = np.zeros_like(np.asarray(params, dtype=float)) + for pair_idx, param_idx in enumerate(valid_indices): + obj_up = objectives[2 * pair_idx] + obj_down = objectives[2 * pair_idx + 1] + fd[param_idx] = float((obj_up - obj_down) / (2.0 * step[param_idx])) + + return fd + + +def make_objective( + make_geometry: Callable[[anp.ndarray, Sequence[float]], td.Geometry], + center: Sequence[float], + tag: str, + base_sim: td.Simulation, + fom: Callable, + tmp_path, + *, + local_gradient: bool, +): + def objective(parameters): + return run_parameter_simulations( + parameters, + make_geometry, + center, + tag, + base_sim, + fom, + tmp_path, + local_gradient=local_gradient, + ) + + return objective + + +@pytest.mark.parametrize("stretch_factor", STRETCH_FACTORS) +@pytest.mark.parametrize("shift_box_center", (True, False)) +@pytest.mark.parametrize("translation_axis", TRANSLATION_AXES) +def test_sphere_triangles_match_fd(stretch_factor, shift_box_center, translation_axis, tmp_path): + base_sim, fom = make_base_simulation(is_3d=True, infinite_dim=None) + + center = [0.0, 0.0, 0.0] + if shift_box_center: + center[0:2] = [0.5 * p for p in PERIODS_UM] + + params0 = anp.array([SPHERE_RADIUS_UM, 0.0]) + + triangle_objective = make_objective( + lambda p, c: make_sphere_triangle_geometry(p, c, stretch_factor, translation_axis), + center, + f"sphere_mesh_{stretch_factor}", + base_sim, + fom, + tmp_path, + local_gradient=LOCAL_GRADIENT, + ) + triangle_objective_fd = make_objective( + lambda p, c: make_sphere_triangle_geometry(p, c, stretch_factor, translation_axis), + center, + f"sphere_mesh_fd_{stretch_factor}", + base_sim, + fom, + tmp_path, + local_gradient=False, + ) + + _, triangle_grad = value_and_grad(triangle_objective)([params0]) + assert triangle_grad is not None + + triangle_grad = np.squeeze(np.asarray(triangle_grad, dtype=float)) + + fd_grad = finite_difference_params(triangle_objective_fd, params0) + + print("stretch", stretch_factor, "shift", shift_box_center, "dimension", translation_axis) + print("triangle_grad\t", triangle_grad.tolist()) + print("fd_grad\t\t", fd_grad.tolist()) + + if COMPARE_TO_FINITE_DIFFERENCE: + mesh_fd_overlap = angled_overlap_deg(triangle_grad, fd_grad) + print( + f"TriangleMesh FD vs. Adjoint angle overlap: {mesh_fd_overlap:.3f}° " + f"(threshold = {ANGLE_OVERLAP_FD_ADJ_THRESH_DEG}°)" + ) + assert mesh_fd_overlap < ANGLE_OVERLAP_FD_ADJ_THRESH_DEG, ( + f"FD–adjoint angle overlap too large: {mesh_fd_overlap:.3f}° " + f"(threshold {ANGLE_OVERLAP_FD_ADJ_THRESH_DEG}°, " + ) + + if SAVE_OUTPUT_DATA: + np.savez( + tmp_path / f"sphere_gradients_stretch_{stretch_factor}_shift_{shift_box_center}.npz", + triangle_grad=triangle_grad, + fd_grad=fd_grad, + ) diff --git a/tests/test_components/autograd/test_autograd.py b/tests/test_components/autograd/test_autograd.py index 8621eaded7..f3e1ae9dab 100644 --- a/tests/test_components/autograd/test_autograd.py +++ b/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") diff --git a/tests/test_components/autograd/test_autograd_triangle_mesh.py b/tests/test_components/autograd/test_autograd_triangle_mesh.py new file mode 100644 index 0000000000..58a7026dc2 --- /dev/null +++ b/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 diff --git a/tidy3d/components/data/data_array.py b/tidy3d/components/data/data_array.py index ea6385a796..ebd3be855e 100644 --- a/tidy3d/components/data/data_array.py +++ b/tidy3d/components/data/data_array.py @@ -78,18 +78,28 @@ class DataArray(xr.DataArray): _data_attrs: dict[str, str] = {} def __init__(self, data, *args: Any, **kwargs: Any) -> None: + # convert numpy object arrays that contain autograd boxes; keep other types as-is + data = self._maybe_convert_object_boxes(data) + # if data is a vanilla autograd box, convert to our box if isbox(data) and not is_tidy_box(data): data = TidyArrayBox.from_arraybox(data) # do the same for xr.Variable or xr.DataArray type - elif ( - isinstance(data, (xr.Variable, xr.DataArray)) - and isbox(data.data) - and not is_tidy_box(data.data) - ): - data.data = TidyArrayBox.from_arraybox(data.data) + elif isinstance(data, (xr.Variable, xr.DataArray)): + if isbox(data.data) and not is_tidy_box(data.data): + data.data = TidyArrayBox.from_arraybox(data.data) super().__init__(data, *args, **kwargs) + @staticmethod + def _maybe_convert_object_boxes(data): + """Convert object arrays of autograd boxes into ArrayBox instances.""" + + if isinstance(data, np.ndarray) and data.dtype == np.object_ and data.size: + # only convert if at least one element is an autograd tracer + if any(isbox(item) for item in data.flat): + return anp.array(data.tolist()) + return data + @classmethod def __get_validators__(cls): """Validators that get run when :class:`.DataArray` objects are added to pydantic models.""" diff --git a/tidy3d/components/geometry/mesh.py b/tidy3d/components/geometry/mesh.py index 5cb4cba58c..39faf5eb5a 100644 --- a/tidy3d/components/geometry/mesh.py +++ b/tidy3d/components/geometry/mesh.py @@ -8,12 +8,15 @@ import numpy as np import pydantic.v1 as pydantic +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 from tidy3d.components.data.dataset import TriangleMeshDataset from tidy3d.components.data.validators import validate_no_nans from tidy3d.components.types import Ax, Bound, Coordinate, MatrixReal4x4, Shapely from tidy3d.components.viz import add_ax_if_none, equal_aspect +from tidy3d.config import config from tidy3d.constants import fp_eps, inf from tidy3d.exceptions import DataError, ValidationError from tidy3d.log import log @@ -45,6 +48,9 @@ class TriangleMesh(base.Geometry, ABC): _no_nans_mesh = validate_no_nans("mesh_dataset") + # cache barycentric sampling patterns keyed by subdivision level + _barycentric_cache: dict[int, np.ndarray] = {} + @pydantic.root_validator(pre=True) @verify_packages_import(["trimesh"]) def _validate_trimesh_library(cls, values: dict[str, Any]) -> dict[str, Any]: @@ -69,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) @@ -695,3 +703,278 @@ def plot( ) return base.Geometry.plot(self, x=x, y=y, z=z, ax=ax, **patch_kwargs) + + def _compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradFieldMap: + """Compute adjoint derivatives for a ``TriangleMesh`` geometry.""" + + vjps: AutogradFieldMap = {} + + if not self.mesh_dataset: + raise DataError("Can't compute derivatives without mesh data.") + + valid_paths = {("mesh_dataset", "surface_mesh")} + for path in derivative_info.paths: + if path not in valid_paths: + raise ValueError(f"No derivative defined w.r.t. 'TriangleMesh' field '{path}'.") + + if ("mesh_dataset", "surface_mesh") not in derivative_info.paths: + return vjps + + triangles = np.asarray(self.triangles, dtype=config.adjoint.gradient_dtype_float) + + # early exit if geometry is completely outside simulation bounds + sim_min, sim_max = map(np.asarray, derivative_info.simulation_bounds) + mesh_min, mesh_max = map(np.asarray, self.bounds) + if np.any(mesh_max < sim_min) or np.any(mesh_min > sim_max): + log.warning( + "'TriangleMesh' lies completely outside the simulation domain.", + log_once=True, + ) + zeros = np.zeros_like(triangles, dtype=config.adjoint.gradient_dtype_float) + vjps[("mesh_dataset", "surface_mesh")] = zeros + return vjps + + # gather surface samples within the simulation bounds + dx = derivative_info.adaptive_vjp_spacing() + samples = self._collect_surface_samples( + triangles=triangles, + spacing=dx, + sim_min=sim_min, + sim_max=sim_max, + ) + + if samples["points"].shape[0] == 0: + zeros = np.zeros_like(triangles, dtype=config.adjoint.gradient_dtype_float) + vjps[("mesh_dataset", "surface_mesh")] = zeros + return vjps + + interpolators = derivative_info.interpolators + if interpolators is None: + interpolators = derivative_info.create_interpolators( + dtype=config.adjoint.gradient_dtype_float + ) + + g = derivative_info.evaluate_gradient_at_points( + samples["points"], + samples["normals"], + samples["perps1"], + samples["perps2"], + interpolators, + ) + + # accumulate per-vertex contributions using barycentric weights + weights = (samples["weights"] * g).real + normals = samples["normals"] + faces = samples["faces"] + bary = samples["barycentric"] + + contrib_vec = weights[:, None] * normals + + triangle_grads = np.zeros_like(triangles, dtype=config.adjoint.gradient_dtype_float) + for vertex_idx in range(3): + scaled = contrib_vec * bary[:, vertex_idx][:, None] + np.add.at(triangle_grads[:, vertex_idx, :], faces, scaled) + + vjps[("mesh_dataset", "surface_mesh")] = triangle_grads + return vjps + + def _collect_surface_samples( + self, + triangles: np.ndarray, + spacing: float, + sim_min: np.ndarray, + sim_max: np.ndarray, + ) -> dict[str, np.ndarray]: + """Collect sample points, normals, tangents, weights, and barycentric data.""" + + dtype = config.adjoint.gradient_dtype_float + tol = config.adjoint.edge_clip_tolerance + + sim_min = np.asarray(sim_min, dtype=dtype) + sim_max = np.asarray(sim_max, dtype=dtype) + + points_list: list[np.ndarray] = [] + normals_list: list[np.ndarray] = [] + perps1_list: list[np.ndarray] = [] + perps2_list: list[np.ndarray] = [] + weights_list: list[np.ndarray] = [] + faces_list: list[np.ndarray] = [] + bary_list: list[np.ndarray] = [] + + spacing = max(float(spacing), np.finfo(float).eps) + triangles_arr = np.asarray(triangles, dtype=dtype) + + # In 2D simulations one simulation dimension is collapsed (size ~ 0) and all + # gradients should be reported per-unit-length along that dimension. Track + # the collapsed axes here so we can normalize the sample weights later. + sim_extents = sim_max - sim_min + collapsed_axes = np.isclose(sim_extents, 0.0, atol=tol) + per_unit_scale = 1.0 + if np.any(collapsed_axes): + coords = triangles_arr.reshape(-1, 3) + geom_min = np.min(coords, axis=0) + geom_max = np.max(coords, axis=0) + geom_extents = geom_max - geom_min + collapsed_extents = np.maximum(geom_extents[collapsed_axes], tol) + per_unit_scale = float(np.prod(collapsed_extents)) + + for face_index, tri in enumerate(triangles_arr): + area, normal = self._triangle_area_and_normal(tri) + if area <= AREA_SIZE_THRESHOLD: + continue + + perps = self._triangle_tangent_basis(tri, normal) + if perps is None: + continue + perp1, perp2 = perps + + edge_lengths = ( + np.linalg.norm(tri[1] - tri[0]), + np.linalg.norm(tri[2] - tri[1]), + np.linalg.norm(tri[0] - tri[2]), + ) + subdivisions = self._subdivision_count(area, spacing, edge_lengths) + barycentric = self._get_barycentric_samples(subdivisions, dtype) + num_samples = barycentric.shape[0] + base_weight = area / num_samples + if per_unit_scale != 1.0: + base_weight /= per_unit_scale + + sample_points = barycentric @ tri + + # Only enforce clipping along axes that have a finite simulation extent. This keeps + # triangles that extend along the "collapsed" axis of a 2D simulation (size == 0) + # from being discarded even though they are physically valid. + valid_axes = np.abs(sim_max - sim_min) > tol + inside_mask = np.all( + sample_points[:, valid_axes] >= (sim_min - tol)[valid_axes], axis=1 + ) & np.all(sample_points[:, valid_axes] <= (sim_max + tol)[valid_axes], axis=1) + if not np.any(inside_mask): + continue + + sample_points = sample_points[inside_mask] + bary_inside = barycentric[inside_mask] + n_samples_inside = sample_points.shape[0] + + normal_tile = np.repeat(normal[None, :], n_samples_inside, axis=0) + perp1_tile = np.repeat(perp1[None, :], n_samples_inside, axis=0) + perp2_tile = np.repeat(perp2[None, :], n_samples_inside, axis=0) + weights_tile = np.full(n_samples_inside, base_weight, dtype=dtype) + faces_tile = np.full(n_samples_inside, face_index, dtype=int) + + points_list.append(sample_points) + normals_list.append(normal_tile) + perps1_list.append(perp1_tile) + perps2_list.append(perp2_tile) + weights_list.append(weights_tile) + faces_list.append(faces_tile) + bary_list.append(bary_inside) + + if not points_list: + return { + "points": np.zeros((0, 3), dtype=dtype), + "normals": np.zeros((0, 3), dtype=dtype), + "perps1": np.zeros((0, 3), dtype=dtype), + "perps2": np.zeros((0, 3), dtype=dtype), + "weights": np.zeros((0,), dtype=dtype), + "faces": np.zeros((0,), dtype=int), + "barycentric": np.zeros((0, 3), dtype=dtype), + } + + return { + "points": np.concatenate(points_list, axis=0), + "normals": np.concatenate(normals_list, axis=0), + "perps1": np.concatenate(perps1_list, axis=0), + "perps2": np.concatenate(perps2_list, axis=0), + "weights": np.concatenate(weights_list, axis=0), + "faces": np.concatenate(faces_list, axis=0), + "barycentric": np.concatenate(bary_list, axis=0), + } + + @staticmethod + def _triangle_area_and_normal(triangle: np.ndarray) -> tuple[float, np.ndarray]: + """Return area and outward normal of the provided triangle.""" + + edge01 = triangle[1] - triangle[0] + edge02 = triangle[2] - triangle[0] + cross = np.cross(edge01, edge02) + norm = np.linalg.norm(cross) + if norm <= 0.0: + return 0.0, np.zeros(3, dtype=triangle.dtype) + normal = (cross / norm).astype(triangle.dtype, copy=False) + area = 0.5 * norm + return area, normal + + @classmethod + def _subdivision_count( + cls, + area: float, + spacing: float, + edge_lengths: Optional[tuple[float, float, float]] = None, + ) -> int: + """Determine the number of subdivisions needed for the given area and spacing.""" + + spacing = max(float(spacing), np.finfo(float).eps) + + target = np.sqrt(max(area, 0.0)) + area_based = np.ceil(np.sqrt(2.0) * target / spacing) + + edge_based = 0.0 + if edge_lengths: + max_edge = max(edge_lengths) + if max_edge > 0.0: + edge_based = np.ceil(max_edge / spacing) + + subdivisions = max(1, int(max(area_based, edge_based))) + return subdivisions + + @classmethod + def _get_barycentric_samples(cls, subdivisions: int, dtype: np.dtype) -> np.ndarray: + """Return barycentric sample coordinates for a subdivision level.""" + + if subdivisions not in cls._barycentric_cache: + cls._barycentric_cache[subdivisions] = cls._build_barycentric_samples(subdivisions) + return cls._barycentric_cache[subdivisions].astype(dtype, copy=False) + + @staticmethod + def _build_barycentric_samples(subdivisions: int) -> np.ndarray: + """Construct barycentric sampling points for a given subdivision level.""" + + if subdivisions <= 1: + return np.array([[1.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0]]) + + bary = [] + for i in range(subdivisions): + for j in range(subdivisions - i): + l1 = (i + 1.0 / 3.0) / subdivisions + l2 = (j + 1.0 / 3.0) / subdivisions + l0 = 1.0 - l1 - l2 + bary.append((l0, l1, l2)) + return np.asarray(bary, dtype=float) + + @staticmethod + def _triangle_tangent_basis( + triangle: np.ndarray, normal: np.ndarray + ) -> Optional[tuple[np.ndarray, np.ndarray]]: + """Compute orthonormal tangential vectors for a triangle.""" + + tol = np.finfo(triangle.dtype).eps + edges = [triangle[1] - triangle[0], triangle[2] - triangle[0], triangle[2] - triangle[1]] + + edge = None + for candidate in edges: + length = np.linalg.norm(candidate) + if length > tol: + edge = (candidate / length).astype(triangle.dtype, copy=False) + break + + if edge is None: + return None + + perp1 = edge + perp2 = np.cross(normal, perp1) + perp2_norm = np.linalg.norm(perp2) + if perp2_norm <= tol: + return None + perp2 = (perp2 / perp2_norm).astype(triangle.dtype, copy=False) + return perp1, perp2