gaussianProcessStan.Rmd

---
title: "Gaussian Process in Stan"
author: "MC"
date: "July 3, 2015"
output: html_document
---

# Intro

The following demonstrates a gaussian process using the Bayesian programming language Stan.  I also have a pure R approach ([1](https://github.com/m-clark/Miscellaneous-R-Code/blob/master/ModelFitting/gp%20Examples/gaussianprocessNoisy.R),
[2](https://github.com/m-clark/Miscellaneous-R-Code/blob/master/ModelFitting/gp%20Examples/gaussianprocessNoiseFree.R)), where you will find a little more context,  and a Matlab implementation ([1](https://github.com/m-clark/Miscellaneous-R-Code/blob/master/ModelFitting/gp%20Examples/gprDemoNoiseFree.m)).  

The R code for this demo is [here](https://github.com/m-clark/Miscellaneous-R-Code/blob/master/ModelFitting/gp%20Examples/gaussianProcessStan.R), the Stan model code [here](https://github.com/m-clark/Miscellaneous-R-Code/blob/master/ModelFitting/gp%20Examples/gpStanModelCode.stan).  A primary reference is Rasmussen & Williams (2006).


# Data and Parameter Setup
To start we will need to generate some data.  We'll have a simple x,y setup with a sample size of 20.  But we'll also create some test data to examine predictions later.

## Data
```{r datasetup}
set.seed(1234)
N = 20
Ntest = 200
x = rnorm(N, sd=5)
y = sin(x) + rnorm(N, sd=.1)
xtest = seq(min(x)-1, max(x)+1, l=Ntest)
plot(x,y, pch=19, col='#ff5500')
```

## Covariance function and parameters
In this demo we'll use the squared exponential kernel, which has three parameters to estimate.

```{r covFunc}
# parameters
eta_sq = 1
rho_sq = 1
sigma_sq = .1


# Covariance function same as implemented in the Stan code.
Kfn <- function (x, eta_sq, rho_sq, sigma_sq) {
  N = length(x)
  Sigma = matrix(NA, N, N)
  
  # off diag elements
  for (i in 1:(N-1)) {
    for (j in (i+1):N) {
      Sigma[i,j] <- eta_sq * exp(-rho_sq * (x[i] - x[j])^2);
      Sigma[j,i] <- Sigma[i,j];
    }
  }
  
  # diagonal elements
  for (k in 1:N)
    Sigma[k,k] <- eta_sq + sigma_sq; # + jitter
  Sigma
}

```

## Visualize the priors
With everything in place we can look at the some draws from the prior distribution.

```{r visPrior}
xinit = seq(-5,5,.2)
xprior = MASS::mvrnorm(3, 
                       mu=rep(0, length(xinit)), 
                       Sigma=Kfn(x=xinit,
                                 eta_sq = eta_sq, 
                                 rho_sq = rho_sq, 
                                 sigma_sq = sigma_sq))

library(reshape2)
gdat = melt(data.frame(x=xinit, y=t(xprior)), id='x')

library(ggvis)
gdat %>% 
  ggvis(~x, ~value) %>% 
  group_by(variable) %>% 
  layer_paths(strokeOpacity:=.5) %>% 
  add_axis('x', grid=F) %>% 
  add_axis('y', grid=F)
```

## Stan Code and Data
Now we are ready for the Stan code.  I've kept to the Stan manual section on Gaussian Processes ([github](https://github.com/stan-dev/example-models/tree/master/misc/gaussian-process])).

```{r stanCode}
gp = "
data {
  int<lower=1> N;                                # initial sample size
  vector[N] x;                                   # covariate
  vector[N] y;                                   # target
  int<lower=0> Ntest;                            # prediction set sample size
  vector[Ntest] xtest;                           # prediction values for covariate
}

transformed data {
  vector[N] mu;
  
  mu <- rep_vector(0, N);                        # mean function
}

parameters {
  real<lower=0> eta_sq;                          # parameters of squared exponential covariance function
  real<lower=0> inv_rho_sq;
  real<lower=0> sigma_sq;
}

transformed parameters {
  real<lower=0> rho_sq;
  rho_sq <- inv(inv_rho_sq);
}

model {
  matrix[N,N] Sigma;

  # off-diagonal elements for covariance matrix
  for (i in 1:(N-1)) {
    for (j in (i+1):N) {
      Sigma[i,j] <- eta_sq * exp(-rho_sq * pow(x[i] - x[j],2));
      Sigma[j,i] <- Sigma[i,j];
    }
  }

  # diagonal elements
  for (k in 1:N)
    Sigma[k,k] <- eta_sq + sigma_sq;             # + jitter for pos def

  # priors
  eta_sq ~ cauchy(0,5);
  inv_rho_sq ~ cauchy(0,5);
  sigma_sq ~ cauchy(0,5);

  # sampling distribution
  y ~ multi_normal(mu,Sigma);
}

generated quantities {
  vector[Ntest] muTest;                          # The following produces the posterior predictive draws
  vector[Ntest] yRep;                            # see GP section of Stan man- 'Analytical Form...'
  matrix[Ntest,Ntest] L;
  {
  matrix[N,N] Sigma;
  matrix[Ntest,Ntest] Omega;
  matrix[N,Ntest] K;
  matrix[Ntest,N] K_transpose_div_Sigma;
  matrix[Ntest,Ntest] Tau;

  # Sigma
  for (i in 1:N)
    for (j in 1:N)
      Sigma[i,j] <- exp(-pow(x[i] - x[j],2)) + if_else(i==j, 0.1, 0.0);

  # Omega
  for (i in 1:Ntest)
    for (j in 1:Ntest)
      Omega[i,j] <- exp(-pow(xtest[i] - xtest[j],2)) + if_else(i==j, 0.1, 0.0);

  # K
  for (i in 1:N)
    for (j in 1:Ntest)
      K[i,j] <- exp(-pow(x[i] - xtest[j],2));

  K_transpose_div_Sigma <- K' / Sigma;
  muTest <- K_transpose_div_Sigma * y;
  Tau <- Omega - K_transpose_div_Sigma * K;

  for (i in 1:(Ntest-1))
    for (j in (i+1):Ntest)
      Tau[i,j] <- Tau[j,i];

  L <- cholesky_decompose(Tau);
  }

  yRep <- multi_normal_cholesky_rng(muTest, L);
}
"
```

# Model Fitting and Summary
## Compile check
When using Stan it's good to do a very brief compile check before getting too far into debugging or the primary model fit.  You don't want to waste any more time than necessary to simply see if the code possibly works at all.

```{r compileCheck, eval=FALSE}
standata = list(N=N, x=x, y=y, xtest=xtest, Ntest=200)

library(rstan)
fit0 = stan(data=standata, model_code = gp, iter = 1, chains=1)
```

## Primary fit
Now we can do the main model fit. With the following setup it took about 2 minutes on my machine.  The fit object is almost 1 gig though, so will not be investigated here except visually in terms of the posterior predictive draws produced in the generated quantities section of the Stan code.

```{r mainFit, eval=FALSE}
iterations = 12000
wu = 2000
th = 20
chains = 4

library(parallel)
cl = makeCluster(4)
clusterExport(cl, c('gp',  'standata', 'fit0','iterations', 'wu', 'th', 'chains'))
clusterEvalQ(cl, library(rstan))


p = proc.time()
fit = parLapply(cl, seq(chains), function(chain) stan(data=standata, model_code = gp, 
                                                      iter = iterations,
                                                      warmup = wu, thin=th, 
                                                      chains=1, chain_id=chain, 
                                                      fit = fit0)
                )
(proc.time() - p)/3600

stopCluster(cl)
```

### Visualize Posterior Predicitive


```{r visPPD, echo=-3, eval=c(-1,-12)}
yRep = extract(fit, 'yRep')$yRep

load('stangp.RData'); suppressPackageStartupMessages(require(rstan))

gdat = data.frame(x,y)
gdat2 = melt(data.frame(x = sort(xtest), y=t(yRep[sample(2000, 3),])), id='x')

gdat2 %>% 
  ggvis(~x, ~value) %>% 
  group_by(variable) %>% 
  layer_paths(strokeOpacity:=.25) %>% 
  layer_points(x=~x, y=~y, fill:='#ff5500', data=gdat) %>% 
  add_axis('x', grid=F) %>% 
  add_axis('y', grid=F)  


# Visualize fit
yRepMean = get_posterior_mean(fit, 'yRep')[,5]
gdat3 = data.frame(x = sort(xtest), y=yRepMean)

gdat3 %>% 
  ggvis(~x, ~y) %>% 
  layer_paths(strokeOpacity:=.5, stroke:='blue') %>% 
  layer_points(x=~x, y=~y, fill:='#ff5500', data=gdat) %>% 
  add_axis('x', grid=F) %>% 
  add_axis('y', grid=F)  

```