Merge 0aeefa0 into e348ca6

.travis.yml

@@ -29,13 +29,12 @@ env:
         # overridden underneath. They are defined here in order to save having
         # to repeat them for all configurations.
         - NUMPY_VERSION=stable
-        # - ASTROPY_VERSION=stable
-        - ASTROPY_VERSION=development # TODO: until v3.0 is released
+        - ASTROPY_VERSION=stable
         - MAIN_CMD='python setup.py'
         - SETUP_CMD='test'
-        # TODO: A lot of the dependencies are because of nbsphinx...we may want
-        #   to drop the nbsphinx component of the build when on travis?
-        - PIP_DEPENDENCIES='emcee schwimmbad nbsphinx twobody'
+        # - PIP_DEPENDENCIES='emcee schwimmbad nbsphinx twobody'
+        # Dev version of twobody!
+        - PIP_DEPENDENCIES='emcee schwimmbad nbsphinx git+https://github.com/adrn/twobody'
         - CONDA_CHANNELS='astropy-ci-extras astropy'
         - CONDA_DEPENDENCIES='cython scipy h5py matplotlib pandoc nbconvert ipykernel'
 
@@ -68,9 +67,9 @@ matrix:
         - python: 3.6
           env: ASTROPY_VERSION=development
 
-        # Try older numpy versions
-        - python: 3.5
-          env: NUMPY_VERSION=1.11
+        # Try numpy development version
+        - python: 3.6
+          env: NUMPY_VERSION=development
 
         # TODO: Test with development version of twobody
         # - python: 3.6

thejoker/plot.py

@@ -1,3 +1,6 @@
+# Standard library
+import warnings
+
 # Third-party
 from astropy.time import Time
 import astropy.units as u
@@ -6,7 +9,7 @@
 __all__ = ['plot_rv_curves']
 
 
-def plot_rv_curves(samples, t_grid, n_plot=None, rv_unit=None, data=None,
+def plot_rv_curves(samples, t_grid=None, rv_unit=None, data=None,
                    ax=None, plot_kwargs=dict(), data_plot_kwargs=dict(),
                    add_labels=True, relative_to_t0=False):
     """
@@ -17,10 +20,10 @@ def plot_rv_curves(samples, t_grid, n_plot=None, rv_unit=None, data=None,
     ----------
     samples : :class:`~thejoker.sampler.JokerSamples`
         Posterior samples from The Joker.
-    t_grid : array_like, `~astropy.time.Time`
-        Array of times. Either in BMJD or as an Astropy time object.
-    n_plot : int, optional
-        The maximum number of samples to plot. Defaults to 128.
+    t_grid : array_like, `~astropy.time.Time`, optional
+        Array of times. Either in BMJD or as an Astropy Time object. If not
+        specified, the time grid will be set to the data range with a small
+        buffer.
     rv_unit : `~astropy.units.UnitBase`, optional
         The units to use when plotting RV's.
     data : `~thejoker.data.RVData`, optional
@@ -49,14 +52,24 @@ def plot_rv_curves(samples, t_grid, n_plot=None, rv_unit=None, data=None,
     else:
         fig = ax.figure
 
-    if not isinstance(t_grid, Time): # Assume BMJD
-        t_grid = Time(t_grid, format='mjd', scale='tcb')
+    if t_grid is None:
+        w = np.ptp(data.t.mjd)
+        dt = samples['P'].to(u.day).value.min() / 10
+        t_grid = np.arange(data.t.mjd.min() - w*0.05,
+                           data.t.mjd.max() + w*0.05 + dt,
+                           dt)
 
-    if n_plot is None:
-        n_plot = len(samples)
-    n_plot = min(n_plot, len(samples))
+        if len(t_grid) > 1e4:
+            warnings.warn("Time grid has more than 10000 grid points, so "
+                          "plotting orbits could be very slow! Set 't_grid' "
+                          "manually to decrease the number of grid points.",
+                          ResourceWarning)
+
+    elif not isinstance(t_grid, Time): # Assume BMJD
+        t_grid = Time(t_grid, format='mjd', scale='tcb')
 
     # scale the transparency of the lines
+    n_plot = len(samples)
     Q = 4. # HACK
     line_alpha = 0.05 + Q / (n_plot + Q)
 
@@ -66,6 +79,7 @@ def plot_rv_curves(samples, t_grid, n_plot=None, rv_unit=None, data=None,
     # default plotting style
     style = plot_kwargs.copy()
     style.setdefault('linestyle', '-')
+    style.setdefault('linewidth', 0.5)
     style.setdefault('alpha', line_alpha)
     style.setdefault('marker', '')
     style.setdefault('color', '#555555')

thejoker/sampler/fast_likelihood.pyx

@@ -72,7 +72,7 @@ cdef void design_matrix(double P, double phi0, double ecc, double omega,
     if n_trend > 1: # only needed if more than constant trend
         for i in range(1, n_trend):
             for j in range(n_times):
-                A_T[1+i, j] = pow(t[j], i)
+                A_T[1+i, j] = pow(t[j] - t0, i)
 
 
 cdef void get_ivar(double[::1] ivar, double s, double[::1] new_ivar):
@@ -99,12 +99,13 @@ cdef void get_ivar(double[::1] ivar, double s, double[::1] new_ivar):
 
 
 cdef double tensor_vector_scalar(double[:,::1] A_T, double[::1] ivar,
-                                 double[::1] y,
+                                 double[::1] y, double[:, ::1] Lambda,
                                  double[::1] linear_pars,
                                  int sample_linear_pars,
                                  double[:, ::1] ATCinvA, double[:, ::1] _A,
-                                 double[::1] p, int[:,::1] ipiv,
-                                 double[:,::1] work, int[:,::1] lwork):
+                                 int[:,::1] ipiv,
+                                 double[:,::1] work, int[:,::1] lwork,
+                                 double[::1] p):
     """Construct objects used to compute the marginal log-likelihood.
 
     Parameters
@@ -115,6 +116,9 @@ cdef double tensor_vector_scalar(double[:,::1] A_T, double[::1] ivar,
         Inverse-variance matrix.
     y : `numpy.ndarray`
         Data (in this case, radial velocities).
+    Lambda : `numpy.ndarray`
+        Inverse variance matrix for a Gaussian prior applied to the linear
+        parameters.
 
     Outputs
     -------
@@ -161,6 +165,7 @@ cdef double tensor_vector_scalar(double[:,::1] A_T, double[::1] ivar,
             for j in range(n_pars):
                 # implicit transpose of first A_T in the below
                 ATCinvA[i, j] += A_T[i, k] * ivar[k] * A_T[j, k]
+                ATCinvA[i, j] += Lambda[i, j]
 
     _A[:, :] = ATCinvA
 
@@ -192,9 +197,9 @@ cdef double tensor_vector_scalar(double[:,::1] A_T, double[::1] ivar,
     if sample_linear_pars == 1:
         cov = np.linalg.inv(ATCinvA)
         # TODO: set np.random.set_state(state) outside of here to make deterministic
-        K, *v_terms = np.random.multivariate_normal(p, cov)
-        linear_pars[0] = K
-        linear_pars[1] = v_terms[0] # TODO: expects just 1 term
+        v_terms = np.random.multivariate_normal(p, cov)
+        for k in range(n_pars):
+            linear_pars[k] = v_terms[k]
 
     return 0.5*logdet - 0.5*chi2
 
@@ -279,14 +284,16 @@ cpdef batch_marginal_ln_likelihood(double[:,::1] chunk,
         int n
         int n_samples = chunk.shape[0]
         int n_times = len(data)
-        int n_pars = 2 # always have K, v0
+        int n_poly = joker_params.poly_trend # polynomial trend terms
+        int n_pars = 1 + n_poly # always have K, v trend terms
 
         double[::1] t = np.ascontiguousarray(data._t_bmjd, dtype='f8')
         double[::1] rv = np.ascontiguousarray(data.rv.value, dtype='f8')
         double[::1] ivar = np.ascontiguousarray(data.ivar.value, dtype='f8')
 
         # inverse variance array with jitter included
         double[::1] jitter_ivar = np.zeros(data.ivar.value.shape)
+        double[:, ::1] Lambda = np.ascontiguousarray(joker_params.linear_par_Vinv)
 
         # transpose of design matrix
         double[:,::1] A_T = np.zeros((n_pars, n_times))
@@ -300,6 +307,8 @@ cpdef batch_marginal_ln_likelihood(double[:,::1] chunk,
         int _fixed_jitter
         double jitter
 
+        double[::1] linear_pars = np.zeros(n_pars)
+
         # needed for temporary storage in tensor_vector_scalar
         double[:, ::1] ATCinvA = np.zeros((n_pars, n_pars))
         double[:, ::1] _A = np.zeros((n_pars, n_pars))
@@ -319,17 +328,16 @@ cpdef batch_marginal_ln_likelihood(double[:,::1] chunk,
         if _fixed_jitter == 0:
             jitter = chunk[n, 4]
 
-        # TODO: n_trend is hard set to 1 here
         design_matrix(chunk[n, 0], chunk[n, 1], chunk[n, 2], chunk[n, 3],
-                      t, t0, A_T, 1)
+                      t, t0, A_T, n_poly)
 
         # Note: jitter must be in same units as the data RV's / ivar!
         get_ivar(ivar, jitter, jitter_ivar)
 
         # compute things needed for the ln(likelihood)
-        ll[n] = tensor_vector_scalar(A_T, jitter_ivar, rv,
-                                     p, 0, # IGNORED PLACEHOLDER
-                                     ATCinvA, _A, p, ipiv, work, lwork)
+        ll[n] = tensor_vector_scalar(A_T, jitter_ivar, rv, Lambda,
+                                     linear_pars, 0, # IGNORED PLACEHOLDER
+                                     ATCinvA, _A, ipiv, work, lwork, p)
 
     return ll
 
@@ -352,17 +360,19 @@ cpdef batch_get_posterior_samples(double[:,::1] chunk,
     """
 
     cdef:
-        int n
+        int n, j
         int n_samples = chunk.shape[0]
         int n_times = len(data)
-        int n_pars = 2 # always have K, v0
+        int n_poly = joker_params.poly_trend # polynomial trend terms
+        int n_pars = 1 + n_poly # always have K, v trend terms
 
         double[::1] t = np.ascontiguousarray(data._t_bmjd, dtype='f8')
         double[::1] rv = np.ascontiguousarray(data.rv.value, dtype='f8')
         double[::1] ivar = np.ascontiguousarray(data.ivar.value, dtype='f8')
 
         # inverse variance array with jitter included
         double[::1] jitter_ivar = np.zeros(data.ivar.value.shape)
+        double[:, ::1] Lambda = np.ascontiguousarray(joker_params.linear_par_Vinv)
 
         # transpose of design matrix
         double[:,::1] A_T = np.zeros((n_pars, n_times))
@@ -373,6 +383,8 @@ cpdef batch_get_posterior_samples(double[:,::1] chunk,
         int _fixed_jitter
         double jitter
 
+        double[::1] linear_pars = np.zeros(n_pars)
+
         # needed for temporary storage in tensor_vector_scalar
         double[:, ::1] ATCinvA = np.zeros((n_pars, n_pars))
         double[:, ::1] _A = np.zeros((n_pars, n_pars))
@@ -391,29 +403,29 @@ cpdef batch_get_posterior_samples(double[:,::1] chunk,
         pars[n, 3] = chunk[n, 3] # omega
         pars[n, 4] = chunk[n, 4] # jitter
 
-        # TODO: hard set n_trend=1 (v0) because removing support for that
-        design_matrix(chunk[n,0], chunk[n,1], chunk[n,2], chunk[n,3],
-                      t, t0, A_T, 1)
+        design_matrix(chunk[n, 0], chunk[n, 1], chunk[n, 2], chunk[n, 3],
+                      t, t0, A_T, n_poly)
 
         # jitter must be in same units as the data RV's / ivar!
         get_ivar(ivar, chunk[n, 4], jitter_ivar)
 
         # compute things needed for the ln(likelihood)
-        ll = tensor_vector_scalar(A_T, jitter_ivar, rv,
-                                  p, 1, # do make samples (1 = True)
-                                  ATCinvA, _A, p, ipiv, work, lwork)
+        ll = tensor_vector_scalar(A_T, jitter_ivar, rv, Lambda,
+                                  linear_pars, 1, # do make samples (1 = True)
+                                  ATCinvA, _A, ipiv, work, lwork, p)
 
-        K = p[0]
+        K = linear_pars[0]
         if K < 0:
-            # log.warning("Swapping K")
+            # Swapping sign of K!
             K = np.abs(K)
             pars[n, 3] += np.pi
             pars[n, 3] = pars[n, 3] % (2*np.pi) # HACK: I think this is safe
 
         pars[n, 5] = K
-        pars[n, 6] = p[1] # TODO: we assume it's just v0
+        for j in range(n_poly):
+            pars[n, 6+j] = linear_pars[1+j]
 
         if return_logprobs:
-            pars[n, 7] = ll
+            pars[n, 6+n_poly] = ll
 
     return np.array(pars)

thejoker/sampler/params.py

@@ -1,6 +1,7 @@
 # Third-party
 from astropy.utils.misc import isiterable
 import astropy.units as u
+import numpy as np
 
 __all__ = ['JokerParams']
 
@@ -31,6 +32,17 @@ class JokerParams:
         If sampling over the jitter as an extra non-linear parameter,
         you must also specify the units of the jitter prior. See note
         above about the ``jitter`` argument.
+    poly_trend : int, optional
+        If specified, sample over a polynomial velocity trend with the specified
+        number of coefficients. For example, ``poly_trend=3`` will sample over
+        parameters of a long-term quadratic velocity trend. Default is 1, just a
+        constant velocity shift.
+    linear_par_Vinv : array_like, optional
+        Inverse variance matrix that specifies the Gaussian prior on the linear
+        parameters, i.e., the semi-amplitude and velocity trend paramters. The
+        units must be in inverse, squared ``v_unit`` and ``v_unit/day^n`` where
+        ``v_unit`` is the jitter velocity unit, and ``day^n`` corresponds to
+        each polynomial trend coefficient.
     anomaly_tol : float (optional)
         Convergence tolerance passed to
         :func:`twobody.eccentric_anomaly_from_mean_anomaly`.
@@ -53,20 +65,30 @@ class JokerParams:
     @u.quantity_input(P_min=u.day, P_max=u.day)
     def __init__(self, P_min, P_max,
                  jitter=None, jitter_unit=None,
-                 trend=None,
+                 poly_trend=1, linear_par_Vinv=None,
                  anomaly_tol=1E-10, anomaly_maxiter=128):
 
-        # TODO: if trend is a class, sample it, if it's an instance, fix vals?
-        # TODO: add trend classes back in...
-
         # the names of the default parameters
-        self.default_params = ['P', 'M0', 'e', 'omega', 'jitter', 'K', 'v0']
+        self.default_params = ['P', 'M0', 'e', 'omega', 'jitter', 'K']
+
+        self.poly_trend = int(poly_trend)
+        self.default_params += ['v{0}'.format(i)
+                                for i in range(self.poly_trend)]
 
         self.P_min = P_min
         self.P_max = P_max
         self.anomaly_tol = float(anomaly_tol)
         self.anomaly_maxiter = int(anomaly_maxiter)
 
+        _n_linear = 1 + self.poly_trend
+        if linear_par_Vinv is None:
+            self.linear_par_Vinv = 1e-10 * np.eye(_n_linear)
+        else:
+            self.linear_par_Vinv = np.array(linear_par_Vinv)
+            if self.linear_par_Vinv.shape != (_n_linear, _n_linear):
+                raise ValueError("Linear parameter inverse variance prior must "
+                                 "have shape ({0}, {0})".format(_n_linear))
+
         # validate the input jitter specification
         if jitter is None:
             jitter = 0 * u.km/u.s

thejoker/sampler/sampler.py

@@ -190,15 +190,18 @@ def _unpack_full_samples(self, result, prior_units, return_logprobs,
 
         n, n_params = samples_arr.shape
 
-        samples = JokerSamples(t0=t0)
+        samples = JokerSamples(t0=t0, poly_trend=self.params.poly_trend)
 
         # TODO: need to keep track of this elsewhere...
         nonlin_params = ['P', 'M0', 'e', 'omega', 'jitter']
         for k, key in enumerate(nonlin_params):
             samples[key] = samples_arr[:, k] * prior_units[k]
 
         samples['K'] = samples_arr[:, k + 1] * prior_units[-1]  # jitter unit
-        samples['v0'] = samples_arr[:, k + 2] * prior_units[-1]  # jitter unit
+
+        for i in range(self.params.poly_trend):
+            _unit = prior_units[-1] / u.day**i  # HACK: jitter unit per day
+            samples['v'+str(i)] = samples_arr[:, k + 2 + i] * _unit
 
         if return_logprobs:
             return samples, ln_prior