feat(adjoint): Add conductivity gradient for CustomMedium

Introduces support for computing gradients with respect to the `conductivity` field in `CustomMedium`.

To achieve this, the derivative computation logic was generalized:
- The `_derivative_field_cmp` method is refactored to accept a `component` parameter ('real', 'imag', 'complex') to compute the VJP for different parts of the complex permittivity.
- For the 'imag' component, the derivative is scaled by `1 / (omega * epsilon_0)` to convert from derivative w.r.t. complex permittivity to derivative w.r.t. conductivity.
- The `_compute_derivatives` dispatcher now handles the new `"conductivity"` parameter path.
- A new autograd test is added to validate the conductivity gradient calculation on a `CustomMedium` with constant permittivity.

Author: yaugenst-flex

CHANGELOG.md

@@ -18,11 +18,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `InternalAbsorber` class for placing first-order absorbing boundary conditions on planes inside the simulation domain. Internal absorbers are automatically wrapped in a PEC frame with a backing PEC plate on the non-absorbing side.
 - Added `absorber` field (default: `True`) to `WavePort` for automatically placing an absorber behind the port.
 - Added `conjugated_dot_product` field in `ModeMonitor` (default: `True`) and `WavePort` (default: `False`) to allow selecting the conjugated or non-conjugated dot product for mode decomposition.
+- Support for gradients with respect to the `conductivity` of a `CustomMedium`.
 
 ### Changed
 - Validate mode solver object for large number of grid points on the modal plane.
 - Adaptive minimum spacing for `PolySlab` integration is now wavelength relative and a minimum discretization is set for computing gradients for cylinders.
 - The `TerminalComponentModeler` defaults to the pseudo wave definition of scattering parameters. The new field `s_param_def` can be used to switch between either pseudo or power wave definitions.
+- Add support for `np.unwrap` in `tidy3d.plugins.autograd`.
 
 ### Fixed
 - Fixed missing amplitude factor and handling of negative normal direction case when making adjoint sources from `DiffractionMonitor`.

tests/test_components/test_autograd.py

@@ -319,30 +319,42 @@ def make_structures(params: anp.ndarray) -> dict[str, td.Structure]:
     eps_arr = 1.01 + 0.5 * (anp.tanh(matrix @ params).reshape(DA_SHAPE) + 1)
 
     nx, ny, nz = eps_arr.shape
+    da_coords = {
+        "x": np.linspace(-0.5, 0.5, nx),
+        "y": np.linspace(-0.5, 0.5, ny),
+        "z": np.linspace(-0.5, 0.5, nz),
+    }
 
     custom_med = td.Structure(
         geometry=box,
         medium=td.CustomMedium(
             permittivity=td.SpatialDataArray(
                 eps_arr,
-                coords={
-                    "x": np.linspace(-0.5, 0.5, nx),
-                    "y": np.linspace(-0.5, 0.5, ny),
-                    "z": np.linspace(-0.5, 0.5, nz),
-                },
+                coords=da_coords,
+            ),
+        ),
+    )
+
+    # custom medium with variable permittivity and conductivity data
+    conductivity_arr = 0.01 * (anp.tanh(matrix @ params).reshape(DA_SHAPE) + 1)
+    custom_med_with_conductivity = td.Structure(
+        geometry=box,
+        medium=td.CustomMedium(
+            permittivity=td.SpatialDataArray(
+                eps_arr,
+                coords=da_coords,
+            ),
+            conductivity=td.SpatialDataArray(
+                conductivity_arr,
+                coords=da_coords,
             ),
         ),
     )
 
     # custom medium with vector valued permittivity data
     eps_ii = td.ScalarFieldDataArray(
         eps_arr.reshape(nx, ny, nz, 1),
-        coords={
-            "x": np.linspace(-0.5, 0.5, nx),
-            "y": np.linspace(-0.5, 0.5, ny),
-            "z": np.linspace(-0.5, 0.5, nz),
-            "f": [td.C_0],
-        },
+        coords=da_coords | {"f": [td.C_0]},
     )
 
     custom_med_vec = td.Structure(
@@ -484,6 +496,7 @@ def make_structures(params: anp.ndarray) -> dict[str, td.Structure]:
         "center_list": center_list,
         "size_element": size_element,
         "custom_med": custom_med,
+        "custom_med_with_conductivity": custom_med_with_conductivity,
         "custom_med_vec": custom_med_vec,
         "polyslab": polyslab,
         "polyslab_dispersive": polyslab_dispersive,
@@ -581,6 +594,7 @@ def plot_sim(sim: td.Simulation, plot_eps: bool = True) -> None:
     "center_list",
     "size_element",
     "custom_med",
+    "custom_med_with_conductivity",
     "custom_med_vec",
     "polyslab",
     "complex_polyslab",
@@ -1738,7 +1752,7 @@ def J(eps):
     monkeypatch.setattr(
         td.CustomPoleResidue,
         "_derivative_field_cmp",
-        lambda self, E_der_map, eps_data, dim: dJ_deps / 3.0,
+        lambda self, E_der_map, spatial_data, dim, freqs, component="real": dJ_deps / 3.0,
     )
 
     import importlib
@@ -2434,3 +2448,52 @@ def objective(x):
 
     with pytest.raises(ValueError):
         g = ag.grad(objective)(1.0)
+
+
+def test_custom_medium_conductivity_only_gradient(rng, use_emulated_run, tmp_path):
+    """Test conductivity gradients for CustomMedium with constant permittivity."""
+
+    monitor, postprocess = make_monitors()["field_point"]
+
+    def objective(params):
+        """Objective function testing only conductivity gradient (constant permittivity)."""
+        len_arr = np.prod(DA_SHAPE)
+        matrix = rng.random((len_arr, N_PARAMS))
+
+        # constant permittivity
+        eps_arr = np.ones(DA_SHAPE) * 2.0
+
+        # variable conductivity
+        conductivity_arr = 0.05 * (anp.tanh(3 * matrix @ params).reshape(DA_SHAPE) + 1)
+
+        nx, ny, nz = DA_SHAPE
+        coords = {
+            "x": np.linspace(-0.5, 0.5, nx),
+            "y": np.linspace(-0.5, 0.5, ny),
+            "z": np.linspace(-0.5, 0.5, nz),
+        }
+
+        custom_med_struct = td.Structure(
+            geometry=td.Box(center=(0, 0, 0), size=(1, 1, 1)),
+            medium=td.CustomMedium(
+                permittivity=td.SpatialDataArray(eps_arr, coords=coords),
+                conductivity=td.SpatialDataArray(conductivity_arr, coords=coords),
+            ),
+        )
+
+        sim = SIM_BASE.updated_copy(
+            structures=[custom_med_struct],
+            monitors=[monitor],
+        )
+
+        data = run(
+            sim,
+            path=str(tmp_path / "sim_test.hdf5"),
+            task_name="conductivity_only_grad_test",
+            verbose=False,
+        )
+        return postprocess(data, data[monitor.name])
+
+    val, grad = ag.value_and_grad(objective)(params0)
+
+    assert anp.all(grad != 0.0), "some gradients are 0 for conductivity-only test"

tidy3d/components/medium.py

@@ -1682,13 +1682,15 @@ def _not_loaded(field):
     def _derivative_field_cmp(
         self,
         E_der_map: ElectromagneticFieldDataset,
-        eps_data: PermittivityDataset,
+        spatial_data: PermittivityDataset,
         dim: str,
     ) -> np.ndarray:
-        coords_interp = {key: val for key, val in eps_data.coords.items() if len(val) > 1}
-        dims_sum = {dim for dim in eps_data.coords.keys() if dim not in coords_interp}
+        coords_interp = {key: val for key, val in spatial_data.coords.items() if len(val) > 1}
+        dims_sum = {dim for dim in spatial_data.coords.keys() if dim not in coords_interp}
 
-        eps_coordinate_shape = [len(eps_data.coords[dim]) for dim in eps_data.dims if dim in "xyz"]
+        eps_coordinate_shape = [
+            len(spatial_data.coords[dim]) for dim in spatial_data.dims if dim in "xyz"
+        ]
 
         # compute sizes along each of the interpolation dimensions
         sizes_list = []
@@ -2898,13 +2900,27 @@ def _compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradField
         vjps = {}
 
         for field_path in derivative_info.paths:
-            if field_path == ("permittivity",):
+            if field_path[0] == "permittivity":
                 vjp_array = 0.0
                 for dim in "xyz":
                     vjp_array += self._derivative_field_cmp(
                         E_der_map=derivative_info.E_der_map,
-                        eps_data=self.permittivity,
+                        spatial_data=self.permittivity,
                         dim=dim,
+                        freqs=np.atleast_1d(derivative_info.frequencies),
+                        component="real",
+                    )
+                vjps[field_path] = vjp_array
+
+            elif field_path[0] == "conductivity":
+                vjp_array = 0.0
+                for dim in "xyz":
+                    vjp_array += self._derivative_field_cmp(
+                        E_der_map=derivative_info.E_der_map,
+                        spatial_data=self.conductivity,
+                        dim=dim,
+                        freqs=np.atleast_1d(derivative_info.frequencies),
+                        component="imag",
                     )
                 vjps[field_path] = vjp_array
 
@@ -2913,10 +2929,11 @@ def _compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradField
                 dim = key[-1]
                 vjps[field_path] = self._derivative_field_cmp(
                     E_der_map=derivative_info.E_der_map,
-                    eps_data=self.eps_dataset.field_components[key],
+                    spatial_data=self.eps_dataset.field_components[key],
                     dim=dim,
+                    freqs=np.atleast_1d(derivative_info.frequencies),
+                    component="complex",
                 )
-
             else:
                 raise NotImplementedError(
                     f"No derivative defined for 'CustomMedium' field: {field_path}."
@@ -2927,14 +2944,18 @@ def _compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradField
     def _derivative_field_cmp(
         self,
         E_der_map: ElectromagneticFieldDataset,
-        eps_data: PermittivityDataset,
+        spatial_data: CustomSpatialDataTypeAnnotated,
         dim: str,
+        freqs: np.ndarray,
+        component: str = "real",
     ) -> np.ndarray:
-        """Compute derivative with respect to the ``dim`` components within the custom medium."""
-        coords_interp = {key: eps_data.coords[key] for key in "xyz"}
+        """Compute the derivative with respect to a material property component."""
+        coords_interp = {key: spatial_data.coords[key] for key in "xyz"}
         coords_interp = {key: val for key, val in coords_interp.items() if len(val) > 1}
 
-        eps_coordinate_shape = [len(eps_data.coords[dim]) for dim in eps_data.dims if dim in "xyz"]
+        eps_coordinate_shape = [
+            len(spatial_data.coords[dim]) for dim in spatial_data.dims if dim in "xyz"
+        ]
 
         E_der_dim_interp = E_der_map[f"E{dim}"]
 
@@ -2972,10 +2993,26 @@ def _derivative_field_cmp(
             # if sizes_list is empty, then reduce() fails
             d_vol = np.array(1.0)
 
-        # TODO: probably this could be more robust. eg if the DataArray has weird edge cases
-        E_der_dim_interp = (
-            E_der_dim_interp.interp(**coords_interp, assume_sorted=True).fillna(0.0).real.sum("f")
-        )
+        E_der_dim_interp_complex = E_der_dim_interp.interp(
+            **coords_interp, assume_sorted=True
+        ).fillna(0.0)
+
+        if component == "imag":
+            # convert from derivative w.r.t. complex permittivity to derivative w.r.t. conductivity
+            E_der_dim_interp = E_der_dim_interp_complex.imag
+            # frequency-dependent scaling must be applied before summing over frequencies
+            for freq in freqs:
+                vjp_imag = Medium.eps_sigma_to_eps_complex(
+                    eps_real=0, sigma=E_der_dim_interp.sel(f=freq), freq=freq
+                ).imag
+                E_der_dim_interp.loc[{"f": freq}] = -vjp_imag
+        elif component == "complex":
+            # for complex permittivity in eps_dataset, return the full complex derivative
+            E_der_dim_interp = E_der_dim_interp_complex
+        else:
+            E_der_dim_interp = E_der_dim_interp_complex.real
+
+        E_der_dim_interp = E_der_dim_interp.sum("f")
 
         try:
             E_der_dim_interp = E_der_dim_interp * d_vol.reshape(E_der_dim_interp.shape)
@@ -3975,15 +4012,39 @@ def _sel_custom_data_inside(self, bounds: Bound):
 
         return self.updated_copy(eps_inf=eps_inf_reduced, poles=poles_reduced)
 
+    def _derivative_field_cmp(
+        self,
+        E_der_map: ElectromagneticFieldDataset,
+        spatial_data: CustomSpatialDataTypeAnnotated,
+        dim: str,
+        freqs=None,
+        component: str = "complex",
+    ) -> np.ndarray:
+        """Compatibility wrapper for derivative computation.
+
+        Accepts the extended signature used by other custom media (
+        e.g., `CustomMedium._derivative_field_cmp`) while delegating the actual
+        computation to the base implementation that only depends on
+        `E_der_map`, `spatial_data`, and `dim`.
+
+        Parameters `freqs` and `component` are ignored for this model since the
+        derivative is taken with respect to the complex permittivity directly.
+        """
+        return super()._derivative_field_cmp(
+            E_der_map=E_der_map, spatial_data=spatial_data, dim=dim
+        )
+
     def _compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradFieldMap:
         """Compute adjoint derivatives by preparing array data and calling the static helper."""
 
         dJ_deps_complex = 0.0
         for dim in "xyz":
             dJ_deps_complex += self._derivative_field_cmp(
                 E_der_map=derivative_info.E_der_map,
-                eps_data=self.eps_inf,
+                spatial_data=self.eps_inf,
                 dim=dim,
+                freqs=np.atleast_1d(derivative_info.frequencies),
+                component="complex",
             )
 
         poles_vals = [