feature[adjoint]: adjust adjoint source fwidth to decay before zero frequency when possible

Author: Gregory Roberts

CHANGELOG.md

@@ -23,6 +23,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Internal adjoint helper methods are now prefixed with an underscore to separate them from the public API.
 - Drop the dependency on `gdspy`, which has been unmaintained for over two years. Interfaces previously relying on `gdspy` now use its maintained successor, `gdstk`, with equivalent functionality.
 - Small (around 1e-4) numerical precision improvements in EME solver.
+- Adjoint source frequency width is adjusted to decay sufficiently before zero frequency when possible to improve accuracy of simulation normalization when using custom current sources.
 
 ## [2.8.4] - 2025-05-15
 

tests/test_components/test_autograd.py

@@ -1225,6 +1225,147 @@ def objective(args):
         ag.grad(objective)(params0)
 
 
+def test_adjoint_src_width():
+    """Test the adjoint source width for single sources decays by f=0."""
+
+    f0 = td.C_0 / 1.55
+    fwidth = f0
+
+    fwidths = f0 * np.linspace(0.1, 1.0, 5)
+
+    adj_srcs = [
+        td.PointDipole(
+            center=(0, 0, 0),
+            source_time=td.GaussianPulse(freq0=f0, fwidth=fwidth),
+            polarization="Ex",
+        )
+        for fwidth in fwidths
+    ]
+
+    adj_srcs_fwidth = td.SimulationData._adjoint_src_width_single(adj_srcs)
+
+    for src in adj_srcs_fwidth:
+        assert np.isclose(
+            (src.source_time.freq0 - f0) / f0, 0.0
+        ), "f0 of adjoint source should be centered on original f0"
+
+        check_fwidth = (
+            src.source_time.freq0
+            - td.components.data.sim_data.NUM_ADJOINT_FWIDTH_TO_ZERO * src.source_time.fwidth
+        ) / src.source_time.freq0
+
+        assert np.isclose(check_fwidth, 0.0) or (
+            check_fwidth > 0.0
+        ), "fwidth of adjoint source should decay sufficiently before f=0"
+
+
+def test_broadband_adjoint_src_width():
+    """Test the broadband adjoint source handling for choosing fwidth."""
+
+    # Test the case where we have a custom current source and a wide adjoint source width that overlaps with zero.
+    # In this case, we want to issue a warning to the user about the adjoint accuracy of this setup.
+    f0_high = td.C_0 / 1.55
+    f0_low = 0.1 * f0_high
+
+    f0_adj_all = [f0_low, f0_high]
+
+    fwidth = 0.1 * f0_high
+
+    adj_srcs = []
+    x = np.array([0.0])
+    y = np.array([0.0])
+    z = np.array([0.0])
+    for f0 in f0_adj_all:
+        f = np.array([f0])
+
+        coords = dict(x=x, y=y, z=z, f=f)
+
+        dataset = td.FieldDataset(Ex=td.ScalarFieldDataArray(np.ones((1, 1, 1, 1)), coords=coords))
+
+        adj_srcs.append(
+            td.CustomCurrentSource(
+                center=(0, 0, 0),
+                size=(0, 0, 0),
+                source_time=td.GaussianPulse(freq0=f0, fwidth=fwidth),
+                current_dataset=dataset,
+            )
+        )
+
+    EXPECTED_WARNING_MSG_PIECE = (
+        "Adjoint source generated with a frequency spectrum that extends to or overlaps with 0 Hz"
+    )
+    with AssertLogLevel("WARNING", contains_str=EXPECTED_WARNING_MSG_PIECE):
+        broadband_f0, broadband_fwidth = td.SimulationData._adjoint_src_width_broadband(adj_srcs)
+
+        f0_expected = 0.5 * (np.max(f0_adj_all) + np.min(f0_adj_all))
+
+        fwidth_expected = (
+            f0_expected - np.min(f0_adj_all)
+        ) / td.components.data.sim_data.NUM_ADJOINT_FWIDTH_TO_FMIN
+
+        assert np.isclose(
+            (f0_expected - broadband_f0) / f0_expected, 0.0
+        ), "Expected freq0 not matching for broadband source"
+        assert np.isclose(
+            (fwidth_expected - broadband_fwidth) / fwidth_expected, 0.0
+        ), "Expected fwidth not matching for broadband source"
+
+    # Test the case where we need a wider pulse to cover all the adjoint frequencies than we would otherwise choose for
+    # each individual adjoint source
+    f0_broadband = np.linspace(f0_low, f0_high, 10)
+    fwidth_broadband = 0.1 * np.mean(f0_broadband)
+
+    adj_srcs = [
+        td.PointDipole(
+            center=(0, 0, 0),
+            source_time=td.GaussianPulse(freq0=f0, fwidth=fwidth_broadband),
+            polarization="Ex",
+        )
+        for f0 in f0_broadband
+    ]
+
+    broadband_f0, broadband_fwidth = td.SimulationData._adjoint_src_width_broadband(adj_srcs)
+
+    f0_expected = 0.5 * (np.max(f0_broadband) + np.min(f0_broadband))
+    fwidth_expected = (
+        f0_expected - np.min(f0_broadband)
+    ) / td.components.data.sim_data.NUM_ADJOINT_FWIDTH_TO_FMIN
+
+    assert np.isclose(
+        (f0_expected - broadband_f0) / f0_expected, 0.0
+    ), "Expected freq0 not matching for broadband source"
+    assert np.isclose(
+        (fwidth_expected - broadband_fwidth) / fwidth_expected, 0.0
+    ), "Expected fwidth not matching for broadband source"
+
+    # Test the case where we have a narrow set of frequencies for the adjoint sources and so we can
+    # choose a wider overall source than is needed for covering those frequencies. This larger pulse width
+    # in frequency will shorten the time pulse.
+    f0_broadband = np.linspace(0.95 * f0_high, 1.05 * f0_high, 10)
+    fwidth_broadband = 0.1 * np.mean(f0_broadband)
+
+    adj_srcs = [
+        td.PointDipole(
+            center=(0, 0, 0),
+            source_time=td.GaussianPulse(freq0=f0, fwidth=fwidth_broadband),
+            polarization="Ex",
+        )
+        for f0 in f0_broadband
+    ]
+
+    broadband_f0, broadband_fwidth = td.SimulationData._adjoint_src_width_broadband(adj_srcs)
+
+    f0_expected = 0.5 * (np.max(f0_broadband) + np.min(f0_broadband))
+    fwidth_expected = f0_expected / td.components.data.sim_data.NUM_ADJOINT_FWIDTH_TO_ZERO
+
+    assert np.isclose(
+        (f0_expected - broadband_f0) / f0_expected, 0.0
+    ), "Expected freq0 not matching for broadband source"
+    assert np.isclose(
+        (fwidth_expected - broadband_fwidth) / fwidth_expected, 0.0
+    ), "Expected fwidth not matching for broadband source"
+
+
 @pytest.mark.parametrize("colocate", [True, False])
 @pytest.mark.parametrize("objtype", ["flux", "intensity"])
 def test_interp_objectives(use_emulated_run, colocate, objtype):

tidy3d/components/data/sim_data.py

@@ -22,6 +22,7 @@
 from ..file_util import replace_values
 from ..monitor import Monitor
 from ..simulation import Simulation
+from ..source.current import CustomCurrentSource
 from ..source.time import GaussianPulse
 from ..source.utils import SourceType
 from ..structure import Structure
@@ -43,6 +44,12 @@
 # residuals below this are considered good fits for broadband adjoint source creation
 RESIDUAL_CUTOFF_ADJOINT = 1e-6
 
+# for adjoint source, the minimum number of FWIDTH between the center frequency and zero
+NUM_ADJOINT_FWIDTH_TO_ZERO = 3
+# for broadband adjoint source, the minimum number of FWIDTH to reach the lowest frequency
+# that is covered by the broadband pulse
+NUM_ADJOINT_FWIDTH_TO_FMIN = 0.5
+
 
 class AdjointSourceInfo(Tidy3dBaseModel):
     """Stores information about the adjoint sources to pass to autograd pipeline."""
@@ -1123,31 +1130,49 @@ def _fwidth_adj(self) -> float:
         normalize_index_fwd = self.simulation.normalize_index or 0
         return self.simulation.sources[normalize_index_fwd].source_time.fwidth
 
+    @staticmethod
+    def _adjoint_src_width_single(adj_srcs: list[SourceType]) -> list[SourceType]:
+        """Ensure the adjoint source sufficiently decays before zero frequency."""
+        adj_srcs_process_fwidth = []
+        for adj_src in adj_srcs:
+            source_time = adj_src.source_time
+            freq0 = source_time.freq0
+
+            fwidth = np.minimum(freq0 / NUM_ADJOINT_FWIDTH_TO_ZERO, source_time.fwidth)
+
+            adj_srcs_process_fwidth.append(
+                adj_src.updated_copy(source_time=source_time.updated_copy(fwidth=fwidth))
+            )
+
+        return adj_srcs_process_fwidth
+
     def _process_adjoint_sources(self, adj_srcs: list[SourceType]) -> list[AdjointSourceInfo]:
         """Compute list of final sources along with a post run normalization for adj fields."""
         # dictionary mapping hash of sources with same freq dependence to list of time-dependencies
         hashes_to_sources = defaultdict(None)
         hashes_to_src_times = defaultdict(list)
 
+        adj_srcs_process_fwidth = self._adjoint_src_width_single(adj_srcs)
+
         tmp_src_time = GaussianPulse(freq0=C_0, fwidth=inf)
-        for src in adj_srcs:
+        for src in adj_srcs_process_fwidth:
             tmp_src = src.updated_copy(source_time=tmp_src_time)
             tmp_src_hash = tmp_src._hash_self()
             hashes_to_sources[tmp_src_hash] = src
             hashes_to_src_times[tmp_src_hash].append(src.source_time)
 
         # Group sources by frequency or port, whichever gives fewer groups
         num_ports = len(hashes_to_src_times)
-        num_unique_freqs = len({src.source_time.freq0 for src in adj_srcs})
+        num_unique_freqs = len({src.source_time.freq0 for src in adj_srcs_process_fwidth})
 
         log.info(f"Found {num_ports} spatial ports and {num_unique_freqs} unique frequencies.")
 
         adjoint_infos = []
         if num_unique_freqs <= num_ports:
             log.info("Grouping adjoint sources by frequency.")
-            unique_freqs = {src.source_time.freq0 for src in adj_srcs}
+            unique_freqs = {src.source_time.freq0 for src in adj_srcs_process_fwidth}
             for freq0 in unique_freqs:
-                group = [src for src in adj_srcs if src.source_time.freq0 == freq0]
+                group = [src for src in adj_srcs_process_fwidth if src.source_time.freq0 == freq0]
                 post_norm = xr.DataArray(data=np.array([1 + 0j]), coords={"f": [freq0]})
                 adjoint_infos.append(
                     AdjointSourceInfo(sources=group, post_norm=post_norm, normalize_sim=True)
@@ -1184,14 +1209,48 @@ def _process_adjoint_sources_broadband(
 
         return [src_broadband], post_norm_amps
 
+    @staticmethod
+    def _adjoint_src_width_broadband(adj_srcs: list[SourceType]) -> float:
+        """Find the adjoint source fwidth that sufficiently covers all adjoint frequencies."""
+
+        adj_srcs_f0 = [adj_src.source_time.freq0 for adj_src in adj_srcs]
+        middle_f0 = 0.5 * (np.max(adj_srcs_f0) + np.min(adj_srcs_f0))
+        min_f0 = np.min(adj_srcs_f0)
+
+        # width of source to sufficiently decay by zero frequency
+        decay_by_f0_fwidth = middle_f0 / NUM_ADJOINT_FWIDTH_TO_ZERO
+        # width of source to sufficiently cover all adjoint frequencies
+        fwidth_to_min_f0 = (middle_f0 - min_f0) / NUM_ADJOINT_FWIDTH_TO_FMIN
+
+        # log warning if the adjoint pulse width is not sufficiently decayed by zero frequency
+        # which may cause some issues in the adjoint accuracy when using field sources
+        if (fwidth_to_min_f0 > decay_by_f0_fwidth) and isinstance(adj_srcs[0], CustomCurrentSource):
+            log.warning(
+                "Adjoint source generated with a frequency spectrum that extends to or overlaps with 0 Hz. "
+                "This can introduce errors into the gradient computation."
+            )
+
+        print(f"source widths: {decay_by_f0_fwidth}, {fwidth_to_min_f0}")
+
+        # Choose a wider pulse width in frequency especially when the min/max frequencies
+        # for the broadband pulse might be very close together
+        adj_src_fwidth = np.maximum(decay_by_f0_fwidth, fwidth_to_min_f0)
+
+        return middle_f0, adj_src_fwidth
+
     def _make_broadband_source(self, adj_srcs: list[SourceType]) -> SourceType:
         """Make a broadband source for a set of adjoint sources."""
 
+        adj_src_f0, adj_src_fwidth = self._adjoint_src_width_broadband(adj_srcs)
+
         source_index = self.simulation.normalize_index or 0
+
         src_time_base = self.simulation.sources[source_index].source_time.updated_copy(
             amplitude=1.0, phase=0.0
         )
-        src_broadband = adj_srcs[0].updated_copy(source_time=src_time_base)
+        src_broadband = adj_srcs[0].updated_copy(
+            source_time=src_time_base.updated_copy(freq0=adj_src_f0, fwidth=adj_src_fwidth)
+        )
 
         return src_broadband