Fix NaN gradients of the CARMA kernel

news/228.bugfix

@@ -0,0 +1 @@
+Fixed NaN gradients of the `CARMA` kernel with respect to its parameters, caused by square roots of terms that are exactly zero.

src/tinygp/kernels/quasisep.py

@@ -773,7 +773,9 @@ def __init__(self, alpha: Any, beta: Any):
         _complex_select = _complex_mask * complex_idx % 2
 
         # Construct the obsservation model => real + complex
-        om_real = jnp.sqrt(jnp.abs(acf.real))
+        # (_safe_sqrt keeps the gradient finite for terms that are exactly zero,
+        # including the ones discarded by jnp.where below)
+        om_real = _safe_sqrt(jnp.where(_real_mask, jnp.abs(acf.real), 0.0))
 
         a, b, c, d = (
             2 * acf.real,
@@ -785,10 +787,11 @@ def __init__(self, alpha: Any, beta: Any):
         d2 = jnp.square(d)
         s2 = c2 + d2
         denom = jnp.where(_real_mask, 1.0, 2 * c * s2)
-        h2_2 = d2 * (a * c - b * d) / denom
-        h2 = jnp.sqrt(h2_2)
+        h2_2 = jnp.where(_real_mask, 0.0, d2 * (a * c - b * d) / denom)
+        h2 = _safe_sqrt(h2_2)
         denom = jnp.where(_real_mask, 1.0, d)
-        h1 = (c * h2 - jnp.sqrt(a * d2 - s2 * h2_2)) / denom
+        rad = jnp.where(_real_mask, 0.0, a * d2 - s2 * h2_2)
+        h1 = (c * h2 - _safe_sqrt(rad)) / denom
         om_complex = jnp.array([h1, h2])
 
         # for complex roots, every conjugate pair match one full celerite term,
@@ -900,6 +903,16 @@ def transition_matrix(self, X1: JAXArray, X2: JAXArray) -> JAXArray:
         return tm_real + tm_complex_diag + -tm_complex_u.T + tm_complex_u
 
 
+def _safe_sqrt(x: JAXArray) -> JAXArray:
+    """Square root that is 0, with a zero derivative, for non-positive input
+
+    The derivative of ``jnp.sqrt`` at 0 is infinite, which turns gradients into
+    NaN when a term is exactly zero or is discarded by ``jnp.where``.
+    """
+    positive = x > 0
+    return jnp.where(positive, jnp.sqrt(jnp.where(positive, x, 1.0)), 0.0)
+
+
 @jax.jit
 def carma_roots(poly_coeffs: JAXArray) -> JAXArray:
     roots = jnp.roots(poly_coeffs[::-1], strip_zeros=False)

tests/test_kernels/test_quasisep.py

@@ -2,6 +2,7 @@
 import jax.numpy as jnp
 import jax.scipy as jsp
 import pytest
+from jax.test_util import check_grads
 from numpy import random as np_random
 
 from tinygp import GaussianProcess
@@ -138,6 +139,34 @@ def loss(params):
     loss(params)
 
 
+@pytest.mark.parametrize(
+    "alpha, beta",
+    [
+        # real root
+        ([0.01], [0.1]),
+        # complex roots, the second entry of the observation model is exactly 0
+        ([0.5, 0.1], [0.1]),
+        # complex roots
+        ([0.1, 1.0], [0.1, 0.05]),
+        # real and complex roots
+        ([0.5, 1.5, 1.2], [0.3, 0.2]),
+    ],
+)
+def test_carma_grad(data, alpha, beta):
+    x, y, _ = data
+
+    def loss(alpha, beta):
+        kernel = quasisep.CARMA(alpha=alpha, beta=beta)
+        return -GaussianProcess(kernel, x, diag=0.01).log_probability(y)
+
+    with jax.enable_x64(True):  # type: ignore
+        alpha = jnp.array(alpha, dtype=jnp.float64)
+        beta = jnp.array(beta, dtype=jnp.float64)
+        grads = jax.grad(loss, argnums=(0, 1))(alpha, beta)
+        assert all(jnp.all(jnp.isfinite(g)) for g in grads)
+        check_grads(loss, (alpha, beta), order=1, modes=["rev"], rtol=1e-4)
+
+
 def test_carma_quads():
     alpha = jnp.array([1.4, 2.3, 1.5])
     beta = jnp.array([0.1, 0.5])