diff --git a/pymc/gp/cov.py b/pymc/gp/cov.py index 0df212aea1..624009bab3 100644 --- a/pymc/gp/cov.py +++ b/pymc/gp/cov.py @@ -614,6 +614,49 @@ 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`. + + 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 + 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) + + # 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)) + + return coeff * term_z + class Matern52(Stationary): r""" diff --git a/tests/gp/test_cov.py b/tests/gp/test_cov.py index 9334d05831..410e345f11 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,55 @@ 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) + + def test_psd_eigenvalues(self): + """Test PSD implementation using Rayleigh quotients.""" + alpha = 1.5 + ls = 0.1 + N = 1000 + L = 10.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() + + freqs = np.fft.fftfreq(N, d=dx) + omegas = 2 * np.pi * freqs + + 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 + + # 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: def test_1d(self):