Skip to content

Improve type hints for the HSGP contrib module #2211

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
juanitorduz merged 1 commit into from
Jul 2, 2026
Merged

Improve type hints for the HSGP contrib module #2211

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 18 additions & 18 deletions numpyro/contrib/hsgp/approximation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,8 +68,8 @@ def linear_approximation(

def hsgp_squared_exponential(
x: ArrayLike,
alpha: float,
length: float,
alpha: ArrayLike,
length: ArrayLike,
ell: float | int | list[float | int],
m: int | list[int],
non_centered: bool = True,
Expand All @@ -90,8 +90,8 @@ def hsgp_squared_exponential(
approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).

:param ArrayLike x: input data
:param float alpha: amplitude of the squared exponential kernel
:param float length: length scale of the squared exponential kernel
:param ArrayLike alpha: amplitude of the squared exponential kernel
:param ArrayLike length: length scale of the squared exponential kernel
: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
Expand All @@ -117,8 +117,8 @@ def hsgp_squared_exponential(
def hsgp_matern(
x: ArrayLike,
nu: float,
alpha: float,
length: float,
alpha: ArrayLike,
length: ArrayLike,
ell: float | int | list[float | int],
m: int | list[int],
non_centered: bool = True,
Expand All @@ -140,8 +140,8 @@ def hsgp_matern(

:param ArrayLike 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 ArrayLike alpha: amplitude of the squared exponential kernel
:param ArrayLike length: length scale of the squared exponential kernel
: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
Expand All @@ -166,9 +166,9 @@ def hsgp_matern(

def hsgp_rational_quadratic(
x: ArrayLike,
alpha: float,
length: float,
scale_mixture: float,
alpha: ArrayLike,
length: ArrayLike,
scale_mixture: ArrayLike,
ell: float | int | list[float | int],
m: int | list[int],
non_centered: bool = True,
Expand Down Expand Up @@ -202,9 +202,9 @@ def hsgp_rational_quadratic(
approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).

:param ArrayLike x: input data
:param float alpha: amplitude of the Rational Quadratic kernel
:param float length: length scale of the Rational Quadratic kernel (scalar, isotropic only)
:param float scale_mixture: scale mixture parameter (α in the RQ kernel formula).
:param ArrayLike alpha: amplitude of the Rational Quadratic kernel
:param ArrayLike length: length scale of the Rational Quadratic kernel (scalar, isotropic only)
:param ArrayLike scale_mixture: scale mixture parameter (α in the RQ kernel formula).
Controls the relative weighting of small-scale and large-scale variations.
As scale_mixture → ∞, the kernel converges to the squared exponential kernel.
:param float | int | list[float | int] ell: positive value that parametrizes the length of the D-dimensional box so
Expand Down Expand Up @@ -235,7 +235,7 @@ def hsgp_rational_quadratic(


def hsgp_periodic_non_centered(
x: ArrayLike, alpha: float, length: float, w0: float, m: int
x: ArrayLike, alpha: ArrayLike, length: ArrayLike, w0: ArrayLike, m: int
) -> Array:
"""
Low rank approximation for the periodic squared exponential kernel in the non-centered parametrization.
Expand All @@ -248,9 +248,9 @@ def hsgp_periodic_non_centered(
approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023).

:param ArrayLike x: input data
:param float alpha: amplitude
:param float length: length scale
:param float w0: frequency of the periodic kernel
:param ArrayLike alpha: amplitude
:param ArrayLike length: length scale
:param ArrayLike w0: frequency of the periodic kernel
:param int m: number of eigenvalues to compute and include in the approximation
:return: the low-rank approximation linear model
:rtype: Array
Expand Down
4 changes: 2 additions & 2 deletions numpyro/contrib/hsgp/laplacian.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,12 +162,12 @@ def eigenfunctions(x: ArrayLike, ell: float | list[float], m: int | list[int]) -
)


def eigenfunctions_periodic(x: ArrayLike, w0: float, m: int) -> tuple[Array, Array]:
def eigenfunctions_periodic(x: ArrayLike, w0: ArrayLike, m: int) -> tuple[Array, Array]:
"""
Basis functions for the approximation of the periodic kernel.

:param ArrayLike x: The points at which to evaluate the eigenfunctions.
:param float w0: The frequency of the periodic kernel.
:param ArrayLike w0: The frequency of the periodic kernel.
:param int m: The number of eigenfunctions to compute.

.. note::
Expand Down
70 changes: 36 additions & 34 deletions numpyro/contrib/hsgp/spectral_densities.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ def align_param(dim, param):


def spectral_density_squared_exponential(
dim: int, w: ArrayLike, alpha: float, length: float | ArrayLike
dim: int, w: ArrayLike, alpha: ArrayLike, length: ArrayLike
) -> Array:
"""
Spectral density of the squared exponential kernel.
Expand All @@ -42,20 +42,20 @@ def spectral_density_squared_exponential(

:param int dim: dimension
:param ArrayLike w: frequency
:param float alpha: amplitude
:param float length: length scale
:param ArrayLike alpha: amplitude
:param ArrayLike length: length scale
:return: spectral density value
:rtype: Array
"""
length = align_param(dim, length)
c = alpha * jnp.prod(jnp.sqrt(2 * jnp.pi) * length, axis=-1)
c = jnp.prod(jnp.sqrt(2 * jnp.pi) * length, axis=-1) * alpha
e = jnp.exp(-0.5 * jnp.sum(w**2 * length**2, axis=-1))
return c * e


def spectral_density_matern(
dim: int, nu: float, w: ArrayLike, alpha: float, length: float | ArrayLike
) -> float:
dim: int, nu: float, w: ArrayLike, alpha: ArrayLike, length: ArrayLike
) -> Array:
"""
Spectral density of the Matérn kernel.

Expand All @@ -78,10 +78,10 @@ def spectral_density_matern(
:param int dim: dimension
:param float nu: smoothness
:param ArrayLike w: frequency
:param float alpha: amplitude
:param float length: length scale
:param ArrayLike alpha: amplitude
:param ArrayLike length: length scale
:return: spectral density value
:rtype: float
:rtype: Array
""" # noqa: E501
length = align_param(dim, length)
c1 = (
Expand All @@ -98,8 +98,8 @@ def spectral_density_matern(


def diag_spectral_density_squared_exponential(
alpha: float,
length: float | list[float],
alpha: ArrayLike,
length: ArrayLike,
ell: float | int | list[float | int],
m: int | list[int],
dim: int,
Expand All @@ -108,8 +108,8 @@ def diag_spectral_density_squared_exponential(
Evaluates the spectral density of the squared exponential kernel at the first :math:`D \\times m^\\star`
square root eigenvalues of the laplacian operator in :math:`[-L_1, L_1] \\times ... \\times [-L_D, L_D]`.

:param float alpha: amplitude of the squared exponential kernel
:param float length: length scale of the squared exponential kernel
:param ArrayLike alpha: amplitude of the squared exponential kernel
:param ArrayLike length: length scale of the squared exponential kernel
:param float | int | list[float | int] ell: The length of the interval divided by 2 in each dimension.
If a float or int, the same length is used in each dimension.
:param int | list[int] m: The number of eigenvalues to compute for each dimension.
Expand All @@ -125,7 +125,7 @@ def _spectral_density(w):
dim=dim,
w=w,
alpha=alpha,
length=length, # ty: ignore[invalid-argument-type]
length=length,
)

sqrt_eigenvalues_ = sqrt_eigenvalues(ell=ell, m=m, dim=dim) # dim x m
Expand All @@ -135,8 +135,8 @@ def _spectral_density(w):
# TODO support length-D kernel hyperparameters
def diag_spectral_density_matern(
nu: float,
alpha: float,
length: float,
alpha: ArrayLike,
length: ArrayLike,
ell: float | int | list[float | int],
m: int | list[int],
dim: int,
Expand All @@ -146,8 +146,8 @@ def diag_spectral_density_matern(
square root eigenvalues of the laplacian operator in :math:`[-L_1, L_1] \\times ... \\times [-L_D, L_D]`.

:param float nu: smoothness parameter
:param float alpha: amplitude of the Matérn kernel
:param float length: length scale of the Matérn kernel
:param ArrayLike alpha: amplitude of the Matérn kernel
:param ArrayLike length: length scale of the Matérn kernel
:param float | int | list[float | int] ell: The length of the interval divided by 2 in each dimension.
If a float or int, the same length is used in each dimension.
:param int | list[int] m: The number of eigenvalues to compute for each dimension.
Expand Down Expand Up @@ -191,9 +191,9 @@ def modified_bessel_second_kind(v, z):
def spectral_density_rational_quadratic(
dim: int,
w: ArrayLike,
alpha: float,
length: float | ArrayLike,
scale_mixture: float,
alpha: ArrayLike,
length: ArrayLike,
scale_mixture: ArrayLike,
) -> Array:
"""
Spectral density of the Rational Quadratic kernel.
Expand Down Expand Up @@ -240,9 +240,9 @@ def spectral_density_rational_quadratic(

:param int dim: dimension
:param ArrayLike w: frequency
:param float alpha: amplitude (σ² in the spectral density formula)
:param float length: length scale (scalar, isotropic)
:param float scale_mixture: scale mixture parameter (α_mix in the RQ kernel formula).
:param ArrayLike alpha: amplitude (σ² in the spectral density formula)
:param ArrayLike length: length scale (scalar, isotropic)
:param ArrayLike scale_mixture: scale mixture parameter (α_mix in the RQ kernel formula).
Controls the relative weighting of small-scale and large-scale variations.
As scale_mixture → ∞, the kernel converges to the squared exponential kernel.
:return: spectral density value
Expand Down Expand Up @@ -304,9 +304,9 @@ def spectral_density_rational_quadratic(


def diag_spectral_density_rational_quadratic(
alpha: float,
length: float | list[float],
scale_mixture: float,
alpha: ArrayLike,
length: ArrayLike,
scale_mixture: ArrayLike,
ell: float | int | list[float | int],
m: int | list[int],
dim: int,
Expand All @@ -315,9 +315,9 @@ def diag_spectral_density_rational_quadratic(
Evaluates the spectral density of the Rational Quadratic kernel at the first :math:`D \\times m^\\star`
square root eigenvalues of the laplacian operator in :math:`[-L_1, L_1] \\times ... \\times [-L_D, L_D]`.

:param float alpha: amplitude of the Rational Quadratic kernel
:param float length: length scale of the Rational Quadratic kernel
:param float scale_mixture: scale mixture parameter (α in the RQ kernel formula).
:param ArrayLike alpha: amplitude of the Rational Quadratic kernel
:param ArrayLike length: length scale of the Rational Quadratic kernel
:param ArrayLike scale_mixture: scale mixture parameter (α in the RQ kernel formula).
Controls the relative weighting of small-scale and large-scale variations.
As scale_mixture → ∞, the kernel converges to the squared exponential kernel.
:param float | int | list[float | int] ell: The length of the interval divided by 2 in each dimension.
Expand All @@ -335,7 +335,7 @@ def _spectral_density(w):
dim=dim,
w=w,
alpha=alpha,
length=length, # ty: ignore[invalid-argument-type]
length=length,
scale_mixture=scale_mixture,
)

Expand All @@ -356,7 +356,9 @@ def modified_bessel_first_kind(v, z):
return jnp.exp(jnp.abs(z)) * tfp.math.bessel_ive(v, z)


def diag_spectral_density_periodic(alpha: float, length: float, m: int) -> Array:
def diag_spectral_density_periodic(
alpha: ArrayLike, length: ArrayLike, m: int
) -> Array:
"""
Not actually a spectral density but these are used in the same
way. These are simply the first `m` coefficients of the low rank
Expand All @@ -367,8 +369,8 @@ def diag_spectral_density_periodic(alpha: float, length: float, m: int) -> Array
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 float alpha: amplitude
:param float length: length scale
:param ArrayLike alpha: amplitude
:param ArrayLike length: length scale
:param int m: number of eigenvalues
:return: "spectral density" vector
:rtype: Array
Expand Down
Loading