diff --git a/CHANGELOG.md b/CHANGELOG.md index a04d01b741..1bc4a30cd1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -42,6 +42,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Added `smoothed_projection` for topology optimization of completely binarized designs. - Added more RF-specific mode characteristics to `MicrowaveModeData`, including propagation constants (alpha, beta, gamma), phase/group velocities, wave impedance, and automatic mode classification with configurable polarization thresholds in `MicrowaveModeSpec`. - Introduce `tidy3d.rf` namespace to consolidate all RF classes. +- Added support of `TriangleMesh` for autograd. ### Breaking Changes - Edge singularity correction at PEC and lossy metal edges defaults to `True`. 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..9ab680efc8 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 @@ -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) 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..93a32d2442 --- /dev/null +++ b/tests/test_components/autograd/numerical/test_autograd_polyslab_trianglemesh_numerical.py @@ -0,0 +1,269 @@ +# 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 +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 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..b455576d39 --- /dev/null +++ b/tests/test_components/autograd/numerical/test_autograd_sphere_trianglemesh_numerical.py @@ -0,0 +1,576 @@ +# Test autograd gradients for scaled spheres represented as TriangleMesh +# geometries, comparing to finite differences for validation. +from __future__ import annotations + +import sys +from collections.abc import Sequence +from pathlib import Path +from typing import Callable + +import autograd.numpy as anp +import numpy as np +import pytest +from autograd import value_and_grad +from matplotlib import pyplot as plt + +import tidy3d as td +import tidy3d.web as web +from tests.test_components.autograd.numerical.test_autograd_box_polyslab_numerical import ( + angled_overlap_deg, +) +from tests.test_components.autograd.test_autograd_triangle_mesh import subdivide_triangles +from tidy3d import config +from tidy3d.components.geometry.primitives import _base_icosahedron + +config.local_cache.enabled = True +config.logging.level = "DEBUG" + +WL_UM = 0.65 +SPHERE_RADIUS_UM = 0.5 * WL_UM +SCALE_FACTORS = (0.2, 1.0, 5.0) +SCALE_AXES = (0, 1, 2) + +FREQ0 = td.C_0 / WL_UM +SRC_OFFSET = -2.5 +MONITOR_OFFSET = 2.5 +N_MAT = 2 +PERMITTIVITY = N_MAT**2 +ICOSAHEDRON_SUBDIVISIONS = 3 +LOCAL_GRADIENT = True +VERBOSE = False +SHOW_PRINT_STATEMENTS = True +SAVE_OUTPUT_DATA = True +ANGLE_OVERLAP_FD_ADJ_THRESH_DEG = 10.0 +VERTEX_FD_STEP = 1e-3 +FINITE_DIFF_STEP = 1e-2 +FINITE_DIFF_STEP_NATIVE = 1e-3 +GRID_STEPS_PER_WVL = 40 +td.config.adjoint.points_per_wavelength = 10 +measure_flux_spec = False + +freqs = td.C_0 / np.linspace(0.6, 0.7, 101) + +if SHOW_PRINT_STATEMENTS: + sys.stdout = sys.stderr + + +def make_base_simulation( + radii: list[float], *, extra_structures: Sequence[td.Structure] | None = None +) -> tuple[td.Simulation, callable]: + sim_size_3d = [ + 2 * radii[0] + 2 * WL_UM, + 2 * radii[1] + 2 * WL_UM, + (MONITOR_OFFSET - SRC_OFFSET) + 2 * WL_UM + 2 * radii[2], + ] + + plane_wave = td.PlaneWave( + center=(0.0, 0.0, SRC_OFFSET), + size=(*sim_size_3d[:2], 0.0), + source_time=td.GaussianPulse(freq0=FREQ0, fwidth=0.2 * FREQ0), + direction="+", + ) + + flux_monitors = [ + td.FieldMonitor( + center=(0.0, 0.0, MONITOR_OFFSET), + size=(*sim_size_3d[:2], 0.0), + freqs=FREQ0, + name="field", + ) + ] + if measure_flux_spec: + flux_monitors.append( + td.FieldMonitor( + center=(0.0, 0.0, MONITOR_OFFSET), + size=(*sim_size_3d[:2], 0.0), + freqs=freqs, + name="field_spectrum", + ) + ) + + boundary_spec_3d = td.BoundarySpec( + x=td.Boundary.pml(), + y=td.Boundary.pml(), + z=td.Boundary.pml(), + ) + + base_sim = td.Simulation( + center=(0.0, 0.0, 0.0), + size=tuple(sim_size_3d), + monitors=flux_monitors, + sources=[plane_wave], + structures=list(extra_structures) if extra_structures else [], + run_time=2e-11, + boundary_spec=boundary_spec_3d, + grid_spec=td.GridSpec( + grid_x=td.UniformGrid(dl=WL_UM / GRID_STEPS_PER_WVL), + grid_y=td.UniformGrid(dl=WL_UM / GRID_STEPS_PER_WVL), + grid_z=td.UniformGrid(dl=WL_UM / GRID_STEPS_PER_WVL), + ), + ) + + def fom(sim_data): + return sim_data["field"].flux.values + + return base_sim, fom + + +def make_overlap_cube_structure(radii: Sequence[float]) -> td.Structure: + radii_arr = np.asarray(radii, dtype=float) + size_x = float(radii_arr[0]) + size_y = float(2.0 * radii_arr[1]) + size_z = float(2.0 * radii_arr[2]) + cube_center = (size_x / 2.0, 0.0, 0.0) + cube = td.Box(center=cube_center, size=(size_x, size_y, size_z)) + cube_medium = td.Medium(permittivity=PERMITTIVITY) + return td.Structure(geometry=cube, medium=cube_medium) + + +def run_parameter_simulations( + parameter_sets: list[anp.ndarray], + make_geometry, + box_center, + tag: str, + base_sim: td.Simulation, + fom, + artifact_dir: Path, + *, + local_gradient: bool, +): + simulation_dict = {} + output_dir = artifact_dir + output_dir.mkdir(parents=True, exist_ok=True) + + for idx, param_values in enumerate(parameter_sets): + geometry = make_geometry(param_values, box_center) + structure = td.Structure( + geometry=geometry, + medium=td.Medium(permittivity=PERMITTIVITY), + ) + + base_structures = list(getattr(base_sim, "structures", ())) + structures = [structure, *base_structures] + grid_spec = td.GridSpec.auto( + min_steps_per_wvl=GRID_STEPS_PER_WVL, override_structures=[structure] + ) + sim = base_sim.updated_copy(structures=structures, grid_spec=grid_spec, validate=True) + task_name = f"{tag}_idx{idx}" + simulation_dict[task_name] = sim + + if len(simulation_dict) == 1: + key, sim = next(iter(simulation_dict.items())) + result_path = output_dir / f"{sim._hash_self()}.hdf5" + sim_data = web.run( + sim, + task_name=key, + path=str(result_path), + local_gradient=local_gradient, + verbose=VERBOSE, + ) + return fom(sim_data) + + sim_data_map = web.run_async( + simulation_dict, + path_dir=str(output_dir), + local_gradient=local_gradient, + verbose=VERBOSE, + ) + + return [fom(sim_data_map[key]) for key in simulation_dict] + + +_ICOSAHEDRON_VERTS, _ICOSAHEDRON_FACES = _base_icosahedron() + + +def make_sphere_triangle_geometry( + params: anp.ndarray, + center: Sequence[float], + scale_factor: float, + scale_axis: int, + subdivisions: int = ICOSAHEDRON_SUBDIVISIONS, + in_plane_subdivisions: int = 0, +) -> td.Geometry: + radii = anp.array(params, dtype=float) + triangles = td.Sphere.unit_sphere_triangles(subdivisions=subdivisions) + if in_plane_subdivisions > 0: + for _ in range(in_plane_subdivisions): + triangles = subdivide_triangles(triangles) + triangles = anp.array(triangles) + triangles = triangles * radii + axis_selector = anp.equal(anp.arange(3), scale_axis) + scale_vec = anp.where(axis_selector, scale_factor, 1.0) + triangles = triangles * scale_vec + center_arr = anp.array(center, dtype=float) + triangles = triangles + center_arr + mesh = td.TriangleMesh.from_triangles(triangles) + return mesh + + +def make_mesh_sphere_from_radius(params: anp.ndarray, center: Sequence[float]) -> td.Geometry: + radius = params[0] + radii = anp.full(3, radius) + return make_sphere_triangle_geometry(radii, center, scale_factor=1.0, scale_axis=0) + + +def finite_difference_params(objective, params: anp.ndarray, finite_diff_step) -> np.ndarray: + step = np.full_like(np.asarray(params, dtype=float), finite_diff_step, dtype=float) + 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 finite_difference_params_step_batch( + objective, params: anp.ndarray, finite_diff_step +) -> np.ndarray: + """Compute central finite-difference gradients for one or more FD step sizes. + + Args: + objective: Callable accepting a batch of parameter vectors and returning objectives + (shape (N,) or (N,1)). + params: Parameter vector (1D array). + finite_diff_step: Either a scalar, array of per-parameter step sizes, or list of such. + + Returns: + np.ndarray: Finite-difference gradient(s). + Shape is (len(steps), n_params) if multiple steps, + or (n_params,) for a single step. + """ + params = anp.asarray(params, dtype=float) + + # Normalize `finite_diff_step` to a list of arrays, one per step value + if np.isscalar(finite_diff_step): + step_list = [anp.full_like(params, float(finite_diff_step), dtype=float)] + else: + finite_diff_step = np.atleast_1d(finite_diff_step) + if finite_diff_step.ndim == 1 and finite_diff_step.size == params.size: + # one per parameter + step_list = [anp.asarray(finite_diff_step, dtype=float)] + else: + # list/array of scalar step values + step_list = [anp.full_like(params, float(s), dtype=float) for s in finite_diff_step] + + grads = [] + + for step in step_list: + perturbations = [] + valid_indices = [] + for idx in range(params.size): + params_up = anp.array(params) + params_down = anp.array(params) + 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).ravel() + + fd = np.zeros_like(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] = (obj_up - obj_down) / (2.0 * step[param_idx]) + + grads.append(fd) + + grads = np.stack(grads, axis=0) + return grads[0] if len(grads) == 1 else grads + + +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.numerical +@pytest.mark.parametrize("scale_factor", (1,)) +@pytest.mark.parametrize("scale_axis", (0,)) +@pytest.mark.parametrize("overlap_cube", (False,)) +def test_sphere_triangles_match_fd( + scale_factor, scale_axis, overlap_cube, tmp_path, numerical_case_dir +): + """ + Compares FD gradients with gradients from _compute_derivatives in TriangleMesh. + Note that FD gradients are very noise which is why there will be some failing tests with a fixed FD-step + """ + if scale_factor == 1 and scale_axis > 0: + pytest.skip("Skipping duplicate test.") + + initial_params = [SPHERE_RADIUS_UM, SPHERE_RADIUS_UM, SPHERE_RADIUS_UM] + params0 = anp.array(initial_params) + + radii = initial_params.copy() + radii[scale_axis] *= scale_factor + extra_structures = [make_overlap_cube_structure(radii)] if overlap_cube else [] + base_sim, fom = make_base_simulation(radii=radii, extra_structures=extra_structures) + + center = [0.0, 0.0, 0.0] + + part_make_geom = lambda p, c: make_sphere_triangle_geometry(p, c, scale_factor, scale_axis) + + triangle_objective = make_objective( + part_make_geom, + center, + f"sphere_mesh_{scale_factor}_axis_{scale_axis}_cube_{overlap_cube}", + base_sim, + fom, + tmp_path, + local_gradient=LOCAL_GRADIENT, + ) + triangle_objective_fd = make_objective( + part_make_geom, + center, + f"sphere_mesh_fd_{scale_factor}_axis_{scale_axis}_cube_{overlap_cube}", + 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, FINITE_DIFF_STEP) + + print("scale", scale_factor, "axis", scale_axis, "overlap_cube", overlap_cube) + print("triangle_grad\t", triangle_grad.tolist()) + print("fd_grad\t\t", fd_grad.tolist()) + + 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( + numerical_case_dir + / f"sphere_gradients_mesh_scale_{scale_factor}_axis_{scale_axis}_cube_{overlap_cube}.npz", + triangle_grad=triangle_grad, + fd_grad=fd_grad, + ) + + +@pytest.mark.skip +def test_grad_insensitive_to_face_splitting(tmp_path, numerical_case_dir): + scale_factor = 1 + scale_axis = 0 + + initial_params = [SPHERE_RADIUS_UM] * 3 + params0 = anp.array(initial_params) + + radii = initial_params.copy() + radii[scale_axis] *= scale_factor + base_sim, fom = make_base_simulation(radii=radii) + + center = [0.0, 0.0, 0.0] + + # clean objective names + obj_name_base = "sphere_mesh_subdiv_0" + obj_name_subdiv_1 = "sphere_mesh_subdiv_1" + obj_name_subdiv_2 = "sphere_mesh_subdiv_2" + + triangle_objective_base = make_objective( + lambda p, c: make_sphere_triangle_geometry( + p, c, scale_factor, scale_axis, subdivisions=0, in_plane_subdivisions=0 + ), + center, + obj_name_base, + base_sim, + fom, + tmp_path, + local_gradient=LOCAL_GRADIENT, + ) + + triangle_objective_subdiv_1 = make_objective( + lambda p, c: make_sphere_triangle_geometry( + p, c, scale_factor, scale_axis, subdivisions=0, in_plane_subdivisions=1 + ), + center, + obj_name_subdiv_1, + base_sim, + fom, + tmp_path, + local_gradient=LOCAL_GRADIENT, + ) + + triangle_objective_subdiv_2 = make_objective( + lambda p, c: make_sphere_triangle_geometry( + p, c, scale_factor, scale_axis, subdivisions=0, in_plane_subdivisions=2 + ), + center, + obj_name_subdiv_2, + base_sim, + fom, + tmp_path, + local_gradient=LOCAL_GRADIENT, + ) + + # ---- Evaluate adjoint gradients for base and subdivided meshes ---- + _, grad_base = value_and_grad(triangle_objective_base)([params0]) + _, grad_subdiv_1 = value_and_grad(triangle_objective_subdiv_1)([params0]) + _, grad_subdiv_2 = value_and_grad(triangle_objective_subdiv_2)([params0]) + + assert grad_base is not None + assert grad_subdiv_1 is not None + assert grad_subdiv_2 is not None + + grad_base = np.squeeze(np.asarray(grad_base, dtype=float)) + grad_subdiv_1 = np.squeeze(np.asarray(grad_subdiv_1, dtype=float)) + grad_subdiv_2 = np.squeeze(np.asarray(grad_subdiv_2, dtype=float)) + + print("grad_base \t", grad_base.tolist()) + print("grad_subdiv_1 \t", grad_subdiv_1.tolist()) + print("grad_subdiv_2 \t", grad_subdiv_2.tolist()) + + # Optional: angles for debugging / log inspection. + angle_base_vs_1 = angled_overlap_deg(grad_base, grad_subdiv_1) + angle_base_vs_2 = angled_overlap_deg(grad_base, grad_subdiv_2) + print( + f"Base vs subdiv-1 angle: {angle_base_vs_1:.3f}°; " + f"Base vs subdiv-2 angle: {angle_base_vs_2:.3f}°" + ) + + np.testing.assert_allclose(grad_base, grad_subdiv_1, rtol=0.1, atol=5e-6) + np.testing.assert_allclose(grad_base, grad_subdiv_2, rtol=0.1, atol=5e-6) + + if SAVE_OUTPUT_DATA: + np.savez( + numerical_case_dir / "sphere_gradients_mesh_subdiv.npz", + grad_base=grad_base, + grad_subdiv_1=grad_subdiv_1, + grad_subdiv_2=grad_subdiv_2, + ) + + +@pytest.mark.skip +@pytest.mark.parametrize("scale_factor", SCALE_FACTORS) +@pytest.mark.parametrize("scale_axis", SCALE_AXES) +@pytest.mark.parametrize("overlap_cube", (False, True)) +def test_triangle_sphere_fd_step_sweep_ref( + tmp_path, scale_factor, scale_axis, overlap_cube, numerical_case_dir +): + global SPHERE_RADIUS_UM + SPHERE_RADIUS_UM = SPHERE_RADIUS_UM * 4 + initial_params = [SPHERE_RADIUS_UM, SPHERE_RADIUS_UM, SPHERE_RADIUS_UM] + params0 = anp.array(initial_params) + + radii = initial_params.copy() + radii[scale_axis] *= scale_factor + extra_structures = [make_overlap_cube_structure(radii)] if overlap_cube else [] + base_sim, fom = make_base_simulation(radii=radii, extra_structures=extra_structures) + + center = [0.0, 0.0, 0.0] + + part_make_geom = lambda p, c: make_sphere_triangle_geometry(p, c, scale_factor, scale_axis) + + # ---- Build objectives ---- + triangle_objective_fd = make_objective( + part_make_geom, + center, + "sphere_mesh_fd_step_sweep", + base_sim, + fom, + tmp_path, + local_gradient=False, + ) + triangle_objective_autograd = make_objective( + part_make_geom, + center, + "sphere_mesh_autograd_ref", + base_sim, + fom, + tmp_path, + local_gradient=True, + ) + + # ---- Compute autograd gradient (for reference) ---- + _, autograd_grad = value_and_grad(triangle_objective_autograd)([params0]) + autograd_grad = np.squeeze(np.asarray(autograd_grad, dtype=float)) + + # ---- Finite difference sweep ---- + steps = np.logspace(-4, -1, num=12) + fd_grads = finite_difference_params_step_batch(triangle_objective_fd, params0, steps) + + fd_grads = np.asarray(fd_grads, dtype=float) + + # ---- Plot ---- + fig, ax = plt.subplots(figsize=(6, 4)) + for idx, label in enumerate(["radius_x", "radius_y", "radius_z"]): + ax.plot(steps, fd_grads[:, idx], marker="o", label=label) + # Add horizontal reference line for autograd + ax.axhline( + autograd_grad[idx], + color=ax.get_lines()[-1].get_color(), + linestyle="--", + alpha=0.7, + label=f"{label} (autograd)", + ) + + ax.set_xscale("log") + ax.set_xlabel("Finite difference step (µm)") + ax.set_ylabel("Gradient value") + ax.set_title("FD gradients vs. step size") + ax.grid(True, which="both", ls=":") + ax.legend() + + fig_path = ( + numerical_case_dir + / f"fd_step_sweep_scale_{scale_factor}_axis_{scale_axis}_cube_{overlap_cube}.png" + ) + + fig.savefig(fig_path, dpi=200) + plt.close(fig) + + np.savez( + numerical_case_dir + / f"fd_step_sweep_scale_{scale_factor}_axis_{scale_axis}_cube_{overlap_cube}.npz", + steps=steps, + gradients=fd_grads, + autograd_grad=autograd_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..0e254525c4 --- /dev/null +++ b/tests/test_components/autograd/test_autograd_triangle_mesh.py @@ -0,0 +1,387 @@ +"""Tests for TriangleMesh autograd derivatives.""" + +from __future__ import annotations + +import numpy as np +import numpy.testing as npt +import pytest + +import tidy3d as td + +from ...utils import AssertLogLevel + +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), +} + +OUT_OF_BOUNDS_WARNING = "Some triangles from the mesh lie outside the simulation bounds - this may lead to inaccurate gradients." + + +@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 subdivide_triangles(triangles: np.ndarray) -> np.ndarray: + """Split each triangle into four smaller ones via ``TriangleMesh.subdivide_faces``.""" + + triangles = np.asarray(triangles, dtype=float) + vertices = triangles.reshape(-1, 3) + face_count = triangles.shape[0] + faces = np.arange(vertices.shape[0], dtype=int).reshape(face_count, 3) + refined_vertices, refined_faces = td.TriangleMesh.subdivide_faces(vertices, faces) + return refined_vertices[refined_faces] + + +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)) + + +@pytest.mark.parametrize( + "simulation_bounds, expect_warning", + [ + (((-0.1, -0.1, -0.1), (0.6, 0.6, 0.6)), True), + (((-2.0, -2.0, -2.0), (2.0, 2.0, 2.0)), False), + ], +) +def test_triangle_mesh_partial_bounds_sampling_warns(simulation_bounds, expect_warning): + """Surface sampling logs a warning only when triangles leave the simulation bounds.""" + + vertices, faces = MESH_DEFINITIONS["tetrahedron"] + mesh = td.TriangleMesh.from_vertices_faces(vertices, faces) + sim_min, sim_max = simulation_bounds + + log_ctx = ( + AssertLogLevel("WARNING", contains_str=OUT_OF_BOUNDS_WARNING) + if expect_warning + else AssertLogLevel(None) + ) + + with log_ctx: + samples = mesh._collect_surface_samples( + triangles=mesh.triangles, + spacing=0.05, + sim_min=np.array(sim_min), + sim_max=np.array(sim_max), + ) + + assert set(samples) == { + "points", + "normals", + "perps1", + "perps2", + "weights", + "faces", + "barycentric", + } + + +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_gradients_insensitive_to_face_splitting(watertight_mesh): + """Refining triangles does not change the directional derivative.""" + + base_mesh = watertight_mesh + base_triangles = base_mesh.triangles + refined_mesh = td.TriangleMesh.from_triangles(subdivide_triangles(base_triangles)) + + coeffs = np.array([0.41, -0.33, 0.27], dtype=float) + offset = 0.12 + grad_func = linear_grad_func_factory(coeffs, offset) + + spacing = 0.006 + derivative_info = DummyDerivativeInfo(grad_func, spacing=spacing) + derivative_info_refined = DummyDerivativeInfo(grad_func, spacing=spacing) + + grad_base = base_mesh._compute_derivatives(derivative_info)[("mesh_dataset", "surface_mesh")] + grad_refined = refined_mesh._compute_derivatives(derivative_info_refined)[ + ("mesh_dataset", "surface_mesh") + ] + + def displacement_field(points: np.ndarray) -> np.ndarray: + x, y, z = points.T + return np.column_stack((0.2 * x - 0.1 * y, 0.15 * y + 0.05 * z, -0.25 * z + 0.12 * x)) + + disp_base = displacement_field(base_triangles.reshape(-1, 3)).reshape(base_triangles.shape) + refined_triangles = refined_mesh.triangles + disp_refined = displacement_field(refined_triangles.reshape(-1, 3)).reshape( + refined_triangles.shape + ) + + dir_base = float(np.sum(grad_base * disp_base)) + dir_refined = float(np.sum(grad_refined * disp_refined)) + + npt.assert_allclose(dir_base, dir_refined, rtol=5e-3, atol=5e-6) diff --git a/tests/test_components/test_geometry.py b/tests/test_components/test_geometry.py index d00a11515b..315205a147 100644 --- a/tests/test_components/test_geometry.py +++ b/tests/test_components/test_geometry.py @@ -838,7 +838,7 @@ def make_polyslabs(gap_size): return td.PolySlab.from_gds(gds_cell=cell, gds_layer=0, axis=2, slab_bounds=(-1, 1)) polyslabs_gap = make_polyslabs(gap_size=0.3) - assert len(polyslabs_gap) == 2, "untouching polylsabs were merged incorrectly." + assert len(polyslabs_gap) == 2, "untouching polyslabs were merged incorrectly." polyslabs_touching = make_polyslabs(gap_size=0) assert len(polyslabs_touching) == 1, "polyslabs didn't merge correctly." @@ -1068,6 +1068,29 @@ def test_custom_surface_geometry(tmp_path): td.TriangleMesh.from_vertices_faces(vertices_small, faces) +@pytest.mark.parametrize("binary", [True, False]) +def test_triangle_mesh_to_stl_roundtrip(tmp_path, binary): + """Exporting with 'to_stl' should be readable via 'from_stl'.""" + + vertices = 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], + ] + ) + faces = np.array([[1, 2, 3], [0, 3, 2], [0, 1, 3], [0, 2, 1]]) + mesh = td.TriangleMesh.from_vertices_faces(vertices, faces) + + export_path = tmp_path / ("mesh_binary.stl" if binary else "mesh_ascii.stl") + mesh.to_stl(str(export_path), binary=binary) + + roundtrip = td.TriangleMesh.from_stl(str(export_path)) + + assert np.allclose(roundtrip.triangles, mesh.triangles) + + def test_geo_group_sim(): geo_grp = td.TriangleMesh.from_stl("tests/data/two_boxes_separate.stl") geos_orig = list(geo_grp.geometries) 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..2b43d5c3b6 100644 --- a/tidy3d/components/geometry/mesh.py +++ b/tidy3d/components/geometry/mesh.py @@ -3,17 +3,23 @@ from __future__ import annotations from abc import ABC +from os import PathLike from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union import numpy as np import pydantic.v1 as pydantic +from autograd import numpy as anp +from numpy.typing import NDArray +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 @@ -44,6 +50,7 @@ class TriangleMesh(base.Geometry, ABC): ) _no_nans_mesh = validate_no_nans("mesh_dataset") + _barycentric_samples: dict[int, NDArray] = pydantic.PrivateAttr(default_factory=dict) @pydantic.root_validator(pre=True) @verify_packages_import(["trimesh"]) @@ -69,7 +76,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) @@ -216,6 +225,28 @@ def process_single(mesh: TrimeshType) -> TriangleMesh: return process_single(meshes[solid_index]) raise ValidationError("No solid found at 'solid_index' in the stl file.") + @verify_packages_import(["trimesh"]) + def to_stl( + self, + filename: PathLike, + *, + binary: bool = True, + ) -> None: + """Export this TriangleMesh to an STL file. + + Parameters + ---------- + filename : str + Output STL filename. + binary : bool = True + Whether to write binary STL. Set False for ASCII STL. + """ + triangles = get_static(self.mesh_dataset.surface_mesh.data) + mesh = self._triangles_to_trimesh(triangles) + + file_type = "stl" if binary else "stl_ascii" + mesh.export(file_obj=filename, file_type=file_type) + @classmethod @verify_packages_import(["trimesh"]) def from_trimesh(cls, mesh: trimesh.Trimesh) -> TriangleMesh: @@ -234,7 +265,7 @@ def from_trimesh(cls, mesh: trimesh.Trimesh) -> TriangleMesh: return cls.from_vertices_faces(mesh.vertices, mesh.faces) @classmethod - def from_triangles(cls, triangles: np.ndarray) -> TriangleMesh: + def from_triangles(cls, triangles: NDArray) -> TriangleMesh: """Create a :class:`.TriangleMesh` from a numpy array containing the triangles of a surface mesh. @@ -268,7 +299,7 @@ def from_triangles(cls, triangles: np.ndarray) -> TriangleMesh: @classmethod @verify_packages_import(["trimesh"]) - def from_vertices_faces(cls, vertices: np.ndarray, faces: np.ndarray) -> TriangleMesh: + def from_vertices_faces(cls, vertices: NDArray, faces: NDArray) -> TriangleMesh: """Create a :class:`.TriangleMesh` from numpy arrays containing the data of a surface mesh. The first array contains the vertices, and the second array contains faces formed from triples of the vertices. @@ -305,12 +336,19 @@ def from_vertices_faces(cls, vertices: np.ndarray, faces: np.ndarray) -> Triangl @classmethod @verify_packages_import(["trimesh"]) def _triangles_to_trimesh( - cls, triangles: np.ndarray + cls, triangles: NDArray ) -> Trimesh: # -> We need to get this out of the classes and into functional methods operating on a class (maybe still referenced to the class) """Convert an (N, 3, 3) numpy array of triangles to a ``trimesh.Trimesh``.""" import trimesh - return trimesh.Trimesh(**trimesh.triangles.to_kwargs(triangles)) + # ``triangles`` may contain autograd ``ArrayBox`` entries when differentiating + # geometry parameters. ``trimesh`` expects plain ``float`` values, so strip any + # tracing information before constructing the mesh. + triangles = np.array(triangles) + if triangles.dtype == np.object_: + triangles = anp.array(triangles.tolist()) + triangles = np.asarray(get_static(triangles), dtype=np.float64) + return trimesh.Trimesh(**trimesh.triangles.to_kwargs(get_static(triangles))) @classmethod def from_height_grid( @@ -319,7 +357,7 @@ def from_height_grid( direction: Literal["-", "+"], base: float, grid: tuple[np.ndarray, np.ndarray], - height: np.ndarray, + height: NDArray, ) -> TriangleMesh: """Construct a TriangleMesh object from grid based height information. @@ -625,20 +663,18 @@ def intersections_plane( log.warning(f"Error encountered: {e}") return self.bounding_box.intersections_plane(x=x, y=y, z=z) - def inside( - self, x: np.ndarray[float], y: np.ndarray[float], z: np.ndarray[float] - ) -> np.ndarray[bool]: + def inside(self, x: NDArray, y: NDArray, z: NDArray) -> np.ndarray[bool]: """For input arrays ``x``, ``y``, ``z`` of arbitrary but identical shape, return an array with the same shape which is ``True`` for every point in zip(x, y, z) that is inside the volume of the :class:`Geometry`, and ``False`` otherwise. Parameters ---------- - x : np.ndarray[float] + x : np.ndarray Array of point positions in x direction. - y : np.ndarray[float] + y : np.ndarray Array of point positions in y direction. - z : np.ndarray[float] + z : np.ndarray Array of point positions in z direction. Returns @@ -649,7 +685,6 @@ def inside( arrays = tuple(map(np.array, (x, y, z))) self._ensure_equal_shape(*arrays) - inside = np.zeros((arrays[0].size,), dtype=bool) arrays_flat = map(np.ravel, arrays) arrays_stacked = np.stack(tuple(arrays_flat), axis=-1) inside = self.trimesh.contains(arrays_stacked) @@ -695,3 +730,303 @@ 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: NDArray, + spacing: float, + sim_min: NDArray, + sim_max: NDArray, + ) -> dict[str, np.ndarray]: + """Deterministic per-triangle sampling used historically.""" + + 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) + + 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)) + + warned = False + 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 + + 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.all(inside_mask) and not warned: + log.warning( + "Some triangles from the mesh lie outside the simulation bounds - this may lead to inaccurate gradients." + ) + warned = True + 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: 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 + + def _get_barycentric_samples(self, subdivisions: int, dtype: np.dtype) -> np.ndarray: + """Return barycentric sample coordinates for a subdivision level.""" + + cache = self._barycentric_samples + if subdivisions not in cache: + cache[subdivisions] = self._build_barycentric_samples(subdivisions) + return 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 subdivide_faces(vertices: NDArray, faces: NDArray) -> tuple[np.ndarray, np.ndarray]: + """Uniformly subdivide each triangular face by inserting edge midpoints.""" + + midpoint_cache: dict[tuple[int, int], int] = {} + verts_list = [np.asarray(v, dtype=float) 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] + vm = 0.5 * (verts_list[i] + verts_list[j]) + verts_list.append(vm) + 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))) + + verts_arr = np.asarray(verts_list, dtype=float) + return verts_arr, np.asarray(new_faces, dtype=int) + + @staticmethod + def _triangle_tangent_basis( + triangle: NDArray, normal: 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 diff --git a/tidy3d/components/geometry/primitives.py b/tidy3d/components/geometry/primitives.py index 667ef5cb1a..a1f9313ab9 100644 --- a/tidy3d/components/geometry/primitives.py +++ b/tidy3d/components/geometry/primitives.py @@ -9,11 +9,15 @@ import numpy as np import pydantic.v1 as pydantic import shapely +from pydantic.v1 import PrivateAttr from shapely.geometry.base import BaseGeometry from tidy3d.components.autograd import AutogradFieldMap, TracedSize1D from tidy3d.components.autograd.derivative_utils import DerivativeInfo from tidy3d.components.base import cached_property, skip_if_fields_missing +from tidy3d.components.geometry import base +from tidy3d.components.geometry.mesh import TriangleMesh +from tidy3d.components.geometry.polyslab import PolySlab from tidy3d.components.types import Axis, Bound, Coordinate, MatrixReal4x4, Shapely from tidy3d.config import config from tidy3d.constants import LARGE_NUMBER, MICROMETER @@ -21,9 +25,6 @@ from tidy3d.log import log from tidy3d.packaging import verify_packages_import -from . import base -from .polyslab import PolySlab - # for sampling conical frustum in visualization _N_SAMPLE_CURVE_SHAPELY = 40 @@ -32,6 +33,61 @@ # Default number of points to discretize polyslab in `Cylinder.to_polyslab()` _N_PTS_CYLINDER_POLYSLAB = 51 +_MAX_ICOSPHERE_SUBDIVISIONS = 7 # this would have 164K vertices and 328K faces +_DEFAULT_EDGE_FRACTION = 0.25 + + +def _base_icosahedron() -> tuple[np.ndarray, np.ndarray]: + """Return vertices and faces of a unit icosahedron.""" + + 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=int, + ) + return vertices, faces + + +_ICOSAHEDRON_VERTS, _ICOSAHEDRON_FACES = _base_icosahedron() class Sphere(base.Centered, base.Circular): @@ -42,6 +98,8 @@ class Sphere(base.Centered, base.Circular): >>> b = Sphere(center=(1,2,3), radius=2) """ + _icosphere_cache: dict[int, tuple[np.ndarray, float]] = PrivateAttr(default_factory=dict) + def inside( self, x: np.ndarray[float], y: np.ndarray[float], z: np.ndarray[float] ) -> np.ndarray[bool]: @@ -178,6 +236,86 @@ def _surface_area(self, bounds: Bound) -> float: return area + @classmethod + def unit_sphere_triangles( + cls, + *, + target_edge_length: Optional[float] = None, + subdivisions: Optional[int] = None, + ) -> np.ndarray: + """Return unit sphere triangles discretized via an icosphere.""" + + unit_tris = UNIT_SPHERE._unit_sphere_triangles( + target_edge_length=target_edge_length, + subdivisions=subdivisions, + copy_result=True, + ) + return unit_tris + + def _unit_sphere_triangles( + self, + *, + target_edge_length: Optional[float] = None, + subdivisions: Optional[int] = None, + copy_result: bool = True, + ) -> np.ndarray: + """Return cached unit-sphere triangles with optional copying.""" + if target_edge_length is not None and subdivisions is not None: + raise ValueError("Specify either target_edge_length OR subdivisions, not both.") + + if subdivisions is None: + subdivisions = self._subdivisions_for_edge(target_edge_length) + + triangles, _ = self._icosphere_data(subdivisions) + return np.array(triangles, copy=copy_result) + + def _subdivisions_for_edge(self, target_edge_length: Optional[float]) -> int: + if target_edge_length is None or target_edge_length <= 0.0: + return 0 + + for subdiv in range(_MAX_ICOSPHERE_SUBDIVISIONS + 1): + _, max_edge = self._icosphere_data(subdiv) + if max_edge <= target_edge_length: + return subdiv + + log.warning( + "Requested sphere mesh edge length %.3e μm requires more than %d subdivisions. " + "Clipping to the finest available mesh.", + target_edge_length, + _MAX_ICOSPHERE_SUBDIVISIONS, + log_once=True, + ) + return _MAX_ICOSPHERE_SUBDIVISIONS + + def _icosphere_data(self, subdivisions: int) -> tuple[np.ndarray, float]: + cache = self._icosphere_cache + if subdivisions in cache: + return cache[subdivisions] + + vertices = np.asarray(_ICOSAHEDRON_VERTS, dtype=float) + faces = np.asarray(_ICOSAHEDRON_FACES, dtype=int) + if subdivisions > 0: + vertices = vertices.copy() + faces = faces.copy() + for _ in range(subdivisions): + vertices, faces = TriangleMesh.subdivide_faces(vertices, faces) + + norms = np.linalg.norm(vertices, axis=1, keepdims=True) + norms = np.where(norms == 0.0, 1.0, norms) + vertices = vertices / norms + + triangles = vertices[faces] + max_edge = self._edge_length(triangles) + cache[subdivisions] = (triangles, max_edge) + return triangles, max_edge + + @staticmethod + def _edge_length(triangles: np.ndarray) -> float: + return float(np.linalg.norm(triangles[0, 0] - triangles[0, 1])) + + +UNIT_SPHERE = Sphere(center=(0.0, 0.0, 0.0), radius=1.0) + class Cylinder(base.Centered, base.Circular, base.Planar): """Cylindrical geometry with optional sidewall angle along axis