Gaussian Process modelling of an intercept+effect model with noisy replicates

gp_multi_effect.R

library(ggplot2)
library(rstan)

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

NX = 100 #number of x-axis grid points
NZ = 10 #number of noisy replicates per condition

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

#define some wiggly functions on x
f_intercept = sin(x)*dnorm(x,5,8)
f_effect = dnorm(x,-5,2)/5

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

#show the effect function
plot(x=x,y=f_effect,type='l')

#combine the intercept and effect functions to create condition functions
condition1 = f_intercept + f_effect
condition2 = f_intercept - f_effect

#show the two condition functions
plot(x=x,y=condition1,type='l',ylim=range(c(condition1,condition2)),ylab='')
lines(x=x,y=condition2)

#generate a number of noisy replicates for condition 1 
#(each row is a replicate)
Y1 = matrix(
    f_intercept + f_effect + rnorm(NX*NZ,0,.02)
    , nrow = NZ
    , byrow = TRUE
)

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

#generate a number of noisy replicates for condition 2
#(each row is a replicate)
Y2 = matrix(
    f_intercept - f_effect + rnorm(NX*NZ,0,.02)
    , nrow = NZ
    , byrow = TRUE
)

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


#compile the stan model
model = stan_model(file="gp_multi_effect.stan")

#sample the posterior given the model & data
post <- sampling(
    model
    , data = list(
        NX = NX
        , X = x
        , NY1 = nrow(Y1)
        , Y1 = Y1
        , NY2 = nrow(Y2)
        , Y2 = Y2
        , fudge = 1e-3
    )
    , iter = 2000
    , chains = 4 
    , cores = 4
    , seed = 1
    , pars = c('L_intercept','L_effect')
    , include = FALSE
)
alarm()

print(post)

#look at the posterior on the intercept function
#  ribbon is the 95% credible interval
#  line is the data-generating true function
intercept = extract(
    post
    , pars = 'f_intercept'
)[[1]]
loci = apply(intercept,2,quantile,prob=.025)
hici = apply(intercept,2,quantile,prob=.975)
ggplot(
    data = data.frame(
        x = x
        , f_intercept = f_intercept
        , lo = loci
        , hi = hici
    )
) +
    geom_ribbon(
        mapping = aes(
            x = x
            , ymin = lo
            , ymax = hi
        )
    )+
    geom_line(
        mapping = aes(
            x = x
            , y = f_intercept
        )
    )

#look at the posterior on the effect function
#  ribbon is the 95% credible interval
#  line is the data-generating true function
effect = extract(
    post
    , pars = 'f_effect'
)[[1]]
loci = apply(effect,2,quantile,prob=.025)
hici = apply(effect,2,quantile,prob=.975)
ggplot(
    data = data.frame(
        x = x
        , f_effect = f_effect
        , lo = loci
        , hi = hici
    )
) +
    geom_ribbon(
        mapping = aes(
            x = x
            , ymin = lo
            , ymax = hi
        )
    )+
    geom_line(
        mapping = aes(
            x = x
            , y = f_effect
        )
    )

gp_multi_effect.stan

data {
    // NX: num unique X values
    int<lower=1> NX ; 
    // X: x values
    vector[NX] X ; 
    // NY1: number of replicates in condition 1
    int<lower=1> NY1 ;
    // Y1: data from condition 1
    matrix[NY1,NX] Y1 ;
    // NY2: number of replicates in condition 2
    int<lower=1> NY2 ;
    // Y2: data from condition 2
    matrix[NY2,NX] Y2 ;
    // fudge: a fudge factor for computation
    real<lower=0> fudge ;
}
transformed data {
    // mu: to be set to zero, used for GPs
    vector[NX] mu ;
    for (nx in 1:NX){
        mu[nx] <- 0 ;
    }
}
parameters {
    // noise: sampling noise for replicates
    real<lower=0> noise ;
    // eta_intercept: GP parameter, scale of output (Y) for intercept
    real<lower=0> eta_intercept ;
    // eta_effect: GP parameter, scale of output (Y) for effect
    real<lower=0> eta_effect ;
    // rho_intercept: GP parameter, x-axis scale of influence between points (wiggliness) for intercept
    real<lower=0> rho_intercept ;
    // rho_effect: GP parameter, x-axis scale of influence between points (wiggliness) for effect
    real<lower=0> rho_effect ;
    // f_intercept: intercept function values
    vector[NX] f_intercept ;
    // f_effect: effect function values
    vector[NX] f_effect ;
}
transformed parameters{
    // L_intercept: cholesky decomposition for intercept
    matrix[NX,NX] L_intercept ;
    // L_effect: cholesky decomposition for effect
    matrix[NX,NX] L_effect ;
    {
        // covMat_intercept: covariance matrix for intercept
        matrix[NX,NX] covMat_intercept ;
        // covMat_intercept: covariance matrix for effect
        matrix[NX,NX] covMat_effect ;
        //  covMat off-diagonal
        for (i in 1:(NX-1)) {
            for (j in (i+1):NX) {
                covMat_intercept[i,j] <- eta_intercept * exp(- pow(X[i] - X[j],2)*rho_intercept) ;
                covMat_intercept[j,i] <- covMat_intercept[i,j] ; // lower mirrors upper
                covMat_effect[i,j] <- eta_effect * exp(- pow(X[i] - X[j],2)*rho_effect) ;
                covMat_effect[j,i] <- covMat_effect[i,j] ; // lower mirrors upper
            }
        }
        // covMat diagonal: adding jitter to avoid computation issues ??
        for (k in 1:NX){
            covMat_intercept[k,k] <- eta_intercept + fudge ; 
            covMat_effect[k,k] <- eta_effect + fudge ;
        }  
        L_intercept <- cholesky_decompose(covMat_intercept) ;
        L_effect <- cholesky_decompose(covMat_effect) ;
    }
}
model {
    // priors on GP kernel parameters
    eta_intercept ~ weibull(2,1) ;
    rho_intercept ~ student_t(4,0,1) ;
    eta_effect ~ weibull(2,1) ;
    rho_effect ~ student_t(4,0,1) ;
    // prior on noise
    noise ~ weibull(2,1) ;
    // assert sampling of f_intercept as GP
    f_intercept ~ multi_normal_cholesky(mu,L_intercept) ;
    // assert sampling of f_effect as GP
    f_effect ~ multi_normal_cholesky(mu,L_effect) ;
    // assert sampling of each replicate of Y1 as the combo (+) of GPs plus noise
    for(ny1 in 1:NY1){
        Y1[ny1] ~ normal(f_intercept+f_effect,noise) ;
    }
    // assert sampling of each replicate of Y1 as the combo (-) of GPs plus noise
    for(ny2 in 1:NY2){
        Y2[ny2] ~ normal(f_intercept-f_effect,noise) ;
    }
}