Gaussian process in greta (Matern covariance)


#1

As mentioned over on Github I am fiddling with spatial models a la: https://becarioprecario.bitbucket.io/spde-gitbook/ch-intro.html#sec:intro

Below is the code I could come up with inspired from: https://mc-stan.org/docs/2_19/stan-users-guide/fit-gp-section.html:

cMatern <- function(h, nu, kappa) {
  intt <- ifelse(h > 0, besselK(h * kappa, nu) * (h * kappa)^nu / 
           (base::gamma(nu) * 2^(nu - 1)), 1) 
  #intt <- intt * sigma
  return(intt)
}

# simulate some data
n <- 200
pts <- cbind(s1 = sample(1:n / n - 0.5 / n)^2,
             s2 = sample(1:n / n - 0.5 / n)^2)
dmat <- as.matrix(dist(pts))

# model parameter 
beta0 <- 10
sigma2e <- 0.3
sigma2u <- 5
kappa <- 7
nu <- 1

mcor <- cMatern(dmat, nu, kappa) 
mcov <- sigma2e * diag(nrow(mcor)) + mcor * sigma2u

R <- chol(mcov)
y1 <- beta0 + drop(crossprod(R, rnorm(n))) 

# greta modelling
## parameters
beta_0 <- normal(0, 1)
sigma_e <- cauchy(0, 5, truncation = c(0, Inf))
sigma_u <- cauchy(0, 5, truncation = c(0, Inf))
kappa <- normal(0, 1)
nu <- normal(0, 1)
eta <- normal(0, 1)

# linear predictor derivation
K <- cMatern(dmat, nu, kappa) 
K <- K * sigma_u
diag(K) <- diag(K) + 1e-9
L_K <- chol(K)
f <- L_K * eta + beta_0
# model
distribution(y1) <- normal(f, sigma_e)

Which does no work as I get the following error: “Error: cannot convert objects with missing or infinite values to greta_arrays”

Browsing through the greta forum this could be an interesting place to start.


#2

Yup, you can definitely do GPs, that’s one of the main reasons I originally wrote greta! There’s a prototype extension package https://github.com/greta-dev/greta.gp to make these models easier, but it’s currently halfway through a refactor and doesn’t work with the current version of greta. Hoping to polish it off soon, because it’s a nice example of the sort of simple interface you can create with an extension package.

The reason your code isn’t working is that you are passing greta arrays into a bunch of functions for which there aren’t yet greta equivalents: ifelse, besselK, and base::gamma. Bessel K and gamma functions could easily be added to greta (once I work out what to call the gamma function without it being confused with the distribution), and I could probably get ifelse to work.

However, I never usually work with the full Matern covariance, because nu is really hard to estimate from data. This full version of the Matern covariance can’t be fitted in INLA either; nu has to specified. It should definitely be possible though, so let me know if you really want that and we can work out a covariance function for it.

For now, here’s a simpler example using a Matern 3/2 kernel (ie. nu fixed at 3/2).

Simulate data:

# Matern 3/2 correlation (no marginal variances added)
matern32 <- function (distances, lengthscale) {
  a <- sqrt(3 * distances ^ 2) / lengthscale
  (1 + a) * exp(-a)
}

# simulate some data
n <- 200
pts <- cbind(s1 = sample(1:n / n - 0.5 / n)^2,
             s2 = sample(1:n / n - 0.5 / n)^2)
dmat <- as.matrix(dist(pts))

# model parameter
beta0 <- 10
sigma2e <- 0.3
sigma2u <- 5
lengthscale <- 1/7

mcor <- matern32(dmat, lengthscale)
mcov <- sigma2e * diag(nrow(mcor)) + mcor * sigma2u

R <- chol(mcov)
y1 <- beta0 + drop(crossprod(R, rnorm(n)))

Fit Matern GP with the linear predictor approach as in your code. Note eta needs to be a length n vector, you need to matrix-multiply it with L_K, and R’s chol() (and therefore greta’s) actually returns the upper triangular matrix, so needs to be transposed to get the lower triangular matrix. I use the one_by_one argument to greta to deal with Cholesky decomposition failures with bad parameter sets.

library(greta)

beta_0 <- normal(0, 1)
sigma_e <- cauchy(0, 5, truncation = c(0, Inf))
sigma_u <- cauchy(0, 5, truncation = c(0, Inf))
lengthscale <- normal(0, 1, truncation = c(0, Inf))
eta <- normal(0, 1, dim = n)

# linear predictor
K <- matern32(dmat, lengthscale)
K <- K * sigma_u ^ 2
K <- K + diag(n) * 1e-9
L_K <- t(chol(K))
f <- L_K %*% eta + beta_0

distribution(y1) <- normal(f, sigma_e)

m <- model(beta_0, lengthscale, sigma_u)
draws <- mcmc(m, one_by_one = TRUE)

This parameterisation has a latent Gaussian process, and then the data are observed with additional error on top of that GP. Since the response is Gaussian in this case, it’s much nicer to reparameterise as having the data drawn from the Gaussian process but with extra marginal variance. Which is actually how you simulated the data. Note this won’t work with non-Gaussian observation models

The following does away with the eta parameter, doesn’t need jitter to be added to the kernel, and samples much more nicely:

beta_0 <- normal(0, 1)
sigma_e <- cauchy(0, 5, truncation = c(0, Inf))
sigma_u <- cauchy(0, 5, truncation = c(0, Inf))
lengthscale <- normal(0, 1, truncation = c(0, Inf))

# covariance matrix
K <- matern32(dmat, lengthscale)
K <- K * sigma_u ^ 2
K <- K + diag(n) * sigma_e ^ 2

# define a multivariate normal on the data *rowwise*, so that observations are
# correlated (each row is an independent MVN)
mu <- beta_0 + zeros(1, n)
ty <- t(y1)
distribution(ty) <- multivariate_normal(mu, K)

m <- model(beta_0, lengthscale, sigma_u)
draws <- mcmc(m, one_by_one = TRUE)

#3

OK, I figured now was the time to get greta.gp working.

There’s a new branch and pull request open which depends on greta 0.3.1, which is itself still a release candidate (waiting on CRAN to accept) and pull request. You can install both by doing:

remotes::install_github("greta-dev/greta@rc_0_3_1")
remotes::install_github("goldingn/greta.gp@roll_forward")

Hopefully these will both be the master branches (and on CRAN, for greta) in a couple of days so those installation instructions will become redundant.

Here’s how you would do the same two models with greta.gp:

linear predictor (latent Gaussian process) version:

library(greta)
library(greta.gp)

beta_0 <- normal(0, 1)
sigma_e <- cauchy(0, 5, truncation = c(0, Inf))
sigma_u <- cauchy(0, 5, truncation = c(0, Inf))
lengthscale <- normal(0, 1, truncation = c(0, Inf))

K <- mat32(lengthscale, sigma_u ^ 2)
f_gp <- gp(pts, K, tol = 1e-9)
f <- f_gp + beta_0

distribution(y1) <- normal(f, sigma_e)

m1 <- model(beta_0, lengthscale, sigma_u)
draws1 <- mcmc(m1, one_by_one = TRUE)

observed Gaussian process version:

beta_0 <- normal(0, 1)
sigma_e <- cauchy(0, 5, truncation = c(0, Inf))
sigma_u <- cauchy(0, 5, truncation = c(0, Inf))
lengthscale <- normal(0, 1, truncation = c(0, Inf))

# covariance matrix
K1 <- mat32(lengthscale, sigma_u ^ 2)
K2 <- white(sigma_e ^ 2)
K <- K1 + K2

# define a multivariate normal on the data *rowwise*, so that observations are
# correlated (each row is an independent MVN)
mu <- beta_0 + zeros(1, n)
sigma <- K(pts)
ty <- t(y1)
distribution(ty) <- multivariate_normal(mu, sigma)

m2 <- model(beta_0, lengthscale, sigma_u)
draws2 <- mcmc(m2, one_by_one = TRUE)

A nice thing about this is that you can use greta.gp::project() to get a greta array for GP values in new locations, and greta::calculate() to get posterior samples for them. So you can easily do Gaussian process prediction to new locations even if you didn’t know the prediction locations before sampling.

s_seq <- seq(0, 1, length.out = 10)
pts_new <- expand.grid(s1 = s_seq,
                       s2 = s_seq)
f_gp_mn_new <- project(f_gp, pts_new)
f_mn_new <- f_gp_mn_new + beta_0
f_mn_new_draws <- calculate(f_mn_new, draws1)
f_mn_new_est <- colMeans(as.matrix(f_mn_new_draws))

That only works with the first one of these examples though. You have to do prediction yourself for the latter case (observed GP with noise):

K_xstar_x <- K1(pts_new, pts)
K_xstar_xstar <- K1(pts_new, pts_new)
i_sigma <- solve(sigma)
f_mn_new <- beta_0 + K_xstar_x %*% (i_sigma %*% as.matrix(y1))

f_mn_new_draws <- calculate(f_mn_new, draws2)
f_mn_new_est <- colMeans(as.matrix(f_mn_new_draws))

Note that both of these only compute posterior estimates of the mean of the GP at these locations, not realisations of the functions drawn from the GP. You’d need to calculate the posterior covariance matrix and draw more samples from a multivariate normal again to get posterior realisations of the function values.


#4

Oops, moving too fast! Accepted on CRAN now, so you can do:

remotes::install_github("greta-dev/greta")
remotes::install_github("goldingn/greta.gp")

#5

Wow thanks a lot for the super fast response and happy to hear that this is integrated in an already developed package. I’ll have a try.


#6

Should the second line rather be:

remotes::install_github("greta-dev/greta.gp")

#7

Yes it should, thanks!