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

@@ -283,7 +283,7 @@ def get_default_weights(cls, eps: ArrayComplex1D) -> Tuple[float, float]:
     def pole_residue(self) -> PoleResidue:
         """Corresponding :class:`.PoleResidue` model."""
         if self.eps_inf is None or self.poles is None:
-            return None
+            return PoleResidue(eps_inf=1, poles=[])
         return PoleResidue(
             eps_inf=self.eps_inf,
             poles=list(
@@ -306,13 +306,47 @@ def values(self) -> ArrayComplex1D:
         """Evaluate model at sample frequencies."""
         return self.evaluate(self.omega)
 
+    @cached_property
+    def loss_omega_pole_samples(self) -> ArrayFloat1D:
+        """Samples to check around each pole for passivity."""
+        ranges_omega = []
+        for pole, res in zip(self.poles, self.residues):
+            cr = np.real(res)
+            ci = np.imag(res)
+            ar = np.real(pole)
+            ai = np.imag(pole)
+            # no extra checking needed for marginally stable poles; these are handled later
+            if ar == 0:
+                continue
+            if cr == 0:
+                pole_extrema = [-ai]
+            else:
+                disc = ci**2 + cr**2
+                pole_extrema = [
+                    -ai + ar * (ci + np.sqrt(disc)) / cr,
+                    -ai + ar * (ci - np.sqrt(disc)) / cr,
+                ]
+
+            ranges_omega.append(np.abs(pole_extrema))
+        if len(ranges_omega) == 0:
+            return []
+        return np.concatenate(ranges_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)
+        range_omega_poles = self.loss_omega_pole_samples
+        return np.concatenate((range_omega, range_omega_poles))
+
     @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 +632,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 +868,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,