[mottopanikeiku]

spectral_density_rational_quadratic multiplies by a * (2 * scale_mixture)**((dim - 1) / 2) where the D-dimensional Fourier transform needs a**dim (with a = sqrt(2 * scale_mixture) * length), so for dim > 1 the density is off by a factor length**(1 - dim). The zero-frequency branch already has a**dim but adds the same extra (2 * scale_mixture)**((dim - 1) / 2) factor. These densities are the basis-coefficient variances in hsgp_rational_quadratic, so the prior itself is wrong. The existing covariance tests all use length=1, where the first error vanishes. I found this by reading and testing the code.

The fix is in _log_spectral_density_rational_quadratic, the log-density helper from #2365: I dropped the extra dimension factor from both branches, use dim * log(a) in the regular branch, and corrected both formulas in the docstring. 1-D values, the tiny floor and the safe argument at w = 0 don't change.

import numpy as np
from numpyro.contrib.hsgp.spectral_densities import spectral_density_rational_quadratic

# alpha=1, length=0.5, scale_mixture=2: k(x) = (1 + ||x||^2)^-2,
# whose integral over R^2 (the spectral density at 0) is pi.
s0 = spectral_density_rational_quadratic(dim=2, w=np.zeros(2), alpha=1.0, length=0.5, scale_mixture=2.0)
print(float(s0), np.pi)

quantity	before	after	reference
spectral density at 0 (2-D, length 0.5)	6.283185	3.141593	pi = 3.141593
2-D HSGP prior variance (alpha 1, length 0.5)	1.9999552	0.9999776	1

Tests:

• test_spectral_density_rational_quadratic_normalization checks the density in 1, 2 and 3 dimensions, for length 0.5/1/2, at zero and nonzero frequency. The reference is independent of the implementation: SciPy quadrature of the RBF spectral density against the Gamma mixing distribution that defines the RQ kernel.
• Two non-unit-length 2-D cases added to the existing RQ covariance test (compared with scikit-learn's RationalQuadratic).
• The new cases fail on master (12 failed) and pass with the fix: pytest test/contrib/hsgp gives 161 passed, and test_spectral_densities.py with JAX_ENABLE_X64=1 gives 70 passed.
• Ruff check/format, ty and the license-header check pass.

I made this PR and directed its development, with AI assistance from Claude and GPT-6 Astra.

[juanitorduz]

I will wait until #2365 is mreged

[mottopanikeiku]

Sounds good. Once #2365 is merged I'll rebase this onto its log-density helper and re-run the tests.

[mottopanikeiku]

rebased onto master now that #2365 is in (8e9543d). the fix now lives in #2365's _log_spectral_density_rational_quadratic helper: i dropped the extra (2 * scale_mixture)**((dim - 1) / 2) factor from both branches and use dim * log(a) in the regular branch, so 1-d values, the tiny floor and the safe argument at w = 0 are unchanged. master still gives 2π instead of π for the 2-d s(0) at length 0.5, scale_mixture 2. the new tests fail there (12 failed) and pass here: pytest test/contrib/hsgp 161 passed, and test_spectral_densities.py with JAX_ENABLE_X64=1 70 passed.

[github-actions]

Benchmark report

this PR fix-hsgp-rq-normalization at a290139c vs baseline master at 3d990ab7

  run time:     unchanged across 32 benchmarks
- compile time: 1 slower, 0 faster

Significant changes (1)

                                ──────── run time ───────     ────── compile time ──────
  benchmark                     baseline   this PR      Δ     baseline   this PR       Δ
────────────────────────────────────────────────────────────────────────────────────────
- predictive_forward_sampling   389.2 ms  388.7 ms  -0.1%      99.6 ms  144.4 ms  +45.0%

_{Red is slower, green is faster; a row is coloured by the worse of its two columns. A delta in parentheses cleared the threshold on a measurement below the resolution floor, so it is shown without being called a change. † marks a benchmark that could not be compared — see below.}

Full results

distributions

                                 ──────── run time ────────     ────── compile time ──────
  benchmark                      baseline  this PR        Δ     baseline   this PR       Δ
──────────────────────────────────────────────────────────────────────────────────────────
  biject_to_constraints            3.0 ms   3.0 ms    -1.1%     164.2 ms  173.4 ms   +5.6%
  categorical_log_prob             1.3 ms   1.3 ms    +4.3%      47.0 ms   47.9 ms   +1.8%
  dirichlet_log_prob               401 µs   407 µs    +1.6%     812.0 ms  832.1 ms   +2.5%
  dirichlet_sample                76.0 ms  75.6 ms    -0.5%     678.7 ms  650.0 ms   -4.2%
  gamma_log_prob                   1.3 ms   1.4 ms    +2.5%       3.56 s    3.55 s   -0.3%
  gamma_sample                    27.9 ms  28.0 ms    +0.4%     585.6 ms  616.0 ms   +5.2%
  lkj_cholesky_sample              4.2 ms   4.3 ms    +2.0%     817.9 ms  827.8 ms   +1.2%
  mixture_same_family_log_prob     566 µs   570 µs    +0.7%      58.6 ms   56.8 ms   -3.0%
  multivariate_normal_log_prob     583 µs   532 µs  (-8.7%)     373.6 ms  369.1 ms   -1.2%
  normal_log_prob                  419 µs   423 µs    +1.0%      41.8 ms   48.6 ms  +16.4%
  normal_sample                   37.1 ms  37.1 ms    +0.0%     134.7 ms  137.1 ms   +1.8%
  stick_breaking_transform         4.7 ms   4.5 ms    -4.0%     164.3 ms  142.8 ms  -13.1%
  student_t_log_prob               2.5 ms   2.5 ms    +0.4%     118.9 ms  115.1 ms   -3.2%
  truncated_normal_log_prob        460 µs   443 µs    -3.8%      52.7 ms   42.9 ms  -18.5%

handlers

                                  ───────── run time ─────────     ────── compile time ──────
  benchmark                       baseline   this PR         Δ     baseline   this PR       Δ
─────────────────────────────────────────────────────────────────────────────────────────────
  initialize_model_hierarchical    31.5 ms   31.7 ms     +0.8%       2.56 s    2.65 s   +3.3%
  log_density_hierarchical          2.4 ms    2.5 ms     +1.5%     851.8 ms  885.9 ms   +4.0%
  nested_handler_stack              1.0 ms    1.0 ms     -0.4%       582 µs    638 µs   +9.6%
  potential_energy_and_grad          19 µs     21 µs  (+13.2%)      73.7 ms   71.9 ms   -2.5%
- predictive_forward_sampling     389.2 ms  388.7 ms     -0.1%      99.6 ms  144.4 ms  +45.0%
  trace_seeded_model                517 µs    520 µs     +0.6%     356.5 ms  364.4 ms   +2.2%

mcmc

                             ──────── run time ───────     ───── compile time ─────
  benchmark                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────
  hmc_logistic_regression    498.2 ms  498.4 ms  +0.0%       2.02 s   2.10 s  +4.0%
  nuts_dense_mass_funnel     751.6 ms  745.0 ms  -0.9%       1.80 s   1.70 s  -6.0%
  nuts_eight_schools         861.0 ms  861.6 ms  +0.1%       1.78 s   1.83 s  +2.4%
  nuts_hierarchical_glm        2.17 s    2.19 s  +0.5%       3.05 s   3.03 s  -0.7%
  nuts_logistic_regression   864.7 ms  865.7 ms  +0.1%       2.08 s   2.10 s  +0.8%
  nuts_vectorized_chains       1.58 s    1.60 s  +1.2%       1.88 s   1.81 s  -3.6%

svi

                                             ──────── run time ───────     ───── compile time ─────
  benchmark                                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────────────────────
  svi_autodelta_map_logistic                 211.6 ms  208.9 ms  -1.3%       2.08 s   2.09 s  +0.4%
  svi_autodiagonalnormal_hierarchical        657.8 ms  669.0 ms  +1.7%       3.31 s   3.36 s  +1.7%
  svi_automultivariatenormal_eight_schools   503.4 ms  510.2 ms  +1.4%       3.26 s   3.37 s  +3.4%
  svi_autonormal_logistic                    583.8 ms  576.8 ms  -1.2%       2.48 s   2.39 s  -3.6%
  svi_multi_particle_elbo                    825.9 ms  856.3 ms  +3.7%       2.45 s   2.47 s  +1.1%
  svi_trace_mean_field_elbo                  763.9 ms  774.1 ms  +1.3%       3.84 s   3.64 s  -5.3%

Methodology and environment

Each benchmark is set up untimed, then called once with the JAX caches cleared and several more times warm. Run is the fastest warm call; compile is the first call minus that, i.e. the tracing, lowering and XLA compilation the warm calls did not have to pay for.

Both refs were measured on the same runner over 2 interleaved round(s), taking the best observation per benchmark. A result is called neutral when it moves less than ±5% (run) or ±25% (compile), or when the measurement itself is under 1 ms (run) / 50 ms (compile) — a shared CI runner cannot resolve changes below that. Compile time gets the looser band because it is measured once per round rather than best-of-N, and swings by roughly 20% even between two runs of identical code. A delta shown in parentheses did clear its threshold, but on a measurement below the resolution floor, so it is reported without being called a change.

	baseline	this PR
ref	master	fix-hsgp-rq-normalization
commit	3d990ab7	a290139c
numpyro	0.22.0	0.22.0
jax	0.12.0.1	0.12.0.1
backend	cpu	cpu
python	3.14.4	3.14.4

Runner: Linux-7.0.0-1017-azure-x86_64-with-glibc2.43, 4 CPUs.

_{Produced by this benchmark run.}

[juanitorduz]

trim the oversized basis counts in the new covariance cases. I found no blocking issues.

[mottopanikeiku]

trimmed the new covariance cases to [48, 48] and [20, 20], without loosening the 1% tolerance. 55 focused tests pass; the covariance cases also pass with x64.

[antonsoo]

Checked a290139 against 3d990ab on CPU, Python 3.14.7 / JAX 0.10.2.

Independent Gamma-mixture quadrature agrees with the spectrum and its first two length derivatives in 36 cases across 1D/2D/3D. The largest scaled error in x64 is 1.1e-13.

I also conditioned the public HSGP model on 96 recorded DWD station temperatures and predicted 64 other stations, comparing both parameterizations with independent dense RQ conditioning. With fixed amplitude 1.7, mixture 3 and length 0.4:

quantity	base	head	dense reference
mean prior variance	4.24999	1.70000	1.7
maximum prediction difference from dense	0.26050	5.3e-7	0
maximum posterior variance difference	0.55981	1.7e-7	0

Length 1.8 also recovers the specified prior and predictions. All three 1D covariance/posterior controls are bit-identical between trees. Float32 checks remain finite: covariance-conditioned means differ from dense by at most 2.2e-4, and public model means match those within 2.1e-4, in standardized temperature units. The head HSGP suite passes 161 tests in x64; 10 of the 18 new normalization cases fail numerically on the base.

Script, pinned data provenance and reports. These are fixed-hyperparameter numerical checks, not forecast-quality or scale-mixture-gradient evidence. Accelerators and the full repository suite were not tested.

[juanitorduz]

The fix is correct. I checked it against a Gamma-mixture derivation and confirmed squared-exponential and Matérn don't have the same bug. Two test trims:

1. Drop length=1.0 from test_spectral_density_rational_quadratic_normalization. At unit length the regular branch can't see the ℓ^{1−D} error, and the 0.5 and 2.0 cases already cover the S(0) error.
2. Drop the 2d,length=0.5 / [48,48] covariance case. The normalization test already pins that density; keep 2d,length=2.0 as the end-to-end check.

Also, a small correction to the PR body: the S(0) branch is never used by hsgp_rational_quadratic, because the HSGP frequencies are all > 0. Only the regular-branch error changes the HSGP prior.