Add support for power spectral densities to covariance function objects

[bwengals]

What is this PR about? Adding support for power spectral density functions to covariance functions. Any stationary kernel may define a psd method that takes in omega as an input.

class ExpQuad(Stationary):
    def __init__(self, ...):
       ...

    def full(self, X, Xs):
        ....

    # allow this method to be defined
    def psd(self, omega):
        ...

This will be useful for some kernel approximation algorithms like the HSGP approximation and random Fourier features. It should also make it easier to work with other things like spectral mixture kernels too.

This also adds the ability for composite kernels to be constructed and their psd to be calculated.

cov = eta1**2 * pm.gp.cov.ExpQuad(1, ls=ls1) + eta2**2 * pm.gp.cov.Matern52(1, ls=ls2) 
cov.psd(omega) # works

It should be OK now for users to define a custom kernel that subclasses Stationary, but doesn't define .full or .diag methods, just .psd, which is all some approximations would need.

Checklist

• [ ] Explain important implementation details 👆
• [ ] Make sure that the pre-commit linting/style checks pass.
• [ ] Link relevant issues (preferably in nice commit messages) No relevant issues, this is a new feature.
• [ ] Are the changes covered by tests and docstrings?
• [ ] Fill out the short summary sections 👇

Major / Breaking Changes

• Shouldn't break anything

Bugfixes / New features

• Add .psd for pm.gp.cov.ExpQuad, pm.gp.cov.Matern52 and pm.gp.cov.Matern32.
• The .psd method should work for any dimensional omega. Meaning, 1 or 2+ dimensional processes.

Docs / Maintenance

• Doesn't add anything major to the docs, just docstrings needed.

[bwengals]

Other than tests and pre-commit, putting this up in draft mode because I need to figure out how to test my translations of the power spectral densities. I got them from here, pg 4.

[codecov[bot]]

Codecov Report

Merging #6036 (dfec72b) into main (3ff4e7a) will decrease coverage by 0.16%. The diff coverage is 76.34%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6036      +/-   ##
==========================================
- Coverage   94.15%   93.98%   -0.17%     
==========================================
  Files         132      133       +1     
  Lines       26726    26921     +195     
==========================================
+ Hits        25163    25301     +138     
- Misses       1563     1620      +57     

Impacted Files	Coverage Δ
pymc/gp/hsgp.py	25.71% <25.71%> (ø)
pymc/gp/cov.py	97.18% <95.14%> (-0.91%)	:arrow_down:
pymc/tests/gp/test_cov.py	99.83% <100.00%> (+0.01%)	:arrow_up:

[fonnesbeck]

Is this ready for review now?

[bwengals]

Sorry I missed your comment @fonnesbeck. Almost, I still need to add tests for the HSGP class, but everything else should be pretty stable.

[bwengals]

I also still need to make a notebook for pymc-examples. Here's a quick usage example though that should run off this branch. Additive covariances work as well as higher dimensional X.

import numpy as np
import pymc as pm
import matplotlib.pyplot as plt
import pymc.sampling_jax

t = np.linspace(0.0, 10.0, 500)
cov = 4**2 * pm.gp.cov.ExpQuad(1, ls=3) + 1**2 * pm.gp.cov.Matern52(1, ls=0.5)
f_true = pm.draw(pm.MvNormal.dist(mu=np.zeros(len(t)), cov=cov(t[:, None])), 1).flatten()

sigma_true = 1.0
y = f_true + sigma_true * np.random.randn(len(f_true))

plt.plot(t, f_true)
plt.plot(t, y, '.');

with pm.Model() as model:
    ell1 = pm.InverseGamma("ell1", mu=2, sigma=4)
    ell2 = pm.InverseGamma("ell2", mu=1, sigma=2)
    cov = 4**2 * pm.gp.cov.ExpQuad(1, ls=ell1) + 1**2 * pm.gp.cov.Matern52(1, ls=ell2)
   
    gp = pm.gp.HSGP(m=[200], c=5.0, cov_func=cov)
    f = gp.prior("f", X=t[:, None])

    sigma = pm.HalfNormal("sigma", sigma=1)
    pm.Normal("y", mu=f, sigma=sigma, observed=y) 
    
    idata = pm.sampling_jax.sample_numpyro_nuts()

[ghost]

Is this ready to be reviewed for merging? Happy to Crete an example notebook if that's what is holding things up.

[bwengals]

Well maybe two things are left, one is need to add tests for the HSGP class. The other is maybe handling the periodic covariance. I was going to punt on that since it's structured a bit differently (doesn't use the PSD) and save it for another PR, but it might not be too hard to include.

Sure! Do you have a example in mind? Whatever the example is, I do want to make sure to include a widget of some sort to help with choosing m, c or L parameters, happy to contribute that part. I think that'd be really handy.

[ghost]

I always have plenty of baseball data (along the lines of what @danhphan has been using).

[danhphan]

Hi, the spin dataset will be on this link in pymc-examples/examples/data folder soon when the MOGPs notebook is merged.

[bwengals]

closing in favor of #6458