Enabling Celerite2 SHO GP kernels.

demos/JWST/S5_fit_par_template.epf

@@ -109,14 +109,18 @@ u2          0.2             'free'          0           1           U
 # Polynomial model variables (c0--c9 for 0th--9th order polynomials in time); Fitting at least c0 is very strongly recommended!
 # Exponential ramp model variables (r0--r1 for one exponential ramp, r2--r3 for a second exponential ramp)
 # Common-mode amplitude parameters (cm1 and cm2)
-# GP model parameters (A, m for the first kernel; A1, m1 for the second kernel; etc.) in log scale
+# GP model parameters (A, m for the first kernel; A1, m1 for the second kernel; etc.) in log scale.
+# The SHO GP kernel takes (in addition to A and m) a third parameter Q, the quality factor of the damped oscillator. 
+#     Various values of Q correspond to different noise sources, (e.g. stellar granulation, pulsations, etc).
 # Step-function model variables (step# and steptime# for step-function model #; e.g. step0 and steptime0)
 # Drift model variables (xpos, ypos, xwidth, ywidth)
 # --------------------
 c0          1.015           'free'          1           0.05        N
 c1          0               'free'          0           0.01        N
+
 # A           -15             'free'          -25         0           U
 # m           -4              'free'          -7          0           U
+# Q           0.707           'free'          0.5         10          U
 
 # -----------
 # ** White noise **

demos/JWST/S5_template.ecf

@@ -70,8 +70,8 @@ target_accept   0.85
 
 #GP inputs
 kernel_inputs   ['time'] #options: time
-kernel_class    ['Matern32'] #options: ExpSquared, Matern32, Exp, RationalQuadratic for george, Matern32 for celerite (sums of kernels possible for george separated by commas)
-GP_package      'celerite' #options: george, celerite
+kernel_class    ['Matern32'] #options: ExpSquared, Matern32, Exp, RationalQuadratic for george and tinygp, Matern32, SHO for celerite (sums of kernels possible for george separated by commas)
+GP_package      'celerite' #options: george, celerite, tinygp
 
 # Plotting controls
 interp          False   # Should astrophysical model be interpolated (useful for uneven sampling like that from HST)

docs/source/ecf.rst

@@ -1154,11 +1154,11 @@ Only used for fits with a GP. A list of the covariates to be used when fitting a
 
 kernel_class
 ^^^^^^^^^^^^
-Only used for fits with a GP. A list of the types of GP kernels to use. For the george GP package, this includes ExpSquared, Matern32, Exp, and RationalQuadratic. For the celerite package, this only includes Matern32. It is possible to sum multiple kernels possible for george by listing multiple kernels.
+Only used for fits with a GP. A list of the types of GP kernels to use. For the george and tinygp GP packages, this includes ExpSquared, Matern32, Exp, and RationalQuadratic. For the celerite package, this includes Matern32 and SHO. It is possible to sum multiple kernels possible for george by listing multiple kernels.
 
 GP_package
 ^^^^^^^^^^
-Only used for fits with a GP. The Python GP package to use, with the options of 'george' or 'celerite'.
+Only used for fits with a GP. The Python GP package to use, with the options of 'george', 'celerite' or 'tinygp'.
 
 useHODLR
 ^^^^^^^^
@@ -1471,6 +1471,8 @@ Available fitting parameters are:
          That said, there are no hard and fast rules about what your priors should be, and you will need to experiment to find what works best.
          If there are multiple kernels that are being added, the second kernel's parameters will be ``A1`` and ``m1``, and so on.
 
+      - ``Q`` - ``SHO`` kernel only (in addition to ``A`` and ``m``)! The quality factor of the damped harmonic oscillator for the ``SHO`` kernel as used in the ``celerite`` GP implementation. Different values of ``Q`` can describe different physical processes (such as stellar granulation and pulsations), and a detailed discussion of these can be found in Section 3. of Foreman-Mackey et al. (2017, arxiv:1703.09710).
+
    - White Noise Parameters - options are ``scatter_mult`` for a multiplier to the expected noise from Stage 3 (recommended), or ``scatter_ppm`` to directly fit the noise level in ppm.
 
 

src/eureka/S5_lightcurve_fitting/models/GPModel.py

@@ -59,12 +59,15 @@ def __init__(self, kernel_types, kernel_input_names, lc,
                 raise AssertionError('Celerite2 cannot compute multi-'
                                      'dimensional GPs. Please choose a '
                                      'different GP code')
-            elif self.kernel_types[0] != 'Matern32':
+            elif self.kernel_types[0] not in ['Matern32', 'SHO']:
                 raise AssertionError('Our celerite2 implementation currently '
-                                     'only supports a Matern32 kernel.')
+                                     'only supports Matern32 or SHO kernels.')
 
         # Setup coefficients
-        self.coeffs = np.zeros((self.nchannel_fitted, self.nkernels, 2))
+        nparam = 2
+        if 'SHO' in self.kernel_types: # SHO kernels need Q as well
+            nparam = 3
+        self.coeffs = np.zeros((self.nchannel_fitted, self.nkernels, nparam))
         self._parse_coeffs()
 
     def _parse_coeffs(self):
@@ -81,8 +84,11 @@ def _parse_coeffs(self):
             else:
                 chankey = f'_ch{chan}'
 
-            for i, par in enumerate(['A', 'm']):
-                for k in range(self.nkernels):
+            for k in range(self.nkernels):
+                paramlist = ['A', 'm']
+                if self.kernel_types[k] == "SHO":
+                    paramlist += 'Q'
+                for i, par in enumerate(paramlist):
                     if k == 0:
                         kernelkey = ''
                     else:
@@ -326,7 +332,7 @@ def get_kernel(self, kernel_name, k, c=0):
                 kernel = amp*kernels.ExpSquaredKernel(
                     metric, ndim=self.nkernels, axes=k)
             elif kernel_name == 'RationalQuadratic':
-                kernel = amp*kernels.RationalQuadraticKernel(
+                kernel = amp*kernels.RationalQuadraticKernel( 
                     log_alpha=1, metric=metric, ndim=self.nkernels, axes=k)
             elif kernel_name == 'Exp':
                 kernel = amp*kernels.ExpKernel(
@@ -338,24 +344,37 @@ def get_kernel(self, kernel_name, k, c=0):
                                      'ExpSquared, RationalQuadratic, Exp.')
         elif self.gp_code_name == 'celerite':
             # get metric and amplitude for the current kernel and channel
+            # "amp" and "metric" broadly correspond to "sigma" and "rho"
+            # which are the specific celerite2 parameter names
             sigma = np.sqrt(np.exp(self.coeffs[c, k, 0]))
-            metric = np.exp(self.coeffs[c, k, 1])
-
-            kernel = celerite2.terms.Matern32Term(sigma=sigma, rho=metric)
+            rho = np.exp(self.coeffs[c, k, 1])
+            if kernel_name == 'Matern32':
+                kernel = celerite2.terms.Matern32Term(sigma=sigma, rho=rho)
+            elif kernel_name == 'SHO':
+                # reparametrize sigma, rho into omega_0 and S_0 as defined in celerite2
+                w0 = np.sqrt(3) / rho
+                S0 = (sigma**2) / w0
+                Q = self.coeffs[c, k, 2]
+                kernel = celerite2.terms.SHOTerm(S0=S0, w0=w0, Q=Q)
+            else:
+                raise AssertionError(f'The kernel {kernel_name} is not in the '
+                                     'currently supported list of kernels for '
+                                     'celerite2 which includes:\nMatern32 '
+                                     'and SHO.')
         elif self.gp_code_name == 'tinygp':
             # get metric and amplitude for the current kernel and channel
             amp = np.exp(self.coeffs[c, k, 0])
-            metric = np.exp(self.coeffs[c, k, 1]*2)
+            rho = np.exp(self.coeffs[c, k, 1])
 
             if kernel_name == 'Matern32':
-                kernel = amp*tinygp.kernels.Matern32(metric)
+                kernel = tinygp.kernels.quasisep.Matern32(scale=rho, sigma=np.sqrt(amp))
             elif kernel_name == 'ExpSquared':
                 kernel = amp*tinygp.kernels.ExpSquared(metric)
             elif kernel_name == 'RationalQuadratic':
                 kernel = amp*tinygp.kernels.RationalQuadratic(alpha=1,
                                                               scale=metric)
             elif kernel_name == 'Exp':
-                kernel = amp*tinygp.kernels.Exp(metric)
+                kernel = tinygp.kernels.quasisep.Exp(scale=metric, sigma=np.sqrt(amp))
             else:
                 raise AssertionError(f'The kernel {kernel_name} is not in the '
                                      'currently supported list of kernels for '