fix(adjoint): fix frequency selection in adjoint postprocessing frequency batching to make sure all data aligns in frequency dimension before passing to geometry or medium for derivative computation

Author: Gregory Roberts

CHANGELOG.md

@@ -7,16 +7,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
-### Added
-
-### Changed
-
-### Fixed
-- Giving opposite boundaries different names no longer causes a symmetry validator failure.
-- Fixed issue with parameters in `InverseDesignResult` sometimes being outside of the valid parameter range.
-- Disallow `EMEFieldMonitor` in EME simulations with `EMELengthSweep`.
-
-
 ## [2.9.0] - 2025-08-04
 
 ### Added
@@ -94,6 +84,10 @@ with fewer layers than recommended.
 - Giving opposite boundaries different names no longer causes a symmetry validator failure.
 - Fixed issue with parameters in `InverseDesignResult` sometimes being outside of the valid parameter range.
 - Fixed performance regression for multi-frequency adjoint calculations.
+- Disallow `EMEFieldMonitor` in EME simulations with `EMELengthSweep`.
+- Fixed bug in adjoint postprocessing frequency batching that was causing gradients to be zero or incorrect. The error was surfacing when selecting a subset of the monitor frequencies in the objective function.
+
+
 
 ## [2.8.5] - 2025-07-07
 

tests/test_components/test_autograd_mode_polyslab_numerical.py

@@ -359,14 +359,34 @@ def test_finite_difference_mode_data_polyslab(
     # due to a multifrequency objective function.
     monitor_top_weights = rng.random(NUM_MODE_MONITOR_FREQUENCIES)
     monitor_bottom_weights = rng.random(NUM_MODE_MONITOR_FREQUENCIES)
+    frequency_selection_mask = np.arange(0, NUM_MODE_MONITOR_FREQUENCIES)
+
+    # sometimes, test what happens when we only use one of the frequencies from the mode monitors
+    # to catch handling of different frequencies being present in the forward and adjoint monitors
+    if rng.random() > 0.5:
+        frequency_selection_mask = rng.integers(1, NUM_MODE_MONITOR_FREQUENCIES)
+        monitor_top_weights = monitor_top_weights[frequency_selection_mask]
+        monitor_bottom_weights = monitor_bottom_weights[frequency_selection_mask]
 
     def eval_fn(sim_data):
         return np.sum(
             monitor_top_weights
-            * np.abs(sim_data["monitor_mode_top"].amps.sel(direction="+").values) ** 2
+            * np.abs(
+                sim_data["monitor_mode_top"]
+                .amps.sel(direction="+")
+                .isel(f=frequency_selection_mask)
+                .data
+            )
+            ** 2
         ) + np.sum(
             monitor_bottom_weights
-            * np.abs(sim_data["monitor_mode_bottom"].amps.sel(direction="+").values) ** 2
+            * np.abs(
+                sim_data["monitor_mode_bottom"]
+                .amps.sel(direction="+")
+                .isel(f=frequency_selection_mask)
+                .data
+            )
+            ** 2
         )
 
     polyslab_height_um = POLYSLAB_HEIGHT_WVL * adj_wvl_um

tidy3d/web/api/autograd/autograd.py

@@ -989,22 +989,25 @@ def _compute_eps_array(medium, frequencies):
     return DataArray(data=np.array(eps_data), dims=("f",), coords={"f": frequencies})
 
 
-def _slice_field_data(field_data: dict, freq_slice: slice) -> dict:
+def _slice_field_data(
+    field_data: dict,
+    freqs: np.ndarray,
+) -> dict:
     """Slice field data dictionary along frequency dimension.
 
     Parameters
     ----------
     field_data : dict
         Dictionary of field components.
-    freq_slice : slice
-        Frequency slice to apply.
+    freqs : np.ndarray
+        Frequencies to select.
 
     Returns
     -------
     dict
         Sliced field data dictionary.
     """
-    return {k: v.isel(f=freq_slice) for k, v in field_data.items()}
+    return {k: v.sel(f=freqs) for k, v in field_data.items()}
 
 
 def postprocess_adj(
@@ -1119,26 +1122,32 @@ def postprocess_adj(
             chunk_end = min(chunk_start + freq_chunk_size, n_freqs)
             freq_slice = slice(chunk_start, chunk_end)
 
+            select_adjoint_freqs = adjoint_frequencies[freq_slice]
+
             # slice field data for current chunk
-            E_der_map_chunk = _slice_field_data(E_der_map.field_components, freq_slice)
-            D_der_map_chunk = _slice_field_data(D_der_map.field_components, freq_slice)
-            E_fwd_chunk = _slice_field_data(E_fwd.field_components, freq_slice)
-            E_adj_chunk = _slice_field_data(E_adj.field_components, freq_slice)
-            D_fwd_chunk = _slice_field_data(D_fwd.field_components, freq_slice)
-            D_adj_chunk = _slice_field_data(D_adj.field_components, freq_slice)
-            eps_data_chunk = _slice_field_data(eps_fwd.field_components, freq_slice)
+            E_der_map_chunk = _slice_field_data(E_der_map.field_components, select_adjoint_freqs)
+            D_der_map_chunk = _slice_field_data(D_der_map.field_components, select_adjoint_freqs)
+            E_fwd_chunk = _slice_field_data(E_fwd.field_components, select_adjoint_freqs)
+            E_adj_chunk = _slice_field_data(E_adj.field_components, select_adjoint_freqs)
+            D_fwd_chunk = _slice_field_data(D_fwd.field_components, select_adjoint_freqs)
+            D_adj_chunk = _slice_field_data(D_adj.field_components, select_adjoint_freqs)
+            eps_data_chunk = _slice_field_data(eps_fwd.field_components, select_adjoint_freqs)
 
             # slice epsilon arrays
-            eps_in_chunk = eps_in.isel(f=freq_slice)
-            eps_out_chunk = eps_out.isel(f=freq_slice)
+            eps_in_chunk = eps_in.sel(f=select_adjoint_freqs)
+            eps_out_chunk = eps_out.sel(f=select_adjoint_freqs)
             eps_background_chunk = (
-                eps_background.isel(f=freq_slice) if eps_background is not None else None
+                eps_background.sel(f=select_adjoint_freqs) if eps_background is not None else None
             )
             eps_no_structure_chunk = (
-                eps_no_structure.isel(f=freq_slice) if eps_no_structure is not None else None
+                eps_no_structure.sel(f=select_adjoint_freqs)
+                if eps_no_structure is not None
+                else None
             )
             eps_inf_structure_chunk = (
-                eps_inf_structure.isel(f=freq_slice) if eps_inf_structure is not None else None
+                eps_inf_structure.sel(f=select_adjoint_freqs)
+                if eps_inf_structure is not None
+                else None
             )
 
             # create derivative info with sliced data