########## R script: ozoneMidwestViaStan.Rs ##########

# For performing MCMC-based fitting of a bivariate spline fit 
# to ozone data from the mid-west the U.S.A.

# Last changed: 06 OCT 2026

# Set tasks:

doMCMC <- TRUE            ;   compileCode <- TRUE
doSummMCMC <- TRUE
doImagePlot <- TRUE
do3Dplot <- TRUE

# Load required packages:

library(rstan)   ;   library(rgl)
library(fields)  ;   library(HRW)

# Load required functions:

source("Ztps.r")
source("rglSetup.R")
source("pointsInPoly.r")

# Set MCMC parameters:

nBurn <- 1000          # Length of burnin.
nKept <- 1000          # Size of the kept sample.

# Read in and extract data:

ozoneMidwest <- read.table("ozoneMidwest.txt",header=TRUE)
yOrig  <- ozoneMidwest$ozone
x1Orig <- ozoneMidwest$longitude
x2Orig <- ozoneMidwest$latitude

# Standardise the data for MCMC:

miny <- min(yOrig)     ;    maxy <- max(yOrig)
minx1 <- min(x1Orig)   ;    maxx1 <- max(x1Orig)
minx2 <- min(x2Orig)   ;    maxx2 <- max(x2Orig)

y <- 2*(yOrig - miny)/(maxy - miny)
x1 <- 2*(x1Orig - minx1)/(maxx1 - minx1)
x2 <- 2*(x2Orig - minx2)/(maxx2 - minx2)

# Set up matrices for bivariate penalized spline model:

ozoneMidwestKnots <- read.table("ozoneMidwestKnots.txt")
X <- cbind(1,x1,x2)
Z <- Ztps(cbind(x1,x2),ozoneMidwestKnots)

# Do plot of (x1,x2) values and show the positions of the knots:

plot(x1,x2,bty="l",xlab=expression(x[1]),ylab=expression(x[2]))
points(ozoneMidwestKnots,cex=3,col="red",lwd=2)
legend("topleft",legend=c("geographical data","knot positions"),lty=rep(0,2),
        col=c("black","red"),pch=rep(1,2),pt.cex=c(1,3),lwd=c(1,2))
readline("Hit Enter to continue.\n")


if (doMCMC)
{
   # Specify model in Stan:

   bivPenSplModel <- 
      'data
      {
         int n;                       int ncX;
         int ncZ;                     vector[n] y;                 
         matrix[n,ncX] X;             matrix[n,ncZ] Z;               
         real<lower = 0> sigmabeta;   real<lower = 0> ssigma;     
      }
      parameters 
      {
         vector[ncX] beta;       real<lower = 0> sigmaeps;   
         vector[ncZ] u;          real<lower = 0> sigmau; 
      }
      model 
      {
         y ~ normal(X*beta + Z*u,sigmaeps);
         u ~ normal(0,sigmau);    
         beta ~ normal(0,sigmabeta); 
         sigmau ~ cauchy(0,ssigma); 
         sigmaeps ~ cauchy(0,ssigma); 
      }'
   
   # Store data in a list in format required by Stan:

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

   # Compile code for model if required:

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

   # Obtain MCMC samples for each parameter using Stan:

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

   betaMCMC <- as.vector(extract(stanObj,"beta[1]",permuted=FALSE))
   betaMCMC <- rbind(betaMCMC,as.vector(extract(stanObj,"beta[2]",permuted=FALSE)))
   betaMCMC <- rbind(betaMCMC,as.vector(extract(stanObj,"beta[3]",permuted=FALSE)))

   sigmaepsMCMC <- as.vector(extract(stanObj,"sigmaeps",permuted=FALSE))
   sigmauMCMC <- as.vector(extract(stanObj,"sigmau",permuted=FALSE))

   uMCMC <- NULL
   for (k in 1:ncol(Z))
   {
      charVar <- paste("u[",as.character(k),"]",sep = "") 
      uMCMC <- rbind(uMCMC,extract(stanObj,charVar,permuted = FALSE))
   }
}

if (doSummMCMC)
{

   CTC <- crossprod(cbind(X,cbind(Z))) 
   Dmat <- diag(c(rep(0,3),rep(1,ncol(Z))))

   lambdaMCMC <- (sigmaepsMCMC/sigmauMCMC)^2
 
   EDFMCMC <- rep(NA,length(lambdaMCMC))
   for (iMCMC in 1:length(lambdaMCMC))
   {
      dVecCurr <- c(rep(0,3),rep(lambdaMCMC[iMCMC],ncol(Z)))

      HatMatrix <- solve(CTC+diag(dVecCurr),CTC,tol=1e-20)

      EDFMCMC[iMCMC] <- sum(diag(HatMatrix))
   }

   parNamesVal <- list(c("effective","degrees of","freedom"),expression(sigma[epsilon]))

   MCMClist <- list(cbind(EDFMCMC,sigmaepsMCMC))
 
   summMCMC(MCMClist,parNames=parNamesVal)
 
   readline("Hit Enter to continue.\n")
}

if (doImagePlot|do3Dplot)
{

   ozoneMidwestBdry <- read.table("ozoneMidwestBdry.txt")

   ng <- 51
   x1grid <- seq(0,2,length=ng)
   x2grid <- seq(0,2,length=ng)
   x1x2Mesh <- as.matrix(expand.grid(x1grid,x2grid))

   Xmesh <- cbind(1,x1x2Mesh)
   Zmesh <- Ztps(x1x2Mesh,ozoneMidwestKnots)

   fitMCMC <- Xmesh%*%betaMCMC + Zmesh%*%uMCMC
   fitMesh <- apply(fitMCMC,1,mean)  ; fitMeshSave <- fitMesh
   innerInds <- pointsInPoly(x1x2Mesh,ozoneMidwestBdry)
   fitMesh[!innerInds] <- NA
   fitMesh <- matrix(fitMesh,ng,ng)
}

if (doImagePlot)
{
   x1gridOrig <- minx1 + 0.5*x1grid*(maxx1 - minx1)   
   x2gridOrig <- minx2 + 0.5*x2grid*(maxx2 - minx2)   
   fitMeshOrig <- miny + 0.5*fitMesh*(maxy - miny)  

   ozoneMidwestBdryOrig <- matrix(NA,nrow(ozoneMidwestBdry),ncol(ozoneMidwestBdry))
    
   ozoneMidwestBdryOrig[,1] <- minx1 + 0.5*ozoneMidwestBdry[,1]*(maxx1 - minx1) 
   ozoneMidwestBdryOrig[,2] <- minx2 + 0.5*ozoneMidwestBdry[,2]*(maxx2 - minx2) 

   image.plot(x1gridOrig,x2gridOrig,fitMeshOrig,
              xlim=range(ozoneMidwestBdryOrig[,1]),
              ylim=range(ozoneMidwestBdryOrig[,2]),
              col=terrain.colors(1000),
              bty="n",xlab="longitude",ylab="latitude",
              legend.args = list(text="ozone",cex = 1.8, adj = 0.75),
              cex.lab=1.8)   
   lines(ozoneMidwestBdryOrig,lwd=2) 
   
   points(x1Orig,x2Orig,col="indianred3",lwd=2)

   readline("Hit Enter to continue.\n")
}

if (do3Dplot)
{
   createAxesandBase("longitude","latitude","ozone level")
   rhcs.spheres(x1,x2,y,col="red",radius=sp.rad)
   
   innerInds <- pointsInPoly(x1x2Mesh,ozoneMidwestBdry)

   al.val <- 0.75
   alvec <- rep(al.val,length(fitMesh))
   alvec[innerInds==FALSE] <- 0

   numCol <- 501
   zcolgrid <- seq(min(fitMeshSave),max(fitMeshSave),length=numCol)

   tz <- (round((numCol-1)*(fitMeshSave - min(fitMeshSave))
          /(max(fitMeshSave)-min(fitMeshSave))) +1)

   colorlut <- c(terrain.colors(numCol))

   colvec <- colorlut[tz] # assign colors to heights for each point
   colvec[innerInds==FALSE] <- bg.col

   rhcs.surface(x1grid,x2grid,fitMesh,color=colvec,front="lines",alpha=alvec)
}

########## End of ozoneMidwestViaStan ##########