Improve passivity enforcement near high-Q poles in FastDispersionFitter

CHANGELOG.md

@@ -8,6 +8,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - Added value_and_grad function to the autograd plugin, importable via `from tidy3d.plugins.autograd import value_and_grad`. Supports differentiating functions with auxiliary data (`value_and_grad(f, has_aux=True)`).
 
+### Fixed
+- Improved passivity enforcement near high-Q poles in `FastDispersionFitter`. Failed passivity enforcement could lead to simulation divergences.
+
 ## [2.7.2] - 2024-08-07
 
 ### Added

tests/test_plugins/test_dispersion_fitter.py

@@ -272,3 +272,16 @@ def test_dispersion_guess(random_data):
     medium, rms = fitter.fit(num_tries=2)
 
     medium_new, rms_new = fitter.fit(num_tries=1, guess=medium)
+
+
+def test_dispersion_loss_samples():
+    wvls = np.array([275e-3, 260e-3, 255e-3])
+    n_nAlGaN = np.array([2.72, 2.68, 2.53])
+
+    nAlGaN_fitter = FastDispersionFitter(wvl_um=wvls, n_data=n_nAlGaN)
+    nAlGaN_mat, _ = nAlGaN_fitter.fit()
+
+    freq_list = nAlGaN_mat.angular_freq_to_Hz(nAlGaN_mat._imag_ep_extrema_with_samples())
+    ep = nAlGaN_mat.eps_model(freq_list)
+    for e in ep:
+        assert e.imag >= 0

tidy3d/plugins/dispersion/fit_fast.py

@@ -37,6 +37,9 @@
 # when poles are close to omega, can cause invalid response function, and we reject model
 OMEGA_POLE_CLOSE_ATOL = 1e-10
 
+# check this many samples evenly spaced around each pole for passivity
+LOSS_CHECK_SAMPLES_PER_POLE = 2
+
 
 class AdvancedFastFitterParam(Tidy3dBaseModel):
     """Advanced fast fitter parameters."""
@@ -306,13 +309,34 @@ def values(self) -> ArrayComplex1D:
         """Evaluate model at sample frequencies."""
         return self.evaluate(self.omega)
 
+    @cached_property
+    def loss_omega_samples(self) -> ArrayFloat1D:
+        """Frequencies to sample loss to ensure it is within bounds."""
+        # let's check a big range in addition to the imag_extrema
+        range_omega = np.logspace(LOSS_CHECK_MIN, LOSS_CHECK_MAX, LOSS_CHECK_NUM)
+
+        # check samples around each pole
+        ranges_omega = [range_omega]
+        for pole in self.poles:
+            center = -np.imag(pole)
+            width = 3 * np.abs(np.real(pole))
+            if width < OMEGA_POLE_CLOSE_ATOL:
+                continue
+            ranges_omega.append(
+                np.linspace(
+                    abs(center - width / 2), abs(center + width / 2), LOSS_CHECK_SAMPLES_PER_POLE
+                )
+            )
+
+        return np.concatenate(ranges_omega)
+
     @cached_property
     def loss_in_bounds_violations(self) -> ArrayFloat1D:
         """Return list of frequencies where model violates loss bounds."""
 
         extrema_list = PoleResidue.imag_ep_extrema(zip(self.poles, self.residues))
-        # let's check a big range in addition to the imag_extrema
-        range_omega = np.logspace(LOSS_CHECK_MIN, LOSS_CHECK_MAX, LOSS_CHECK_NUM)
+        range_omega = self.loss_omega_samples
+
         omega = np.concatenate((range_omega, extrema_list))
         loss = self.evaluate(omega).imag
         bmin, bmax = self.loss_bounds
@@ -598,7 +622,7 @@ def enforce_passivity(
 
         model = self.updated_copy(passivity_optimized=True)
         violations = model.loss_in_bounds_violations
-        range_omega = np.logspace(LOSS_CHECK_MIN, LOSS_CHECK_MAX, LOSS_CHECK_NUM)
+        range_omega = model.loss_omega_samples
         violations = np.unique(np.concatenate((violations, range_omega)))
 
         # only need one iteration since poles are fixed
@@ -834,7 +858,6 @@ def make_configs():
                 "Unweighted RMS error %.3g",
                 best_model.unweighted_rms_error,
             )
-
         return (
             best_model.pole_residue.updated_copy(frequency_range=self.frequency_range),
             best_model.rms_error,