adding faster conditional mean predictions for celerite

.pre-commit-config.yaml

@@ -21,6 +21,7 @@ repos:
   hooks:
   - id: isort
     args: []
+    additional_dependencies: [toml]
 
 - repo: https://github.com/psf/black
   rev: 19.3b0

.vscode/settings.json

@@ -23,6 +23,6 @@
     "files.trimTrailingWhitespace": true,
     "C_Cpp.autoAddFileAssociations": false,
     "editor.formatOnSave": true,
-    "python.pythonPath": "${workspaceFolder}/env/bin/python",
+    "python.pythonPath": "env/bin/python",
     "git.ignoreLimitWarning": true
 }

docs/notebooks/gp.ipynb

@@ -348,9 +348,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.8"
+   "version": "3.7.5"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }

requirements-dev.txt

@@ -1,5 +1,6 @@
 -r requirements-test.txt
 black
 isort
+toml
 flake8
 nbstripout

requirements-test.txt

@@ -2,6 +2,7 @@
 scipy
 nose
 parameterized
+arviz
 pytest
 pytest-cov>=2.6.1
 pytest-env

src/exoplanet/gp/celerite.py

@@ -6,6 +6,7 @@
 import theano.tensor as tt
 
 from ..citations import add_citations_to_model
+from ..theano_ops.celerite.conditional_mean import ConditionalMeanOp
 from ..theano_ops.celerite.diag_dot import DiagDotOp
 from ..theano_ops.celerite.factor import FactorOp
 from ..theano_ops.celerite.solve import SolveOp
@@ -56,6 +57,7 @@ def __init__(self, kernel, x, diag, J=-1, model=None):
                 J = -1
         self.J = J
 
+        self.z = None
         self.x = tt.as_tensor_variable(x)
         self.diag = tt.as_tensor_variable(diag)
         self.a, self.U, self.V, self.P = self.kernel.get_celerite_matrices(
@@ -69,6 +71,8 @@ def __init__(self, kernel, x, diag, J=-1, model=None):
         self.vector_solve_op = SolveOp(J=self.J, n_rhs=1)
         self.general_solve_op = SolveOp(J=self.J)
 
+        self.conditional_mean_op = ConditionalMeanOp(J=self.J)
+
     def log_likelihood(self, y):
         self.y = y
         self.z, _, _ = self.vector_solve_op(
@@ -85,7 +89,17 @@ def log_likelihood(self, y):
     def apply_inverse(self, rhs):
         return self.general_solve_op(self.U, self.P, self.d, self.W, rhs)[0]
 
-    def predict(self, t=None, return_var=False, return_cov=False, kernel=None):
+    def predict(
+        self,
+        t=None,
+        return_var=False,
+        return_cov=False,
+        kernel=None,
+        fast_mean=True,
+    ):
+        if self.z is None:
+            raise RuntimeError("log_likelihood must be called before predict")
+
         mu = None
         if t is None and kernel is None:
             mu = self.y - self.diag * self.z[:, 0]
@@ -107,7 +121,15 @@ def predict(self, t=None, return_var=False, return_cov=False, kernel=None):
             Kss = kernel.value(t[:, None] - t[None, :])
 
         if mu is None:
-            mu = tt.dot(Kxs, self.z)[:, 0]
+            if fast_mean:
+                U_star, V_star, inds = kernel.get_conditional_mean_matrices(
+                    self.x, t
+                )
+                mu = self.conditional_mean_op(
+                    self.U, self.V, self.P, self.z[:, 0], U_star, V_star, inds
+                )
+            else:
+                mu = tt.dot(Kxs, self.z)[:, 0]
 
         if not (return_var or return_cov):
             return mu

src/exoplanet/gp/celerite_test.py

@@ -98,7 +98,6 @@ def _get_theano_kernel(celerite_kernel):
 @pytest.mark.parametrize(
     "celerite_kernel",
     [
-        "cterms.RealTerm(log_a=0.1, log_c=0.5)",
         "cterms.RealTerm(log_a=0.1, log_c=0.5) + "
         "cterms.RealTerm(log_a=-0.1, log_c=0.7)",
         "cterms.ComplexTerm(log_a=0.1, log_c=0.5, log_d=0.1)",
@@ -119,6 +118,7 @@ def test_gp(celerite_kernel, seed=1234):
     celerite_kernel = eval(celerite_kernel)
     np.random.seed(seed)
     x = np.sort(np.random.rand(100))
+    t = np.sort(np.random.rand(50))
     yerr = np.random.uniform(0.1, 0.5, len(x))
     y = np.sin(x)
     diag = yerr ** 2
@@ -129,6 +129,11 @@ def test_gp(celerite_kernel, seed=1234):
     celerite_mu, celerite_cov = celerite_gp.predict(y)
     _, celerite_var = celerite_gp.predict(y, return_cov=False, return_var=True)
 
+    celerite_mu_t, celerite_cov_t = celerite_gp.predict(y, t=t)
+    _, celerite_var_t = celerite_gp.predict(
+        y, t=t, return_cov=False, return_var=True
+    )
+
     kernel = _get_theano_kernel(celerite_kernel)
     gp = GP(kernel, x, diag)
     loglike = gp.log_likelihood(y).eval()
@@ -142,6 +147,13 @@ def test_gp(celerite_kernel, seed=1234):
     assert np.allclose(var.eval(), celerite_var)
     assert np.allclose(cov.eval(), celerite_cov)
 
+    mu = gp.predict(t)
+    _, var = gp.predict(t, return_var=True)
+    _, cov = gp.predict(t, return_cov=True)
+    assert np.allclose(mu.eval(), celerite_mu_t)
+    assert np.allclose(var.eval(), celerite_var_t)
+    assert np.allclose(cov.eval(), celerite_cov_t)
+
 
 def test_integrated_diag(seed=1234):
     np.random.seed(seed)

src/exoplanet/gp/terms.py

@@ -137,17 +137,27 @@ def get_celerite_matrices(self, x, diag):
         )
 
         dx = x[1:] - x[:-1]
-        P = tt.concatenate(
-            (
-                tt.exp(-cr[None, :] * dx[:, None]),
-                tt.exp(-cc[None, :] * dx[:, None]),
-                tt.exp(-cc[None, :] * dx[:, None]),
-            ),
-            axis=1,
-        )
+        c = tt.concatenate((cr, cc, cc))
+        P = tt.exp(-c[None, :] * dx[:, None])
 
         return a, U, V, P
 
+    def get_conditional_mean_matrices(self, x, t):
+        ar, cr, ac, bc, cc, dc = self.coefficients
+        x = tt.as_tensor_variable(x)
+
+        inds = tt.extra_ops.searchsorted(x, t)
+        _, U_star, V_star, _ = self.get_celerite_matrices(t, t)
+
+        c = tt.concatenate((cr, cc, cc))
+        dx = t - x[tt.minimum(inds, x.size - 1)]
+        U_star *= tt.exp(-c[None, :] * dx[:, None])
+
+        dx = x[tt.maximum(inds - 1, 0)] - t
+        V_star *= tt.exp(-c[None, :] * dx[:, None])
+
+        return U_star, V_star, inds
+
     def to_dense(self, x, diag):
         K = self.value(x[:, None] - x[None, :])
         K += tt.diag(diag)

src/exoplanet/theano_ops/celerite/conditional_mean.cc

@@ -0,0 +1,62 @@
+#section support_code_apply
+
+// Apply-specific main function
+int APPLY_SPECIFIC(conditional_mean)(PyArrayObject* input0, PyArrayObject* input1,
+                                     PyArrayObject* input2, PyArrayObject* input3,
+                                     PyArrayObject* input4, PyArrayObject* input5,
+                                     PyArrayObject* input6, PyArrayObject** output0) {
+  using namespace exoplanet;
+
+  // Read in U and extract the N and J dimensions
+  int success = 0;
+  npy_intp N, J;
+  auto U_in = get_matrix_input<DTYPE_INPUT_0>(&N, &J, input0, &success);
+  if (success) return 1;
+  if (CELERITE_J != Eigen::Dynamic && J != CELERITE_J) {
+    PyErr_Format(PyExc_ValueError, "runtime value of J does not match compiled value");
+    return 1;
+  }
+
+  npy_intp M, J2;
+  auto U_star_in = get_matrix_input<DTYPE_INPUT_4>(&M, &J2, input4, &success);
+  if (success) return 1;
+  if (J2 != J) {
+    PyErr_Format(PyExc_ValueError, "dimension mismatch");
+    return 1;
+  }
+
+  npy_intp input1_shape[] = {N, J};
+  npy_intp input2_shape[] = {N - 1, J};
+  npy_intp input3_shape[] = {N};
+  npy_intp input5_shape[] = {M, J};
+  npy_intp input6_shape[] = {M};
+  auto V_in = get_input<DTYPE_INPUT_1>(2, input1_shape, input1, &success);
+  auto P_in = get_input<DTYPE_INPUT_2>(2, input2_shape, input2, &success);
+  auto z_in = get_input<DTYPE_INPUT_3>(1, input3_shape, input3, &success);
+  auto V_star_in = get_input<DTYPE_INPUT_5>(2, input5_shape, input5, &success);
+  auto inds_in = get_input<DTYPE_INPUT_6>(1, input6_shape, input6, &success);
+  if (success) return 1;
+
+  auto mu_out =
+      allocate_output<DTYPE_OUTPUT_0>(1, input6_shape, TYPENUM_OUTPUT_0, output0, &success);
+  if (success) return 1;
+
+  Eigen::Map<Eigen::Matrix<DTYPE_INPUT_0, Eigen::Dynamic, CELERITE_J, CELERITE_J_ORDER>> U(U_in, N,
+                                                                                           J);
+  Eigen::Map<Eigen::Matrix<DTYPE_INPUT_1, Eigen::Dynamic, CELERITE_J, CELERITE_J_ORDER>> V(V_in, N,
+                                                                                           J);
+  Eigen::Map<Eigen::Matrix<DTYPE_INPUT_2, Eigen::Dynamic, CELERITE_J, CELERITE_J_ORDER>> P(
+      P_in, N - 1, J);
+  Eigen::Map<Eigen::Matrix<DTYPE_INPUT_3, Eigen::Dynamic, 1>> z(z_in, N);
+  Eigen::Map<Eigen::Matrix<DTYPE_INPUT_4, Eigen::Dynamic, CELERITE_J, CELERITE_J_ORDER>> U_star(
+      U_star_in, M, J);
+  Eigen::Map<Eigen::Matrix<DTYPE_INPUT_5, Eigen::Dynamic, CELERITE_J, CELERITE_J_ORDER>> V_star(
+      V_star_in, M, J);
+  Eigen::Map<Eigen::Matrix<DTYPE_INPUT_6, Eigen::Dynamic, 1>> inds(inds_in, M);
+
+  Eigen::Map<Eigen::Matrix<DTYPE_OUTPUT_0, Eigen::Dynamic, 1>> mu(mu_out, N);
+
+  celerite::conditional_mean(U, V, P, z, U_star, V_star, inds, mu);
+
+  return 0;
+}

src/exoplanet/theano_ops/celerite/conditional_mean.py

@@ -0,0 +1,24 @@
+# -*- coding: utf-8 -*-
+
+__all__ = ["ConditionalMeanOp"]
+
+import theano.tensor as tt
+from theano import gof
+
+from .base_op import CeleriteBaseOp
+
+
+class ConditionalMeanOp(CeleriteBaseOp):
+
+    func_file = "./conditional_mean.cc"
+    func_name = "APPLY_SPECIFIC(conditional_mean)"
+    num_input = 7
+    output_ndim = (1,)
+
+    def make_node(self, *args):
+        in_args = [tt.as_tensor_variable(a) for a in args]
+        out_args = [in_args[3].type()]
+        return gof.Apply(self, in_args, out_args)
+
+    def infer_shape(self, node, shapes):
+        return (shapes[-1],)