Add power_spectral_density to RatQuad covariance kernel for HSGP support

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

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