## Simple linear regression model using JAGS
## Author: Ben Saville
## Date: Oct 11, 2012

setwd("~/Berry/Conferences/Vandy2018/SavilleJAGSintro/Rjags")
library('rjags')


################################################################
# Data: allergy medication
# y contains the N hours of relief experienced
# x contains the N dosages.

    ## Input data as a list
df = list(
  y = c(9.1, 5.5, 12.3, 9.2, 14.2, 16.8, 22, 18.3, 
        24.5, 22.7), 
  x = c(3, 3, 4, 5, 6, 6, 7, 8, 8, 9), 
  N = 10
)
    
################## Classical Reg results ###################

freq.model <- lm(y~x,df)
summary(freq.model)
confint(freq.model)
  ## Variance estimate
(summary(freq.model)$sigma)^2
  
with(df,plot(x,y,xlab="dose",ylab="hours",main="Hours of relief by medication dose"))
abline(freq.model$coefficients)

################# Or input data from data set from a data frame (convert to list)
mydata =data.frame(cbind(df$y,df$x))
names(mydata) = c("y","x")    
    ## i.e., starting here
df2 = list(  'y'  = with(mydata, y),
            'x' = with(mydata, x),
            'N'  = nrow(mydata))

as.list(as.data.frame(mydata))

## You can fit model using either df or df2

############################# Fitting model with JAGS ############################################

################# Inital values
inits1 = list( beta0 = 0, beta1 = 0, tau = 1)                                                     # No seed
#inits1 = list( beta0 = 0, beta1 = 0, tau = 1, .RNG.seed=1, .RNG.name='base::Wichmann-Hill')      # Sets seed

################# Conjugate model ############################
# n.chains = number of chains to monitor (with different starting values)
# n.adapt = initial sample phase to optimize the sampling mechanism (e.g. Metropolis Hastings step size)
# inits = initial values from above
# "allergy-bugConjugate.txt" is the filename which contains the model

jags.Conj <- jags.model("allergy-bugConjugate.txt",data=df,n.chains=1, n.adapt=1000, inits=inits1)
  ## burn-in
update( jags.Conj , n.iter=2000 )
  ## Store MCMC samples after burn-in
  ## Arugments are the parameters to be monitored and the number of samples to take
mcmc.samples.Conj <- coda.samples(jags.Conj,c('beta0','beta1','sig'),10000)
summary(mcmc.samples.Conj)
plot(mcmc.samples.Conj)

    ### Calculate probability that the slope is greater than 1
post.sam = as.matrix(mcmc.samples.Conj[1])
mean(post.sam[,'beta1'] > 0 )

################### Non-conjugate model ###################################

jags <- jags.model("allergy-bug.txt",data=df,n.chains=1,n.adapt=1000, inits=inits1)
  ## burn-in
update(jags, n.iter=2000 )
mcmc.samples <- coda.samples(jags,c('beta0','beta1','sig'),10000)
summary(mcmc.samples)
#png('PostGraphs.png')
plot(mcmc.samples)
#dev.off()


######################### 3 chains & convergence diagnostics, Non-conjugate model ############

inits2 = list( beta0 = 0.087, beta1 = -3.914, tau = 1.976 )
inits3 = list( beta0 = 0.568, beta1 = 2.440, tau = 0.841 )

jags3 <- jags.model("allergy-bug.txt",data=df,n.chains=3,n.adapt=1000, inits=list(inits1,inits2,inits3))
  ## burn-in
update(jags3, n.iter=2000 )
  ## Samples
mcmc.samples3 <- coda.samples(jags3,c('beta0','beta1','sig'),10000)
summary(mcmc.samples3)
plot(mcmc.samples3)
gelman.plot(mcmc.samples3)
gelman.diag(mcmc.samples3)
HPDinterval(mcmc.samples3)

############ Two Parameters, Non-conjugate model ###################################
    ## Create a txt file for use by JAGS
fileForJAGS <- "allergy-bug2.txt"
cat("
# This file is created dynamically by the main knitr file for this project.
# Make any changes there, not here!
model {
for(i in 1:N)  {
    y[i] ~ dnorm(mu[i], tau);
    mu[i] <- beta0 + beta1*x[i] + beta2*tmt[i];
    }

  beta0 ~ dnorm(0, 0.001);
  beta1 ~ dnorm(0, 0.001);
  beta2 ~ dnorm(0, 0.001);
  tau ~ dgamma(0.01, 0.01);
  sig <- 1/tau;
}
", file= fileForJAGS)

    ## Additional variable: Treatment status (binary)
df$tmt <- rep(c(1,0),5)
freq.model <- lm(y~x + tmt,df)
summary(freq.model)
confint(freq.model)
with(df,plot(x,y,xlab="dose",ylab="hours", ylim=c(0,25),
      main="Hours of relief by medication dose & tmt", col=c("red","blue")[df$tmt+1]))
legend("topleft",c("Tmt","PBO"), pch=c(1,1), col=c("blue","red"))



inits2param = list( beta0 = 0, beta1 = 0, tau = 1)   
jags <- jags.model("allergy-bug2.txt",data=df,n.chains=1,n.adapt=1000, inits=inits2param)
  ## burn-in
update(jags, n.iter=2000 )
mcmc.samples <- coda.samples(jags,c('beta0','beta1','beta2','sig'),10000)
summary(mcmc.samples)
plot(mcmc.samples)
as.matrix(mcmc.samples)
mat.mcmc <- as.matrix(mcmc.samples)
mean(mat.mcmc[,'beta2'])  ## Mean Treatment effect
mean(mat.mcmc[,'beta2'] < 0 )  ### Probability treatment is superior to PBO