```python
import numpy as np
import matplotlib.pyplot as plt
plt.style.use('apw-notebook.mplstyle')
%matplotlib inline
np.random.seed(1234)
```

$$
\newcommand{\mat}[1]{\boldsymbol{\rm #1}}
\newcommand{\vec}[1]{\boldsymbol{#1}}
\newcommand{\T}[1]{{#1}^{\mathsf{T}}}
\newcommand{\inv}[1]{{#1}^{-1}}
$$

# Gaussian processes

GPs are a type of very flexible models that are primarily useful for prediction, discovery, and making models robust to unknown error properties. A lot of the GP literature is focused around time-series analysis, but they are useful in other contexts as well in which one variable (the analog to 'time') is very precisely measured or perfectly known. They are pretty slow to compute, but people like Dan Foreman-Mackey take speed very seriously (see [george](http://github.com/dfm/george) and [celerite](http://celerite.readthedocs.io/)).

The default resource for Gaussian processes is [Rasmussen & Williams](http://www.gaussianprocess.org/gpml/).

In previous classes we talked about the problem of estimating the mean flux of a star given noisy measurements. Recall that we assumed that the data were independent, uncorrelated, and Gaussian, and with those assumptions the likelihood over all data was separable into a product over individual (per–datum) Gaussians. For $N$ data points labeled $n$, we wrote:

$$
p(\{f_n\} \,|\, f_0, \sigma_n^2) = \prod_n^N \mathcal{N}(f_n \,|\, f_0, \sigma_n^2) \\
\ln p(\{f_n\} \,|\, f_0, \sigma_n^2) = -\frac{1}{2}\sum_n^N \left[\frac{(f_n - f_0)^2}{\sigma_n^2} + \ln\left(2\pi \,\sigma_n^2\right)\right]
$$

Remember that we can also write this log-likelihood using matrix notation by defining a _residual vector_, $\vec{r} = \vec{f} - f_0$, and a _covariance matrix_, $\mat{C}$:

$$
\begin{align}
\vec{f} &= \begin{bmatrix}
             f_{1} \\
             f_{2} \\
             \vdots \\
             f_{N}
           \end{bmatrix}\\
\mat{C} &= \begin{bmatrix}
             \sigma_1^2 & 0 & \cdots & 0 \\
             0 & \sigma_2^2 & \cdots & 0  \\
             \vdots & \vdots & \ddots & \vdots \\
             0 & 0 & \cdots & \sigma_{N}^2
           \end{bmatrix}\\
\end{align}
$$

With these objects, the log-likelihood is:

$$
\ln p(\vec{f} \,|\, f_0, \mat{C}) = -\frac{1}{2} \T{\vec{r}} \, \inv{\mat{C}} \, \vec{r} 
    - \frac{1}{2} \ln\det\mat{C} - \frac{N}{2}\ln 2\pi
$$

Writing it like this you can interpret the data vector $\vec{f}$ as being a single draw from an $N$-dimensional Gaussian with mean $f_0$ and (diagonal) covariance matrix $\mat{C}$. The diagonality of the covariance matrix is a result of our assumption that the data are uncorrelated or independent. But what if the data are correlated in some way? Or, what if there are off-diagonal elements in the covariance matrix? Let's look at both cases:


```python
# for simplicity, set f_0 to 0
f0 = 0.
n_data = 64

times = np.random.uniform(0, 100, n_data)
times.sort()

y_err = np.random.uniform(0.1, 1., n_data)

# first construct a purely diagonal covariance matrix
C1 = np.diag(y_err**2)

# now create a covariance matrix with correlated noise
C2 = 8 * np.exp(-0.5 * (times[:,None]-times[None,:])**2 / 10.**2)
C2[np.diag_indices_from(C2)] += y_err**2
```


```python
fig,axes = plt.subplots(2, 1, figsize=(5,10))
axes[0].imshow(C1, cmap='Greys_r')
axes[1].imshow(C2, cmap='Greys_r')
for ax in axes:
    ax.xaxis.set_visible(False)
    ax.yaxis.set_visible(False)
fig.tight_layout()
```


```python
f1 = np.random.multivariate_normal(np.full(n_data, f0), C1)
f2 = np.random.multivariate_normal(np.full(n_data, f0), C2)

fig,axes = plt.subplots(2, 1, figsize=(8,8), sharex=True, sharey=True)

axes[0].errorbar(times, f1, yerr=y_err, marker='.', linestyle='none')
axes[1].errorbar(times, f2, yerr=y_err, marker='.', linestyle='none')

axes[0].set_ylim(-8,8)
```

If we knew this covariance matrix for the data, then doing least-squares fitting would be as simple as ever:


```python
def solve_linalg(times, flux, Cov):
    # construct design matrix and invert covariance matrix
    X = np.vander(times, N=1)
    Cinv = np.linalg.inv(Cov)

    best_f0 = np.linalg.inv(X.T @ Cinv @ X) @ (X.T @ Cinv @ flux)
    cov_f0 = np.linalg.inv(X.T @ Cinv @ X)
    
    return best_f0, cov_f0
```


```python
x0 = np.linspace(times.min(), times.max(), 1024)

fig,ax = plt.subplots(1, 1, figsize=(8,6))
ax.errorbar(times, f2, yerr=y_err, marker='.', linestyle='none')

for color, C in zip(['#31a354', '#de2d26'], [C2, C1]):
    best_f0, cov_f0 = solve_linalg(times, f2, C)
    samples = np.random.multivariate_normal(best_f0, cov_f0, 10000)

    lines = np.dot(np.vander(x0, 1), samples.T)
    q = np.percentile(lines, [16, 84], axis=1)

    ax.plot(x0, np.full_like(x0, best_f0), linewidth=2, alpha=0.8, color=color)
    ax.fill_between(x0, q[0], q[1], alpha=0.25, color=color)

ax.set_ylim(-8,8)
```

But what if we don't know the noise structure? For example:


```python
from astropy.io import fits
data = fits.getdata('../data/k2.fits')
```


```python
data = data[np.isfinite(data['PDCSAP_FLUX'])]
raw_time = data['TIME'].astype(np.float64)
raw_flux = data['PDCSAP_FLUX'].astype(np.float64)
raw_flux_err = data['PDCSAP_FLUX_ERR'].astype(np.float64)
```


```python
med = np.median(raw_flux)
robust_std = 1.48 * np.median(np.abs(raw_flux - med))
```


```python
idx = (raw_flux < (med + 5*robust_std)) & (raw_flux > (med - 5*robust_std))

trim = 512
time = raw_time[idx][:trim] # too much data
flux = raw_flux[idx][:trim]
flux_err = raw_flux_err[idx][:trim]

time = time - time.min()
```


```python
plt.figure(figsize=(12,6))
plt.errorbar(time, flux, flux_err, marker='.', linestyle='none')
plt.xlabel('time')
plt.ylabel('flux')
```

We're only given single error-bars for each flux measurement, i.e. just the diagonal terms of the covariance matrix. But there is likely some very correlated noise in this light curve! What do we do? One (very bad) idea would be to add all off-diagonal elements to your model as parameters! For an $N\times N$ symmetric matrix, that amounts to adding $\frac{N\,(N-1)}{2}$ parameters to your model (and scales like the number of data points _squared_!). No one would seriously consider doing that...

Instead: what if we can write a simple parametrized _kernel function_, $k$, that generates the off-diagonal elements as a function of location in the matrix? That is:

$$
C_{ij} = \sigma_i^2 \, \delta_{ij} + k(t_i, t_j)
$$

We now just have to specify this function, $k$. This is where the magic comes in...the function must produce a symmetric, positive, semi-definite function, so there aren't too many options. We'll introduce a few others later, but for now let's use the _Squared Exponential Kernel_:

$$
k(x_i, x_j) = a^2 \, \exp\left(-\frac{(x_i - x_j)^2}{2\,b^2}\right)
$$

(yes it's just another Gaussian...). This kernel says that data points near $x_i$ are correlated with a length scale set by the parameter $b$ and a correlation amplitude set by the parameter $a$. 

__Exercise__: Implement this kernel function below


```python
def sq_exp_kernel(x_i, x_j, a, b):
    # TODO: Delete the line below and fill in this function
    pass
```


```python
a = 4.
b = 1.
Cov = sq_exp_kernel(time[:,None], time[None], a, b)

plt.figure(figsize=(6,6))
plt.imshow(Cov, cmap='Greys_r')
plt.gca().xaxis.set_visible(False)
plt.gca().yaxis.set_visible(False)
plt.tight_layout()
```

We'll now implement a log-likelihood function for a Gaussian process with a simple constant mean model (same as we did above for inferring the mean flux):

$$
\ln p(\vec{f} \,|\, f_0, \mat{C}) = -\frac{1}{2} \T{\vec{r}} \, \inv{\mat{C}} \, \vec{r} 
    - \frac{1}{2} \ln\det\mat{C} - \frac{N}{2}\ln 2\pi
$$

Replacing $\vec{C}$ with the sum of the diagonal covariance matrix with that generated by the kernel function.

The parameters that come in will be $(f_0, \ln a, \ln b)$. You'll need to invert a large matrix in the likelihood call. Rather than using `numpy.linalg.inv`, I recommend using `numpy.linalg.solve` in the following way. As a part of the likelihood expression, you need to compute the vector $\vec{y}$ in:

$$
\vec{y} = \inv{\mat{C}}\,\vec{r}\\
\vec{r} = \mat{C}\,\vec{y}
$$

Instead of explicitly doing the inversion, you can do:
$$
\vec{y} = {\rm solve}\left(\mat{C}, \vec{r}\right)
$$

Also, instead of taking the log of the determinant, use `numpy.linalg.slogdet` (which is much more numerically stable).

__Exercise__: Implement the likelihood function below


```python
def ln_likelihood(pars, t, flux, flux_err):   
    # TODO: Delete the line below and fill in this function
    pass
```

We now need to implement priors for the parameters. We'll assume a uniform prior on $f_0$ and uniform priors on the log hyper-parameters $\ln a$ and $\ln b$. Choose wide and sensible ranges for the parameters.

__Exercise:__ Implement the prior function below


```python
def ln_prior(pars):
    # TODO: Delete the line below and fill in this function
    pass
```

Now we can run MCMC to infer the parameters:


```python
import emcee
```


```python
def ln_posterior(pars, t, flux, flux_err):
    lp = ln_prior(pars)
    if not np.isfinite(lp):
        return -np.inf
    
    ll = ln_likelihood(pars, t, flux, flux_err)
    if not np.isfinite(ll):
        return -np.inf
    
    return lp + ll
```


```python
nwalkers = 32
p0 = [np.median(flux), 4., -2.]

sampler = emcee.EnsembleSampler(nwalkers, dim=len(p0), lnpostfn=ln_posterior,
                                args=(time, flux, flux_err))
```

__Warning__: this could take a few minutes...


```python
_p0 = emcee.utils.sample_ball(p0, std=[1E-3]*len(p0), size=nwalkers)

pos,_,_ = sampler.run_mcmc(_p0, 256)
sampler.reset()
_ = sampler.run_mcmc(pos, 256)
```


```python
for k in range(len(p0)):
    plt.figure()
    for walker in sampler.chain[...,k]:
        plt.plot(walker, marker='', drawstyle='steps-mid', color='k', alpha=0.1)
```

We'll use [@dfm](https://github.com/dfm)'s Gaussian process package [George](https://github.com/dfm/george) to plot samples from the GP at each choice of the hyper-parameters $a, b$:


```python
import george
```


```python
def plot_gp(grid, time, flux, flux_err, f0, a, b):
    K = a**2 * george.kernels.ExpSquaredKernel(b)
    gp = george.GP(K, mean=f0)
    gp.compute(time, flux_err)
    return gp.predict(flux, grid, mean_only=True)
```


```python
plt.figure(figsize=(12,6))
plt.errorbar(time, flux, flux_err, marker='.', linestyle='none')

t_grid = np.linspace(time.min(), time.max(), 1024)
flatchain = np.vstack(sampler.chain[:,::32])

for i in np.random.choice(flatchain.shape[0], size=64, replace=False):
    f0,ln_a,ln_b = flatchain[i]
    mean = plot_gp(t_grid, time, flux, flux_err, f0, np.exp(ln_a), np.exp(ln_b))    
    plt.plot(t_grid, mean, marker='', zorder=100, color='#2b8cbe', alpha=0.25)

plt.xlabel('time')
plt.ylabel('flux')
```

Now that we know how they work, let's use `George` to compute the likelihood (it's much faster!) and re-do the analysis above:


```python
from george import kernels
```


```python
def ln_likelihood2(pars, t, flux, flux_err):       
    f_0 = pars[0]
    a, b = np.exp(pars[1:])
    
    # residual vector - you could replace with a more complicated mean model
    r = flux - f_0 
    
    K = a**2 * kernels.ExpSquaredKernel(b)
    gp = george.GP(K)
    gp.compute(t, flux_err)
    return gp.lnlikelihood(flux - f_0)
```


```python
def ln_prior2(pars):    
    f_0 = pars[0]
    
    if f_0 < 0 or f_0 > 1E10:
        return -np.inf
    
    for ln_p in pars[1:]:
        if ln_p < -10 or ln_p > 10:
            return -np.inf
    
    return 0.
```


```python
def ln_posterior2(pars, t, flux, flux_err):
    lp = ln_prior2(pars)
    if not np.isfinite(lp):
        return -np.inf
    
    ll = ln_likelihood2(pars, t, flux, flux_err)
    if not np.isfinite(ll):
        return -np.inf
    
    return lp + ll
```


```python
nwalkers = 32
p0 = [np.median(flux), 4., -2.]

sampler2 = emcee.EnsembleSampler(nwalkers, dim=len(p0), lnpostfn=ln_posterior2,
                                 args=(time, flux, flux_err))
```


```python
_p0 = emcee.utils.sample_ball(p0, std=[1E-3]*len(p0), size=nwalkers)

pos,_,_ = sampler2.run_mcmc(_p0, 256)
sampler2.reset()
_ = sampler2.run_mcmc(pos, 256)
```


```python
for k in range(len(p0)):
    plt.figure()
    for walker in sampler2.chain[...,k]:
        plt.plot(walker, marker='', drawstyle='steps-mid', color='k', alpha=0.1)
```


```python
flatchain = np.vstack(sampler2.chain[:,::32])
```


```python
plt.figure(figsize=(12,6))
plt.errorbar(time, flux, flux_err, marker='.', linestyle='none')

t_grid = np.linspace(time.min(), time.max(), 1024)
flatchain = np.vstack(sampler2.chain[:,::32])

for i in np.random.choice(flatchain.shape[0], size=64, replace=False):
    f0, ln_a, ln_b = flatchain[i]
    a = np.exp(ln_a)
    b = np.exp(ln_b)
    
    K = a**2 * kernels.ExpSquaredKernel(b)
    gp = george.GP(K, mean=f0)
    gp.compute(time, flux_err)
    mean = gp.predict(flux, t_grid, mean_only=True)
    plt.plot(t_grid, mean, marker='', zorder=100, color='#2b8cbe', alpha=0.25)

plt.xlabel('time')
plt.ylabel('flux')
```

We've been using an exponential-squared kernel function but I never really justified its use. That's a generic problem of GP models: there are a handful of popular kernel functions, but sometimes there's no real motivation to use one over another. In this case, however, the variability in our data (the "noise") clearly has contributions over a large range of frequencies. The kernel we're using has a fixed correlation length-scale, meaning it won't necessarily adapt well when there are sharp features and long-term trends. 

__Exercise__: Re-do the analysis (using `George`) we did above but instead with the following kernel, known as the "Matérn 3/2":

$$
k(x_i, x_j) = a^2 \, \left(1 + \frac{\sqrt{3}\,\left|x_i-x_j\right|}{b}\right) \, \exp\left(-\frac{\sqrt{3}\,\left|x_i-x_j\right|}{b}\right)
$$

(see `george.kernels.Matern32Kernel`).


```python
# TODO: fill in here
```


```python

```

97779e69b300f525fa2f8377d731590f74d543f9
ipynb
Jupyter Notebook
notebooks/Gaussian processes - workbook.ipynb
adrn/AST542
3a633cde68235cae95093e9f080dc0f3429705cd
MIT
2017-03-30T20:13:38.000Z
2021-07-12T00:55:13.000Z
notebooks/Gaussian processes - workbook.ipynb
adrn/AST542
3a633cde68235cae95093e9f080dc0f3429705cd
MIT
notebooks/Gaussian processes - workbook.ipynb
adrn/AST542
3a633cde68235cae95093e9f080dc0f3429705cd
MIT
2021-05-28T15:26:52.000Z
2021-07-09T15:34:41.000Z
Qwen/Qwen-72B
1. YES
2. YES