Basic support for simulating random effects

2 files changed, 86 insertions, 25 deletions

diff --git a/R/power.R b/R/power.R
index 80af473..ad1558a 100644
--- a/R/power.R
+++ b/R/power.R
@@ -151,16 +151,38 @@ bootstrap_data <- function(data, nsims, design){
   lapply(1:nsims, function(i) data[M_samps[[i]], J_samps[[i]]])
 }
 
-check_coefs <- function(coefs, fit, template=FALSE){
+check_coefs <- function(coefs, fit, formulas, template=FALSE){
   required_subs <- names(fit@estimates@estimates)
   required_coefs <- lapply(fit@estimates@estimates, function(x) names(x@estimates))
   required_lens <- lapply(required_coefs, length)
+
+  # If there are random effects, adjust the expected coefficient names
+  # to remove the b vector and add the grouping covariate name
+  rand <- lapply(formulas, lme4::findbars)
+  if(!all(sapply(rand, is.null))){
+    stopifnot(all(required_subs %in% names(formulas)))
+    rvar <- lapply(rand, function(x) unlist(lapply(x, all.vars)))
+    if(!all(sapply(rvar, length)<2)){
+      stop("Only 1 random effect per parameter is supported", call.=FALSE)
+    }
+    for (i in required_subs){
+      if(!is.null(rand[[i]][[1]])){
+        signame <- rvar[[i]]
+        old_coefs <- required_coefs[[i]]
+        new_coefs <- old_coefs[!grepl("b_", old_coefs, fixed=TRUE)]
+        new_coefs <- c(new_coefs, signame)
+        required_coefs[[i]] <- new_coefs
+      }
+    }
+  }
+
   dummy_coefs <- lapply(required_coefs, function(x){
                     out <- rep(0, length(x))
                     x <- gsub("(Intercept)", "intercept", x, fixed=TRUE)
                     names(out) <- x
                     out
                   })
+
   if(template) return(dummy_coefs)
 
   if(is.null(coefs)){
@@ -226,6 +248,7 @@ check_coefs <- function(coefs, fit, template=FALSE){
   }
   coefs[required_subs]
 }
+
 wald <- function(est, se, null_hyp=NULL){
   if(is.null(null_hyp) || is.na(null_hyp)) null_hyp <- 0
   Z <- (est-null_hyp)/se
diff --git a/R/simulate.R b/R/simulate.R
index 3503ba5..7868948 100644
--- a/R/simulate.R
+++ b/R/simulate.R
@@ -59,6 +59,68 @@ blank_umFit <- function(fit_function){
 }
 
 
+setMethod("simulate", "character",
+  function(object, nsim=1, seed=NULL, formulas, coefs=NULL, design, guide=NULL, ...){
+  model <- blank_umFit(object)
+  fit <- suppressWarnings(simulate_fit(model, formulas, guide, design, ...))
+  coefs <- check_coefs(coefs, fit, formulas)
+  coefs <- generate_random_effects(coefs, fit, formulas)
+  fit <- replace_estimates(fit, coefs)
+  ysims <- simulate(fit, nsim)
+  umf <- fit@data
+  # fix this
+  umfs <- lapply(ysims, function(x){
+    if(object=="occuMulti"){
+      umf@ylist <- x
+    } else if(object=="gdistremoval"){
+      umf@yDistance=x$yDistance
+      umf@yRemoval=x$yRemoval
+    } else {
+      umf@y <- x
+    }
+    umf
+  })
+  if(length(umfs)==1) umfs <- umfs[[1]]
+  umfs
+})
+
+
+generate_random_effects <- function(coefs, fit, formulas){
+  required_subs <- names(fit@estimates@estimates)
+  rand <- lapply(formulas, lme4::findbars)
+  if(!all(sapply(rand, is.null))){
+    rvar <- lapply(rand, function(x) unlist(lapply(x, all.vars)))
+    for (i in required_subs){
+      if(!is.null(rand[[i]][[1]])){
+        signame <- rvar[[i]]
+        old_coefs <- coefs[[i]]
+        new_coefs <- old_coefs[names(old_coefs)!=signame]
+
+        # Find levels of factor variable
+        if(signame %in% names(siteCovs(fit@data))){
+          lvldata <- siteCovs(fit@data)[[signame]]
+        } else if(signame %in% names(obsCovs(fit@data))){
+          lvldata <- obsCovs(fit@data)[[signame]]
+        } else if(methods::.hasSlot(fit@data, "yearlySiteCovs") && signame %in% names(yearlySiteCovs(fit@data))){
+          lvldata <- yearlySiteCovs(fit@data)[[signame]]
+        } else {
+          stop("Random effect covariate missing from data", call.=FALSE)
+        }
+
+        if(!is.factor(lvldata)){
+          stop("Random effect covariates must be specified as factors with guide argument", call.=FALSe)
+        }
+        b <- rnorm(length(levels(lvldata)), 0, old_coefs[signame])
+        names(b) <- rep(paste0("b_",i), length(b))
+        new_coefs <- c(new_coefs, b)
+        coefs[[i]] <- new_coefs
+      }
+    }
+  }
+  coefs
+}
+
+
 setGeneric("get_umf_components", function(object, ...) standardGeneric("get_umf_components"))
 
 setMethod("get_umf_components", "unmarkedFit",
@@ -103,30 +165,6 @@ setMethod("simulate_fit", "unmarkedFitOccuRN",
 })
 
 
-setMethod("simulate", "character",
-  function(object, nsim=1, seed=NULL, formulas, coefs=NULL, design, guide=NULL, ...){
-  model <- blank_umFit(object)
-  fit <- suppressWarnings(simulate_fit(model, formulas, guide, design, ...))
-  coefs <- check_coefs(coefs, fit)
-  fit <- replace_estimates(fit, coefs)
-  ysims <- simulate(fit, nsim)
-  umf <- fit@data
-  # fix this
-  umfs <- lapply(ysims, function(x){
-    if(object=="occuMulti"){
-      umf@ylist <- x
-    } else if(object=="gdistremoval"){
-      umf@yDistance=x$yDistance
-      umf@yRemoval=x$yRemoval
-    } else {
-      umf@y <- x
-    }
-    umf
-  })
-  if(length(umfs)==1) umfs <- umfs[[1]]
-  umfs
-})
-
 setMethod("get_umf_components", "unmarkedFitMPois",
           function(object, formulas, guide, design, ...){
   args <- list(...)