Refine a zero-dispersion point with array samples, not Brent

changelog.d/zdw-array-refinement.changed.md

@@ -0,0 +1 @@
+Made `capillary_params` and the other design tools that report a zero-dispersion wavelength about 2.5 times faster by refining the zero with array samples, and stopped them raising `ValueError` when the search met a gas resonance (xenon in a 14 um capillary at 15 bar). Zeros clear of the gas resonances are unchanged within 1e-8; a sign change reported beside a resonance, which only a pump in the vacuum ultraviolet selects, can move within the search's 3.6e13 rad/s sample interval.

src/solprop/tools.py

@@ -148,6 +148,14 @@
 _ZDW_SAMPLES = 512
 """Samples, evenly spaced in frequency, on which the dispersion's sign is read."""
 
+_ZDW_REFINE_ROUNDS = 4
+"""Rounds of array samples that refine a sign change of the dispersion. Four take a
+bracket of 3.6e13 rad/s to 3.4e7 rad/s, where interpolating linearly reaches the noise
+of the differenced ``beta2``."""
+
+_ZDW_REFINE_SAMPLES = 33
+"""Samples of each refinement round: 32 intervals, so each round narrows 32-fold."""
+
 _POLE_FACTOR = 2.0
 """How far the group delay at a refined sign change of ``beta2`` may sit from its value
 at the bracket's ends, in units of the change a bounded ``beta2`` could produce across
@@ -1277,6 +1285,44 @@ def _on_axis_intensity(mode: Mode, peak_power: float, z: float) -> float:
     return per_watt * peak_power
 
 
+def _refine_zero(mode: Mode, low: float, high: float, z: float) -> float:
+    """Return a sign change of the dispersion inside a bracket, by array samples.
+
+    A mode's index models evaluate an array at about the cost of one point, so each
+    round narrows the bracket 32-fold for the cost of one Brent step, where a scalar
+    search costs a glass model evaluation per iterate. The rounds stop short of a pole
+    of the differenced ``beta2``, which a scalar search converges onto until the model
+    returns NaN.
+
+    Parameters
+    ----------
+    mode : Mode
+        The mode.
+    low, high : float
+        The bracket, rad/s, whose ends have ``beta2`` of opposite signs.
+    z : float
+        Position along the waveguide, m.
+
+    Returns
+    -------
+    float
+        The first sign change upwards in frequency, rad/s, linearly interpolated in the
+        last bracket; or NaN when samples that are not finite separate the signs.
+    """
+    first = last = math.nan
+    for _ in range(_ZDW_REFINE_ROUNDS):
+        grid = np.linspace(low, high, _ZDW_REFINE_SAMPLES)
+        sampled = base.beta2(mode, grid, z)
+        crossing = np.flatnonzero(np.sign(sampled[:-1]) * np.sign(sampled[1:]) <= 0.0)
+        if crossing.size == 0:
+            return math.nan
+        i = int(crossing[0])
+        low, high = float(grid[i]), float(grid[i + 1])
+        # The first change has a nonzero sample below it, so the two cannot be equal.
+        first, last = float(sampled[i]), float(sampled[i + 1])
+    return low - first * (high - low) / (last - first)
+
+
 def _zdw_near(mode: Mode, omega: float, z: float) -> float:
     """Return the zero-dispersion wavelength nearest a pump frequency.
 
@@ -1309,19 +1355,15 @@ def _zdw_near(mode: Mode, omega: float, z: float) -> float:
         _ZDW_SAMPLES,
     )
     zeros = []
-    # Brent's method converges onto a pole as readily as onto a zero, and a model may
+    # The refinement converges onto a pole as readily as onto a zero, and a model may
     # divide by zero exactly there; the group-delay test below tells the two apart.
     with _probing(), np.errstate(divide="ignore", invalid="ignore"):
         sampled = base.beta2(mode, grid, z)
         for i in np.flatnonzero(np.sign(sampled[:-1]) * np.sign(sampled[1:]) < 0.0):
             low, high = float(grid[i]), float(grid[i + 1])
-            root = float(
-                brentq(
-                    lambda value: float(base.beta2(mode, np.array([value]), z)[0]),
-                    low,
-                    high,
-                )
-            )
+            root = _refine_zero(mode, low, high, z)
+            if math.isnan(root):
+                continue
             delay = base.beta1(mode, np.array([low, root, high]), z)
             bounded = max(abs(sampled[i]), abs(sampled[i + 1])) * (high - low)
             if abs(delay[1] - delay[0]) <= _POLE_FACTOR * bounded:

tests/test_tools.py

@@ -15,6 +15,7 @@
 from hypothesis import given, settings
 from hypothesis import strategies as st
 from scipy.integrate import quad
+from scipy.optimize import brentq
 from scipy.special import jn_zeros, jv
 
 from solprop import tools
@@ -826,6 +827,62 @@ def test_a_sign_change_through_a_pole_is_not_a_zero_dispersion_point():
     assert math.isnan(tools._zdw_near(pole_mode(pole), OMEGA0, 0.0))
 
 
+def test_a_sign_change_across_samples_that_are_not_finite_is_not_a_zero():
+    # beta1 is undefined in a narrow band about its minimum, which falls between two
+    # samples of the first search: they bracket the only sign change, and every round
+    # of the refinement finds samples that are not finite between the two signs.
+    grid = np.linspace(
+        2.0 * math.pi * C / tools._ZDW_RANGE[1],
+        2.0 * math.pi * C / tools._ZDW_RANGE[0],
+        tools._ZDW_SAMPLES,
+    )
+    centre = 0.5 * (grid[20] + grid[21])
+
+    def delay(w):
+        offset = w - centre
+        # Wider than the step of beta2 (4e12 rad/s here), so beta2 has no finite gap.
+        return np.where(np.abs(offset) < 6e12, np.nan, 1e-30 * offset**2)
+
+    mode = SpecifiedMode(
+        dispersion=lambda w: 1e-30 * (w - centre) ** 3 / 3.0,
+        dispersion_derivative=delay,
+        loss=0.0,
+        effective_area=1.0,
+    )
+    assert np.all(np.isfinite(base.beta2(mode, grid)))
+    assert math.isnan(tools._zdw_near(mode, OMEGA0, 0.0))
+
+
+@pytest.mark.parametrize(
+    ("gas", "radius", "pressure"),
+    [
+        ("He", 14e-6, 15.0),
+        ("Ne", 14e-6, 15.0),
+        ("Ar", 14e-6, 15.0),
+        ("Kr", 14e-6, 15.0),
+        ("He", 50e-6, 1.0),
+        ("Ar", 125e-6, 5.0),
+        # A scalar search of the whole window meets a resonance here and stops at NaN.
+        ("Xe", 14e-6, 15.0),
+    ],
+)
+def test_the_array_refinement_of_a_zero_agrees_with_a_brent_search(
+    gas, radius, pressure
+):
+    # 200 nm to 3 um holds one sign change, clear of the resonances of each gas.
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        mode = MarcatiliMode(radius, gas, pressure)
+        reference = brentq(
+            lambda w: float(base.beta2(mode, np.array([w]))[0]),
+            2.0 * math.pi * C / 3e-6,
+            2.0 * math.pi * C / 200e-9,
+        )
+        wavelength = tools._zdw_near(mode, OMEGA0, 0.0)
+    # beta2 is a difference, good to about 5e-9 (modes.base); both stop on that noise.
+    assert wavelength == pytest.approx(2.0 * math.pi * C / reference, rel=1e-8)
+
+
 def test_a_fill_without_kerr_data_is_refused_by_name():
     inert = replace(species("He"), gamma3=None, kerr=None)
     fill = Mixture.fill({inert: 1.0}, 1.0)