broadband source feature PR comments

tests/test_components/test_source.py

@@ -151,3 +151,81 @@ def get_pol_dir(axis, pol_angle=0, angle_theta=0, angle_phi=0):
     assert np.allclose(
         get_pol_dir(axis=2, angle_theta=np.pi / 4), (-1 / np.sqrt(2), 0, +1 / np.sqrt(2))
     )
+
+
+def test_broadband_source():
+    g = td.GaussianPulse(freq0=1, fwidth=0.1)
+    mode_spec = td.ModeSpec(num_modes=2)
+    fmin, fmax = g.frequency_range(num_fwidth=4)
+    fdiff = (fmax - fmin) / 2
+    fmean = (fmax + fmin) / 2
+
+    def check_freq_grid(freq_grid, num_freqs):
+        """Test that chebyshev polynomials are orthogonal on provided grid."""
+        cheb_grid = (freq_grid - fmean) / fdiff
+        poly = np.polynomial.chebyshev.chebval(cheb_grid, np.ones(num_freqs))
+        dot_prod_theory = num_freqs + num_freqs * (num_freqs - 1) / 2
+        # print(len(freq_grid), num_freqs)
+        # print(abs(dot_prod_theory - np.dot(poly, poly)))
+        assert len(freq_grid) == num_freqs
+        assert abs(dot_prod_theory - np.dot(poly, poly)) < 1e-10
+
+    # test we can make a broadband gaussian beam
+    num_freqs = 3
+    s = td.GaussianBeam(
+        size=(0, 1, 1), source_time=g, pol_angle=np.pi / 2, direction="+", num_freqs=num_freqs
+    )
+    freq_grid = s.frequency_grid
+    check_freq_grid(freq_grid, num_freqs)
+
+    # test we can make a broadband astigmatic gaussian beam
+    num_freqs = 10
+    s = td.AstigmaticGaussianBeam(
+        size=(0, 1, 1),
+        source_time=g,
+        pol_angle=np.pi / 2,
+        direction="+",
+        waist_sizes=(0.2, 0.4),
+        waist_distances=(0.1, 0.3),
+        num_freqs=num_freqs,
+    )
+    freq_grid = s.frequency_grid
+    check_freq_grid(freq_grid, num_freqs)
+
+    # test we can make a broadband mode source
+    num_freqs = 20
+    s = td.ModeSource(
+        size=(0, 1, 1),
+        direction="+",
+        source_time=g,
+        mode_spec=mode_spec,
+        mode_index=0,
+        num_freqs=num_freqs,
+    )
+    freq_grid = s.frequency_grid
+    check_freq_grid(freq_grid, num_freqs)
+
+    # check validators for num_freqs
+    with pytest.raises(pydantic.ValidationError) as e_info:
+        s = td.GaussianBeam(
+            size=(0, 1, 1), source_time=g, pol_angle=np.pi / 2, direction="+", num_freqs=200
+        )
+    with pytest.raises(pydantic.ValidationError) as e_info:
+        s = td.AstigmaticGaussianBeam(
+            size=(0, 1, 1),
+            source_time=g,
+            pol_angle=np.pi / 2,
+            direction="+",
+            waist_sizes=(0.2, 0.4),
+            waist_distances=(0.1, 0.3),
+            num_freqs=100,
+        )
+    with pytest.raises(pydantic.ValidationError) as e_info:
+        s = td.ModeSource(
+            size=(0, 1, 1),
+            direction="+",
+            source_time=g,
+            mode_spec=mode_spec,
+            mode_index=0,
+            num_freqs=-10,
+        )

tidy3d/components/source.py

@@ -312,7 +312,7 @@ def plot(  #  pylint:disable=too-many-arguments
         z: float = None,
         ax: Ax = None,
         sim_bounds: Bound = None,
-        **patch_kwargs
+        **patch_kwargs,
     ) -> Ax:
 
         # call the `Source.plot()` function first.
@@ -466,24 +466,42 @@ def _dir_vector(self) -> Tuple[float, float, float]:
 class BroadbandSource(Source, ABC):
     """A source with frequency dependent field distributions."""
 
-    num_freqs: pydantic.NonNegativeInt = pydantic.Field(
+    num_freqs: int = pydantic.Field(
         1,
         title="Number of Frequency Points",
-        description="Number of points to approximate the frequency dependence of injected field.",
+        description="Number of points used to approximate the frequency dependence of injected "
+        "field. A Chebyshev interpolation is used, thus, only a small number of points, i.e., less "
+        "than 20, is typically sufficient to obtain converged results.",
+        ge=1,
+        le=99,
     )
 
     @cached_property
     def frequency_grid(self) -> np.ndarray:
         """A Chebyshev grid used to approximate frequency dependence."""
-        fmin, fmax = self.source_time.frequency_range(4)
-        return 0.5 * (fmin + fmax) + 0.5 * (fmax - fmin) * np.cos(
-            0.5 * np.pi * (2 * np.flip(np.arange(self.num_freqs)) + 1) / self.num_freqs
-        )
-
-    def map_frequencies(self, freq_grid) -> np.ndarray:
-        """Map frequency values into (-1,1) interval."""
-        fmin, fmax = self.source_time.frequency_range(4)
-        return (freq_grid - 0.5 * (fmin + fmax)) / (0.5 * (fmax - fmin))
+        freq_min, freq_max = self.source_time.frequency_range(num_fwidth=4)
+        freq_avg = 0.5 * (freq_min + freq_max)
+        freq_diff = 0.5 * (freq_max - freq_min)
+        uni_points = (2 * np.arange(self.num_freqs) + 1) / (2 * self.num_freqs)
+        cheb_points = np.cos(np.pi * np.flip(uni_points))
+        return freq_avg + freq_diff * cheb_points
+
+    @pydantic.validator("num_freqs", always=True, allow_reuse=True)
+    def _warn_if_large_number_of_freqs(cls, val):
+        """Warn if a large number of frequency points is requested."""
+
+        if val is None:
+            return val
+
+        if val >= 20:
+            logging.warning(
+                f"A large number ({val}) of frequency points is used in a broadband source. "
+                "This can slow down simulation time and is only needed if the mode fields are "
+                "expected to have a very sharp frequency dependence. 'num_freqs' < 20 is "
+                "sufficient in most cases."
+            )
+
+        return val
 
 
 """ Source current profiles defined by (1) angle or (2) desired mode. Sets theta and phi angles."""