ENH: Spectral density of periodic kernels for HSGPs

[theorashid]

Before

No response

After

No response

Context for the issue:

I was looking to port the birthdays model example for HSGPs (see numpyro and Aki).

This model builds a long term trend, smooth year seasonality, slowly varying day of week effect, day of the year effect and special floating days effects. The long term trend and seasonal effect are built using GPs, the former with a SE kernel and the latter with a periodic kernel.

The power spectral density for the SE (ExpQuad) kernel is implemented, but there is no spectral density for the periodic kernel. (This numpyro example also calls it a periodic SE kernel, although the Stan page says just "periodic", so something is wrong).

@bwengals is the spectral density for the periodic kernel something worth implementing or is it too niche? I think it's also a low rank approximation rather than exact

PS: There should be a HSGP example! I would be willing to contribute this as an example, although looking at the numpyro example, there might be a bit too much with the different holiday effects. Perhaps Bill's/McElreath's cherry blossom example might be better.

[twiecki]

@theorashid Thanks for your issue. I think a PSD for the periodic kernel as well as a HSGP example would be great contributions!

[theorashid]

Hey Thomas, hey Bill,

Looking at this a bit more closely, the paper appendix B says "A GP model with a periodic covariance function does not fit in the framework of the HSGP approximation covered in this study, but it has also a low-rank representation" It then goes on to explain the low-rank approximation to the periodic squared exponential kernel, which I'm guessing is the same as "periodic" as they both have the form $\exp{(- \sin^2(\text{stuff}))}$.

The problem is, the expansion relies on Bessel functions of the first kind. The author of the numpyro example calls tfp.math.bessel_ive, presumably because it isn't in jnp. So because this function isn't in pytensor, I don't think it can be done.

[twiecki]

We could borrow them like we do here: https://github.com/pymc-devs/pytensor/pull/403/files

[theorashid]

Thanks! Okay that's one problem fixed. I've run into a new one that might need @bwengals help.

The expansion doesn't use the eigenvectors of the Laplacian (i.e, not just the sinusoidal eigenfunctions then phi @ (pt.sqrt(psd) * beta)). It actually uses periodic eigenfunctions. See the difference in these lines from Stan. So although we can get the power spectral density, the eigenvectors would need a more flexible API to allow different eignenvectors, which may or may not be worth the hassle.

On another note, where are you getting the notation from for the power spectral density in the docs? It's different to the original paper, and it would be nice to know what D means rather than guessing the link with the paper's notation.

[bwengals]

Def worth implementing imo, feel free - either the class or an example. I wanted to implement the periodic case when I was doing HSGP, but the implementations really are pretty different and it would have been a mess. One possibility is making a separate HSGPPeriodic class for just that special case. HSGP could dispatch to it if the user passes in the Periodic covariance. Not sure if it's worth figuring out how to do it in > 1 dimension.

D should be the same as the paper, it's in eqs 1, 2, 3. The code it's called n_dims.

RE Periodic vs. Periodic ExpQuad, @jahall recently added support for Periodic versions of Matern or any other stationary kernel, but it's unrelated to HSGP.

[theorashid]

Ahhh okay that makes sense. So the periodic HSGP is only for $D=1$ in the paper. I did a quick search and I haven't found anyone extending it to more than 1 dimensions.

I have a few more design questions, then I'll open a draft PR and we can work on it there. I'll probably need a reasonable amount of help because I'm unfamiliar with the pymc.gp API.

I haven't tested this code or the maths, but here's the code that I've sketched out and let me know if this is how you viewed it happening:

1. in cov.py

class Periodic(Stationary):
     ...
    
    def power_spectral_density(self, m: TensorLike) -> TensorVariable:
        """
        Not actually a spectral density but these are used in the same
        way. These are the first `m` coefficients of the low rank
        approximation for the periodic kernel.
        """
        a = 1 / (self.ls ** 2)
        J = pt.arange(0, m)
        c = pt.where(J > 0, 2, 1)

        # modified bessel first kind
        bessel = pt.iv(J, a)

        q2 = (c * bessel / pt.exp(a))

        return q2

2. In hsgp_approx.py

def calc_eigenvectors_periodic(
    Xs: TensorLike,
    period: TensorLike,
    m: int,
    tl: ModuleType = np,
):
    """
    The eigenvectors do not depend on the covariance hyperparameters.
    m is just an int here because we are not in multiple dimensions.
    If I'm right, I don't think we need the eigenvalues
    """
    w0 = (2 * tl.pi) / period # angular frequency defining the periodicity
    m1 = tl.tile(w0 * Xs[:, None], m)
    m2 = tl.diag(tl.arange(m))
    mw0x = m1 @ m2
    phi_cos = tl.cos(mw0x)
    phi_sin = tl.sin(mw0x)
    return phi_cos, phi_sin


class HSGPPeriodic(Base):
    def __init__():
        # assert cov_func is periodic, D=1 etc
        pass

    def prior_linearized(self, Xs: TensorLike):
        # Index Xs using input_dim and active_dims of covariance function
        Xs, _ = self.cov_func._slice(Xs)

        phi_cos, phi_sin = calc_eigenvectors_periodic(Xs, self.cov_func.period, self._m, tl=pt)

        psd = self.cov_func.power_spectral_density(self._m)

        # will be called as:
        # f = phi_cos @ (psd * beta_cos) + phi_sin @ (psd * beta_sin)

        return (phi_cos, phi_sin), psd

I don't think we need to calculate eigenvalues at all, but let me know if I've interpreted the maths incorrectly.

In the paper, they have the variable alpha, which I think is the variance of the effect. Do you know where we should put this in?

We'll also need to write something so users can just call HSGP rather than HSGPPeriodic

[theorashid]

Closed by #6877