feat(tidy3d): FXC-3296-autograd-support-for-anisotropic-medium-and-custom-anisotropic-medium

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
 

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

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)

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,

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.