CARMA Kernel - Githubissues

I found a couple issues with the CARMA kernel. Sorry to bomb you with all these issues, but tinpygp is so great that I can't resist playing with it. :)

First: two places throw errors when I tried to jax.jit the loss function:

https://github.com/dfm/tinygp/blob/992eb3db76587c62a27330835f91541b00acde8f/src/tinygp/kernels/quasisep.py#L669 Adding strip_zeros=False fixes the problem for me.
https://github.com/dfm/tinygp/blob/992eb3db76587c62a27330835f91541b00acde8f/src/tinygp/kernels/quasisep.py#L698 Changing np.roll to jnp.roll fixes the issue

Second: the CARMA kernel returns -inf for a reasonable kernel. A minimum testing kernel is the CARMA(1,0)/DRW kernel, which is equivalent to the Exp kernel with scale = 1/alpha_1, sigma_exp^2 = sigma_carma^2/(2*alpha_1). That said, I think all CARMA(1,0) kernel should not return -inf for the log_likelihood.

Here is a testing example:

import numpy as np
import jax
import jax.numpy as jnp
from tinygp import kernels, GaussianProcess
from tinygp.kernels.quasisep import CARMA
jax.config.update("jax_enable_x64", True)

# simulate white noise time series
npts = 200
t = np.sort(np.random.uniform(0, 3650, npts))
y = np.random.normal(loc=0, scale=0.3, size=(npts))
yerr = np.random.normal(loc=0, scale=0.03, size=(npts))

# build_gp and loss functions
def build_gp(params):
    kernel = CARMA.init(
        alpha=params["alpha"], beta=params["beta"], sigma=params["sigma"]
    )
    diag = yerr**2
    return GaussianProcess(kernel, t, diag=diag, mean=params["mean"])

@jax.jit
def loss(params):
    gp = build_gp(params)
    return -gp.log_probability(y)

# evaluate log-likelihood
params = {
    "mean": 0.0,
    "alpha": np.array([1 / 100]),
    "beta": np.array([1]),
    "sigma": .1,
}
loss(params)

what I got:

DeviceArray(inf, dtype=float64)

Here is a different implementation of the CARMA kernel, which seems to give a likelihood value and I tested it against eztao (A CARMA modeling tool relying on the celerite package). Note that the beta parameters in my version is simply the beta parameters in yours multiplied by your sigma parameter. My implementation is obviously a lot uglier, but I am happy to do a PR if makes sense.

dfm / tinygp

CARMA Kernel #89