Merge pull request #191 from California-Planet-Search/next-release

Version 1.2.0

example_planets/k2-131_celerite.py

@@ -3,10 +3,8 @@
 import os
 import radvel
 
-# This setup file is provided to illustrate how to set up
-# a config file for use with Dan Foreman-Mackey's `celerite`
-# package. The results it produces are not meant to be
-# compared with those of Dai et al. (2017).
+# This setup file is provided to illustrate how to set up a config
+# file for use with Dan Foreman-Mackey's `celerite` package.
 
 # Data from Dai+ 2017
 instnames = ['harps-n','pfs'] 
@@ -31,59 +29,51 @@
 fitting_basis = 'per tc secosw sesinw k'
 
 # Set initial parameter value guesses.
-params = radvel.Parameters(nplanets,basis=fitting_basis, planet_letters=planet_letters)
+params = radvel.Parameters(
+  nplanets,basis=fitting_basis, planet_letters=planet_letters
+)
 
 params['per1'] = radvel.Parameter(value=Porb)
 params['tc1'] = radvel.Parameter(value=Tc)
-params['sesinw1'] = radvel.Parameter(value=0.,vary=False) # hold eccentricity fixed at 0
+params['sesinw1'] = radvel.Parameter(value=0.,vary=False) 
 params['secosw1'] = radvel.Parameter(value=0.,vary=False)
 params['k1'] = radvel.Parameter(value=6.55)
 params['dvdt'] = radvel.Parameter(value=0.,vary=False)
 params['curv'] = radvel.Parameter(value=0.,vary=False)
 time_base = np.median(t)
 
 # Define Celerite GP hyperparameters.
-
-# First Celerite Term. AC>=BD must be true to ensure 
-# positive-definiteness. (A == params['1_logA'].value, etc.)
-params['1_logA'] = radvel.Parameter(value=np.log(26.))
-params['1_logB'] = radvel.Parameter(value=np.log(.0005))
-params['1_logC'] = radvel.Parameter(value=np.log(.5))
-params['1_logD'] = radvel.Parameter(value=np.log(.0005))
-
-# Second Celerite Term (real). Setting vary=False for 1_logB and 
-# 1_logD ensures that this term remains real throughout the fitting process.
-params['2_logA'] = radvel.Parameter(value=np.log(26.)) 
-params['2_logB'] = radvel.Parameter(value=np.log(0.005), vary=False)
-params['2_logC'] = radvel.Parameter(value=np.log(0.1)) 
-params['2_logD'] = radvel.Parameter(value=np.log(0.005), vary=False)
+params['gp_B'] = radvel.Parameter(value=30**2., vary=True)
+params['gp_C'] = radvel.Parameter(value=1., vary=True) 
+params['gp_L'] = radvel.Parameter(value=9., vary=True)
+params['gp_Prot'] = radvel.Parameter(value=9., vary=True)
 
 hnames = {}
 for tel in instnames:
-  hnames[tel] = ['1_logA','1_logB','1_logC','1_logD',
-                 '2_logA','2_logB','2_logC','2_logD'
-                 ]
+  hnames[tel] = ['gp_B','gp_C','gp_L','gp_Prot']
 
-kernel_name = {'harps-n':"Celerite", 
-               'pfs':"Celerite"}
+kernel_name = {'harps-n':"Celerite", 'pfs':"Celerite"}
 
 jit_guesses = {'harps-n':2.0, 'pfs':5.0}
-
+  
 def initialize_instparams(tel_suffix):
 
-    indices = telgrps[tel_suffix]
+  indices = telgrps[tel_suffix]
 
-    params['gamma_'+tel_suffix] = radvel.Parameter(value=np.mean(vel[indices]))
-    params['jit_'+tel_suffix] = radvel.Parameter(value=jit_guesses[tel_suffix]) 
+  params['gamma_'+tel_suffix] = radvel.Parameter(value=np.mean(vel[indices]))
+  params['jit_'+tel_suffix] = radvel.Parameter(value=jit_guesses[tel_suffix]) 
 
 
 for tel in instnames:
-    initialize_instparams(tel)
-
-priors = [radvel.prior.Gaussian('per1', Porb, Porb_unc),
-          radvel.prior.Gaussian('tc1', Tc, Tc_unc),
-          radvel.prior.Jeffreys('k1', 0.01, 10.),
-          radvel.prior.Jeffreys('jit_pfs', 0.01, 10.),
-          radvel.prior.Jeffreys('jit_harps-n', 0.01,10.),
-          radvel.prior.Jeffreys('1_logA', 0.005, 75.0)
-          ]
+  initialize_instparams(tel)
+  priors = [
+    radvel.prior.Gaussian('per1', Porb, Porb_unc),
+    radvel.prior.Gaussian('tc1', Tc, Tc_unc),
+    radvel.prior.Jeffreys('k1', 0.01, 10.),
+    radvel.prior.Jeffreys('jit_pfs', 0.01, 10.),
+    radvel.prior.Jeffreys('jit_harps-n', 0.01,10.),
+    radvel.prior.Gaussian('gp_L',9.5,1.), # constraints from photometry on hyperparams (Dai et al. 2017)
+    radvel.prior.Gaussian('gp_Prot',9.64,0.12),
+    radvel.prior.HardBounds('gp_C',0.,1e10),  # keep other two hyperparams positive
+    radvel.prior.HardBounds('gp_B',0.,1e10)
+  ]

radvel/__init__.py

@@ -19,7 +19,7 @@ def _custom_warningfmt(msg, *a, **b):
 __all__=['model', 'likelihood', 'posterior', 'mcmc', 'prior', 'utils',
          'fitting', 'report', 'cli', 'driver', 'gp']
 
-__version__ = '1.1.12'
+__version__ = '1.2.0'
 
 MODULEDIR, filename = os.path.split(__file__)
 DATADIR = os.path.join(sys.prefix, 'radvel_example_data')

radvel/driver.py

@@ -19,6 +19,7 @@
 
 import radvel
 from radvel.plot import orbit_plots, mcmc_plots
+from radvel.mcmc import statevars
 from astropy import constants as c
 from numpy import inf
 
@@ -168,6 +169,9 @@ def mcmc(args):
             post, nwalkers=args.nwalkers, nrun=args.nsteps, ensembles=args.ensembles, burnGR=args.burnGR,
             maxGR=args.maxGR, minTz=args.minTz, minsteps=args.minsteps, thin=args.thin, serial=args.serial)
 
+    mintz = statevars.mintz
+    maxgr = statevars.maxgr
+
     # Convert chains into synth basis
     synthchains = chains.copy()
     for par in post.params.keys():
@@ -241,7 +245,9 @@ def mcmc(args):
                  'chainfile': os.path.abspath(csvfn),
                  'summaryfile': os.path.abspath(saveto),
                  'nwalkers': args.nwalkers,
-                 'nsteps': args.nsteps}
+                 'nsteps': args.nsteps,
+                 'minTz': mintz,
+                 'maxGR': maxgr}
     save_status(statfile, 'mcmc', savestate)
 
 

radvel/gp.py

@@ -8,12 +8,12 @@
 
 warnings.simplefilter('once')
 
-# implemented kernels & examples of possible names for associated hyperparameters
+# implemented kernels & examples of their associated hyperparameters
 KERNELS = {
     "SqExp": ['gp_length','gp_amp'],
     "Per": ['gp_per','gp_length','gp_amp'],
     "QuasiPer": ['gp_per','gp_perlength','gp_explength','gp_amp'],
-    "Celerite": ['1_logA','1_logB','1_logC','1_logD']
+    "Celerite": ['gp_B','gp_C','gp_L','gp_Prot']
 }
 
 if sys.version_info[0] < 3:
@@ -236,8 +236,8 @@ class QuasiPerKernel(Kernel):
 
     .. math::
 
-        C_{ij} = \\eta_1^2 * exp( \\frac{ -|t_i - t_j|^2 }{ \\eta_2^2 } 
-                           - \\frac{ \\sin^2(\\frac{ \\pi|t_i-t_j| }{ \\eta_3^2 } ) }{ 2\\eta_4^2 } )     
+        C_{ij} = B/(2+C) * exp( -|t_i - t_j| / L) *  
+                           (\\cos(\\frac{ 2\\pi|t_i-t_j| }{ P_{rot} }) + (1+C) )     
 
     Args:
         hparams (dict of radvel.Parameter): dictionary containing
@@ -327,44 +327,31 @@ def compute_covmatrix(self, errors):
 
         return self.covmatrix
 
+
 class CeleriteKernel(Kernel):
-    """ A wrapper for the celerite.solver.CholeskySolver() methods and attributes.
+    """
+    Class that computes and stores a matrix approximating the quasi-periodic
+    kernel.
 
-    Class that computes and stores a kernel that can be modeled as the sum of 
-    celerite terms, or kernels of the following form:
+    See `radvel/example_planets/k2-131_celerite.py` for an example of a setup
+    file that uses this Kernel object.
+
+    See celerite.readthedocs.io and Foreman-Mackey et al. 2017. AJ, 154, 220
+    (equation 56) for more details.
+
+    An arbitrary element, :math:`C_{ij}`, of the matrix is:
 
     .. math::
 
-        C_{ij} = \\sum_{n=1}^{J} \\frac{1}{2}(a_n + ib_n)e^{-(c_n + id_n)\\tau_{nm}} 
-                               + \\frac{1}{2}(a_n - ib_n)e^{-(c_n - id_n)\\tau_{nm}}
-
-    where :math:`J` is the number of terms in the sum, and :math:`\\tau_{nm}=|t_i - t_j|`. 
-    
-    The hyperparameters of this kernel are :math:`a_{n}` , :math:`b_{n}` , :math:`c_{n}` , and :math:`d_n` 
-    for each term :math:`C_{ij}` in the sum. 
-
-    See celerite.readthedocs.io for more information about celerite kernels and
-    computation.
-
-    Note: 
-        For this kernel to be positive-definite, we must have :math:`a_nc_n \\ge b_nd_n`
-        at all times. The CeleriteLikelihood object will raise a warning
-        if it ever detects a non-positive-definite kernel.
-
-    :param: hparams (dict of radvel.Parameter): dictionary containing
-        radvel.Parameter objects that are GP hyperparameters
-        of this kernel. Must contain a multiple of 4 Parameter object
-        with the following names:
-                - `k_logA*`: the natural log of :math:`a_{k}`. 
-                - `k_logB*`: the natural log of :math:`b_{k}`.
-                - `k_logC*`: the natural log of :math:`c_{k}`. 
-                - `k_logD*`: the natural log of :math:`d_{k}`. 
-        (where k is a 1-indexed integer identifying coefficients of a 
-        particular term, :math:`1 <= k <= N_{terms}`, and * is an optional 
-        suffix, e.g. 'hires'). Suffix is useful when fitting several individual 
-        GPLikelihoods with different hyperparameters using the 
-        CompositeLikelihood object.
+        C_{ij} = \\eta_1^2 * exp( \\frac{ -|t_i - t_j|^2 }{ \\eta_2^2 } 
+                           - \\frac{ \\sin^2(\\frac{ \\pi|t_i-t_j| }{ \\eta_3^2 } ) }{ 2\\eta_4^2 } )     
 
+    Args:
+        hparams (dict of radvel.Parameter): dictionary containing
+            radvel.Parameter objects that are GP hyperparameters
+            of this kernel. Must contain exactly four objects, 'gp_B*',
+            'gp_C*', 'gp_L*', and 'gp_Prot*', where * is a suffix 
+            identifying these hyperparameters with a likelihood object.
     """
 
     @property
@@ -373,57 +360,64 @@ def name(self):
 
     def __init__(self, hparams):
 
-        assert len(hparams) > 0 and len(hparams) % 4 == 0, \
-            "CeleriteKernel requires a positive integer number of terms, each" \
-             + "with 4 coefficients. See CeleriteKernel documentation."
-        self.num_terms = int(len(hparams) / 4)
-        self.hparams = np.zeros((self.num_terms, 4))
+        self.hparams = {}
+        for par in hparams:
+            if par.startswith('gp_B'):
+                self.hparams['gp_B'] = hparams[par]
+            if par.startswith('gp_C'):
+                self.hparams['gp_C'] = hparams[par]
+            if par.startswith('gp_L'):
+                self.hparams['gp_L'] = hparams[par]
+            if par.startswith('gp_Prot'):
+                self.hparams['gp_Prot'] = hparams[par]
+
+        assert len(self.hparams) == 4, """
+CeleriteKernel requires exactly 4 hyperparameters with names 'gp_B', 'gp_C', 'gp_L', and 'gp_Prot'.
+        """
 
-        # set up hyperparameter arrays
         try:
-            for par in hparams:
-                index = int(par[0]) - 1
-                if 'logA' in par:
-                    self.hparams[index,0] = hparams[par].value
-                if 'logB' in par:
-                    self.hparams[index,1] = hparams[par].value
-                if 'logC' in par:
-                    self.hparams[index,2] = hparams[par].value
-                if 'logD' in par:
-                    self.hparams[index,3] = hparams[par].value
+            self.hparams['gp_Prot'].value
+            self.hparams['gp_C'].value
+            self.hparams['gp_B'].value
+            self.hparams['gp_L'].value
+        except KeyError:
+            raise KeyError("""
+CeleriteKernel requires hyperparameters 'gp_B*', 'gp_C*', 'gp_L', and 'gp_Prot*'.
+                """)
         except AttributeError:
-            raise AttributeError("CeleriteKernel requires dictionary of" \
-                                 + " radvel.Parameter objects as input.")
-        except IndexError:
-            raise IndexError("CeleriteKernel hyperparameter indices (k in k_logA*)"
-                             + " were named incorrectly. See CeleriteKernel documentation.")
-        except ValueError:
-            raise ValueError("CeleriteKernel hyperparameter indices (k in k_logA*)"
-                             + " were named incorrectly. See CeleriteKernel documentation.")
+            raise AttributeError("CeleriteKernel requires dictionary of radvel.Parameter objects as input.")
 
     # get arrays of real and complex parameters
     def compute_real_and_complex_hparams(self):
-        real_indices = []
-        complex_indices = []
-        for col in np.arange(self.num_terms):
-            if np.exp(self.hparams[col,1])==0 \
-            and np.exp(self.hparams[col,3])==0: 
-                real_indices.append(col)
-            else:
-                complex_indices.append(col)
-        self.real = self.hparams[real_indices]
-        self.complex = self.hparams[complex_indices]
+
+        self.real = np.zeros((1, 4))
+        self.complex = np.zeros((1, 4))
+
+        B = self.hparams['gp_B'].value
+        C = self.hparams['gp_C'].value
+        L = self.hparams['gp_L'].value
+        Prot = self.hparams['gp_Prot'].value
+
+        # Foreman-Mackey et al. (2017) eq 56
+        self.real[0,0] = B*(1+C)/(2+C)
+        self.real[0,2] = 1/L
+        self.complex[0,0] = B/(2+C)
+        self.complex[0,1] = 0.
+        self.complex[0,2] = 1/L
+        self.complex[0,3] = 2*np.pi/Prot
 
     def __repr__(self):
+
+        B = self.hparams['gp_B'].value
+        C = self.hparams['gp_C'].value
+        L = self.hparams['gp_L'].value
+        Prot = self.hparams['gp_Prot'].value
+
         msg = (
-            "Celerite Kernel with log(a) = {}, log(b) = {}, log(c) = {}, log(d) = {}."
-        ).format(
-            self.hparams[:,0], self.hparams[:,1], 
-            self.hparams[:,2], self.hparams[:,3]
-          )
+            "Celerite Kernel with B = {}, C = {}, L = {}, Prot = {}."
+        ).format(B, C, L, Prot)
         return msg
 
-
     def compute_distances(self, x1, x2):
         """
         The celerite.solver.CholeskySolver object does 
@@ -438,6 +432,7 @@ def compute_distances(self, x1, x2):
         self.U = np.empty((0,0))
         self.V = self.U
 
+
     def compute_covmatrix(self, errors):
         """ Compute the Cholesky decomposition of a celerite kernel
 
@@ -453,8 +448,6 @@ def compute_covmatrix(self, errors):
         solver = CholeskySolver()
 
         self.compute_real_and_complex_hparams()
-        self.real = np.exp(self.real) # (celerite hyperparameters are fit in log-space)
-        self.complex = np.exp(self.complex)
         solver.compute(
             0., self.real[:,0], self.real[:,2], 
             self.complex[:,0], self.complex[:,1], 
@@ -463,4 +456,4 @@ def compute_covmatrix(self, errors):
             self.x, errors**2
         )
 
-        return solver
+        return solver