Repository: flexcompute/tidy3d
Pull request #2455: adjust adjoint source fwidth to decay before zero frequency when possible
Author: login: groberts-flex
Labels: name: 2.9
description: will go into version 2.9.*
Created: 2025-05-13T23:12:29Z | Merged: 2025-05-23T11:29:45Z

For sources that run a single frequency simulation, this chooses a source width according to decay naturally before f=0. For sources that run a broadband simulation and are custom current sources, we attempt to adjust the source width to decay before f=0 while covering all of the adjoint frequencies in the pulse. When this is not possible, a warning is generated.

The results for a custom current source before and after the change are below for when we have a forward simulation with fwidth=f0

<img width="1058" alt="image" src="https://github.com/user-attachments/assets/80e251d5-7ed0-43d4-a163-33fb230b8d25" />



Diff:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9601c0e518..f2aba2f42a 100644
--- a/CHANGELOG.md
+++ b/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
 
diff --git a/tests/test_components/test_autograd.py b/tests/test_components/test_autograd.py
index 67ae10fedd..e035da0903 100644
--- a/tests/test_components/test_autograd.py
+++ b/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):
diff --git a/tidy3d/components/data/sim_data.py b/tidy3d/components/data/sim_data.py
index 1bba53a6d7..d0c1142b44 100644
--- a/tidy3d/components/data/sim_data.py
+++ b/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,14 +1130,32 @@ 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
@@ -1138,16 +1163,16 @@ def _process_adjoint_sources(self, adj_srcs: list[SourceType]) -> list[AdjointSo
 
         # 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
 


Reviews:
user: login: yaugenst-flex
body: Thanks for this fix @groberts-flex, code looks good overall. Can we add a test for this behavior? I think it should be relatively simple to just construct a couple of adjoint sources and run them through `_make_broadband_source`?
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-20T08:37:30Z
user: login: yaugenst-flex
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-23T03:50:50Z
user: login: yaugenst-flex
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-23T03:51:06Z
user: login: yaugenst-flex
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-23T03:51:45Z
user: login: groberts-flex
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T11:16:28Z
user: login: groberts-flex
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T11:16:47Z
user: login: yaugenst-flex
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-23T11:22:12Z

Review comments:
diff_hunk: @@ -1187,9 +1206,32 @@ def process_adjoint_sources_broadband(
     def _make_broadband_source(self, adj_srcs: list[SourceType]) -> SourceType:
         """Make a broadband source for a set of adjoint sources."""
 
+        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."
path: tidy3d/components/data/sim_data.py
user: login: yaugenst-flex
body: Space missing: `Hz.` -> `Hz. `
created_at: 2025-05-20T08:33:09Z
author_association: COLLABORATOR
diff_hunk: @@ -35,6 +35,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Warn if more than 20 frequencies are used in EME, as this may lead to slower or more expensive simulations.
 - EME now supports 2D simulations.
 - 'EMESimulation' now supports 'PermittivityMonitor'.
+- Change `VisualizationSpec` validator for checking validity of user specified colors to only issue a warning if matplotlib is not installed instead of an error.
path: CHANGELOG.md
user: login: yaugenst-flex
body: Did this slip in?
created_at: 2025-05-23T03:50:50Z
author_association: COLLABORATOR
diff_hunk: @@ -35,6 +35,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Warn if more than 20 frequencies are used in EME, as this may lead to slower or more expensive simulations.
 - EME now supports 2D simulations.
 - 'EMESimulation' now supports 'PermittivityMonitor'.
+- Change `VisualizationSpec` validator for checking validity of user specified colors to only issue a warning if matplotlib is not installed instead of an error.
+- 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.
path: CHANGELOG.md
user: login: yaugenst-flex
body: Changelog needs to go under "Unreleased"
created_at: 2025-05-23T03:51:06Z
author_association: COLLABORATOR
path: tests/test_components/test_autograd.py
user: login: yaugenst-flex
body: Looks great!
created_at: 2025-05-23T03:51:45Z
author_association: COLLABORATOR
diff_hunk: @@ -35,6 +35,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Warn if more than 20 frequencies are used in EME, as this may lead to slower or more expensive simulations.
 - EME now supports 2D simulations.
 - 'EMESimulation' now supports 'PermittivityMonitor'.
+- Change `VisualizationSpec` validator for checking validity of user specified colors to only issue a warning if matplotlib is not installed instead of an error.
path: CHANGELOG.md
user: login: groberts-flex
body: Hmm yeah I'm not sure what happened here because I'm not seeing this VisualizationSpec changelog entry on develop but the PR with it was merged a few days ago.  Looking back at that PR, I think I had put that entry in the wrong location as well (2.8.4 instead of unreleased). 
created_at: 2025-05-23T11:16:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -35,6 +35,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Warn if more than 20 frequencies are used in EME, as this may lead to slower or more expensive simulations.
 - EME now supports 2D simulations.
 - 'EMESimulation' now supports 'PermittivityMonitor'.
+- Change `VisualizationSpec` validator for checking validity of user specified colors to only issue a warning if matplotlib is not installed instead of an error.
+- 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.
path: CHANGELOG.md
user: login: groberts-flex
body: thank you, good catch!
created_at: 2025-05-23T11:16:47Z
author_association: CONTRIBUTOR

PR comments:
user: login: groberts-flex
created_at: 2025-05-13T23:14:07Z
body: There is also minor improvement for mode sources when fwidth=f0 as seen below:

<img width="1061" alt="image" src="https://github.com/user-attachments/assets/661058fe-71f7-4335-8163-c0407eb7eed6" />

author_association: CONTRIBUTOR
user: login: groberts-flex
created_at: 2025-05-22T19:35:46Z
body: > Thanks for this fix @groberts-flex, code looks good overall. Can we add a test for this behavior? I think it should be relatively simple to just construct a couple of adjoint sources and run them through `_make_broadband_source`?

Yes, definitely, great point! Added a couple of tests for this to make sure the single source widths were being set correctly as well as the broadband source (and that we issue a warning for custom field sources that still end up wide bandwidth based on the adjoint frequencies being computed for)
author_association: CONTRIBUTOR