pyro-ppl / numpyro

Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.
https://num.pyro.ai
Apache License 2.0
2.09k stars 227 forks source link

fix(contrib.hsgp): convert matern spectral density from frequency domain #1811

Closed brendancooley closed 4 weeks ago

brendancooley commented 4 weeks ago

As I was implementing these tests (to support development in #1805) I noticed that the Matérn spectral density implementation was not giving results consistent with scikit-learn. After staring at Rasmussen and Williams Chapter 4 for a bit I finally noticed that their spectral densities are in the frequency domain, with the inputs scaled by $2 \pi$. Riutort-Mayol 2023 adopt Rasmussen and Williams' parameterization initially but note that the input vectors are frequencies and convert appropriately for the specific $\nu = \infty$, $\nu =3/2$ and $\nu =5/2$ examples (see screenshots below).

After dropping the $4 \pi^2$ term from the Matern spectral density function I am able to replicate the exact covariance functions from scikit-learn.

image image

@juanitorduz let me know if I'm missing anything or misunderstanding here -- still trying to wrap my head around some of this stuff and get the terminology right. Additionally, let me know if you have any suggestions on how to make the test suite stronger. The goal is to add vector-valued alpha and length examples to the tests and use those to guide development on #1805.

juanitorduz commented 4 weeks ago

@brendancooley, do we also need to fix the formula in the LaTeX docstring?

brendancooley commented 4 weeks ago

@brendancooley, do we also need to fix the formula in the LaTeX docstring?

Good call! Just pushed fix