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

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 - Add support for `np.unwrap` in `tidy3d.plugins.autograd`.
+- Support for gradients with respect to the `conductivity` of a `CustomMedium`.
 
 ### Fixed
 - Arrow lengths are now scaled consistently in the X and Y directions, and their lengths no longer exceed the height of the plot window.
@@ -100,8 +101,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 - Fixed `reverse` property of `td.Scene.plot_structures_property()` to also reverse the colorbar.
-
-### Fixed
 - Fixed bug in surface gradient computation where fields, instead of gradients, were being summed in frequency.
 
 ## [2.8.2] - 2025-04-09

tests/test_components/test_autograd.py

@@ -315,30 +315,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(
@@ -480,6 +492,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,
@@ -577,6 +590,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",
@@ -1661,7 +1675,7 @@ def J(eps):
     monkeypatch.setattr(
         td.CustomPoleResidue,
         "_derivative_field_cmp",
-        lambda self, E_der_map, eps_data, dim, freqs: dJ_deps,
+        lambda self, E_der_map, spatial_data, dim, freqs, component="real": dJ_deps,
     )
 
     import importlib
@@ -2351,3 +2365,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 gonductivity 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

@@ -1663,14 +1663,16 @@ def _not_loaded(field):
     def _derivative_field_cmp(
         self,
         E_der_map: ElectromagneticFieldDataset,
-        eps_data: PermittivityDataset,
+        spatial_data: PermittivityDataset,
         dim: str,
         freqs: NDArray,
     ) -> 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 = []
@@ -2835,14 +2837,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.frequency),
+                        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.frequency),
+                        component="imag",
                     )
                 vjps[field_path] = vjp_array
 
@@ -2851,11 +2866,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.frequency),
+                    component="complex",
                 )
-
             else:
                 raise NotImplementedError(
                     f"No derivative defined for 'CustomMedium' field: {field_path}."
@@ -2866,15 +2881,18 @@ def _compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradField
     def _derivative_field_cmp(
         self,
         E_der_map: ElectromagneticFieldDataset,
-        eps_data: PermittivityDataset,
+        spatial_data: PermittivityDataset,
         dim: str,
         freqs: 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}"].sel(f=freqs)
 
@@ -2912,10 +2930,25 @@ 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
+            # dJ/d(sigma) = Im[dJ/d(eps_complex)] / (omega * epsilon_0)
+            omegas = 2 * np.pi * freqs
+            E_der_dim_interp = E_der_dim_interp_complex.imag
+            for i, omega in enumerate(omegas):
+                E_der_dim_interp.loc[{"f": freqs[i]}] = E_der_dim_interp.sel(f=freqs[i]) / (
+                    omega * EPSILON_0
+                )
+            E_der_dim_interp = E_der_dim_interp.sum("f")
+        elif component == "complex":
+            # for complex permittivity in eps_dataset, return the full complex derivative
+            E_der_dim_interp = E_der_dim_interp_complex.sum("f")
+        else:
+            E_der_dim_interp = E_der_dim_interp_complex.real.sum("f")
 
         try:
             E_der_dim_interp = E_der_dim_interp * d_vol.reshape(E_der_dim_interp.shape)
@@ -3909,7 +3942,7 @@ def _compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradField
         for dim in "xyz":
             dJ_deps += 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.frequency),
             )