Source code for celerite2.pymc3.distribution
# -*- coding: utf-8 -*-
__all__ = ["CeleriteNormal"]
import numpy as np
import pymc3 as pm
import theano.tensor as tt
from pymc3.distributions.distribution import Continuous
import celerite2.driver
[docs]class CeleriteNormal(Continuous):
"""A multivariate normal distribution with a celerite covariance matrix
Args:
gp (celerite2.theano.GaussianProcess): The Gaussian Process computation
engine.
"""
def __init__(self, gp, *args, **kwargs):
super().__init__(*args, **kwargs)
self.gp = gp
self.mean = (
self.median
) = self.mode = self.gp.mean_value + tt.zeros_like(self.gp._t)
def random(self, point=None, size=None):
if size is None:
size = tuple()
else:
if not isinstance(size, tuple):
try:
size = tuple(size)
except TypeError:
size = (size,)
mu, U, P, d, W = pm.distributions.distribution.draw_values(
[self.mean, self.gp._U, self.gp._P, self.gp._d, self.gp._W],
point=point,
size=size,
)
n = np.random.randn(*(size + tuple([d.shape[-1]])))
func = np.vectorize(
celerite2.driver.dot_tril,
signature="(n,j),(m,j),(n),(n,j),(n)->(n)",
)
return func(U, P, d, W, n) + mu
def logp(self, y):
return self.gp.log_likelihood(y)