Add power_spectral_density to RatQuad covariance kernel for HSGP support

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"""

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