########## R script: othordontModel2.Rs ##########

# For fitting a random intercepts and slopes linear mixed model
# to the `Orthodont data frame within the 'nlme' package,
# with a fixed effect for gender.

# Last changed: 01 AUG 2024

# Set flags:

doMCMC <- TRUE              ;      compileCode <- TRUE
doSummMCMC <- TRUE       
doFitPlot <- FALSE
doResidPlot <- FALSE

# Load in data:

library(nlme)
xOrig <- Orthodont$age
yOrig <- Orthodont$distance
idnumOrig <- as.numeric(Orthodont$Subject)
idnum <- match(idnumOrig,unique(idnumOrig))
xlabVal <- "age in years"
ylabVal <- "distance from the pituitary to the pterygomaxillary fissure"
layoutVec <- c(9,3)

# Standardise for Bayesian analysis:

mean.x <- mean(xOrig)  ;    sd.x <- sd(xOrig)
mean.y <- mean(yOrig)  ;    sd.y <- sd(yOrig)

x <- (xOrig - mean.x)/sd.x   ;    y <- (yOrig - mean.y)/sd.y

# Obtain X matrix:

isMale <- as.numeric(as.character(Orthodont$Sex)=="Male") 
X <- cbind(1,x,isMale)

# Set dimension values:

numObs <- length(y)
m <- length(unique(idnum))

# Do Markov chain Monte Carlo if requested:

if (doMCMC)
{
   library(rstan)

   # Set MCMC sample size paramaters:

   nBurn <- 1000
   nKept <- 1000
   nThin <- 1

   randIntAndSlpModel <-
   'data
   {
      int<lower=1> numObs;     int<lower=1> m;
      matrix[numObs,3] X;      real y[numObs];                  
      int idnum[numObs];
   }
   transformed data
   {
      vector[2] zeroVec;           zeroVec = rep_vector(0,2);
   }
   parameters
   {
      vector[3] beta;            vector[2] uLin[m];
      cov_matrix[2] SigmaU;      real<lower=0> sigmaEps;     
      vector[2] aSigma;
   }
   transformed parameters
   {
      vector[numObs] meanVec;  cov_matrix[2] SigmaUInv;
      meanVec = X*beta + to_vector(uLin[idnum,1])
                + to_vector(uLin[idnum,2]).*col(X,2);
      SigmaUInv = inverse(SigmaU);
   }
   model
   {
      vector[2] scaleSigma;

      y ~ normal(meanVec,sigmaEps);
      uLin ~ multi_normal(zeroVec,SigmaU);

      beta ~ normal(0,100000);

      aSigma ~ inv_gamma(0.5,0.0000000001);
      for (k in 1:2) scaleSigma[k] = 4/aSigma[k];
      SigmaU ~ inv_wishart(3,diag_matrix(scaleSigma));

      sigmaEps ~ cauchy(0,10000);
   }'

   # Set up data list:

   allData <- list(y=y,X=X,idnum=idnum,numObs=numObs,m=m)

   # Compile code for model if required:

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

   # Perform MCMC:

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

   # Save and extract relevant MCMC samples:

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

   SigmaU11MCMC <- as.vector(extract(stanObj,"SigmaU[1,1]",permuted=FALSE))
   SigmaU12MCMC <- as.vector(extract(stanObj,"SigmaU[1,2]",permuted=FALSE))
   SigmaU22MCMC <- as.vector(extract(stanObj,"SigmaU[2,2]",permuted=FALSE))

   sigmaU0MCMC <- sqrt(SigmaU11MCMC)
   sigmaU1MCMC <- sqrt(SigmaU22MCMC)
   rhoUMCMC <- SigmaU12MCMC/(sigmaU0MCMC*sigmaU1MCMC)

   sigmaEpsMCMC <- as.vector(extract(stanObj,"sigmaEps",permuted=FALSE))
   u0MCMC <- NULL   ; u1MCMC <- NULL
   for (i in 1:m)
   {
      charVar <- paste("uLin[",as.character(i),",1]",sep = "") 
      u0MCMC <- rbind(u0MCMC,extract(stanObj,charVar,permuted = FALSE))
      charVar <- paste("uLin[",as.character(i),",2]",sep = "") 
      u1MCMC <- rbind(u1MCMC,extract(stanObj,charVar,permuted = FALSE))
   }

   # Convert to original units:

   beta1OrigMCMC <- beta1MCMC*(sd.y/sd.x)
   beta2OrigMCMC <- beta2MCMC*sd.y
   beta0OrigMCMC <- mean.y + sd.y*beta0MCMC - beta1MCMC*(mean.x*sd.y/sd.x)
   sigmaU0OrigMCMC <- sigmaU0MCMC*sd.y
   sigmaU1OrigMCMC <- sigmaU1MCMC*sd.y/sd.x
   rhoUOrigMCMC <- rhoUMCMC
   sigmaEpsOrigMCMC <- sigmaEpsMCMC*sd.y
}

if (doSummMCMC)
{
   parNamesVal <- list(c(expression(beta[0])),c(expression(beta[1])),
                       c(expression(beta[2])),
                       c(expression(sigma[u1])),c(expression(sigma[u2])),
                       c(expression(rho[u])),c(expression(sigma[epsilon])))

   MCMClist <- list(cbind(beta0OrigMCMC,beta1OrigMCMC,beta2OrigMCMC, 
                          sigmaU0OrigMCMC,sigmaU1OrigMCMC,rhoUOrigMCMC,sigmaEpsOrigMCMC))
 
   summMCMC(MCMClist,parNames=parNamesVal)

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


######### End of orthodontModel2.Rs ##########