add prior for linear parameters

thejoker/sampler/fast_likelihood.pyx

@@ -99,7 +99,7 @@ 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,
@@ -116,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
     -------
@@ -162,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
 
@@ -289,6 +293,7 @@ cpdef batch_marginal_ln_likelihood(double[:,::1] chunk,
 
         # 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))
@@ -330,7 +335,7 @@ cpdef batch_marginal_ln_likelihood(double[:,::1] chunk,
         get_ivar(ivar, jitter, jitter_ivar)
 
         # compute things needed for the ln(likelihood)
-        ll[n] = tensor_vector_scalar(A_T, jitter_ivar, rv,
+        ll[n] = tensor_vector_scalar(A_T, jitter_ivar, rv, Lambda,
                                      linear_pars, 0, # IGNORED PLACEHOLDER
                                      ATCinvA, _A, ipiv, work, lwork, p)
 
@@ -367,6 +372,7 @@ cpdef batch_get_posterior_samples(double[:,::1] chunk,
 
         # 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))
@@ -404,7 +410,7 @@ cpdef batch_get_posterior_samples(double[:,::1] chunk,
         get_ivar(ivar, chunk[n, 4], jitter_ivar)
 
         # compute things needed for the ln(likelihood)
-        ll = tensor_vector_scalar(A_T, jitter_ivar, rv,
+        ll = tensor_vector_scalar(A_T, jitter_ivar, rv, Lambda,
                                   linear_pars, 1, # do make samples (1 = True)
                                   ATCinvA, _A, ipiv, work, lwork, p)
 

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']
 
@@ -36,6 +37,12 @@ class JokerParams:
         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`.
@@ -58,7 +65,7 @@ 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,
-                 poly_trend=1,
+                 poly_trend=1, linear_par_Vinv=None,
                  anomaly_tol=1E-10, anomaly_maxiter=128):
 
         # the names of the default parameters
@@ -73,6 +80,15 @@ def __init__(self, P_min, 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