Edits to rjags interface code to make it more readable

1 files changed, 180 insertions, 158 deletions

diff --git a/R/run_rjags.R b/R/run_rjags.R
index 02a0ec9..2cbe51a 100644
--- a/R/run_rjags.R
+++ b/R/run_rjags.R
@@ -36,196 +36,218 @@ run_rjags <- function(data, inits, params, modfile, mcmc_info,
 }
 
 # Setup and run rjags----------------------------------------------------------
-run.model <- function(model.file=NULL,data=NULL,inits=NULL,parameters.to.save,n.chains=NULL,
-                      n.iter,n.burnin,n.thin,n.adapt,verbose=TRUE,model.object=NULL,
-                      update=FALSE,parallel=FALSE,na.rm=TRUE){
-
-if(verbose){pb="text"} else {pb="none"}
-
-if(update){
-  #Recompile model
-  m <- model.object
-  if(verbose | parallel==TRUE){
-  m$recompile()
-  } else {null <- capture.output(
-          m$recompile()
-          )}
-
-} else {
-  #Compile model
-  if(verbose | parallel==TRUE){
-    m <- jags.model(file=model.file,data=data,inits=inits,n.chains=n.chains,n.adapt=0)
+run.model <- function(model.file=NULL, data=NULL, inits=NULL, parameters.to.save,
+                      n.chains=NULL, n.iter, n.burnin, n.thin, n.adapt,
+                      verbose=TRUE, model.object=NULL, update=FALSE,
+                      parallel=FALSE, na.rm=TRUE){
+  
+  if(verbose){
+    pb <- "text"
   } else {
-    null <- capture.output(
-    m <- jags.model(file=model.file,data=data,inits=inits,n.chains=n.chains,n.adapt=0,quiet=TRUE)
-    )
+    pb <- "none"
   }
-}
 
-#Adaptive phase using adapt()
-total.adapt <- 0
+  if(update){
+    #Recompile model
+    m <- model.object
+    if(verbose | parallel==TRUE){
+      m$recompile()
+    } else {
+      null <- capture.output(m$recompile())
+    }
 
-if(!is.null(n.adapt)){
-  if(n.adapt>0){
-    if(verbose){
-      cat('Adaptive phase,',n.adapt,'iterations x',n.chains,'chains','\n')
-      cat('If no progress bar appears JAGS has decided not to adapt','\n','\n')
-      sufficient.adapt <- adapt(object=m,n.iter=n.adapt,progress.bar=pb,end.adaptation=TRUE)
+  } else {
+    #Compile model
+    if(verbose | parallel==TRUE){
+      m <- jags.model(file=model.file,data=data,inits=inits,n.chains=n.chains,n.adapt=0)
     } else {
       null <- capture.output(
-      sufficient.adapt <- adapt(object=m,n.iter=n.adapt,progress.bar=pb,end.adaptation=TRUE)
-      )}
-    total.adapt <- total.adapt + n.adapt
-  } else{
-    if(verbose){cat('No adaptive period specified','\n','\n')}
-    #If no adaptation period specified:
-    #Force JAGS to not adapt (you have to allow it to adapt at least 1 iteration)
-    if(!update){
+      m <- jags.model(file=model.file,data=data,inits=inits,n.chains=n.chains,n.adapt=0,quiet=TRUE)
+      )
+    }
+  }
+
+  #Adaptive phase using adapt()
+  total.adapt <- 0
+
+  if(!is.null(n.adapt)){
+    if(n.adapt>0){
       if(verbose){
-        sufficient.adapt <- adapt(object=m,n.iter=1,end.adaptation=TRUE)
+        cat('Adaptive phase,',n.adapt,'iterations x',n.chains,'chains','\n')
+        cat('If no progress bar appears JAGS has decided not to adapt','\n','\n')
+        sufficient.adapt <- adapt(object=m, n.iter=n.adapt, progress.bar=pb, end.adaptation=TRUE)
       } else {
         null <- capture.output(
-          sufficient.adapt <- adapt(object=m,n.iter=1,end.adaptation=TRUE)
+        sufficient.adapt <- adapt(object=m, n.iter=n.adapt, progress.bar=pb, end.adaptation=TRUE)
         )}
+      total.adapt <- total.adapt + n.adapt
+    } else{
+      if(verbose){cat('No adaptive period specified','\n','\n')}
+      #If no adaptation period specified:
+      #Force JAGS to not adapt (you have to allow it to adapt at least 1 iteration)
+      if(!update){
+        if(verbose){
+          sufficient.adapt <- adapt(object=m, n.iter=1, end.adaptation=TRUE)
+        } else {
+          null <- capture.output(
+            sufficient.adapt <- adapt(object=m, n.iter=1, end.adaptation=TRUE)
+          )}
+      }
+      total.adapt <- 0
     }
-    total.adapt <- 0
-  }
-} else {
-
-  maxloops <- 100
-  n.adapt.iter <- 100
-
-  for (i in 1:maxloops){
-    if(verbose){cat('Adaptive phase.....','\n')}
-    sufficient.adapt <- adapt(object=m,n.iter=n.adapt.iter,progress.bar='none')
-    total.adapt <- total.adapt + n.adapt.iter
-    if(i==maxloops){
-      if(verbose){warning(paste("Reached max of",maxloops*n.adapt.iter,"adaption iterations; set n.adapt to > 10000"))}
-      null <- adapt(object=m,n.iter=1,end.adaptation = TRUE)
-      break
-    }
-    if(sufficient.adapt){
-      null <- adapt(object=m,n.iter=1,end.adaptation = TRUE)
-      if(verbose){cat('Adaptive phase complete','\n','\n')}
-      break
+  } else {
+
+    maxloops <- 100
+    n.adapt.iter <- 100
+
+    for (i in 1:maxloops){
+      if(verbose) cat('Adaptive phase.....','\n')
+      sufficient.adapt <- adapt(object=m, n.iter=n.adapt.iter, progress.bar='none')
+      total.adapt <- total.adapt + n.adapt.iter
+      if(i==maxloops){
+        if(verbose){
+          warning(paste("Reached max of",maxloops*n.adapt.iter,"adaption iterations; set n.adapt to > 10000"))
+        }
+        null <- adapt(object=m,n.iter=1,end.adaptation = TRUE)
+        break
+      }
+      if(sufficient.adapt){
+        null <- adapt(object=m,n.iter=1,end.adaptation = TRUE)
+        if(verbose) cat('Adaptive phase complete','\n','\n')
+        break
+      }
     }
   }
 
-}
-if(!sufficient.adapt&total.adapt!=0&verbose){warning("JAGS reports adaptation was incomplete. Consider increasing n.adapt")}
+  if(!sufficient.adapt & total.adapt!=0){
+    warning("JAGS reports adaptation was incomplete. Consider increasing n.adapt")
+  }
+
+  #Burn-in phase using update()
+  if(n.burnin>0){
+    if(verbose){
+      cat('\n','Burn-in phase,',n.burnin,'iterations x',n.chains,'chains','\n','\n')
+      update(object=m,n.iter=n.burnin,progress.bar=pb)
+      cat('\n')
+    } else {
+      null <- capture.output(
+      update(object=m,n.iter=n.burnin,progress.bar=pb)
+    )}
+  } else if(verbose){
+    cat('No burn-in specified','\n','\n')
+  }
 
-#Burn-in phase using update()
-if(n.burnin>0){
+  #Sample from posterior using coda.samples()
   if(verbose){
-    cat('\n','Burn-in phase,',n.burnin,'iterations x',n.chains,'chains','\n','\n')
-    update(object=m,n.iter=n.burnin,progress.bar=pb)
+    cat('Sampling from joint posterior,',(n.iter-n.burnin),
+        'iterations x',n.chains,'chains','\n','\n')
+    samples <- coda.samples(model=m, variable.names=parameters.to.save,
+                            n.iter=(n.iter-n.burnin), thin=n.thin,
+                            na.rm=na.rm, progress.bar=pb)
     cat('\n')
   } else {
     null <- capture.output(
-    update(object=m,n.iter=n.burnin,progress.bar=pb)
-  )}
-} else if(verbose){cat('No burn-in specified','\n','\n')}
-
-#Sample from posterior using coda.samples()
-if(verbose){
-  cat('Sampling from joint posterior,',(n.iter-n.burnin),'iterations x',n.chains,'chains','\n','\n')
-  samples <- coda.samples(model=m,variable.names=parameters.to.save,n.iter=(n.iter-n.burnin),thin=n.thin,
-                          na.rm=na.rm, progress.bar=pb)
-  cat('\n')
-} else {
-  null <- capture.output(
-  samples <- coda.samples(model=m,variable.names=parameters.to.save,n.iter=(n.iter-n.burnin),thin=n.thin,
-                          na.rm=na.rm, progress.bar=pb)
-  )}
-
-return(list(m=m,samples=samples,total.adapt=total.adapt,sufficient.adapt=sufficient.adapt))
+    samples <- coda.samples(model=m, variable.names=parameters.to.save,
+                            n.iter=(n.iter-n.burnin), thin=n.thin,
+                            na.rm=na.rm, progress.bar=pb)
+    )
+  }
+
+  return(list(m=m,samples=samples,
+              total.adapt=total.adapt,
+              sufficient.adapt=sufficient.adapt))
 }
 
 # Setup parallel and run rjags-------------------------------------------------
-run.parallel <- function(data=NULL,inits=NULL,parameters.to.save,model.file=NULL,n.chains,n.adapt,n.iter,n.burnin,n.thin,
-                         modules,factories,DIC,model.object=NULL,update=FALSE,verbose=TRUE,n.cores=NULL) {
-
-#Save current library paths
-current.libpaths <- .libPaths()
-
-#Set up clusters
-cl = makeCluster(n.cores)
-on.exit(stopCluster(cl))
-clusterExport(cl = cl, ls(), envir = environment())
-clusterEvalQ(cl,.libPaths(current.libpaths))
-
-if(verbose){
-cat('Beginning parallel processing using',n.cores,'cores. Console output will be suppressed.\n')}
+run.parallel <- function(data=NULL, inits=NULL, parameters.to.save, model.file=NULL,
+                         n.chains, n.adapt, n.iter, n.burnin, n.thin,
+                         modules, factories, DIC,
+                         model.object=NULL, update=FALSE, 
+                         verbose=TRUE, n.cores=NULL) {
 
-#Function called in each core
-jags.clust <- function(i){
+  #Save current library paths
+  current.libpaths <- .libPaths()
 
-#Load modules
-set.modules(modules,DIC)
-set.factories(factories)
+  #Set up clusters
+  cl <- makeCluster(n.cores)
+  on.exit(stopCluster(cl))
+  clusterExport(cl = cl, ls(), envir = environment())
+  clusterEvalQ(cl,.libPaths(current.libpaths))
 
-if(update){
-  #Recompile model
-  cluster.mod <- model.object[[i]]
-
-  #Run model
-  rjags.output <- run.model(model.file=NULL,data=NULL,inits=NULL,parameters.to.save,n.chains=1,n.iter,n.burnin=0,n.thin,n.adapt,
-                            verbose=FALSE,model.object=cluster.mod,update=TRUE,parallel=TRUE, na.rm=FALSE)
-
-} else {
-
-  #Set initial values for cluster
-  cluster.inits <- inits[[i]]
-
-  #Run model
+  if(verbose){
+    cat('Beginning parallel processing using',n.cores,
+        'cores. Console output will be suppressed.\n')
+  }
 
-  rjags.output <- run.model(model.file,data,inits=cluster.inits,parameters.to.save,n.chains=1,n.iter,
-                  n.burnin,n.thin,n.adapt,verbose=FALSE,parallel=TRUE, na.rm=FALSE)
+  #Function called in each core
+  jags.clust <- function(i){
 
-}
+    #Load modules
+    set.modules(modules,DIC)
+    set.factories(factories)
 
-return(list(samp=rjags.output$samples[[1]],mod=rjags.output$m,total.adapt=rjags.output$total.adapt,sufficient.adapt=rjags.output$sufficient.adapt))
+    if(update){
+      #Recompile model
+      cluster.mod <- model.object[[i]]
+      #Run model
+      rjags.output <- run.model(model.file=NULL,data=NULL,inits=NULL,parameters.to.save,
+                                n.chains=1,n.iter,n.burnin=0,n.thin,n.adapt,
+                                verbose=FALSE,model.object=cluster.mod,update=TRUE,
+                                parallel=TRUE, na.rm=FALSE)
+    } else {
+      #Set initial values for cluster
+      cluster.inits <- inits[[i]]
+      #Run model
+      rjags.output <- run.model(model.file,data,inits=cluster.inits,parameters.to.save,
+                                n.chains=1,n.iter,n.burnin,n.thin,n.adapt,verbose=FALSE,
+                                parallel=TRUE, na.rm=FALSE)
+    }
 
-}
+    return(list(samp=rjags.output$samples[[1]],
+                mod=rjags.output$m,
+                total.adapt=rjags.output$total.adapt,
+                sufficient.adapt=rjags.output$sufficient.adapt))
+    }
 
-#Do analysis
-par <- clusterApply(cl=cl,x=1:n.chains,fun=jags.clust)
-
-#Create empty lists
-out <- samples <- model <- list()
-total.adapt <- sufficient.adapt <- vector(length=n.chains)
-
-#Save samples and model objects from each cluster
-total.adapt <- sapply(par, function(x) x[[3]])
-starts <- sapply(par, function(x) stats::start(x$samp))
-ends <- sapply(par, function(x) stats::end(x$samp))
-nsamp <- sapply(par, function(x) coda::niter(x$samp))
-thins <- sapply(par, function(x) coda::thin(x$samp))
-#mc_start <- max(total.adapt) + n.burnin + n.thin
-stopifnot(all(nsamp == nsamp[1]))
-stopifnot(all(thins == n.thin))
-for (i in 1:n.chains){
-  samples[[i]] <- coda::mcmc(par[[i]]$samp,start=max(starts),thin=n.thin)
-  model[[i]] <- par[[i]]$m
-  sufficient.adapt[i] <- par[[i]]$sufficient.adapt
-}
-out$samples <- as.mcmc.list(samples)
-# Remove columns with all NA
-try({
-  all_na <- apply(as.matrix(out$samples),2, function(x) all(is.na(x)))
-  out$samples <- out$samples[,!all_na,drop=FALSE]
-})
-out$model <- model
-out$total.adapt <- total.adapt
-out$sufficient.adapt <- sufficient.adapt
-names(out$model) <- sapply(1:length(out$model),function(i){paste('cluster',i,sep="")})
-
-if(verbose){
-cat('\nParallel processing completed.\n\n')
-}
+  #Do parallel analysis
+  par <- clusterApply(cl=cl, x=1:n.chains, fun=jags.clust)
+
+  #Create empty lists
+  out <- samples <- model <- list()
+  total.adapt <- sufficient.adapt <- vector(length=n.chains)
+
+  #Save samples from each cluster into mcmc.list
+  total.adapt <- sapply(par, function(x) x[[3]])
+  starts <- sapply(par, function(x) stats::start(x$samp))
+  ends <- sapply(par, function(x) stats::end(x$samp))
+  nsamp <- sapply(par, function(x) coda::niter(x$samp))
+  thins <- sapply(par, function(x) coda::thin(x$samp))
+  stopifnot(all(nsamp == nsamp[1]))
+  stopifnot(all(thins == n.thin))
+  for (i in 1:n.chains){
+    samples[[i]] <- coda::mcmc(par[[i]]$samp,start=max(starts),thin=n.thin)
+    model[[i]] <- par[[i]]$m
+    sufficient.adapt[i] <- par[[i]]$sufficient.adapt
+  }
+  out$samples <- as.mcmc.list(samples)
+  
+  # Remove columns with all NA
+  try({
+    all_na <- apply(as.matrix(out$samples),2, function(x) all(is.na(x)))
+    out$samples <- out$samples[,!all_na,drop=FALSE]
+  })
+  
+  # Save other info in output object
+  out$model <- model
+  out$total.adapt <- total.adapt
+  out$sufficient.adapt <- sufficient.adapt
+  names(out$model) <- sapply(1:length(out$model),function(i){paste('cluster',i,sep="")})
 
-return(out)
+  if(verbose){
+    cat('\nParallel processing completed.\n\n')
+  }
 
+  return(out)
 }