Skip to content

fix(state_space): avoid overflow in the Miller Bessel recurrence (TruncatedPeriodic zeros/NaNs) - #827

Open
flyingdoubleG wants to merge 1 commit into
QuantClimate:mainfrom
flyingdoubleG:fix-bessel-miller-overflow
Open

flyingdoubleG wants to merge 1 commit into
QuantClimate:mainfrom
flyingdoubleG:fix-bessel-miller-overflow

Conversation

@flyingdoubleG

Copy link
Copy Markdown

Checklist

  • I've formatted the new code by running uv run poe format before committing (ruff format + ruff check on the changed files).
  • I've added tests for new code.
  • I've added docstrings for the new code.

Description

_stable_scaled_ive computes Ĩ_k(c) = e^{-c} I_k(c), with c = 1/(4ℓ²), for TruncatedPeriodic and TruncatedPeriodicSDE. For c <= truncation_order it uses Miller's downward recurrence. That recurrence works on unnormalised values, starting from 1 at k_max = 2K + 20, and only rescales at the end. Those values grow like (2 k_max / c)^{k_max}, so they overflow to inf. The rescale i0e(c) / ratios[0] is then 0, and every returned value is 0 or NaN.

Where the overflow happens, compared with scipy.special.ive:

  • float32: the default kernel already fails. TruncatedPeriodic(lengthscale=1.0, variance=1.0, period=1.0) uses truncation_order=6, so c = 0.25:

    >>> TruncatedPeriodic().gram(jnp.array([[0.0], [0.3]])).as_matrix()
    [[nan nan] [nan nan]]                  # Periodic gives [[1, 0.721], [0.721, 1]]

    With truncation_order=6, it fails for every lengthscale above about 0.4.

  • float64:

    truncation_order fails for lengthscale above
    10 ~800
    12 ~330
    20 ~30
    25 ~12

    Example: _stable_scaled_ive(1.8e-3, 25) returns all zeros, but the true Ĩ_0 is 0.998.

The existing tests stop at c = 1e-3 with truncation_order = 20, which is just inside the safe region.

Fix

The recurrence now runs on the ratios r_k = I_k / I_{k-1}, from r_{k_max+1} = 0:

r_k = c / (2k + c r_{k+1})

This is the same Miller recurrence I_{k-1} = (2k/c) I_k + I_{k+1}, divided through by I_k. It has the same starting point k_max = 2K + 20 and the same truncation as before. The result is Ĩ_k = i0e(c) · r_1 ⋯ r_k. Each ratio is in (0, 1), so no step can overflow. Very small values underflow to zero, which is correct. The function stays jit- and grad-friendly (lax.scan plus cumprod). The other two branches are unchanged.

Tests

New tests in tests/test_state_space/test_stable_bessel.py:

  • test_stable_scaled_ive_small_c_high_order_does_not_overflow: float64, the (c, K) pairs in the overflow region, values against scipy, and finite gradients.
  • test_stable_scaled_ive_float32: float32 inputs, including the default kernel's c = 0.25.

Results:

  • On main, 9 of the new cases fail.
  • With this change, tests/test_state_space passes: 354 passed, 1 skipped.
  • I also compared all c in [1e-8, 1e4] for K in {1, 3, 6, 12, 25} against scipy.special.ive. On main, the failures were the overflow cases above. With the fix, the only differences left come from the unchanged forward-recurrence branch, at c just above K = 25 (relative error 3e-7 on the highest orders).

A CHANGELOG entry is added under Unreleased → Fixed.

Issue Number: N/A

The downward recurrence in _stable_scaled_ive formed unnormalised values that
grow like (2 k_max / c)^k_max and overflow for small c = 1/(4 l^2) and high
truncation orders, so TruncatedPeriodic and TruncatedPeriodicSDE got zero or
NaN variances. In float32 this already happens for the default kernel.

Run the recurrence on the ratios I_k / I_{k-1} = c / (2k + c r_{k+1}), which
stay in (0, 1), and take cumulative products scaled by i0e(c).
@github-actions github-actions Bot added documentation Improvements or additions to documentation tests release size/s ci Continuous Integration labels Nov 4, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ci Continuous Integration documentation Improvements or additions to documentation release size/s tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant