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):