GPs with many replicates per grid point

many_reps.r

library(ggplot2)
library(rstan)

N = 100 #number of x-axis grid points
M = 10 #number of noisy replicates per grid point
noise = 1 #amount of noise added to each replicate

#define the x-axis grid
x = seq(-10,10,length.out=N)

#define a wiggly function on x
f = scale(sin(x)*dnorm(x,5,8))[,1]

#show the intercept function
plot(x=x,y=f,type='l')


#set random seed for reproducibility
set.seed(1) 

#generate a number of noisy replicates
#(each row is a replicate)
Y = matrix(
  f + rnorm(N*M,0,noise)
  , nrow = M
  , byrow = TRUE
)

#show a couple of the noisy replicates
plot(x=x,y=Y[1,],type='l')
plot(x=x,y=Y[2,],type='l')

#scale x so it's 0-1
x_scaled = (x-min(x))/max(x-min(x))

#scale y to mean 0 and sd 1
Y_scaled = (Y-mean(Y))/sd(Y)

#compile the stan model
# This model requires that each replicate has all the values of x
model = stan_model(file="many_reps.stan")

#sample the posterior given the model & data
post <- sampling(
  model
  , data = list(
    N = N
    , X = x_scaled
    , M = M
    , Y = Y_scaled
    , fudge = 1e-3
  )
  , iter = 2e3
  , chains = 4
  , cores = 4
  , seed = 1
)
alarm()

print(
  post
  , pars = c('noise','eta','rho')
  , digits = 2
  , probs = c(.025,.975)
)
get_elapsed_time(post)/60
mean(rowSums(get_elapsed_time(post)/60))
sd(rowSums(get_elapsed_time(post)/60))

#look at the posterior on the intercept function
#  ribbon is the 95% credible interval
#  line is the data-generating true function
f_post = extract(
  post
  , pars = 'f'
)[[1]]
f_post = f_post*sd(Y)+mean(Y) #re-scale back to original data scale
loci = apply(f_post,2,quantile,prob=.025)
hici = apply(f_post,2,quantile,prob=.975)
ggplot(
  data = data.frame(
    x = x
    , f = f
    , lo = loci
    , hi = hici
  )
) +
  geom_point(
    data.frame(
      x = rep(x,times=M)
      , y = as.vector(t(Y))
    )
    , mapping = aes(
      x = x
      , y = y
    )
    , alpha = .5
  )+
  geom_ribbon(
    mapping = aes(
      x = x
      , ymin = lo
      , ymax = hi
    )
    , alpha = .5
    , fill = 'green'
  )+
  geom_line(
    mapping = aes(
      x = x
      , y = f
    )
    , colour = 'blue'
  )



#compile the stan model
# This model is OK with replicates having different values of x
model2 = stan_model(file="many_unequal_reps.stan")

#sample the posterior given the model & data
post2 <- sampling(
  model2
  , data = list(
    nX = N
    , X = x_scaled #unique values of x observed
    , nY = M*N
    , Y = as.vector(t(Y_scaled))
    , XY = rep(1:N,times=M)
    , fudge = 1e-3
  )
  , iter = 2e3
  , chains = 4
  , cores = 4
  , seed = 1
)
alarm()

print(
  post2
  , pars = c('noise','eta','rho')
  , digits = 2
  , probs = c(.025,.975)
)
get_elapsed_time(post2)/60
mean(rowSums(get_elapsed_time(post2)/60))
sd(rowSums(get_elapsed_time(post2)/60))

#look at the posterior on the intercept function
#  ribbon is the 95% credible interval
#  line is the data-generating true function
f_post2 = extract(
  post2
  , pars = 'f'
)[[1]]
f_post2 = f_post2*sd(Y)+mean(Y) #re-scale back to original data scale
loci2 = apply(f_post2,2,quantile,prob=.025)
hici2 = apply(f_post2,2,quantile,prob=.975)
ggplot(
  data = data.frame(
    x = x
    , f = f
    , lo = loci2
    , hi = hici2
  )
) +
  geom_point(
    data.frame(
      x = rep(x,times=M)
      , y = as.vector(t(Y))
    )
    , mapping = aes(
      x = x
      , y = y
    )
    , alpha = .5
  )+
  geom_ribbon(
    mapping = aes(
      x = x
      , ymin = lo
      , ymax = hi
    )
    , alpha = .5
    , fill = 'green'
  )+
  geom_line(
    mapping = aes(
      x = x
      , y = f
    )
    , colour = 'blue'
  )

many_reps.stan

data {
  // N: num unique X values
  int<lower=1> N ; 
  // X: x values
  vector[N] X ; 
  // M: number of replicates in per x
  int<lower=1> M ;
  // Y: data 
  matrix[M,N] Y ;
  // fudge: a fudge factor for computation
  real<lower=0> fudge ;
}
transformed data {
  // dx: computing the pairwise distances between x values
  matrix[N,N] dx ;
  //off-diagonal
  for (i in 1:(N-1)) {
    for (j in (i+1):N) {
      dx[i,j] = -( X[i] - X[j] )^2  ;
      dx[j,i] = dx[i,j] ; // lower mirrors upper
    }
  }
  //diagonal
  for (n in 1:N){
    dx[n,n] = 0 ;
  }
}
parameters {
  // noise: sampling noise for replicates
  real<lower=0> noise ;
  // eta_intercept: GP parameter, scale of output (Y)
  real<lower=0> eta ;
  // rho_intercept: GP parameter, x-axis scale of influence between points (wiggliness)
  real<lower=0> rho ;
  // z: dummy variable for quicker multivariate normal computation for GP
  vector[N] z ;
}
transformed parameters{
  // f: function values
  vector[N] f ;
  {
    // L: cholesky decomposition
    matrix[N,N] L ;
    // covMat: covariance matrix
    matrix[N,N] covMat ;
    //  covMat off-diagonal
    for (i in 1:(N-1)) {
      for (j in (i+1):N) {
        covMat[i,j] = eta * exp(dx[i,j]*rho) ;
        covMat[j,i] = covMat[i,j] ; // lower mirrors upper
      }
    }
    // covMat diagonal: adding jitter to avoid computation issues ??
    for (n in 1:N){
      covMat[n,n] = eta + fudge ; 
    }  
    L = cholesky_decompose(covMat) ;
    // assert sampling of f as GP
    f = L * z ;
  }
}
model {
  // priors on GP kernel parameters
  eta ~ weibull(2,1) ;
  rho ~ student_t(4,0,1) ;
  // prior on noise
  noise ~ weibull(2,1) ;
  //assert standard normal sampling of dummy variable z
  z ~ normal(0,1) ;
  // assert sampling of each replicate of Y as GP plus noise
  for(m in 1:M){
    Y[m] ~ normal(f,noise) ;
  }
}

many_unequal_reps.stan

data {
  // nX: num unique X values
  int<lower=1> nX ; 
  // X: unique x values
  vector[nX] X ; 
  // nY: number of observations total
  int<lower=1> nY ;
  // Y: data 
  vector[nY] Y ;
  // XY: *index* of x value per observation in Y 
  int XY[nY] ;
  // fudge: a fudge factor for computation
  real<lower=0> fudge ;
}
transformed data {
  // dx: computing the pairwise distances between x values
  matrix[nX,nX] dx ;
  //off-diagonal
  for (i in 1:(nX-1)) {
    for (j in (i+1):nX) {
      dx[i,j] = -( X[i] - X[j] )^2  ;
      dx[j,i] = dx[i,j] ; // lower mirrors upper
    }
  }
  //diagonal
  for (nx in 1:nX){
    dx[nx,nx] = 0 ;
  }
}
parameters {
  // noise: sampling noise for replicates
  real<lower=0> noise ;
  // eta_intercept: GP parameter, scale of output (Y)
  real<lower=0> eta ;
  // rho_intercept: GP parameter, x-axis scale of influence between points (wiggliness)
  real<lower=0> rho ;
  // z: dummy variable for quicker multivariate normal computation for GP
  vector[nX] z ;
}
transformed parameters{
  // f: function values
  vector[nX] f ;
  {
    // L: cholesky decomposition
    matrix[nX,nX] L ;
    // covMat: covariance matrix
    matrix[nX,nX] covMat ;
    //  covMat off-diagonal
    for (i in 1:(nX-1)) {
      for (j in (i+1):nX) {
        covMat[i,j] = eta * exp(dx[i,j]*rho) ;
        covMat[j,i] = covMat[i,j] ; // lower mirrors upper
      }
    }
    // covMat diagonal: adding jitter to avoid computation issues ??
    for (nx in 1:nX){
      covMat[nx,nx] = eta + fudge ; 
    }  
    L = cholesky_decompose(covMat) ;
    // assert sampling of f as zero-centered GP
    f = L * z ;
  }
}
model {
  // priors on GP kernel parameters
  eta ~ weibull(2,1) ;
  rho ~ student_t(4,0,1) ;
  // prior on noise
  noise ~ weibull(2,1) ;
  z ~ normal(0,1) ;
  // assert sampling of each replicate of Y as GP plus noise
  Y ~ normal(f[XY],noise) ;
}

Hi Mike, was working through this Gist. In case it matters, just noticed that line 36 of the first model spec, ‘many_reps.stan’, should probably read vector[N] z ;.

Thanks! And sorry for the delayed reply, I somehow failed to see your comment.