PyMC3 GP -> Celerite

docs/dot/forwardmodel.rst

@@ -68,7 +68,7 @@ We can now call our model on the start dictionary, and plot it like so:
 
     import matplotlib.pyplot as plt
 
-    forward_model_start = m(start_params)
+    forward_model_start, var = m(start_params)
 
     plt.plot(m.lc.time[m.mask], m.lc.flux[m.mask], m.lc.flux_err[m.mask],
              color='k', fmt='.', ecolor='silver')
@@ -114,7 +114,7 @@ We can now call our model on the start dictionary, and plot it like so:
 
     import matplotlib.pyplot as plt
 
-    forward_model_start = m(start_params)
+    forward_model_start, var = m(start_params)
 
     plt.plot(m.lc.time[m.mask][::m.skip_n_points], forward_model_start,
              color='DodgerBlue')
@@ -191,7 +191,7 @@ rate. This time when we plot the result we'll see a more complicated model:
 
     import matplotlib.pyplot as plt
 
-    forward_model_two = m(two_spot_params)
+    forward_model_two, var = m(two_spot_params)
 
     plt.plot(m.lc.time[m.mask][::m.skip_n_points], forward_model_two,
              color='DodgerBlue')

docs/dot/gettingstarted.rst

@@ -124,10 +124,12 @@ a posteriori solution:
     with m:
         map_soln = pm.find_MAP()
 
+    fit, var = m(map_soln)
+
     plt.errorbar(m.lc.time[m.mask], m.lc.flux[m.mask],
                  m.scale_errors * m.lc.flux_err[m.mask], color='k',
                  ecolor='silver', fmt='.')
-    plt.plot(m.lc.time[m.mask][::m.skip_n_points], m(map_soln),
+    plt.plot(m.lc.time[m.mask][::m.skip_n_points], fit,
              color='DodgerBlue')
     plt.gca().set(xlabel='Time [d]', ylabel='Flux')
     plt.show()
@@ -158,7 +160,7 @@ plots, not to give converged publication-ready results. Always make the
 ``draws`` parameter as large as you can tolerate!
 
 The ``init`` keyword argument is set to ``'jitter+adapt_full'``, and this is
-very important. This uses `Daniel Foreman-Mackey's dense mass matrix setting
+important. This uses `Daniel Foreman-Mackey's dense mass matrix setting
 <https://dfm.io/posts/pymc3-mass-matrix/>`_ which is critical for getting fast
 results from highly degenerate model parameterizations (like this one).
 

docs/dot/plots.rst

@@ -16,7 +16,6 @@ the one below:
 
     m, trace_nuts, summary = load_results(results_dir)
 
-    posterior_shear(m, trace_nuts)
     movie(results_dir, m, trace_nuts, xsize=250)
 
 .. raw:: html

dot/model.py

@@ -2,41 +2,42 @@
 
 import numpy as np
 import pymc3 as pm
+import theano.tensor as tt
+from celerite2.theano import terms, GaussianProcess
 
 __all__ = ['Model']
 
 
-class MeanModel(pm.gp.mean.Mean):
+class MeanModel(object):
     """
     Mean model for Gaussian process regression on photometry with starspots
     """
     def __init__(self, light_curve, rotation_period, n_spots, contrast, t0,
                  latitude_cutoff=10, partition_lon=True):
-        pm.gp.mean.Mean.__init__(self)
 
         if contrast is None:
             contrast = pm.TruncatedNormal("contrast", lower=0.01, upper=0.99,
                                           testval=0.4, mu=0.5, sigma=0.5)
         self.contrast = contrast
 
-        self.f0 = pm.TruncatedNormal("f0", mu=0, sigma=1,
-                                     testval=0,
-                                     lower=-1, upper=2)
-
+        self.f0 = pm.TruncatedNormal("f0", mu=light_curve.flux.max(), sigma=1,
+                                     testval=light_curve.flux.max(),
+                                     lower=-2, upper=2)
         self.eq_period = pm.TruncatedNormal("P_eq",
                                             lower=0.8 * rotation_period,
                                             upper=1.2 * rotation_period,
                                             mu=rotation_period,
                                             sigma=0.2 * rotation_period,
                                             testval=rotation_period)
 
-        self.ln_shear = pm.Uniform("ln_shear", lower=-5, upper=np.log(0.8),
-                                   testval=np.log(0.1))
+        eps = 1e-5  # Small but non-zero number
+        BoundedHalfNormal = pm.Bound(pm.HalfNormal, lower=eps, upper=0.8)
+        self.shear  = BoundedHalfNormal("shear", testval=0.2)
 
-        self.comp_inclination = pm.Uniform("comp_inc",
-                                           lower=0,
-                                           upper=np.pi/2,
-                                           testval=np.radians(1))
+        self.comp_inclination = 0  #pm.Uniform("comp_inc",
+                                   #        lower=0,
+                                   #        upper=np.pi/2,
+                                   #        testval=np.radians(1))
 
         if partition_lon:
             lon_lims = 2 * np.pi * np.arange(n_spots + 1) / n_spots
@@ -59,42 +60,40 @@ def __init__(self, light_curve, rotation_period, n_spots, contrast, t0,
                               shape=(1, n_spots),
                               testval=np.pi/2)
 
-        eps = 1e-5  # Small but non-zero number
-        BoundedHalfNormal = pm.Bound(pm.HalfNormal, lower=eps, upper=0.8)
         self.rspot = BoundedHalfNormal("R_spot",
-                                       sigma=0.2,
+                                       sigma=0.4,
                                        shape=(1, n_spots),
                                        testval=0.3)
 
-        self.spot_period = self.eq_period / (1 - pm.math.exp(self.ln_shear) *
-                                             pm.math.sin(self.lat - np.pi / 2) ** 2)
-        self.sin_lat = pm.math.sin(self.lat)
-        self.cos_lat = pm.math.cos(self.lat)
-        self.sin_c_inc = pm.math.sin(self.comp_inclination)
-        self.cos_c_inc = pm.math.cos(self.comp_inclination)
+        self.spot_period = self.eq_period / (1 - self.shear *
+                                             tt.sin(self.lat - np.pi / 2) ** 2)
+        self.sin_lat = tt.sin(self.lat)
+        self.cos_lat = tt.cos(self.lat)
+        self.sin_c_inc = tt.sin(self.comp_inclination)
+        self.cos_c_inc = tt.cos(self.comp_inclination)
         self.t0 = t0
 
     def __call__(self, X):
-        phi = 2 * np.pi / self.spot_period * (X - self.t0) - self.lon
+        phi = 2 * np.pi / self.spot_period * (X[:, None] - self.t0) - self.lon
 
-        spot_position_x = (pm.math.cos(phi - np.pi / 2) *
+        spot_position_x = (tt.cos(phi - np.pi / 2) *
                            self.sin_c_inc *
                            self.sin_lat +
                            self.cos_c_inc *
                            self.cos_lat)
-        spot_position_y = -(pm.math.sin(phi - np.pi/2) *
+        spot_position_y = -(tt.sin(phi - np.pi/2) *
                             self.sin_lat)
         spot_position_z = (self.cos_lat *
                            self.sin_c_inc -
-                           pm.math.sin(phi) *
+                           tt.sin(phi) *
                            self.cos_c_inc *
                            self.sin_lat)
         rsq = spot_position_x ** 2 + spot_position_y ** 2
-        spot_model = self.f0 - pm.math.sum(self.rspot ** 2 * (1 - self.contrast) *
-                                           pm.math.where(spot_position_z > 0,
-                                                         pm.math.sqrt(1 - rsq),
-                                                         0),
-                                           axis=1)
+        spot_model = self.f0 - tt.sum(self.rspot ** 2 *
+                                      (1 - self.contrast) *
+                                      tt.where(spot_position_z > 0,
+                                               tt.sqrt(1 - rsq), 0),
+                                      axis=1)
 
         return spot_model
 
@@ -104,7 +103,7 @@ class DisableLogger(object):
     Simple logger disabler to minimize info-level messages during PyMC3
     integration
     """
-    def __init__(self, verbose):
+    def __init__(self, verbose=False):
         self.verbose = verbose
 
     def __enter__(self):
@@ -120,7 +119,7 @@ class Model(object):
     def __init__(self, light_curve, rotation_period, n_spots, scale_errors=1,
                  skip_n_points=1, latitude_cutoff=10, rho_factor=250,
                  verbose=False, min_time=None, max_time=None, contrast=0.7,
-                 partition_lon=True):
+                 partition_lon=True, use_gp=False, name=None):
         """
         Construct a new instance of `~dot.Model`.
 
@@ -166,15 +165,15 @@ def __init__(self, light_curve, rotation_period, n_spots, scale_errors=1,
         self.scale_errors = scale_errors
 
         self.pymc_model = None
-        self.pymc_gp = None
-        self.pymc_gp_white = None
-        self.pymc_gp_matern = None
+        self.gp = None
 
         self._initialize_model(latitude_cutoff=latitude_cutoff,
                                rho_factor=rho_factor,
-                               partition_lon=partition_lon)
+                               partition_lon=partition_lon,
+                               use_gp=use_gp, name=name)
 
-    def _initialize_model(self, latitude_cutoff, partition_lon, rho_factor):
+    def _initialize_model(self, latitude_cutoff, partition_lon, rho_factor,
+                          use_gp=False, name=None):
         """
         Construct a PyMC3 model instance for use with samplers.
 
@@ -188,7 +187,7 @@ def _initialize_model(self, latitude_cutoff, partition_lon, rho_factor):
             Scale up the GP length scale by a factor `rho_factor`
             larger than the estimated `rotation_period`
         """
-        with pm.Model(name='dot') as model:
+        with pm.Model(name=name) as model:
             mean_func = MeanModel(
                   self.lc,
                   self.rotation_period,
@@ -199,28 +198,27 @@ def _initialize_model(self, latitude_cutoff, partition_lon, rho_factor):
                   partition_lon=partition_lon
             )
 
-            x = self.lc.time[self.mask][::self.skip_n_points]
-            y = self.lc.flux[self.mask][::self.skip_n_points]
-            yerr = self.scale_errors * self.lc.flux_err[self.mask][::self.skip_n_points]
+            x = np.ascontiguousarray(self.lc.time[self.mask][::self.skip_n_points], dtype='float64')
+            y = np.ascontiguousarray(self.lc.flux[self.mask][::self.skip_n_points], dtype='float64')
+            yerr = np.ascontiguousarray(self.scale_errors * self.lc.flux_err[self.mask][::self.skip_n_points], dtype='float64')
 
             ls = rho_factor * self.rotation_period
             mean_err = yerr.mean()
-            gp_white = pm.gp.Marginal(mean_func=mean_func,
-                                      cov_func=pm.gp.cov.WhiteNoise(mean_err))
-            gp_matern = pm.gp.Marginal(cov_func=mean_err ** 2 *
-                                       pm.gp.cov.Matern32(1, ls=ls))
+            sigma = pm.HalfNormal("sigma", sigma=mean_err)
 
-            gp = gp_white + gp_matern
+            if use_gp:
+                # Set up the kernel an GP
+                kernel = terms.Matern32Term(sigma=sigma, rho=ls)
+                gp = GaussianProcess(kernel, t=x, yerr=yerr, mean=mean_func)
+                gp.marginal("gp", observed=y)
+            else:
+                pm.Normal("obs", mu=mean_func(x), sigma=yerr, observed=y)
 
-            gp.marginal_likelihood("y", X=x[:, None], y=y, noise=yerr)
+        if use_gp:
+            self.gp = gp
 
         self.pymc_model = model
-        self.pymc_gp = gp
-        self.pymc_gp_white = gp_white
-        self.pymc_gp_matern = gp_matern
-        self.mean_model = mean_func(x[:, None])
-
-        return self.pymc_model
+        self.mean_model = mean_func(x)
 
     def __enter__(self):
         """
@@ -243,7 +241,7 @@ def _check_model(self):
         if self.pymc_model is None:
             raise ValueError('Must first call `Model._initialize_model` first.')
 
-    def __call__(self, point=None, **kwargs):
+    def __call__(self, point=None, use_gp=False, **kwargs):
         """
         Evaluate the model with input parameters at ``point``
 
@@ -252,9 +250,16 @@ def __call__(self, point=None, **kwargs):
         from exoplanet import eval_in_model
 
         with self.pymc_model:
-            result = eval_in_model(
-                self.mean_model,
-                point=point,
-                **kwargs
+            if use_gp:
+                mu, var = eval_in_model(
+                    self.gp.predict(self.lc.time[self.mask][::self.skip_n_points],
+                                    return_var=True), point
+                )
+
+            mean_eval = eval_in_model(
+                self.mean_model, point
             )
-        return result
+        if use_gp:
+            return mu + mean_eval, var
+        else:
+            return mean_eval