Gaussian Process Sample in the Stan User's Guide

gp.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import print_function, absolute_import, division
from kernel import expand_dims
import numpy as np
import scipy as sci


def GPmodel(x, y, kernel, **options):
  x,y = expand_dims(x),expand_dims(y)
  s = options.get('sigma', 1e-8)
  def func(p):
    p = expand_dims(p)
    Kx = kernel(x,x) + s*s*np.eye(x.size)
    Kp = kernel(p,x)
    Kq = kernel(p,p) + s*s*np.eye(p.size)
    Cx = np.linalg.solve(Kx, np.eye(x.size))
    q = Kp.T.dot(Cx).dot(y.T).flatten()
    qe = Kq - Kp.T.dot(Cx.dot(Kp))
    return q,qe
  return func

gp_model.stan

data {
  int<lower=1> N;
  real X[N];
  vector[N] Y;
}

transformed data {
  vector[N] mu = rep_vector(0, N);
}

parameters {
  real<lower=0> rho;
  real<lower=0> scale;
  real<lower=0> sigma;
}

model {
  matrix[N, N] L_K;
  matrix[N, N] K = cov_exp_quad(X, scale, rho);
  real sq_sigma = square(sigma);

  for (n in 1:N) K[n,n] = K[n,n] + sq_sigma;
  L_K = cholesky_decompose(K);

  rho ~ inv_gamma(5, 5);
  scale ~ std_normal();
  sigma ~ std_normal();

  Y ~ multi_normal_cholesky(mu, L_K);
}

gp_stan.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import print_function, absolute_import, division
from argparse import ArgumentParser as ap
import matplotlib.pyplot as plt
import numpy as np
import pystan, pickle, sys, os, gzip
import gp, kernel


if __name__ == '__main__':
  parser = ap()
  parser.add_argument(
    '-n',dest='n',metavar='n_sample',type=int,default=20)
  parser.add_argument(
    '-s',dest='s',metavar='sigma',type=float,default=0.2)
  parser.add_argument(
    '--iter',dest='niter',metavar='n_iter',type=int,default=5000)
  parser.add_argument(
    '--chain',dest='nchain',metavar='n_chain',type=int,default=4)
  args = parser.parse_args()

  x = np.random.uniform(0,12,args.n)
  y = np.sin(x**1.3)+np.random.normal(scale=args.s,size=args.n)
  w = np.linspace(0,14,1000)
  z = np.sin(w**1.3)

  stan_data = {
    'N': args.n,
    'X': x,
    'Y': y,
  }
  pklfile = 'gp_model.pkl'
  if not os.path.exists(pklfile):
    sm = pystan.StanModel(file='gp_model.stan')
    with open(pklfile, 'wb') as f: pickle.dump(sm, f)
  else:
    with open(pklfile, 'r') as f: sm = pickle.load(f)

  fit = sm.sampling(data=stan_data, iter=args.niter, chains=args.nchain)
  fit.plot()
  print(fit)
  plt.show()

  par = fit.extract(permuted=True)
  dumpfile = 'gp_model.result.gz'
  with gzip.open(dumpfile, 'wb') as f: pickle.dump(par, f)

  rho = np.mean(par['rho'])
  scale = np.mean(par['scale'])
  sigma = np.mean(par['sigma'])
  K = kernel.RBFKernel(rho, scale)
  M = gp.GPmodel(x, y, K, sigma=sigma)
  f, fe = M(w)
  err = fe.diagonal()

  fig, ax = plt.subplots(1,1)
  ax.fill_between(w,f-3*err,f+3*err,facecolor=[.9,.9,.9])
  ax.fill_between(w,f-err,f+err,facecolor=[.8,.8,.8])
  ax.plot(w,z)
  ax.plot(x,y,'o')
  ax.plot(w,f)
  plt.show()

kernel.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import print_function, absolute_import, division
import numpy as np
import scipy as sci
import scipy.special as sp


def expand_dims(arr):
  if arr.ndim==1: arr = np.expand_dims(arr, axis=0)
  return arr


def RBFKernel(rho, scale):
  def func(x1,x2):
    x1,x2 = expand_dims(x1),expand_dims(x2)
    return scale*np.exp(-(x1-x2.T)*(x1-x2.T)/2.0/rho/rho)
  return func


def OUKernel(rho, scale):
  def func(x1,x2):
    x1,x2 = expand_dims(x1),expand_dims(x2)
    return scale*np.exp(-np.abs(x1-x2.T)/rho)
  return func


def MaternKernel(nu, scale):
  def func(x1,x2):
    x1,x2 = expand_dims(x1),expand_dims(x2)
    d = np.abs(x1-x2.T)
    x = np.sqrt(2*nu)*d/scale
    x[x==0] = 1e-9
    return (2**(1-nu))/sp.gamma(nu) * x**nu * sp.kv(nu, x)
  return func