From 64bc5a9ab57e458586b5fdfeb4049d9f3601c239 Mon Sep 17 00:00:00 2001
From: Ken Kellner <ken@kenkellner.com>
Date: Sun, 29 Mar 2020 17:44:53 -0400
Subject: [PATCH] Return all samples instead of summarizing

---
 R/ranef.R                    | 27 +++++++++++++++++++++------
 inst/unitTests/runit.ranef.R | 15 +++++++--------
 man/predict-methods.Rd       | 12 ++++++++----
 man/ranef-methods.Rd         | 28 ++++++++++++++++++----------
 4 files changed, 54 insertions(+), 28 deletions(-)

diff --git a/R/ranef.R b/R/ranef.R
index 81512af6..e7963f42 100644
--- a/R/ranef.R
+++ b/R/ranef.R
@@ -1135,13 +1135,28 @@ setMethod("plot", c("unmarkedRanef", "missing"), function(x, y, ...)
 setMethod("predict", "unmarkedRanef", function(object, func, nsims=100, ...)
 {
 
-  ps <- posteriorSamples(object, nsims)@samples
+  ps <- posteriorSamples(object, nsims=nsims)@samples
+  s1 <- func(ps[,,1])
+  nm <- names(s1)
+  row_nm <- rownames(s1)
+  col_nm <- colnames(s1)
+
+  if(is.vector(s1)){
+    out_dim <- c(length(s1), nsims)
+  } else{
+    out_dim <- c(dim(s1), nsims)
+  }
+
   param <- apply(ps, 3, func)
 
-  pr <- rowMeans(param, na.rm=TRUE)
-  se <- sqrt(apply(param, 1, function(x) stats::var(x, na.rm=TRUE)))
-  lower <- apply(param, 1, function(x) stats::quantile(x, 0.025, na.rm=TRUE))
-  upper <- apply(param, 1, function(x) stats::quantile(x, 0.975, na.rm=TRUE))
+  out <- array(param, out_dim)
   
-  data.frame(Predicted=pr, SE=se, lower=lower, upper=upper)
+  if(is.vector(s1)){
+    rownames(out) <- nm
+  } else {
+    rownames(out) <- row_nm
+    colnames(out) <- col_nm
+  }
+
+  drop(out)
 })
diff --git a/inst/unitTests/runit.ranef.R b/inst/unitTests/runit.ranef.R
index 09b0b6c9..65f32d8b 100644
--- a/inst/unitTests/runit.ranef.R
+++ b/inst/unitTests/runit.ranef.R
@@ -565,14 +565,14 @@ test.ranef.predict <- function(){
   checkEqualsNumeric(dim(ps@samples), c(9,1,10))
 
   myfunc <- function(x){
-    c(mean(x[1:4]), mean(x[5:9]))
+    c(gr1=mean(x[1:4]), gr2=mean(x[5:9]))
   }
 
   pr <- predict(re, fun=myfunc, nsim=10)
-  checkEqualsNumeric(nrow(pr), 2)
-  checkEqualsNumeric(as.numeric(pr[1,]), c(5.275,0.7115,4.1125,6.1937),
-                     tol=1e-4)
-
+  checkEqualsNumeric(dim(pr), c(2,10))
+  checkEquals(rownames(pr), c("gr1","gr2"))
+  checkEqualsNumeric(as.numeric(pr[1,1:3]), c(6.0,4.0,5.75))
+  
   #Dynamic model
   set.seed(7)
   M <- 10
@@ -611,9 +611,8 @@ test.ranef.predict <- function(){
   }
 
   pr <- predict(re1, fun=myfunc, nsim=10)
-  checkEqualsNumeric(nrow(pr), 2*T)
-  checkEqualsNumeric(as.numeric(pr[1,]), c(3.525,0.2189,3.2500,3.7500),
-                     tol=1e-4)
+  checkEqualsNumeric(dim(pr), c(2,5,10))
+  checkEqualsNumeric(pr[1,1:3,1], c(3.5,2.5,1.5))
 
 
 }
diff --git a/man/predict-methods.Rd b/man/predict-methods.Rd
index e757b192..fcc3b21a 100644
--- a/man/predict-methods.Rd
+++ b/man/predict-methods.Rd
@@ -45,13 +45,17 @@ with dimensions M x T, where M is the number of sites and T is the number of
 primary periods (T=1 for single-season models). The output of this function should 
 be a vector or matrix containing the derived parameters of interest. 
 
-The output of predict is the mean value of each derived parameter, in 
-column-major order, along with the standard error and 95\% confidence interval.
-See \code{\link{ranef}} for an example.
-
 You may also manually set the number of draws from the posterior predictive
 distribution with argument \code{nsims}; the default is 100.
 
+The output of \code{predict} will be a vector or array with one more dimension
+than the output of the function supplied \code{func}, corresponding to the number
+of draws requested \code{nsims}. For example, if \code{func}
+outputs a scalar, the output of \code{predict} will be a vector with length
+equal to \code{nsims}. If \code{func} outputs a 3x2 matrix, the output of
+\code{predict} will be an array with dimensions 3x2x\code{nsims}.
+See \code{\link{ranef}} for an example.
+
 Alternatively, you can use the \code{\link{posteriorSamples}} function on the
 \code{ranef} output object to obtain the full posterior predictive distribution. 
 This is useful if you are having trouble designing your custom function or if 
diff --git a/man/ranef-methods.Rd b/man/ranef-methods.Rd
index 2d2548bb..1cb62184 100644
--- a/man/ranef-methods.Rd
+++ b/man/ranef-methods.Rd
@@ -152,23 +152,31 @@ post.arr <- as(re, "array")
 #Generate posterior predictive distribution for a function
 #of random variables using predict()
 
-#First, create a function that operates on a MxT matrix where
-#M = nsites and T = n primary periods (1 in this case)
+#First, create a function that operates on a vector of 
+#length M (if you fit a single-season model) or a matrix of 
+#dimensions MxT (if a dynamic model), where
+#M = nsites and T = n primary periods
 #Our function will generate mean abundance for sites 1-10 and sites 11-20
-myfunc <- function(x){ #x is 20 x 1
+myfunc <- function(x){ #x will be length 20 since M=20
   
   #Mean of first 10 sites
-  group1 <- mean(x[1:10,1])
+  group1 <- mean(x[1:10])
   #Mean of sites 11-20
-  group2 <- mean(x[11:20,1])
-
-  return(c(group1, group2))
+  group2 <- mean(x[11:20])
+  
+  #Naming elements of the output is optional but helpful
+  return(c(group1=group1, group2=group2))
 
 }
 
-#Get mean, SE, and 95% confidence interval for the two derived parameters
-#Using 100 draws from posterior predictive distribution
-(predict(re, func=myfunc, nsims=100))
+#Get 100 samples of the values calculated in your function
+(pr <- predict(re, func=myfunc, nsims=100))
+
+#Summarize posterior
+data.frame(mean=rowMeans(pr),
+           se=apply(pr, 1, stats::sd),
+           lower=apply(pr, 1, stats::quantile, 0.025),
+           upper=apply(pr, 1, stats::quantile, 0.975))
 
 #Alternatively, you can return the posterior predictive distribution
 #and run operations on it separately