From 4840b885277c466fbafa6cce27db51baab4680ae Mon Sep 17 00:00:00 2001 From: Chris Fonnesbeck Date: Fri, 21 Nov 2025 21:59:34 -0500 Subject: [PATCH 1/6] Add power_spectral_density to RatQuad covariance kernel for HSGP support --- pymc/gp/cov.py | 21 +++++++++++++++++++++ tests/gp/test_cov.py | 22 +++++++++++++++++++++- 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index 0df212aea1..ba6c9b66d4 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -614,6 +614,27 @@ def full_from_distance(self, dist: TensorLike, squared: bool = False) -> TensorV -1.0 * self.alpha, ) + def power_spectral_density(self, omega: TensorLike) -> TensorVariable: + r""" + Power spectral density for the Rational Quadratic kernel. + + .. math:: + S(\boldsymbol\omega) = \frac{2 (2\pi\alpha)^{D/2} \prod_{i=1}^D \ell_i}{\Gamma(\alpha)} + \left(\frac{z}{2}\right)^{\nu} + K_{\nu}(z) + where :math:`z = \sqrt{2\alpha} \sqrt{\sum \ell_i^2 \omega_i^2}` and :math:`\nu = \alpha - D/2`. + """ + ls = pt.ones(self.n_dims) * self.ls + alpha = self.alpha + D = self.n_dims + nu = alpha - D / 2.0 + + z = pt.sqrt(2 * alpha) * pt.sqrt(pt.dot(pt.square(omega), pt.square(ls))) + coeff = 2.0 * pt.power(2.0 * np.pi * alpha, D / 2.0) * pt.prod(ls) / pt.gamma(alpha) + term_z = pt.power(z / 2.0, nu) * pt.kv(nu, z) + + return coeff * term_z + class Matern52(Stationary): r""" diff --git a/tests/gp/test_cov.py b/tests/gp/test_cov.py index 9334d05831..e3b485c140 100644 --- a/tests/gp/test_cov.py +++ b/tests/gp/test_cov.py @@ -18,7 +18,7 @@ import pytensor.tensor as pt import pytest -from scipy.special import iv +from scipy.special import gamma, iv, kv import pymc as pm @@ -533,6 +533,26 @@ def test_1d(self): Kd = cov(X, diag=True).eval() npt.assert_allclose(np.diag(K), Kd, atol=1e-5) + def test_psd(self): + omega = np.linspace(0.1, 2, 10) + ell = 0.5 + alpha = 5.0 + D = 1 + + z = np.sqrt(2 * alpha) * ell * np.abs(omega) + nu = alpha - D / 2.0 + + coeff = 2.0 * (2.0 * np.pi * alpha) ** (D / 2.0) * ell / gamma(alpha) + true_1d_psd = coeff * np.power(z / 2.0, nu) * kv(nu, z) + + test_1d_psd = ( + pm.gp.cov.RatQuad(1, ls=ell, alpha=alpha) + .power_spectral_density(omega[:, None]) + .flatten() + .eval() + ) + npt.assert_allclose(true_1d_psd, test_1d_psd, atol=1e-5) + class TestExponential: def test_1d(self): From 99d461624f570cd0c701d94cc92bb5a776bbff31 Mon Sep 17 00:00:00 2001 From: Chris Fonnesbeck Date: Mon, 24 Nov 2025 09:27:10 -0600 Subject: [PATCH 2/6] Added derivation info to docstring --- pymc/gp/cov.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index ba6c9b66d4..226c5e782a 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -623,6 +623,26 @@ def power_spectral_density(self, omega: TensorLike) -> TensorVariable: \left(\frac{z}{2}\right)^{\nu} K_{\nu}(z) where :math:`z = \sqrt{2\alpha} \sqrt{\sum \ell_i^2 \omega_i^2}` and :math:`\nu = \alpha - D/2`. + + Derivation + ---------- + The Rational Quadratic kernel can be expressed as a scale mixture of Squared Exponential kernels: + + .. math:: + k_{RQ}(r) = \int_0^\infty k_{SE}(r; \lambda) p(\lambda) d\lambda + + where :math:`k_{SE}(r; \lambda) = \exp\left(-\frac{\lambda r^2}{2}\right)` and the mixing distribution + on the precision parameter :math:`\lambda` is :math:`\lambda \sim \text{Gamma}(\alpha, \beta)` + with rate parameter :math:`\beta = \alpha \ell^2`. + + By the linearity of the Fourier transform, the PSD of the Rational Quadratic kernel is the expectation + of the PSD of the Squared Exponential kernel with respect to the mixing distribution: + + .. math:: + S_{RQ}(\omega) = \int_0^\infty S_{SE}(\omega; \lambda) p(\lambda) d\lambda + + Substituting the known PSD of the Squared Exponential kernel and evaluating the integral yields + the expression involving the modified Bessel function of the second kind, :math:`K_{\nu}(z)`. """ ls = pt.ones(self.n_dims) * self.ls alpha = self.alpha From 2a3e42084d8bde4ef7e0de1e17efe3e8c6f7fae2 Mon Sep 17 00:00:00 2001 From: Chris Fonnesbeck Date: Mon, 24 Nov 2025 17:05:48 -0600 Subject: [PATCH 3/6] Add test for covariance eigenvalue equivalence --- pymc/gp/cov.py | 8 ++++++-- tests/gp/test_cov.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index 226c5e782a..720fced7ce 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -651,8 +651,12 @@ def power_spectral_density(self, omega: TensorLike) -> TensorVariable: z = pt.sqrt(2 * alpha) * pt.sqrt(pt.dot(pt.square(omega), pt.square(ls))) coeff = 2.0 * pt.power(2.0 * np.pi * alpha, D / 2.0) * pt.prod(ls) / pt.gamma(alpha) - term_z = pt.power(z / 2.0, nu) * pt.kv(nu, z) - + + # Handle singularity at z=0 + safe_z = pt.switch(pt.eq(z, 0), 1.0, z) + term_z = pt.power(safe_z / 2.0, nu) * pt.kv(nu, safe_z) + term_z = pt.switch(pt.eq(z, 0), pt.gamma(nu) / 2.0, term_z) + return coeff * term_z diff --git a/tests/gp/test_cov.py b/tests/gp/test_cov.py index e3b485c140..90da1c4afc 100644 --- a/tests/gp/test_cov.py +++ b/tests/gp/test_cov.py @@ -17,6 +17,7 @@ import pytensor import pytensor.tensor as pt import pytest +import scipy.linalg from scipy.special import gamma, iv, kv @@ -553,6 +554,36 @@ def test_psd(self): ) npt.assert_allclose(true_1d_psd, test_1d_psd, atol=1e-5) + def test_psd_eigenvalues(self): + # Test PSD implementation using Szegő’s Theorem + alpha = 1.5 + ls = 0.5 + N = 500 + L = 50.0 + dx = L / N + X = np.linspace(0, L, N)[:, None] + + with pm.Model(): + cov = pm.gp.cov.RatQuad(1, alpha=alpha, ls=ls) + + K = cov(X).eval() + + evals = scipy.linalg.eigvalsh(K) + evals = np.sort(evals)[::-1] + + freqs = np.fft.fftfreq(N, d=dx) + omegas = 2 * np.pi * freqs + + psd = cov.power_spectral_density(omegas[:, None]).eval() + + psd_scaled = psd.flatten() / dx + psd_sorted = np.sort(psd_scaled)[::-1] + + rel_err = np.abs((evals - psd_sorted) / psd_sorted) + med_rel_err = np.median(rel_err) + + assert med_rel_err < 0.1 + class TestExponential: def test_1d(self): From dccd030bdb1eb1b7933ef6e63560cd7738877b17 Mon Sep 17 00:00:00 2001 From: Chris Fonnesbeck Date: Tue, 25 Nov 2025 20:10:56 -0600 Subject: [PATCH 4/6] Update pymc/gp/cov.py Co-authored-by: Jesse Grabowski <48652735+jessegrabowski@users.noreply.github.com> --- pymc/gp/cov.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index 720fced7ce..d130218a45 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -653,9 +653,9 @@ def power_spectral_density(self, omega: TensorLike) -> TensorVariable: coeff = 2.0 * pt.power(2.0 * np.pi * alpha, D / 2.0) * pt.prod(ls) / pt.gamma(alpha) # Handle singularity at z=0 - safe_z = pt.switch(pt.eq(z, 0), 1.0, z) - term_z = pt.power(safe_z / 2.0, nu) * pt.kv(nu, safe_z) - term_z = pt.switch(pt.eq(z, 0), pt.gamma(nu) / 2.0, term_z) + term_z = pt.switch(pt.eq(z, 0), + pt.gamma(nu) / 2.0, + pt.power(z / 2.0, nu) * pt.kv(nu, z)) return coeff * term_z From 971f6ceb3e1d40bcb2235c8b09038f351fce7c64 Mon Sep 17 00:00:00 2001 From: Chris Fonnesbeck Date: Tue, 25 Nov 2025 20:11:42 -0600 Subject: [PATCH 5/6] Better eigenvalue varname Co-authored-by: Jesse Grabowski <48652735+jessegrabowski@users.noreply.github.com> --- tests/gp/test_cov.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/gp/test_cov.py b/tests/gp/test_cov.py index 90da1c4afc..85055a6860 100644 --- a/tests/gp/test_cov.py +++ b/tests/gp/test_cov.py @@ -568,7 +568,7 @@ def test_psd_eigenvalues(self): K = cov(X).eval() - evals = scipy.linalg.eigvalsh(K) + eigs = scipy.linalg.eigvalsh(K) evals = np.sort(evals)[::-1] freqs = np.fft.fftfreq(N, d=dx) From 46343e232828d2840354b71f94ac607b84ba0daa Mon Sep 17 00:00:00 2001 From: Chris Fonnesbeck Date: Tue, 25 Nov 2025 21:15:14 -0600 Subject: [PATCH 6/6] Replace eigenvalue decomp with Rayleigh quotients --- pymc/gp/cov.py | 8 +++----- tests/gp/test_cov.py | 26 ++++++++++++-------------- 2 files changed, 15 insertions(+), 19 deletions(-) diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index d130218a45..624009bab3 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -651,12 +651,10 @@ def power_spectral_density(self, omega: TensorLike) -> TensorVariable: z = pt.sqrt(2 * alpha) * pt.sqrt(pt.dot(pt.square(omega), pt.square(ls))) coeff = 2.0 * pt.power(2.0 * np.pi * alpha, D / 2.0) * pt.prod(ls) / pt.gamma(alpha) - + # Handle singularity at z=0 - term_z = pt.switch(pt.eq(z, 0), - pt.gamma(nu) / 2.0, - pt.power(z / 2.0, nu) * pt.kv(nu, z)) - + term_z = pt.switch(pt.eq(z, 0), pt.gamma(nu) / 2.0, pt.power(z / 2.0, nu) * pt.kv(nu, z)) + return coeff * term_z diff --git a/tests/gp/test_cov.py b/tests/gp/test_cov.py index 85055a6860..410e345f11 100644 --- a/tests/gp/test_cov.py +++ b/tests/gp/test_cov.py @@ -17,7 +17,6 @@ import pytensor import pytensor.tensor as pt import pytest -import scipy.linalg from scipy.special import gamma, iv, kv @@ -555,11 +554,11 @@ def test_psd(self): npt.assert_allclose(true_1d_psd, test_1d_psd, atol=1e-5) def test_psd_eigenvalues(self): - # Test PSD implementation using Szegő’s Theorem + """Test PSD implementation using Rayleigh quotients.""" alpha = 1.5 - ls = 0.5 - N = 500 - L = 50.0 + ls = 0.1 + N = 1000 + L = 10.0 dx = L / N X = np.linspace(0, L, N)[:, None] @@ -568,21 +567,20 @@ def test_psd_eigenvalues(self): K = cov(X).eval() - eigs = scipy.linalg.eigvalsh(K) - evals = np.sort(evals)[::-1] - freqs = np.fft.fftfreq(N, d=dx) omegas = 2 * np.pi * freqs - psd = cov.power_spectral_density(omegas[:, None]).eval() + j = np.arange(N) + modes = np.exp(2j * np.pi * np.outer(np.arange(N), j) / N) + numerator = np.diag(modes @ K @ modes.conj().T).real + rayleigh_quotient = numerator / N + psd = cov.power_spectral_density(omegas[:, None]).eval() psd_scaled = psd.flatten() / dx - psd_sorted = np.sort(psd_scaled)[::-1] - - rel_err = np.abs((evals - psd_sorted) / psd_sorted) - med_rel_err = np.median(rel_err) - assert med_rel_err < 0.1 + # Trim boundaries where numerical error concentrates + trim = N // 10 + npt.assert_allclose(psd_scaled[trim:-trim], rayleigh_quotient[trim:-trim], atol=1e-2) class TestExponential: