diff --git a/CHANGELOG.md b/CHANGELOG.md index 9f84be6c38..f60a0bf8bf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- Added `symmetrize_mirror`, `symmetrize_rotation`, `symmetrize_diagonal` functions to the autograd plugin. They can be used for enforcing symmetries in topology optimization. +- Klayout plugin automatically finds klayout installation path at common locations. +- Added autograd support for `TriangleMesh`, allowing gradient computation with respect to mesh vertices for inverse design. +- Added `EMESimulation.store_coeffs` to store coefficients from the EME solver, including mode overlaps, interface S matrices, and effective propagation indices. +- Get cell-related information from violation markers in `DRCResults` and `DRCViolation` to the klayout plugin: Use for example `DRCResults.violations_by_cell` to group them. +- Added autograd support for `Sphere`. ### Changed diff --git a/schemas/EMESimulation.json b/schemas/EMESimulation.json index 9a8c933b07..0c756cb12a 100644 --- a/schemas/EMESimulation.json +++ b/schemas/EMESimulation.json @@ -11264,8 +11264,15 @@ ] }, "radius": { - "minimum": 0, - "type": "number" + "anyOf": [ + { + "minimum": 0, + "type": "number" + }, + { + "type": "autograd.tracer.Box" + } + ] }, "type": { "default": "Sphere", diff --git a/schemas/HeatChargeSimulation.json b/schemas/HeatChargeSimulation.json index e1514a8e0d..1b5a38b470 100644 --- a/schemas/HeatChargeSimulation.json +++ b/schemas/HeatChargeSimulation.json @@ -7567,8 +7567,15 @@ ] }, "radius": { - "minimum": 0, - "type": "number" + "anyOf": [ + { + "minimum": 0, + "type": "number" + }, + { + "type": "autograd.tracer.Box" + } + ] }, "type": { "default": "Sphere", diff --git a/schemas/HeatSimulation.json b/schemas/HeatSimulation.json index f1167d86f1..3f02a416e1 100644 --- a/schemas/HeatSimulation.json +++ b/schemas/HeatSimulation.json @@ -7567,8 +7567,15 @@ ] }, "radius": { - "minimum": 0, - "type": "number" + "anyOf": [ + { + "minimum": 0, + "type": "number" + }, + { + "type": "autograd.tracer.Box" + } + ] }, "type": { "default": "Sphere", diff --git a/schemas/ModeSimulation.json b/schemas/ModeSimulation.json index dda93ce073..0f4d386e11 100644 --- a/schemas/ModeSimulation.json +++ b/schemas/ModeSimulation.json @@ -10984,8 +10984,15 @@ ] }, "radius": { - "minimum": 0, - "type": "number" + "anyOf": [ + { + "minimum": 0, + "type": "number" + }, + { + "type": "autograd.tracer.Box" + } + ] }, "type": { "default": "Sphere", diff --git a/schemas/Simulation.json b/schemas/Simulation.json index 9b86e9efdb..ed7a1af676 100644 --- a/schemas/Simulation.json +++ b/schemas/Simulation.json @@ -15334,8 +15334,15 @@ ] }, "radius": { - "minimum": 0, - "type": "number" + "anyOf": [ + { + "minimum": 0, + "type": "number" + }, + { + "type": "autograd.tracer.Box" + } + ] }, "type": { "default": "Sphere", diff --git a/schemas/TerminalComponentModeler.json b/schemas/TerminalComponentModeler.json index 64e8bfabf9..690fc4024f 100644 --- a/schemas/TerminalComponentModeler.json +++ b/schemas/TerminalComponentModeler.json @@ -16456,8 +16456,15 @@ ] }, "radius": { - "minimum": 0, - "type": "number" + "anyOf": [ + { + "minimum": 0, + "type": "number" + }, + { + "type": "autograd.tracer.Box" + } + ] }, "type": { "default": "Sphere", 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 index b455576d39..1086fa46bd 100644 --- a/tests/test_components/autograd/numerical/test_autograd_sphere_trianglemesh_numerical.py +++ b/tests/test_components/autograd/numerical/test_autograd_sphere_trianglemesh_numerical.py @@ -20,6 +20,7 @@ ) from tests.test_components.autograd.test_autograd_triangle_mesh import subdivide_triangles from tidy3d import config +from tidy3d.components.autograd import get_static from tidy3d.components.geometry.primitives import _base_icosahedron config.local_cache.enabled = True @@ -55,7 +56,10 @@ def make_base_simulation( - radii: list[float], *, extra_structures: Sequence[td.Structure] | None = None + radii: list[float], + *, + extra_structures: Sequence[td.Structure] | None = None, + is_2d: bool = False, ) -> tuple[td.Simulation, callable]: sim_size_3d = [ 2 * radii[0] + 2 * WL_UM, @@ -63,12 +67,23 @@ def make_base_simulation( (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="+", - ) + if is_2d: + sim_size_3d[1] = 0.0 + + source_time = td.GaussianPulse(freq0=FREQ0, fwidth=0.2 * FREQ0) + if is_2d: + primary_source = td.PointDipole( + center=(0.0, 0.0, SRC_OFFSET), + source_time=source_time, + polarization="Ez", + ) + else: + primary_source = td.PlaneWave( + center=(0.0, 0.0, SRC_OFFSET), + size=(*sim_size_3d[:2], 0.0), + source_time=source_time, + direction="+", + ) flux_monitors = [ td.FieldMonitor( @@ -90,7 +105,7 @@ def make_base_simulation( boundary_spec_3d = td.BoundarySpec( x=td.Boundary.pml(), - y=td.Boundary.pml(), + y=td.Boundary.pml() if not is_2d else td.Boundary.periodic(), z=td.Boundary.pml(), ) @@ -98,7 +113,7 @@ def make_base_simulation( center=(0.0, 0.0, 0.0), size=tuple(sim_size_3d), monitors=flux_monitors, - sources=[plane_wave], + sources=[primary_source], structures=list(extra_structures) if extra_structures else [], run_time=2e-11, boundary_spec=boundary_spec_3d, @@ -110,7 +125,14 @@ def make_base_simulation( ) def fom(sim_data): - return sim_data["field"].flux.values + dataset = sim_data["field"] + if is_2d: + ex_vals = dataset.Ex.values + ey_vals = dataset.Ey.values + ez_vals = dataset.Ez.values + intensity = anp.abs(ex_vals) ** 2 + anp.abs(ey_vals) ** 2 + anp.abs(ez_vals) ** 2 + return anp.real(anp.mean(intensity)) + return dataset.flux.values return base_sim, fom @@ -206,6 +228,10 @@ def make_sphere_triangle_geometry( return mesh +def make_native_sphere_geometry(params: anp.ndarray, center: Sequence[float]) -> td.Geometry: + return td.Sphere(center=tuple(center), radius=params[0]) + + def make_mesh_sphere_from_radius(params: anp.ndarray, center: Sequence[float]) -> td.Geometry: radius = params[0] radii = anp.full(3, radius) @@ -322,11 +348,13 @@ def objective(parameters): @pytest.mark.numerical -@pytest.mark.parametrize("scale_factor", (1,)) -@pytest.mark.parametrize("scale_axis", (0,)) +@pytest.mark.parametrize("scale_factor", SCALE_FACTORS) +@pytest.mark.parametrize("scale_axis", SCALE_AXES) @pytest.mark.parametrize("overlap_cube", (False,)) +@pytest.mark.parametrize("parametrization", ("center",)) +@pytest.mark.parametrize("is_2d", (False, True)) def test_sphere_triangles_match_fd( - scale_factor, scale_axis, overlap_cube, tmp_path, numerical_case_dir + scale_factor, scale_axis, overlap_cube, parametrization, is_2d, tmp_path, numerical_case_dir ): """ Compares FD gradients with gradients from _compute_derivatives in TriangleMesh. @@ -336,21 +364,33 @@ def test_sphere_triangles_match_fd( pytest.skip("Skipping duplicate test.") initial_params = [SPHERE_RADIUS_UM, SPHERE_RADIUS_UM, SPHERE_RADIUS_UM] - params0 = anp.array(initial_params) + radius_params = 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) + base_sim, fom = make_base_simulation( + radii=radii, extra_structures=extra_structures, is_2d=is_2d + ) center = [0.0, 0.0, 0.0] + center_params = anp.array(center) + dim_tag = "2d" if is_2d else "3d" - part_make_geom = lambda p, c: make_sphere_triangle_geometry(p, c, scale_factor, scale_axis) + if parametrization == "radius": + params0 = radius_params + part_make_geom = lambda p, c: make_sphere_triangle_geometry(p, c, scale_factor, scale_axis) + else: + params0 = center_params + fixed_radii = anp.array(radii) + part_make_geom = lambda p, _c, fixed_radii=fixed_radii: make_sphere_triangle_geometry( + fixed_radii, p, scale_factor, scale_axis + ) triangle_objective = make_objective( part_make_geom, center, - f"sphere_mesh_{scale_factor}_axis_{scale_axis}_cube_{overlap_cube}", + f"sphere_mesh_{scale_factor}_axis_{scale_axis}_cube_{overlap_cube}_param_{parametrization}_dim_{dim_tag}", base_sim, fom, tmp_path, @@ -359,7 +399,7 @@ def test_sphere_triangles_match_fd( triangle_objective_fd = make_objective( part_make_geom, center, - f"sphere_mesh_fd_{scale_factor}_axis_{scale_axis}_cube_{overlap_cube}", + f"sphere_mesh_fd_{scale_factor}_axis_{scale_axis}_cube_{overlap_cube}_param_{parametrization}_dim_{dim_tag}", base_sim, fom, tmp_path, @@ -373,26 +413,42 @@ def test_sphere_triangles_match_fd( fd_grad = finite_difference_params(triangle_objective_fd, params0, FINITE_DIFF_STEP) - print("scale", scale_factor, "axis", scale_axis, "overlap_cube", overlap_cube) + print( + "scale", + scale_factor, + "axis", + scale_axis, + "overlap_cube", + overlap_cube, + "parametrization", + parametrization, + "is_2d", + is_2d, + ) 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) + grad_angle_deg = angled_overlap_deg(triangle_grad, fd_grad) print( - f"TriangleMesh FD vs. Adjoint angle overlap: {mesh_fd_overlap:.3f}° " + f"TriangleMesh FD vs. Adjoint angle overlap: {grad_angle_deg:.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}° " + assert grad_angle_deg < ANGLE_OVERLAP_FD_ADJ_THRESH_DEG, ( + f"FD–adjoint angle overlap too large: {grad_angle_deg:.3f}° " f"(threshold {ANGLE_OVERLAP_FD_ADJ_THRESH_DEG}°, " ) if SAVE_OUTPUT_DATA: + dim_tag = "2d" if is_2d else "3d" np.savez( numerical_case_dir - / f"sphere_gradients_mesh_scale_{scale_factor}_axis_{scale_axis}_cube_{overlap_cube}.npz", + / ( + f"sphere_gradients_mesh_scale_{scale_factor}_axis_{scale_axis}_cube_{overlap_cube}" + f"_param_{parametrization}_dim_{dim_tag}.npz" + ), triangle_grad=triangle_grad, fd_grad=fd_grad, + grad_angle_deg=np.array([grad_angle_deg], dtype=float), ) @@ -488,7 +544,7 @@ def test_grad_insensitive_to_face_splitting(tmp_path, numerical_case_dir): ) -@pytest.mark.skip +@pytest.mark.numerical @pytest.mark.parametrize("scale_factor", SCALE_FACTORS) @pytest.mark.parametrize("scale_axis", SCALE_AXES) @pytest.mark.parametrize("overlap_cube", (False, True)) @@ -574,3 +630,299 @@ def test_triangle_sphere_fd_step_sweep_ref( gradients=fd_grads, autograd_grad=autograd_grad, ) + + +@pytest.mark.numerical +@pytest.mark.parametrize("radius_scale", (0.25, 0.5, 1)) +@pytest.mark.parametrize("overlap_cube", (False, True)) +@pytest.mark.parametrize("parametrization", ("radius", "center")) +@pytest.mark.parametrize("is_2d", (False,), ids=lambda x: "2D" if x else "3D") +def test_native_sphere_match_fd( + radius_scale, overlap_cube, parametrization, is_2d, tmp_path, numerical_case_dir +): + """ + Compares FD gradients with gradients from _compute_derivatives in Sphere. + Note that FD gradients are very noise which is why there will be some failing tests with a fixed FD-step. + Currently, numerical tests fail for 2D as FD gradients are very shaky. + """ + radius = SPHERE_RADIUS_UM * radius_scale + radii = [radius, radius, radius] + extra_structures = ( + [make_overlap_cube_structure([radius, radius / 2, radius / 2])] if overlap_cube else [] + ) + base_sim, fom = make_base_simulation( + radii=radii, extra_structures=extra_structures, is_2d=is_2d + ) + + center = [0.0, 0.0, 0.0] + center_params = anp.array(center) + + if parametrization == "radius": + params0 = anp.array([radius]) + geometry_factory = make_native_sphere_geometry + objective_suffix = "radius" + else: + params0 = center_params + + def geometry_factory(params, _unused_center, radius_fixed=radius): + return td.Sphere(center=tuple(params), radius=radius_fixed) + + objective_suffix = "center" + + native_objective = make_objective( + geometry_factory, + center, + f"native_sphere_scale_{radius_scale}_cube_{overlap_cube}_param_{objective_suffix}_dim_{'2d' if is_2d else '3d'}", + base_sim, + fom, + tmp_path, + local_gradient=LOCAL_GRADIENT, + ) + native_objective_fd = make_objective( + geometry_factory, + center, + f"native_sphere_fd_scale_{radius_scale}_cube_{overlap_cube}_param_{objective_suffix}_dim_{'2d' if is_2d else '3d'}", + base_sim, + fom, + tmp_path, + local_gradient=False, + ) + + _, native_grad = value_and_grad(native_objective)([params0]) + native_grad = np.squeeze(np.asarray(native_grad, dtype=float)) + + fd_grad = finite_difference_params(native_objective_fd, params0, FINITE_DIFF_STEP_NATIVE) + + print( + "native radius scale", + radius_scale, + "overlap_cube", + overlap_cube, + "parametrization", + parametrization, + "is_2d", + is_2d, + ) + print("native_grad\t", native_grad.tolist()) + print("fd_grad\t\t", fd_grad.tolist()) + + if parametrization == "radius": + abs_diff = float(np.abs(native_grad - fd_grad)) + rel_err = abs_diff / max(np.abs(native_grad), np.abs(fd_grad), 1e-12) + print( + f"Native sphere FD vs. Adjoint absolute diff: {abs_diff:.3e}, " + f"relative error: {float(get_static(rel_err)):.3e}" + ) + assert rel_err < 1e-1, ( + f"Native sphere gradients mismatch: abs_diff={abs_diff:.3e}, " + f"rel_err={float(get_static(rel_err)):.3e}, native_grad={native_grad.tolist()}, " + f"fd_grad={fd_grad.tolist()}" + ) + grad_metric = rel_err + else: + grad_angle_deg = angled_overlap_deg(native_grad, fd_grad) + print( + f"Native sphere FD vs. Adjoint angle overlap: {grad_angle_deg:.3f}° " + f"(threshold = {ANGLE_OVERLAP_FD_ADJ_THRESH_DEG}°)", + ) + assert grad_angle_deg < ANGLE_OVERLAP_FD_ADJ_THRESH_DEG, ( + f"FD–adjoint angle overlap too large: {grad_angle_deg:.3f}° " + f"(threshold {ANGLE_OVERLAP_FD_ADJ_THRESH_DEG}°)" + ) + grad_metric = grad_angle_deg + + if SAVE_OUTPUT_DATA: + np.savez( + numerical_case_dir + / ( + f"native_sphere_gradients_scale_{radius_scale}_cube_{overlap_cube}" + f"_param_{parametrization}_dim_{'2d' if is_2d else '3d'}.npz" + ), + native_grad=native_grad, + fd_grad=fd_grad, + grad_metric=np.array([grad_metric], dtype=float), + ) + + +@pytest.mark.numerical +@pytest.mark.parametrize( + "radius_factor", + ( + 0.25, + 0.5, + 1, + ), +) +def test_sphere_cylinder_grads_match_2d(radius_factor, tmp_path, numerical_case_dir): + """Ensure 2D sphere gradients equal those from an equivalent Cylinder cross section.""" + radius = radius_factor * SPHERE_RADIUS_UM + params0 = anp.array([radius, 0.0, 0.0, 0.0]) + radii = [radius, radius, radius] + base_sim, fom = make_base_simulation(radii=radii, extra_structures=None, is_2d=True) + + def make_param_sphere_geometry(params, _center_unused): + rad = params[0] + center = tuple(params[1:4]) + return td.Sphere(center=center, radius=rad) + + def make_param_cylinder_geometry(params, _center_unused): + rad = params[0] + center_x, center_y, center_z = params[1:4] + plane_value = 0.0 + rad_plane_sq = rad**2 - (center_y - plane_value) ** 2 + rad_plane = anp.sqrt(anp.maximum(rad_plane_sq, 1e-15)) + cyl_center = (center_x, plane_value, center_z) + return td.Cylinder(center=cyl_center, radius=rad_plane, length=1.0, axis=1) + + sphere_objective = make_objective( + make_param_sphere_geometry, + [0.0, 0.0, 0.0], + "sphere_parametric_2d", + base_sim, + fom, + tmp_path, + local_gradient=LOCAL_GRADIENT, + ) + cylinder_objective = make_objective( + make_param_cylinder_geometry, + [0.0, 0.0, 0.0], + "cylinder_parametric_2d", + base_sim, + fom, + tmp_path, + local_gradient=LOCAL_GRADIENT, + ) + + sphere_val, sphere_grad = value_and_grad(sphere_objective)([params0]) + cylinder_val, cylinder_grad = value_and_grad(cylinder_objective)([params0]) + sphere_grad = np.squeeze(np.asarray(sphere_grad, dtype=float)) + cylinder_grad = np.squeeze(np.asarray(cylinder_grad, dtype=float)) + + center_indices = [1, 3] + center_angle = angled_overlap_deg(sphere_grad[center_indices], cylinder_grad[center_indices]) + + print("sphere_val\t", sphere_val.tolist()) + print("cylinder_val\t", cylinder_val.tolist()) + + print("sphere_grad\t", sphere_grad.tolist()) + print("cylinder_grad\t", cylinder_grad.tolist()) + print(f"Sphere vs Cylinder center gradient angle: {center_angle:.6f}°") + np.testing.assert_allclose( + sphere_grad[center_indices], cylinder_grad[center_indices], rtol=8e-2, atol=5e-3 + ) + + radius_rel_err = abs(sphere_grad[0] - cylinder_grad[0]) / max(abs(cylinder_grad[0]), 1e-12) + assert radius_rel_err < 0.08, ( + f"Sphere vs cylinder radius gradient mismatch: " + f"{sphere_grad[0]} vs {cylinder_grad[0]} (relative error {radius_rel_err:.3f})" + ) + + if SAVE_OUTPUT_DATA: + np.savez( + numerical_case_dir / "sphere_vs_cylinder_gradients_2d.npz", + sphere_grad=sphere_grad, + cylinder_grad=cylinder_grad, + center_angle_deg=np.array([center_angle], dtype=float), + ) + + +@pytest.mark.skip +@pytest.mark.parametrize("radius_scale", (0.5, 1, 2)) +@pytest.mark.parametrize("overlap_cube", (True, False)) +@pytest.mark.parametrize("is_2d", (False, True)) +def test_native_sphere_fd_step_sweep_ref( + tmp_path, radius_scale, overlap_cube, is_2d, numerical_case_dir +): + """FD step sweep for native sphere with autograd reference.""" + radius = SPHERE_RADIUS_UM * radius_scale + params0 = anp.array([radius]) + + radii = [radius, radius, radius] + extra_structures = [make_overlap_cube_structure(radii)] if overlap_cube else [] + base_sim, fom = make_base_simulation( + radii=radii, extra_structures=extra_structures, is_2d=is_2d + ) + + center = [0.0, 0.0, 0.0] + + # FD objective + native_objective_fd = make_objective( + make_native_sphere_geometry, + center, + f"native_sphere_fd_step_sweep_{radius_scale}_cube_{overlap_cube}_dim_{'2d' if is_2d else '3d'}", + base_sim, + fom, + tmp_path, + local_gradient=False, + ) + + # Autograd objective (reference) + native_objective_autograd = make_objective( + make_native_sphere_geometry, + center, + f"native_sphere_autograd_ref_{radius_scale}_cube_{overlap_cube}_dim_{'2d' if is_2d else '3d'}", + base_sim, + fom, + tmp_path, + local_gradient=True, + ) + + # ---- Autograd gradient ---- + _, autograd_grad = value_and_grad(native_objective_autograd)([params0]) + autograd_grad = float(np.squeeze(np.asarray(autograd_grad, dtype=float))) + print( + f"native autograd gradient (radius_scale={radius_scale}, overlap_cube={overlap_cube}, " + f"is_2d={is_2d}): {autograd_grad}" + ) + + # ---- Finite-difference sweep ---- + min_log = -4 + max_log = -1 + n = (max_log - min_log + 1) * 2 + 1 + steps = np.logspace(min_log, max_log, num=n) + + fd_grads = finite_difference_params_step_batch(native_objective_fd, params0, steps) + fd_grads = np.asarray(fd_grads, dtype=float) + + for step, grad in zip(steps, fd_grads): + print( + f"native finite difference step {step:.1e}: gradient {grad.tolist()} " + f"cube={overlap_cube} radius_scale={radius_scale} is_2d={is_2d}" + ) + + # ---- Plot ---- + fig, ax = plt.subplots(figsize=(6, 4)) + # FD curve + ax.plot(steps, fd_grads[:, 0], marker="o", label="radius (FD)") + # Autograd reference line in same color + ax.axhline( + autograd_grad, + color=ax.get_lines()[-1].get_color(), + linestyle="--", + alpha=0.7, + label="radius (autograd)", + ) + + ax.set_xscale("log") + ax.set_xlabel("Finite difference step (µm)") + ax.set_ylabel("Gradient value") + ax.set_title( + f"Native sphere FD vs autograd (radius_scale={radius_scale}, overlap_cube={overlap_cube})" + ) + ax.grid(True, which="both", ls=":") + ax.legend() + + fig_path = ( + numerical_case_dir + / f"rad_{radius_scale}_cube_{overlap_cube}_dim_{'2d' if is_2d else '3d'}.png" + ) + fig.savefig(fig_path, dpi=200) + plt.close(fig) + + np.savez( + numerical_case_dir + / f"rad_{radius_scale}_cube_{overlap_cube}_dim_{'2d' if is_2d else '3d'}.npz", + steps=steps, + gradients=fd_grads, + autograd_grad=np.array([autograd_grad], dtype=float), + ) diff --git a/tests/test_components/autograd/test_autograd.py b/tests/test_components/autograd/test_autograd.py index b56a6345cb..191a6c1a24 100644 --- a/tests/test_components/autograd/test_autograd.py +++ b/tests/test_components/autograd/test_autograd.py @@ -28,6 +28,7 @@ from tidy3d.components.autograd.utils import get_static, is_tidy_box from tidy3d.components.base import TRACED_FIELD_KEYS_ATTR from tidy3d.components.data.data_array import DataArray +from tidy3d.components.geometry.primitives import discretization_wavelength from tidy3d.config import config from tidy3d.exceptions import AdjointError from tidy3d.plugins.polyslab import ComplexPolySlab @@ -558,6 +559,13 @@ def make_structures(params: anp.ndarray) -> dict[str, td.Structure]: medium=td.Medium(permittivity=mesh_eps), ) + # use first 4 params for radius/center, rest for eps + sphere_geom = td.Sphere(radius=params[0], center=params[1:4]) + sphere = td.Structure( + geometry=sphere_geom, + medium=td.Medium(permittivity=anp.mean(params[4:])), + ) + return { "medium": medium, "center_list": center_list, @@ -573,6 +581,7 @@ def make_structures(params: anp.ndarray) -> dict[str, td.Structure]: "custom_pole_res": custom_pole_res, "cylinder": cylinder, "triangle_mesh": triangle_mesh, + "sphere": sphere, } @@ -1957,10 +1966,8 @@ def test_cylinder_discretization(eps_real): with AssertLogLevel( "WARNING", contains_str="The minimum wavelength inside the cylinder material" ): - cylinder = td.Cylinder(axis=2, length=info.wavelength_min, radius=2 * info.wavelength_min) - expected_wvl_mat = info.wavelength_min * config.adjoint.min_wvl_fraction - wvl_mat = cylinder._discretization_wavelength(derivative_info=info) + wvl_mat = discretization_wavelength(info, "cylinder") assert np.isclose(expected_wvl_mat, wvl_mat), ( "Unexpected wavelength for discretizing cylinder!" @@ -3202,3 +3209,33 @@ def record_method(self, derivative_info): assert object.__getattribute__(big_box, "recorded_bounds_intersect") == group.bounds, ( f"got {object.__getattribute__(big_box, 'recorded_bounds_intersect')} and {group.bounds}" ) + + +@pytest.mark.parametrize("monitor_key", ("mode",)) +def test_autograd_sphere_0_radius(use_emulated_run, monitor_key): + """Integration test that Sphere gradients are non-zero (mirrors cylinder check).""" + + monitor, postprocess = make_monitors()[monitor_key] + + def make_sphere(radius, x0, y0, z0): + return td.Sphere(center=(x0, y0, z0), radius=radius) + + def make_sim(params): + geometry = make_sphere(*params) + structure = td.Structure(geometry=geometry, medium=td.Medium(permittivity=2)) + return SIM_BASE.updated_copy(structures=[structure], monitors=[monitor]) + + p0 = [0.0, 0.0, 0.0, 0.0] + + def objective(params): + sim = make_sim(params) + if PLOT_SIM: + plot_sim(sim, plot_eps=True) + data = run(sim, task_name="autograd_test", verbose=False) + return anp.sum(anp.abs(data[monitor.name].amps)).item() + + with AssertLogLevel("WARNING", contains_str="cannot be computed"): + val_sphere, grad_sphere = ag.value_and_grad(objective)(p0) + # first 4 parameters are related to the geometry + geom_grad = np.asarray(get_static(grad_sphere[:4]), dtype=float) + assert np.allclose(geom_grad, 0.0) diff --git a/tidy3d/components/geometry/primitives.py b/tidy3d/components/geometry/primitives.py index f6916b535d..3fb7473c81 100644 --- a/tidy3d/components/geometry/primitives.py +++ b/tidy3d/components/geometry/primitives.py @@ -12,7 +12,7 @@ from pydantic.v1 import PrivateAttr from shapely.geometry.base import BaseGeometry -from tidy3d.components.autograd import AutogradFieldMap, TracedSize1D +from tidy3d.components.autograd import AutogradFieldMap, TracedSize1D, get_static 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 @@ -90,6 +90,25 @@ def _base_icosahedron() -> tuple[np.ndarray, np.ndarray]: _ICOSAHEDRON_VERTS, _ICOSAHEDRON_FACES = _base_icosahedron() +def discretization_wavelength(derivative_info: DerivativeInfo, geometry_label: str) -> float: + """Choose reference wavelength for surface discretization.""" + wvl0_min = derivative_info.wavelength_min + wvl_mat = wvl0_min / np.max([1.0, np.max(np.sqrt(abs(derivative_info.eps_in)))]) + + grid_cfg = config.adjoint + + min_wvl_mat = grid_cfg.min_wvl_fraction * wvl0_min + if wvl_mat < min_wvl_mat: + log.warning( + f"The minimum wavelength inside the {geometry_label} material is {wvl_mat:.3e} μm, which would " + f"create a large number of discretization points for computing the gradient. " + f"To prevent performance degradation, the discretization wavelength has " + f"been clipped to {min_wvl_mat:.3e} μm.", + log_once=True, + ) + return max(wvl_mat, min_wvl_mat) + + class Sphere(base.Centered, base.Circular): """Spherical geometry. @@ -98,8 +117,42 @@ class Sphere(base.Centered, base.Circular): >>> b = Sphere(center=(1,2,3), radius=2) """ + radius: TracedSize1D = pydantic.Field( + ..., + title="Radius", + description="Radius of geometry.", + units=MICROMETER, + ) + _icosphere_cache: dict[int, tuple[np.ndarray, float]] = PrivateAttr(default_factory=dict) + @verify_packages_import(["trimesh"]) + def to_triangle_mesh( + self, + *, + max_edge_length: Optional[float] = None, + subdivisions: Optional[int] = None, + ) -> TriangleMesh: + """Approximate the sphere surface with a ``TriangleMesh``. + + Parameters + ---------- + max_edge_length : float = None + Maximum edge length for triangulation in micrometers. + subdivisions : int = None + Number of subdivisions for icosphere generation. + + Returns + ------- + TriangleMesh + Triangle mesh approximation of the sphere surface. + """ + + triangles, _ = self._triangulated_surface( + max_edge_length=max_edge_length, subdivisions=subdivisions + ) + return TriangleMesh.from_triangles(triangles) + def inside( self, x: np.ndarray[float], y: np.ndarray[float], z: np.ndarray[float] ) -> np.ndarray[bool]: @@ -278,6 +331,225 @@ def unit_sphere_triangles( ) return unit_tris + def _compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradFieldMap: + """Compute adjoint derivatives using smooth sphere surface samples.""" + valid_paths = {("radius",), *{("center", i) for i in range(3)}} + for path in derivative_info.paths: + if path not in valid_paths: + raise ValueError( + f"No derivative defined w.r.t. 'Sphere' field '{path}'. " + "Supported fields are 'radius' and 'center'." + ) + + if not derivative_info.paths: + return {} + + grid_cfg = config.adjoint + radius = float(get_static(self.radius)) + if radius == 0.0: + log.warning( + "Sphere gradients cannot be computed for zero radius; gradients are zero.", + log_once=True, + ) + return dict.fromkeys(derivative_info.paths, 0.0) + + wvl_mat = discretization_wavelength(derivative_info, "sphere") + target_edge = max(wvl_mat / grid_cfg.points_per_wavelength, np.finfo(float).eps) + triangles, _ = self._triangulated_surface(max_edge_length=target_edge) + triangles = triangles.astype(grid_cfg.gradient_dtype_float, copy=False) + + sim_min, sim_max = ( + np.asarray(arr, dtype=grid_cfg.gradient_dtype_float) + for arr in derivative_info.simulation_bounds + ) + tol = config.adjoint.edge_clip_tolerance + + sim_extents = sim_max - sim_min + collapsed_indices = np.flatnonzero(np.isclose(sim_extents, 0.0, atol=tol)) + if collapsed_indices.size: + if collapsed_indices.size > 1: + return dict.fromkeys(derivative_info.paths, 0.0) + axis_idx = int(collapsed_indices[0]) + plane_value = float(sim_min[axis_idx]) + return self._compute_derivatives_collapsed_axis( + derivative_info=derivative_info, + axis_idx=axis_idx, + plane_value=plane_value, + ) + + trimesh_obj = TriangleMesh._triangles_to_trimesh(triangles) + vertices = np.asarray(trimesh_obj.vertices, dtype=grid_cfg.gradient_dtype_float) + center = np.asarray(self.center, dtype=grid_cfg.gradient_dtype_float) + verts_centered = vertices - center + norms = np.linalg.norm(verts_centered, axis=1, keepdims=True) + norms = np.where(norms == 0, 1, norms) + normals = verts_centered / norms + + if vertices.size == 0: + return dict.fromkeys(derivative_info.paths, 0.0) + + # get vertex weights + faces = np.asarray(trimesh_obj.faces, dtype=int) + face_areas = np.asarray(trimesh_obj.area_faces, dtype=grid_cfg.gradient_dtype_float) + weights = np.zeros(len(vertices), dtype=grid_cfg.gradient_dtype_float) + np.add.at(weights, faces[:, 0], face_areas / 3.0) + np.add.at(weights, faces[:, 1], face_areas / 3.0) + np.add.at(weights, faces[:, 2], face_areas / 3.0) + + perp1, perp2 = self._tangent_basis_from_normals(normals) + + valid_axes = np.abs(sim_max - sim_min) > tol + inside_mask = np.all( + vertices[:, valid_axes] >= (sim_min - tol)[valid_axes], axis=1 + ) & np.all(vertices[:, valid_axes] <= (sim_max + tol)[valid_axes], axis=1) + + if not np.any(inside_mask): + return dict.fromkeys(derivative_info.paths, 0.0) + + points = vertices[inside_mask] + normals_sel = normals[inside_mask] + perp1_sel = perp1[inside_mask] + perp2_sel = perp2[inside_mask] + weights_sel = weights[inside_mask] + + interpolators = derivative_info.interpolators + if interpolators is None: + interpolators = derivative_info.create_interpolators( + dtype=grid_cfg.gradient_dtype_float + ) + + g = derivative_info.evaluate_gradient_at_points( + points, + normals_sel, + perp1_sel, + perp2_sel, + interpolators, + ) + + weighted = (weights_sel * g).real + grad_center = np.sum(weighted[:, None] * normals_sel, axis=0) + grad_radius = np.sum(weighted) + + vjps: AutogradFieldMap = {} + for path in derivative_info.paths: + if path == ("radius",): + vjps[path] = float(grad_radius) + else: + _, idx = path + vjps[path] = float(grad_center[idx]) + + return vjps + + def _compute_derivatives_collapsed_axis( + self, + derivative_info: DerivativeInfo, + axis_idx: int, + plane_value: float, + ) -> AutogradFieldMap: + """Delegate collapsed-axis gradients to a Cylinder cross section.""" + tol = config.adjoint.edge_clip_tolerance + radius = float(self.radius) + center = np.asarray(self.center, dtype=float) + delta = plane_value - center[axis_idx] + radius_sq = radius**2 - delta**2 + if radius_sq <= tol**2: + return dict.fromkeys(derivative_info.paths, 0.0) + + radius_plane = float(np.sqrt(max(radius_sq, 0.0))) + if radius_plane <= tol: + return dict.fromkeys(derivative_info.paths, 0.0) + + cyl_paths: set[tuple[str, int | None]] = set() + need_radius = False + for path in derivative_info.paths: + if path == ("radius",) or path == ("center", axis_idx): + cyl_paths.add(("radius",)) + need_radius = True + elif path[0] == "center" and path[1] != axis_idx: + cyl_paths.add(("center", path[1])) + + if not cyl_paths: + return dict.fromkeys(derivative_info.paths, 0.0) + + cyl_center = center.copy() + cyl_center[axis_idx] = plane_value + cylinder = Cylinder( + center=tuple(cyl_center), + radius=radius_plane, + length=discretization_wavelength(derivative_info, "sphere") * 2.0, + axis=axis_idx, + ) + + bounds_min = list(cyl_center) + bounds_max = list(cyl_center) + for dim in range(3): + if dim == axis_idx: + continue + bounds_min[dim] = center[dim] - radius_plane + bounds_max[dim] = center[dim] + radius_plane + + bounds = (tuple(bounds_min), tuple(bounds_max)) + sim_min_arr, sim_max_arr = ( + np.asarray(arr, dtype=float) for arr in derivative_info.simulation_bounds + ) + intersect_min = tuple(max(bounds[0][i], sim_min_arr[i]) for i in range(3)) + intersect_max = tuple(min(bounds[1][i], sim_max_arr[i]) for i in range(3)) + if any(lo > hi for lo, hi in zip(intersect_min, intersect_max)): + return dict.fromkeys(derivative_info.paths, 0.0) + + derivative_info_cyl = derivative_info.updated_copy( + paths=list(cyl_paths), + bounds=bounds, + bounds_intersect=(intersect_min, intersect_max), + ) + + vjps_cyl = cylinder._compute_derivatives(derivative_info_cyl) + result = dict.fromkeys(derivative_info.paths, 0.0) + vjp_radius = float(vjps_cyl.get(("radius",), 0.0)) if need_radius else 0.0 + + for path in derivative_info.paths: + if path == ("radius",): + result[path] = vjp_radius * (radius / radius_plane) + elif path == ("center", axis_idx): + result[path] = vjp_radius * (delta / radius_plane) + elif path[0] == "center" and path[1] != axis_idx: + result[path] = float(vjps_cyl.get(("center", path[1]), 0.0)) + + return result + + def _edge_length_on_unit_sphere( + self, max_edge_length: Optional[float] = _DEFAULT_EDGE_FRACTION + ) -> Optional[float]: + """Convert ``max_edge_length`` in μm to unit-sphere coordinates.""" + radius = float(self.radius) + if radius <= 0.0: + return None + return max_edge_length / radius + + def _triangulated_surface( + self, + *, + max_edge_length: Optional[float] = None, + subdivisions: Optional[int] = None, + ) -> tuple[np.ndarray, np.ndarray]: + """Return physical and unit triangles for the surface discretization. Pass either max_edge_length or subdivisions.""" + max_edge_length_unit = None + if subdivisions is None: + max_edge_length_unit = self._edge_length_on_unit_sphere(max_edge_length) + + unit_tris = self._unit_sphere_triangles( + target_edge_length=max_edge_length_unit, + subdivisions=subdivisions, + copy_result=False, + ) + + radius = float(get_static(self.radius)) + center = np.asarray(self.center, dtype=float) + dtype = config.adjoint.gradient_dtype_float + + physical = radius * unit_tris + center + return physical.astype(dtype, copy=False), unit_tris.astype(dtype, copy=False) + def _unit_sphere_triangles( self, *, @@ -285,7 +557,7 @@ def _unit_sphere_triangles( subdivisions: Optional[int] = None, copy_result: bool = True, ) -> np.ndarray: - """Return cached unit-sphere triangles with optional copying.""" + """Return cached unit-sphere triangles with optional copying. Pass either target_edge_length or subdivisions.""" if target_edge_length is not None and subdivisions is not None: raise ValueError("Specify either target_edge_length OR subdivisions, not both.") @@ -312,6 +584,32 @@ def _subdivisions_for_edge(self, target_edge_length: Optional[float]) -> int: ) return _MAX_ICOSPHERE_SUBDIVISIONS + @staticmethod + def _tangent_basis_from_normals(normals: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Construct orthonormal tangential bases for each normal vector (vectorized).""" + + dtype = normals.dtype + tol = np.finfo(dtype).eps + + # Normalize normals (in case they are not perfectly unit length). + n_norm = np.linalg.norm(normals, axis=1) + n = normals / np.maximum(n_norm, tol)[:, None] + + # Pick a reference axis least aligned with each normal: argmin(|nx|,|ny|,|nz|). + ref_idx = np.argmin(np.abs(n), axis=1) + ref = np.zeros_like(n) + ref[np.arange(n.shape[0]), ref_idx] = 1.0 + + basis1 = np.cross(n, ref) + b1_norm = np.linalg.norm(basis1, axis=1) + basis1 = basis1 / np.maximum(b1_norm, tol)[:, None] + + basis2 = np.cross(n, basis1) + b2_norm = np.linalg.norm(basis2, axis=1) + basis2 = basis2 / np.maximum(b2_norm, tol)[:, None] + + return basis1, basis2 + def _icosphere_data(self, subdivisions: int) -> tuple[np.ndarray, float]: cache = self._icosphere_cache if subdivisions in cache: @@ -454,31 +752,11 @@ def _points_unit_circle( ys = np.sin(angles) return np.stack((xs, ys), axis=0) - def _discretization_wavelength(self, derivative_info: DerivativeInfo) -> float: - """Choose a reference wavelength for discretizing the cylinder into a `PolySlab`.""" - wvl0_min = derivative_info.wavelength_min - wvl_mat = wvl0_min / np.max([1.0, np.max(np.sqrt(abs(derivative_info.eps_in)))]) - - grid_cfg = config.adjoint - - min_wvl_mat = grid_cfg.min_wvl_fraction * wvl0_min - if wvl_mat < min_wvl_mat: - log.warning( - f"The minimum wavelength inside the cylinder material is {wvl_mat:.3e} μm, which would " - f"create a large number of discretization points for computing the gradient. " - f"To prevent performance degradation, the discretization wavelength has " - f"been clipped to {min_wvl_mat:.3e} μm.", - log_once=True, - ) - wvl_mat = max(wvl_mat, min_wvl_mat) - - return wvl_mat - def _compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradFieldMap: """Compute the adjoint derivatives for this object.""" # compute circumference discretization - wvl_mat = self._discretization_wavelength(derivative_info=derivative_info) + wvl_mat = discretization_wavelength(derivative_info, "cylinder") circumference = 2 * np.pi * self.radius wvls_in_circumference = circumference / wvl_mat