Enabling Celerite2 SHO GP kernels.

[jbrande]

Summary

This PR enables users to select SHO kernels for their GPs when using Celerite2. The SHO kernel takes in the same log-amplitude A and log-timescale m, as well as the quality factor of the oscillator Q.

Why

The existing Matern 3/2 kernel is useful and extensible, but does not necessarily give any physical intuition about the system being observed. A stochastically driven damped harmonic oscillator can reasonably model physical processes commonly seen in transit lightcurves, with the appropriate choice of quality factor. E.g. Q = 1/sqrt(2) can describe stellar granulation in lightcurves.

Key Changes

GPModel.py now allows for SHO kernel choice for celerite GPs. and automatically parses the EPF for all three parameters A, m, Q.

docs/source/ecf.rst also now describes the parameters and explains the motivation for the new GP kernel.

demos/JWST/S5_template.ecf and demos/JWST/S5_fit_par_template.epf also give useful examples for the new GP kernel.

Testing

I've run some tests using the SHO kernel on some very small datasets, but haven't yet exhaustively tried different Q values, or tried to get equivalent results by setting Q=0.5 (which should collapse to the Matern 3/2 kernel). I also haven't tested multi-channel datasets or shared GP parameters yet.

[taylorbell57]

I'm pretty sure sigma shouldn't be square-rooted, right?

[jbrande]

Looking at the george API, the amplitude (sigma^2) that the kernel is multiplied by is just exp(A). In tinygp, we can also do this the same way, but the kernels.quasisep package supposedly speeds things up slightly by letting you pass in sigma on its own to the kernel function. So, sigma should be sqrt(exp(A)), but it may be clearer to just keep the amp = np.exp(self.coeffs[c, k, 0]) and then do sigma = np.sqrt(amp).

[taylorbell57]

Didn't we decide that the *2 on the line:

metric = np.exp(self.coeffs[c, k, 1]*2)

in the george code was a typo?

Or does the metric used by george use a different unit than the other codes?

[jbrande]

Didn't we decide that the *2 on the line:

metric = np.exp(self.coeffs[c, k, 1]*2)

in the george code was a typo?

Or does the metric used by george use a different unit than the other codes?

Yeah, the george API specifies that the metric is squared, where celerite2 and tinygp don't. The typo was in the original placeholder code for the tinygp kernels, where it was copied as-is from the george implementation and as such also squared the metric, when it shouldn't have. That's been fixed in this PR.

[taylorbell57]

Thanks @jbrande! I've totally messed up the target repo (jaxoplanet) while trying to re-work things, so at present here are merge conflicts that are blocking me from accepting the PR. I'll take care of those myself though, so from you're end you're all good to go and this PR is conditionally accepted!

[jbrande]

Thanks for the heads up! Sounds good.