diff --git a/CHANGELOG.md b/CHANGELOG.md index 52b5a6debe..5b772e5017 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - Added priority attribute to `TopologyDesignRegion` to enable manual control of overlapping structures. - `to_mat_file()` method is now available on `ModeSimulationData` and `HeatChargeSimulationData` for exporting results to MATLAB `.mat` files. +- Added autograd support for diagonal `AnisotropicMedium` and `CustomAnisotropicMedium` with diagonal permittivity tensor. ### Changed diff --git a/tests/test_components/autograd/numerical/test_autograd_medium_numerical.py b/tests/test_components/autograd/numerical/test_autograd_medium_numerical.py index df866bcf77..4fb504af6f 100644 --- a/tests/test_components/autograd/numerical/test_autograd_medium_numerical.py +++ b/tests/test_components/autograd/numerical/test_autograd_medium_numerical.py @@ -58,6 +58,60 @@ "polarization": np.pi / 3, "medium_type": "custom", }, + { + "name": "opt_flux_aniso", + "wavelength": 1.0, + "permittivities": (2.1, 2.4, 2.7), + "objective_kind": "flux", + "monitor_size": (np.inf, np.inf, 0.0), + "polarization": 0.0, + "medium_type": "anisotropic", + }, + { + "name": "opt_int_point_aniso", + "wavelength": 1.55, + "permittivities": (1.6, 2.0, 2.5), + "objective_kind": "intensity", + "monitor_size": (0.0, 0.0, 0.0), + "polarization": np.pi / 2, + "medium_type": "anisotropic", + }, + { + "name": "mw_flux_aniso", + "wavelength": 0.8, + "permittivities": (2.2, 2.4, 1.5), + "objective_kind": "flux", + "monitor_size": (np.inf, np.inf, 0.0), + "polarization": np.pi / 8, + "medium_type": "anisotropic", + }, + { + "name": "mw_int_plane_aniso", + "wavelength": 2.1, + "permittivities": (2.6, 2.0, 3.1), + "objective_kind": "intensity", + "monitor_size": (0.2, 0.2, 0.0), + "polarization": np.pi / 4, + "medium_type": "anisotropic", + }, + { + "name": "opt_flux_custom_aniso", + "wavelength": 1.2, + "permittivities": (1.8, 2.3, 2.9), + "objective_kind": "flux", + "monitor_size": (np.inf, np.inf, 0.0), + "polarization": 0.0, + "medium_type": "custom_anisotropic", + }, + { + "name": "mw_int_custom_aniso", + "wavelength": 1.9, + "permittivities": (1.4, 2.0, 1.7), + "objective_kind": "intensity", + "monitor_size": (0.5, 0.5, 0.0), + "polarization": np.pi / 6, + "medium_type": "custom_anisotropic", + }, ] @@ -163,7 +217,7 @@ def _custom_isotropic(val): def _metric_value(case, dataset, freq0): if case["objective_kind"] == "flux": - return dataset.flux.values + return dataset.flux.values.item() ex_vals = dataset.Ex.values ey_vals = dataset.Ey.values ez_vals = dataset.Ez.values diff --git a/tests/test_components/autograd/test_autograd_anisotropic_medium.py b/tests/test_components/autograd/test_autograd_anisotropic_medium.py new file mode 100644 index 0000000000..929fd3b73c --- /dev/null +++ b/tests/test_components/autograd/test_autograd_anisotropic_medium.py @@ -0,0 +1,175 @@ +from __future__ import annotations + +import numpy as np + +import tidy3d as td +from tidy3d.components.autograd.derivative_utils import DerivativeInfo +from tidy3d.components.medium import AnisotropicMedium +from tidy3d.constants import EPSILON_0 + + +def _scalar_field(value: complex, coords: dict[str, list[float]]): + """Helper to build a ``ScalarFieldDataArray`` with a single sample.""" + + data = np.array([[[[value]]]], dtype=np.complex128) + return td.ScalarFieldDataArray(data, coords=coords) + + +def _derivative_info(paths: list[tuple[str, ...]], freq: float): + """Construct a ``DerivativeInfo`` with synthetic field products.""" + + coords = {"x": [0.0], "y": [0.0], "z": [0.0], "f": [freq]} + ex_val = 1.2 + 0.4j + ey_val = -0.8 + 0.1j + ez_val = 0.5 - 0.3j + + field_map = { + "Ex": _scalar_field(ex_val, coords), + "Ey": _scalar_field(ey_val, coords), + "Ez": _scalar_field(ez_val, coords), + } + + eps_no = _scalar_field(1.0, coords) + eps_inf = _scalar_field(2.0, coords) + + return DerivativeInfo( + paths=paths, + E_der_map=field_map, + D_der_map={}, + E_fwd={}, + D_fwd={}, + E_adj={}, + D_adj={}, + eps_data={}, + eps_in=2.0, + eps_out=1.0, + frequencies=np.array([freq]), + bounds=((0.0, 0.0, 0.0), (1.0, 1.0, 1.0)), + bounds_intersect=((0.0, 0.0, 0.0), (1.0, 1.0, 1.0)), + simulation_bounds=((0.0, 0.0, 0.0), (1.0, 1.0, 1.0)), + eps_no_structure=eps_no, + eps_inf_structure=eps_inf, + ) + + +def test_anisotropic_medium_axis_gradients_uniform(): + """Gradients respect axis-specific contributions for diagonal anisotropy.""" + + freq = 2.0e14 + paths = [("xx", "permittivity"), ("yy", "conductivity"), ("zz", "permittivity")] + info = _derivative_info(paths=paths, freq=freq) + + medium = td.AnisotropicMedium( + xx=td.Medium(permittivity=2.0), + yy=td.Medium(permittivity=2.5), + zz=td.Medium(permittivity=3.0), + ) + + grads = medium._compute_derivatives(info) + + ex_real = np.real(info.E_der_map["Ex"].values).item() + ey_imag = np.imag(info.E_der_map["Ey"].values).item() + ez_real = np.real(info.E_der_map["Ez"].values).item() + + assert np.isclose(grads[("xx", "permittivity")], ex_real) + assert np.isclose(grads[("zz", "permittivity")], ez_real) + + expected_sigma = -ey_imag / (2.0 * np.pi * freq * EPSILON_0) + assert np.isclose(grads[("yy", "conductivity")], expected_sigma) + + +def _linear_variation_coords(shape: tuple[int, int, int]): + coords = { + "x": np.linspace(-1.0, 1.0, shape[0]), + "y": np.linspace(-1.0, 1.0, shape[1]), + "z": np.linspace(-1.0, 1.0, shape[2]), + } + X, Y, Z = np.meshgrid(coords["x"], coords["y"], coords["z"], indexing="ij") + factors = 1 + 0.2 * (X + Y + Z) / 3.0 + return coords, factors + + +def test_custom_anisotropic_medium_axis_gradients(): + """Custom anisotropic medium delegates gradients per axis with spatial data.""" + + freq = 1.5e14 + paths = [("xx", "permittivity"), ("yy", "permittivity"), ("zz", "permittivity")] + info = _derivative_info(paths=paths, freq=freq) + + coords, factors = _linear_variation_coords((3, 4, 5)) + + def _make_custom_medium(base_val): + values = base_val * factors + data = td.SpatialDataArray(values, coords=coords) + return td.CustomMedium(permittivity=data) + + medium = td.CustomAnisotropicMedium( + xx=_make_custom_medium(1.8), + yy=_make_custom_medium(2.0), + zz=_make_custom_medium(2.2), + ) + + grads = medium._compute_derivatives(info) + + for comp_name in ("xx", "yy", "zz"): + comp_info = AnisotropicMedium._component_derivative_info(info, comp_name) + comp_grad = medium.components[comp_name]._compute_derivatives(comp_info) + np.testing.assert_allclose(grads[(comp_name, "permittivity")], comp_grad[("permittivity",)]) + + +def test_anisotropic_medium_conductivity_uses_projected_d_map(): + """Conductivity gradients keep only the axis D component.""" + + freq = 2.0e14 + paths = [("yy", "conductivity")] + + coords = {"x": [0.0], "y": [0.0], "z": [0.0], "f": [freq]} + e_map = { + "Ex": _scalar_field(1.0 + 0.0j, coords), + "Ey": _scalar_field(2.0 - 0.5j, coords), + "Ez": _scalar_field(3.0 + 0.0j, coords), + } + d_map = { + "Ex": _scalar_field(1.1 + 0.0j, coords), + "Ey": _scalar_field(2.2 + 0.0j, coords), + "Ez": _scalar_field(3.3 + 0.0j, coords), + } + + eps_no = _scalar_field(1.0, coords) + eps_inf = _scalar_field(2.0, coords) + + info = DerivativeInfo( + paths=paths, + E_der_map=e_map, + D_der_map=d_map, + E_fwd={}, + D_fwd={}, + E_adj={}, + D_adj={}, + eps_data={}, + eps_in=2.0, + eps_out=1.0, + frequencies=np.array([freq]), + bounds=((0.0, 0.0, 0.0), (1.0, 1.0, 1.0)), + bounds_intersect=((0.0, 0.0, 0.0), (1.0, 1.0, 1.0)), + simulation_bounds=((0.0, 0.0, 0.0), (1.0, 1.0, 1.0)), + eps_no_structure=eps_no, + eps_inf_structure=eps_inf, + ) + + medium = td.AnisotropicMedium( + xx=td.Medium(permittivity=2.0), + yy=td.Medium(permittivity=2.5, conductivity=0.1), + zz=td.Medium(permittivity=3.0), + ) + + comp_info = AnisotropicMedium._component_derivative_info(info, "yy") + + assert np.allclose(comp_info.D_der_map["Ex"].values, 0.0) + assert np.allclose(comp_info.D_der_map["Ez"].values, 0.0) + assert np.allclose(comp_info.D_der_map["Ey"].values, d_map["Ey"].values) + + grads = medium.components["yy"]._compute_derivatives(comp_info) + expected_sigma = -np.imag(e_map["Ey"].values).item() / (2.0 * np.pi * freq * EPSILON_0) + + assert np.isclose(grads[("conductivity",)], expected_sigma) diff --git a/tidy3d/components/autograd/derivative_utils.py b/tidy3d/components/autograd/derivative_utils.py index 7c36444687..fe76fde923 100644 --- a/tidy3d/components/autograd/derivative_utils.py +++ b/tidy3d/components/autograd/derivative_utils.py @@ -9,7 +9,8 @@ import xarray as xr from tidy3d.components.data.data_array import FreqDataArray, ScalarFieldDataArray -from tidy3d.components.types import ArrayLike, Bound, tidycomplex +from tidy3d.components.data.utils import _zeros_like +from tidy3d.components.types import ArrayLike, Bound, tidycomplex, xyz from tidy3d.config import config from tidy3d.constants import C_0, EPSILON_0, LARGE_NUMBER, MU_0 from tidy3d.log import log @@ -714,6 +715,42 @@ def _project_in_basis( field_matrix = np.transpose(field_matrix, (1, 0, 2)) return np.einsum("ij...,ij->i...", field_matrix, basis_vector) + def project_der_map_to_axis( + self, axis: xyz, field_type: str = "E" + ) -> dict[str, ScalarFieldDataArray] | None: + """Return a copy of the selected derivative map with only one axis kept. + + Parameters + ---------- + axis: + Axis to keep (``"x"``, ``"y"``, ``"z"``, case-insensitive). + field_type: + Map selector: ``"E"`` (``self.E_der_map``) or ``"D"`` (``self.D_der_map``). + + Returns + ------- + dict[str, ScalarFieldDataArray] | None + Copied map where non-selected components are replaced by zeros, or ``None`` + if the requested map is unavailable. + """ + field_map = {"E": self.E_der_map, "D": self.D_der_map}.get(field_type) + if field_map is None: + raise ValueError("field type must be 'D' or 'E'.") + + axis = axis.lower() + projected = dict(field_map) + if not field_map: + return projected + for dim in "xyz": + key = f"E{dim}" + if key not in field_map: + continue + if dim != axis: + projected[key] = _zeros_like(field_map[key]) + else: + projected[key] = field_map[key] + return projected + def adaptive_vjp_spacing( self, wl_fraction: Optional[float] = None, diff --git a/tidy3d/components/medium.py b/tidy3d/components/medium.py index dd44197d91..1e5d75e74c 100644 --- a/tidy3d/components/medium.py +++ b/tidy3d/components/medium.py @@ -6118,6 +6118,49 @@ def sel_inside(self, bounds: Bound): return self.updated_copy(**dict(zip(["xx", "yy", "zz"], new_comps))) + # --- shared autograd helpers --- + @staticmethod + def _component_derivative_info( + derivative_info: DerivativeInfo, component: str + ) -> DerivativeInfo | None: + """Build ``DerivativeInfo`` filtered to a single anisotropic component.""" + + component_paths = [ + tuple(path[1:]) for path in derivative_info.paths if path and path[0] == component + ] + if not component_paths: + return None + + axis = component[0] # f.e. xx -> x + projected_E = derivative_info.project_der_map_to_axis(axis, "E") + projected_D = derivative_info.project_der_map_to_axis(axis, "D") + return derivative_info.updated_copy( + paths=component_paths, E_der_map=projected_E, D_der_map=projected_D + ) + + def _compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradFieldMap: + """Delegate derivatives for each diagonal component of an anisotropic medium.""" + + components = self.components + for field_path in derivative_info.paths: + if len(field_path) < 2 or field_path[0] not in components: + raise NotImplementedError( + f"No derivative defined for '{type(self).__name__}' field: {field_path}." + ) + + vjps: AutogradFieldMap = {} + for comp_name, component in components.items(): + comp_info = self._component_derivative_info( + derivative_info=derivative_info, component=comp_name + ) + if comp_info is None: + continue + comp_vjps = component._compute_derivatives(comp_info) + for sub_path, value in comp_vjps.items(): + vjps[(comp_name, *sub_path)] = value + + return vjps + class AnisotropicMediumFromMedium2D(AnisotropicMedium): """The same as ``AnisotropicMedium``, but converted from Medium2D.