HSGP periodic (#6877)

* implement periodic psd and pass test

* simplify hsgp test, set up for mark other cov_func

* attempt at if-else logic in HSGP class API

* np.pi in w0

* check arguments for period and nonperiod case

* tests run but failing

* fix one of the mypy errors

* periodic cov pytest uses J not m because mpyp

* avoid AttributeError by init parameterization

* get rid of mypy "None" error

* self.prior_linearized inside if avoids double call

* need to add the mean function in prior

* run pre-commit locally

* add docs to test about why test_prior fails

* freshen up the docs

* improve coverage of args

* mypy ignore type

* fix parameterization test

* use power_spectral_density_approx because cant sum

* tbe -> the typo docs

* expand warning and change the triggering clause

* separate into HSGP and HSGPPeriodic

* remove commented periodic in HSGP test

* appease mypy

* add scale parameter for HSGP

Author: theorashid

pymc/gp/__init__.py

@@ -22,4 +22,4 @@
     MarginalKron,
     MarginalSparse,
 )
-from pymc.gp.hsgp_approx import HSGP
+from pymc.gp.hsgp_approx import HSGP, HSGPPeriodic

pymc/gp/cov.py

@@ -790,6 +790,28 @@ def full_from_distance(self, dist: TensorLike, squared: bool = False) -> TensorV
         r2 = dist if squared else dist**2
         return pt.exp(-0.5 * r2)
 
+    def power_spectral_density_approx(self, J: TensorLike) -> TensorVariable:
+        """
+        Technically, this is not a spectral density but these are the first `m` coefficients of
+        the low rank approximation for the periodic kernel, which are used in the same way.
+        `J` is a vector of `np.arange(m)`.
+
+        The coefficients of the HSGP approximation for the `Periodic` kernel are given by:
+
+        .. math::
+
+            \tilde{q}_j^2 = \frac{2 \text{I}_j (\\ell^{-2})}{\\exp(\\ell^{-2})} \\
+            \tilde{q}_0^2 = \frac{\text{I}_0 (\\ell^{-2})}{\\exp(\\ell^{-2})}
+
+        where $\text{I}_j$ is the modified Bessel function of the first kind.
+        """
+        a = 1 / pt.square(self.ls)
+        c = pt.where(J > 0, 2, 1)
+
+        q2 = c * pt.iv(J, a) / pt.exp(a)
+
+        return q2
+
 
 class Linear(Covariance):
     r"""