########## R script: wolfcampLOOana.Rs ##########

# For performing MCMC-based fitting of a 
# bivariate spline fit to the Wolfcamp aquifer data.

# Last changed: 14 OCT 2024

# Set tasks:

doMCMC <- TRUE     ;  compileCode <- TRUE   
obtainLOOvalue <- TRUE

# Load required packages:

library(rstan)   ;   library(loo)

# Set MCMC parameters:

nBurn <- 1000          # Length of warmup.
nKept <- 1000          # Size of the kept sample.
nThin <- 1             # Thinning factor. 

# Read in and extract data:

wolfcamp <- read.table("wolfcamp.txt",header=TRUE)
yOrig  <- wolfcamp$piezometricHead
x1Orig <- wolfcamp$xCoordinate
x2Orig <- wolfcamp$yCoordinate

# Standardised data for MCMC:

y <- 2*(yOrig - min(yOrig))/(max(yOrig) - min(yOrig))
x1 <- 2*(x1Orig - min(x1Orig))/(max(x1Orig) - min(x1Orig))
x2 <- 2*(x2Orig - min(x2Orig))/(max(x2Orig) - min(x2Orig))

# Set up matrices for bivariate penalized spline model:

wolfcampKnots <- read.table("wolfcampKnots.txt")
X <- cbind(1,x1,x2)

if (doMCMC)
{
   # Specify model in Stan:

   planeModel <- 
      'data
      {
         int n;                       int ncX;
         vector[n] y;                 matrix[n,ncX] X;                     
         real<lower = 0> sigmabeta;   real<lower = 0> ssigma;     
      }
      parameters 
      {
         vector[ncX] beta;       real<lower = 0> sigmaeps;    
      }
      model 
      {
         y ~ normal(X*beta,sigmaeps);
         beta ~ normal(0,sigmabeta);
         sigmaeps ~ cauchy(0,ssigma); 
      }
      generated quantities
      {
         vector[n] log_lik;
         for (i in 1:n)
         {
            log_lik[i] = normal_lpdf(y[i]|X[i]*beta,sigmaeps);
         }
      }'
     
   # Store data in a list in format required by Stan:

   allData <- list(n=length(y),ncX=ncol(X),y=y,X=X,
                   sigmabeta=1e5,ssigma=1e5)

   # Compile code for model if required:

   if (compileCode)
      stanCompilObj <- stan(model_code=planeModel,data=allData,
                               iter=1,chains=1)

   # Obtain MCMC samples for each parameter using Stan:

   stanObj <-  stan(model_code=planeModel,data=allData,
                       warmup=nBurn,iter=(nBurn+nKept),
                       chains=1,thin=nThin,refresh=250,fit=stanCompilObj)
}

if (obtainLOOvalue)
{

   currLOOobj <- loo(extract_log_lik(stanObj))

   currEst <- currLOOobj$estimates

   LOOval <- currEst[3,1]

   cat("\n\n")
   cat("The leave-one-out information criterion for the plane model is:",LOOval)
   cat("\n\n")
}

########## End of wolfcampLOOana ##########