Support multidimensional HSGP (numpyro.contrib.hsgp) (#1803)

* wip(hsgp_nd): add eigen_index function, allow m to take Sequence[int]

wip(hsgp_nd): test list input for m, make dim required arg

wip(hsgp_nd): support batched eigenfunctions, swap dims for m and d

* wip(hsgp_nd): support/test n-d approximations

* doc(hsgp_nd): format ordered lists

doc(hsgp_nd): polish contrib documentation

rm TODO

fix eigenindices docstring

cleanup nb

* fix(hsgp_nd): make python 3.9 happy

* fix(hsgp_nd): make linter happy (space between license and imports)

* fix(stochastic_volatility): legendHandles -> legend_handles

looks like matplotlib snuck in a deprecation here

* test(hsgp_nd): add tests for `_convert_ell` and `eigenindices`

* test(hsgp_nd): specific test for eq 10

docs/source/contrib.rst

@@ -106,8 +106,7 @@ This module contains helper functions for use in the Hilbert Space Gaussian Proc
 described in [1] and [2].
 
 .. warning::
-    This module is experimental. Currently, it only supports Gaussian processes with one-dimensional inputs.
-
+    This module is experimental.
 
 **Why do we need an approximation?** 
 
@@ -123,30 +122,34 @@ a practical approach using NumPyro and PyMC.
 
 Here we provide the main steps and ingredients of the approximation method:
 
-1. Each stationary kernel :math:`k` has an associated spectral density :math:`S(\omega)`. There are closed formulas for the most common kernels. These formulas depend on the hyperparameters of the kernel (e.g. amplitudes and length scales).
-2. We can approximate the spectral density :math:`S(\omega)` as a polynomial series in :math:`||\omega||`. We call :math:`\omega` the frequency.
-3. We can interpret these polynomial terms as "powers" of the Laplacian operator. The key observation is that the Fourier transform of the Laplacian operator is :math:`||\omega||^2`.
-4. Next, we impose Dirichlet boundary conditions on the Laplacian operator which makes it self-adjoint and with discrete spectrum.
-5. We identify the expansion in (2) with the sum of powers of the Laplacian operator in the eigenbasis of (4).
+    1. Each stationary kernel :math:`k` has an associated spectral density :math:`S(\omega)`. There are closed formulas for the most common kernels. These formulas depend on the hyperparameters of the kernel (e.g. amplitudes and length scales).
+
+    2. We can approximate the spectral density :math:`S(\omega)` as a polynomial series in :math:`||\omega||`. We call :math:`\omega` the frequency.
+
+    3. We can interpret these polynomial terms as "powers" of the Laplacian operator. The key observation is that the Fourier transform of the Laplacian operator is :math:`||\omega||^2`.
+
+    4. Next, we impose Dirichlet boundary conditions on the Laplacian operator which makes it self-adjoint and with discrete spectrum.
+
+    5. We identify the expansion in (2) with the sum of powers of the Laplacian operator in the eigenbasis of (4).
 
-For the one dimensional case the approximation formula, in the non-centered parameterization, is:
+Let :math:`m^\star = \prod_{d=1}^D m_d` be the total number of terms of the approximation, where :math:`m_d` is the number of basis functions used in the approximation for the :math:`d`-th dimension. Then, the approximation formula, in the non-centered parameterization, is:
 
 .. math::
 
-    f(x) \approx \sum_{j = 1}^{m} 
-    \overbrace{\color{red}{\left(S(\sqrt{\lambda_j})\right)^{1/2}}}^{\text{all hyperparameters are here!}} 
+    f(x) \approx \sum_{j = 1}^{m^\star} 
+    \overbrace{\color{red}{\left(S(\sqrt{\boldsymbol{\lambda}_j})\right)^{1/2}}}^{\text{all hyperparameters are here!}} 
     \times
-    \underbrace{\color{blue}{\phi_{j}(x)}}_{\text{easy to compute!}}
+    \underbrace{\color{blue}{\phi_{j}(\boldsymbol{x})}}_{\text{easy to compute!}}
     \times
     \overbrace{\color{green}{\beta_{j}}}^{\sim \: \text{Normal}(0,1)}
 
-where :math:`\lambda_j` are the eigenvalues of the Laplacian operator, :math:`\phi_{j}(x)` are the eigenfunctions of the
+where :math:`\boldsymbol{x}` is a :math:`D` vector of inputs, :math:`\boldsymbol{\lambda}_j` are the eigenvalues of the Laplacian operator, :math:`\phi_{j}(\boldsymbol{x})` are the eigenfunctions of the
 Laplacian operator, and :math:`\beta_{j}` are the coefficients of the expansion (see Eq. (8) in [2]). We expect this
-to be a good approximation for a finite number of :math:`m` terms in the series as long as the inputs values :math:`x`
-are not too close to the boundaries `ell` amd `-ell`.
+to be a good approximation for a finite number of :math:`m^\star` terms in the series as long as the inputs values :math:`x`
+are not too close to the boundaries :math:`-L_d` and :math:`L_d`.
 
 .. note::
-    Even though the periodic kernel is not stationary, one can still adapt and find a similar approximation formula. 
+    Even though the periodic kernel is not stationary, one can still adapt and find a similar approximation formula. However, these kernels are not supported for multidimensional inputs.
     See Appendix B in [2] for more details.
 
 **Example:**
@@ -240,22 +243,26 @@ Other kernels can be used similarly.
 
 **References:**
 
-1. Solin, A., Särkkä, S. Hilbert space methods for reduced-rank Gaussian process regression.
-Stat Comput 30, 419-446 (2020).
+    1. Solin, A., Särkkä, S. Hilbert space methods for reduced-rank Gaussian process regression.
+       Stat Comput 30, 419-446 (2020).
 
-2. Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space
-approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).
-
-3. `Orduz, J., A Conceptual and Practical Introduction to Hilbert Space GPs Approximation Methods <https://juanitorduz.github.io/hsgp_intro>`_.
-
-4. `Example: Hilbert space approximation for Gaussian processes <https://num.pyro.ai/en/stable/examples/hsgp.html>`_.
-
-5. `Gelman, Vehtari, Simpson, et al., Bayesian workflow book - Birthdays <https://avehtari.github.io/casestudies/Birthdays/birthdays.html>`_.
+    2. Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space
+       approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).
+    
+    3. `Orduz, J., A Conceptual and Practical Introduction to Hilbert Space GPs Approximation Methods <https://juanitorduz.github.io/hsgp_intro>`_.
+    
+    4. `Example: Hilbert space approximation for Gaussian processes <https://num.pyro.ai/en/stable/examples/hsgp.html>`_.
+    
+    5. `Gelman, Vehtari, Simpson, et al., Bayesian workflow book - Birthdays <https://avehtari.github.io/casestudies/Birthdays/birthdays.html>`_.
 
 .. note::
     The code of this module is based on the code of the example
     `Example: Hilbert space approximation for Gaussian processes <https://num.pyro.ai/en/stable/examples/hsgp.html>`_ by `Omar Sosa Rodríguez <https://github.com/omarfsosa>`_.
 
+eigenindices
+----------------
+.. autofunction:: numpyro.contrib.hsgp.laplacian.eigenindices
+
 sqrt_eigenvalues
 ----------------
 .. autofunction:: numpyro.contrib.hsgp.laplacian.sqrt_eigenvalues

examples/stochastic_volatility.py

@@ -115,7 +115,7 @@ def main(args):
 
     ax.plot(dates, jnp.exp(hmc_states["s"].T), "r", alpha=0.01)
     legend = ax.legend(["returns", "volatility"], loc="upper right")
-    legend.legendHandles[1].set_alpha(0.6)
+    legend.legend_handles[1].set_alpha(0.6)
     ax.set(xlabel="time", ylabel="returns", title="Volatility of S&P500 over time")
 
     plt.savefig("stochastic_volatility_plot.pdf")

numpyro/contrib/hsgp/approximation.py

@@ -5,6 +5,8 @@
 This module contains the low-rank approximation functions of the Hilbert space Gaussian process.
 """
 
+from __future__ import annotations
+
 from jaxlib.xla_extension import ArrayImpl
 
 import jax.numpy as jnp
@@ -34,21 +36,22 @@ def _centered_approximation(phi: ArrayImpl, spd: ArrayImpl, m: int) -> ArrayImpl
 
 
 def linear_approximation(
-    phi: ArrayImpl, spd: ArrayImpl, m: int, non_centered: bool = True
+    phi: ArrayImpl, spd: ArrayImpl, m: int | list[int], non_centered: bool = True
 ) -> ArrayImpl:
     """
     Linear approximation formula of the Hilbert space Gaussian process.
 
     See Eq. (8) in [1].
 
     **References:**
-    1. Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space
-       approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).
+
+        1. Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space
+           approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).
 
     :param ArrayImpl phi: laplacian eigenfunctions
     :param ArrayImpl spd: square root of the diagonal of the spectral density evaluated at square
         root of the first `m` eigenvalues.
-    :param int m: number of eigenfunctions in the approximation
+    :param int | list[int] m: number of eigenfunctions in the approximation
     :param bool non_centered: whether to use a non-centered parameterization
     :return: The low-rank approximation linear model
     :rtype: ArrayImpl
@@ -62,8 +65,8 @@ def hsgp_squared_exponential(
     x: ArrayImpl,
     alpha: float,
     length: float,
-    ell: float,
-    m: int,
+    ell: float | int | list[float | int],
+    m: int | list[int],
     non_centered: bool = True,
 ) -> ArrayImpl:
     """
@@ -75,38 +78,44 @@ def hsgp_squared_exponential(
 
     **References:**
 
-    1. Solin, A., Särkkä, S. Hilbert space methods for reduced-rank Gaussian process regression.
-    Stat Comput 30, 419-446 (2020).
+        1. Solin, A., Särkkä, S. Hilbert space methods for reduced-rank Gaussian process regression.
+           Stat Comput 30, 419-446 (2020).
 
-    2. Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space
-    approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).
+        2. Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space
+           approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).
 
     :param ArrayImpl x: input data
     :param float alpha: amplitude of the squared exponential kernel
     :param float length: length scale of the squared exponential kernel
-    :param float ell: positive value that parametrizes the length of the one-dimensional box so that the input data
-        lies in the interval [-ell, ell]. We expect the approximation to be valid within this interval
-    :param int m: number of eigenvalues to compute and include in the approximation
+    :param float | int | list[float | int] ell: positive value that parametrizes the length of the D-dimensional box so
+        that the input data lies in the interval :math:`[-L_1, L_1] \\times ... \\times [-L_D, L_E]`.
+        We expect the approximation to be valid within this interval
+    :param int | list[m] m: number of eigenvalues to compute and include in the approximation for each dimension
+        (:math:`\\left\\{1, ..., D\\right\\}`).
+        If an integer, the same number of eigenvalues is computed in each dimension.
     :param bool non_centered: whether to use a non-centered parameterization. By default, it is set to True
     :return: the low-rank approximation linear model
     :rtype: ArrayImpl
     """
+    dim = x.shape[-1] if x.ndim > 1 else 1
     phi = eigenfunctions(x=x, ell=ell, m=m)
     spd = jnp.sqrt(
         diag_spectral_density_squared_exponential(
-            alpha=alpha, length=length, ell=ell, m=m
+            alpha=alpha, length=length, ell=ell, m=m, dim=dim
         )
     )
-    return linear_approximation(phi=phi, spd=spd, m=m, non_centered=non_centered)
+    return linear_approximation(
+        phi=phi, spd=spd, m=phi.shape[-1], non_centered=non_centered
+    )
 
 
 def hsgp_matern(
     x: ArrayImpl,
     nu: float,
     alpha: float,
     length: float,
-    ell: float,
-    m: int,
+    ell: float | int | list[float | int],
+    m: int | list[int],
     non_centered: bool = True,
 ):
     """
@@ -118,28 +127,36 @@ def hsgp_matern(
 
     **References:**
 
-    1. Solin, A., Särkkä, S. Hilbert space methods for reduced-rank Gaussian process regression.
-    Stat Comput 30, 419-446 (2020).
+        1. Solin, A., Särkkä, S. Hilbert space methods for reduced-rank Gaussian process regression.
+           Stat Comput 30, 419-446 (2020).
 
-    2. Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space
-    approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).
+        2. Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space
+           approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).
 
     :param ArrayImpl x: input data
     :param float nu: smoothness parameter
     :param float alpha: amplitude of the squared exponential kernel
     :param float length: length scale of the squared exponential kernel
-    :param float ell: positive value that parametrizes the length of the one-dimensional box so that the input data
-        lies in the interval [-ell, ell]. We expect the approximation to be valid within this interval.
-    :param int m: number of eigenvalues to compute and include in the approximation
+    :param float | int | list[float | int] ell: positive value that parametrizes the length of the D-dimensional box so
+        that the input data lies in the interval :math:`[-L_1, L_1] \\times ... \\times [-L_D, L_D]`.
+        We expect the approximation to be valid within this interval
+    :param int | list[m] m: number of eigenvalues to compute and include in the approximation for each dimension
+        (:math:`\\left\\{1, ..., D\\right\\}`).
+        If an integer, the same number of eigenvalues is computed in each dimension.
     :param bool non_centered: whether to use a non-centered parameterization. By default, it is set to True.
     :return: the low-rank approximation linear model
     :rtype: ArrayImpl
     """
+    dim = x.shape[-1] if x.ndim > 1 else 1
     phi = eigenfunctions(x=x, ell=ell, m=m)
     spd = jnp.sqrt(
-        diag_spectral_density_matern(nu=nu, alpha=alpha, length=length, ell=ell, m=m)
+        diag_spectral_density_matern(
+            nu=nu, alpha=alpha, length=length, ell=ell, m=m, dim=dim
+        )
+    )
+    return linear_approximation(
+        phi=phi, spd=spd, m=phi.shape[-1], non_centered=non_centered
     )
-    return linear_approximation(phi=phi, spd=spd, m=m, non_centered=non_centered)
 
 
 def hsgp_periodic_non_centered(
@@ -152,8 +169,8 @@ def hsgp_periodic_non_centered(
 
     **References:**
 
-    1. Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space
-    approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).
+        1. Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space
+           approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).
 
     :param ArrayImpl x: input data
     :param float alpha: amplitude