require(lme4)
require(coda)
if (!interactive())
    pdf(onefile = FALSE)                # plots are much smaller files than for postscript

sessionInfo()

## Illustrate problem with mcmcsamp for two variance components mixed-effects model.

make.data = function(N.batch)
{
  ## Create example data set
  ## N.batch Batches, 15 units per batch = (5 Trt A, 5 Trt B, 5 Trt C)


  ## Trt means, Batch SD, and residual SD
  theta = c(5, 10, 15)
  sd.batch = .5
  sd.resid = 1.5

  d = expand.grid( Batch=factor(1:N.batch), Trt=factor(c("A", "B", "C")), Units=1:5 )
  d = d[order(d$Batch, d$Trt),][,-3]
  d$y = theta[ unclass(d$Trt) ] + rep( rnorm( N.batch, sd=sd.batch ), each=15 ) + rnorm( nrow(d), sd=sd.resid )
  
  d
}


## Create the response variable
set.seed(991)
d = make.data(N.batch=4)

## Plot the data
xyplot( y~Batch|Trt, data=d, layout=c(3, 1))

## Fit linear mixed effects model
fit = lmer2( y ~ Trt + (1|Batch), data=d )

## MCMC sampling from the posterio distribution
M = 25000
samp = mcmcsamp(fit, n=M)

xyplot(samp, scales = list(x = list(axs = 'i')))    ## Uh oh...What's going on with the sd.batch ?
print(summary(samp))
print(autocorr.diag(samp))  ## Uh oh.  Very high autocorrelation for sd.batch!

## Four chains:  MCMC sampling from the posterior distribution
M = 25000
samp = lapply(1:4, function(i){ mcmcsamp( fit, n=M ) })

all.samp = data.frame( Chain=rep(1:4, each=M), Sample=rep(1:M, 4),
              log.sd.resid=as.vector(sapply(samp, function(sa){ sa[,4] })),
              log.sd.batch=as.vector(sapply(samp, function(sa){ sa[,5] })) )

## Plot the chains for the two variance components
xyplot(log.sd.resid+log.sd.batch~Sample, groups=Chain, data=all.samp, outer=TRUE,
       layout = c(1,2), type="l", strip = FALSE, strip.left = TRUE,
       scales = list(y = list(relation="free")))

if (require(BRugs)) {
    ## No problems using BRugs
    ## OpenBugs Model
    model <- function()
    {
        for ( i in 1:N.batch )
        {
            m[i] ~ dnorm(0, tau.batch)
            for ( j in 1:N.unit )
            {
                m.unit[ N.unit*(i-1) +j ] <- m[i] + theta[ Trt[N.unit*(i-1) +j] ]
                y [ N.unit*(i-1) +j ] ~ dnorm( m.unit[ N.unit*(i-1) +j ], tau.err )
            }
        }
        sigma.batch <- 1/sqrt(tau.batch)
        sigma.err <- 1/sqrt(tau.err)
        log.sigma.batch <- log( sigma.batch )
        
        tau.batch ~ dgamma(0.001, 0.001)
        tau.err ~ dgamma(0.001, 0.001)
        for ( k in 1:3 )
        {
            theta[k] ~ dnorm( 0.0, 0.000001 )
        }
    }
    setwd("C:/")
    writeModel(model, "model.txt")
    data = list( y=d$y, Trt=as.vector( unclass(d$Trt) ), N.batch=4, N.unit=15 )
    inits = list( list( theta=c( 6, 8, 10 ), tau.batch=1, tau.err=1, m=rep(0, 4) ),
    list( theta=c( 7, 5, 12 ), tau.batch=.8, tau.err=1.2, m=rep(0, 4) ),
    list( theta=c( 15, 6, 15 ), tau.batch=1.2, tau.err=.8, m=rep(0, 4) ) )
    
    
    BRugsFit("model.txt", data=data, inits=inits, numChains=3,
             parametersToSave=c("theta", "sigma.batch", "sigma.err", "log.sigma.batch"),
             nBurnin=5000, nIter=25000)
    samplesStats("*")
    plotHistory( "log.sigma.batch" )
    plotAutoC( "log.sigma.batch" )
}


################################################
## Repeat the exercise, but with 10 batches   ##
################################################
## Still doesn't look good

## Create the response variable
set.seed(991)
d = make.data(N.batch=10)

fit = lmer2( y ~ Trt + (1|Batch), data=d )
M = 25000
samp = mcmcsamp(fit, n=M)

xyplot(samp, scales = list(x = list(axs = 'i')))
print(summary(samp))
print(autocorr.diag(samp))

samp = lapply(1:4, function(i){ mcmcsamp( fit, n=M ) })

all.samp <-
    data.frame(Chain = rep(1:4, each=M), Sample=rep(1:M, 4),
               log.sd.resid = as.vector(sapply(samp, function(sa){ sa[,4] })),
               log.sd.batch=as.vector(sapply(samp, function(sa){ sa[,5] })))

## Plot the chains for the two variance components
xyplot(log.sd.resid+log.sd.batch~Sample, groups=Chain, data=all.samp, outer=TRUE,
       layout = c(1,2), type="l", strip = FALSE, strip.left = TRUE,
       scales = list(y = list(relation="free")))


################################################
## Repeat the exercise, but with 25 batches   ##
################################################
## This one looks good

## Create the response variable
set.seed(991)
d = make.data(N.batch=25)

fit = lmer2( y ~ Trt + (1|Batch), data=d )
M = 25000
samp = mcmcsamp(fit, n=M)

xyplot(samp, scales = list(x = list(axs = 'i')))
print(summary(samp))
print(autocorr.diag(samp))

samp = lapply(1:4, function(i){ mcmcsamp( fit, n=M ) })

all.samp = data.frame( Chain=rep(1:4, each=M), Sample=rep(1:M, 4),
              log.sd.resid=as.vector(sapply(samp, function(sa){ sa[,4] })),
              log.sd.batch=as.vector(sapply(samp, function(sa){ sa[,5] })) )

## Plot the chains for the two variance components
xyplot(log.sd.resid+log.sd.batch~Sample, groups=Chain, data=all.samp, outer=TRUE,
       layout = c(1,2), type="l", strip = FALSE, strip.left = TRUE,
       scales = list(y = list(relation="free")))