add support for variable coefficients in celerite models

[john-livingston]

I'm interested in being able to model different subsets of the data with different kernels, i.e. K = K1 + K2, where K2 could have zeros everywhere except blocks corresponding to the target data subset. As you mentioned this isn't trivial with celerite, but could work in principle, i.e. by generalizing the kernel to be k(tau=|ti-tj|) = a_1(ti) * a_2(tj) * exp(-ctau) * cos(dtau) + b_1(ti) * b_2(tj) * exp(-ctau) * sin(dtau)

[ericagol]

And easy way to do this would be to simply compute the likelihood for the two data sets separately. Eric Agol Astronomy Professor University of Washington

…

On Sep 19, 2019, at 12:09 AM, John Livingston ***@***.***> wrote: I'm interested in being able to model different subsets of the data with different kernels, i.e. K = K1 + K2, where K2 could have zeros everywhere except blocks corresponding to the target data subset. As you mentioned this isn't trivial with celerite, but could work in principle, i.e. by generalizing the kernel to be k(tau=|ti-tj|) = a_1(ti) * a_2(tj) * exp(-ctau) * cos(dtau) + b_1(ti) * b_2(tj) * exp(-ctau) * sin(dtau) — You are receiving this because you are subscribed to this thread. Reply to this email directly, view it on GitHub, or mute the thread.

[dfm]

@ericagol: The catch here is actually that K1 will be dense (in the celerite sense), but then K2 will be block diagonal (with celerite kernels in each block). The idea here is that you might want to model in transit data using an extra kernel for spot crossings, etc. I think that this could be implemented using variable a and b and then setting them to zero if either of the times are out of transit.

The same machinery will be needed by exoplanet-dev/exoplanet#51.

[ericagol]

Yes, I see!

[dfm]

So - I worked out how to do this. Basically, if you want a kernel with the form:

k(tn, tm) = (
    a(tn) * a_prime(tm) * cos(d * (tn - tm))
    + b(tn) * b_prime(tm) * sin(dj * (tn - tm))
) * exp(-c * (tn - tm))

you can just use:

u_1 = a(tn) * sin(d*tn)
u_2 = a(tn) * cos(d*tn)
u_3 = -b(tn) * cos(d*tn)
u_4 = b(tn) * sin(d*tn)

and:

v_1 = a_prime(tm) * sin(d*tm)
v_2 = a_prime(tm) * cos(d*tm)
v_3 = b_prime(tm) * sin(d*tm)
v_4 = b_prime(tm) * cos(d*tm)

Which means that J will be larger by a factor of 2 (I can't see how to improve on that), but something like this should work for any problem where we need a variable coefficient. We'll need to think a bit about positive definiteness, etc. but this could be useful.

Here's a simple example in Python for a masked GP (ie. where the GP only acts at specific times). I'm still thinking about the best interface to actually combine this with the existing terms but I wanted to write this here so I didn't lose it.

import numpy as np
import matplotlib.pyplot as plt

import exoplanet as xo

import pymc3 as pm
import theano.tensor as tt


class VariableKernel:
    
    def __init__(self, *, a1, a2, b1, b2, c, d):
        self.a1 = a1
        self.a2 = a2
        self.b1 = b1
        self.b2 = b2
        self.c = tt.as_tensor_variable(c)
        self.d = tt.as_tensor_variable(d)
        self.J = 4
        
    def get_celerite_matrices(self, x, diag):
        x = tt.as_tensor_variable(x)
        diag = tt.as_tensor_variable(diag)
        
        a1 = self.a1(x)
        a2 = self.a2(x)
        b1 = self.b1(x)
        b2 = self.b2(x)
        
        a = diag + a1 * a2
        
        cos = tt.cos(self.d * x)
        sin = tt.sin(self.d * x)
        U = tt.stack([
            a1 * sin,
            a1 * cos,
            -b1 * cos,
            b1 * sin,
        ], axis=-1)
        V = tt.stack([
            a2 * sin,
            a2 * cos,
            b2 * sin,
            b2 * cos,
        ], axis=-1)
        
        dx = x[1:] - x[:-1]
        p0 = tt.exp(-self.c * dx)
        P = tt.stack([
            p0,
            p0,
            p0,
            p0,
        ], axis=-1)
        
        return a, U, V, P


x = np.sort(np.random.uniform(0, 10, 500))
diag = 0.01 ** 2 * np.ones_like(x)

a = lambda x: 1.0 * (x < 5.0) * (2.0 < x)
b = lambda x: 1.0 * (x < 5.0) * (2.0 < x)

kernel = VariableKernel(a1=a, a2=a, b1=b, b2=b, c=0.3, d=0.1)
gp = xo.gp.GP(kernel, x, diag)

y = gp.dot_l(np.random.randn(len(x), 1))[:, 0].eval()

plt.plot(x, y)