fix(kernels): parameterise the spectral measure by lengthscale and dimension

[thomaspinder]

Checklist

• I've formatted the new code by running uv run poe format before committing.
• I've added tests for new code.
• I've added docstrings for the new code.

Description

Fixes the spectral density reported by @ahwillia in #612.

StationaryKernel.spectral_density returned a standardised distribution — Normal(0, 1) for RBF, StudentT(2ν, 0, 1) for Matérn — that ignored the lengthscale and was hard-coded to one dimension. Every lengthscale gave the same curve, which is exactly the superimposed plot in the issue.

Why RFF still worked

BasisFunctionComputation.compute_features divided the sampled frequencies by the lengthscale, and scaling() applied the variance:

scaling_factor = _val(kernel.base_kernel.lengthscale)
z = jnp.matmul(x, (frequencies / scaling_factor).T)

So the approximation compensated downstream and was numerically correct. That compensation is what kept the bug latent — the only in-repo consumer of spectral_density was the one place that cancelled the error out. Anyone using the property directly got a wrong answer silently.

The fix

The measure is now D-dimensional and carries diag(ℓ)⁻¹ as its scale (ARD lengthscales included), and the compensating division is removed from compute_features. It satisfies Bochner's theorem

$$k(\tau) = \sigma^2 , \mathbb{E}_{p(\omega)}\big[e^{i \omega^\top \tau}\big]$$

The variance deliberately stays out of the measure: σ² is its total mass, which normalisation divides out, and it re-enters as the explicit prefactor above. A numpyro Distribution cannot represent the unnormalised S(ω) — log_prob is normalised by construction — so the docstring pins the relation to R&W §4.2.1 instead:

$$S(\omega) = \sigma^2 (2\pi)^D p(\omega), \qquad k(\tau) = (2\pi)^{-D}!\int S(\omega) e^{i\omega^\top\tau},d\omega$$

I verified that formula reproduces R&W's closed form for the RBF in 1-D and 3-D.

Compatibility

RFF Gram and cross-covariance values are bit-identical to before. Scaling by a diagonal matrix is exactly elementwise division, so the frequencies are unchanged — test_rff_effective_frequencies_unchanged asserts jnp.array_equal against the old two-step route across dimensions and lengthscales. The lengthscale simply moved from the feature map into the measure it is drawn from.

Two API changes, both noted in the CHANGELOG:

• spectral_density now returns MultivariateNormal / MultivariateStudentT with event_shape == (D,). Callers doing .sample(key, (M, D)) should use .sample(key, (M,)).
• RFF with an explicitly supplied frequencies= now treats them as the frequencies ω directly. They were previously divided by the lengthscale, silently rescaling user-supplied values — arguably a second latent bug.

Tests

New tests/test_kernels/test_spectral_density.py (92 cases). The load-bearing one is a Monte-Carlo Bochner check across RBF/Matérn12/32/52 × D ∈ {1,2,3} × ℓ ∈ {0.5,1,2} × σ² ∈ {1,3} — it pins the parameterisation without reference to any Fourier normalisation convention, so it constrains the measure independently of how it's implemented. Plus the reporter's reproducer, dimensionality, ARD anisotropy, and variance-independence.

uv run poe all-tests: 2582 passed. (The one failure in my working tree is test_root_good picking up an untracked local scratch .md in the repo root — unrelated, and not present in CI.)

Follow-ups this unblocks

• Add spectral densities for RationalQuadratic, PoweredExponential, and Periodic (unlock RFF) #672 — adding spectral densities for RationalQuadratic / PoweredExponential / Periodic now extends a correct contract rather than replicating the bug three more times.
• Hsgp #608 (HSGP) — Hilbert-space methods evaluate S(ω) rather than sampling it, so they need this parameterisation to be right. Worth deciding there whether to add an explicit unnormalised S(ω) accessor; I kept this PR to the reported bug.

Issue Number: #612

🤖 Generated with Claude Code

https://claude.ai/code/session_01Bj9k5fnAZ8JzD4Rg3HMDMj