From 53141021aeb9f326b0dfe3592a13fa7aa413997f Mon Sep 17 00:00:00 2001
From: erik-ringen <ringen.erik@gmail.com>
Date: Tue, 8 Oct 2024 18:45:53 +0200
Subject: [PATCH 1/5] first draft of 'coev_pred_series()'

---
 NAMESPACE               |   1 +
 R/coev_pred_series.R    | 128 ++++++++++++++++++++++++++++++++++++++++
 man/coev_pred_series.Rd |  72 ++++++++++++++++++++++
 3 files changed, 201 insertions(+)
 create mode 100644 R/coev_pred_series.R
 create mode 100644 man/coev_pred_series.Rd

diff --git a/NAMESPACE b/NAMESPACE
index a8bdf60..f21bd6b 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -15,6 +15,7 @@ export(coev_plot_delta_theta)
 export(coev_plot_flowfield)
 export(coev_plot_predictive_check)
 export(coev_plot_selection_gradient)
+export(coev_pred_series)
 export(coev_simulate_coevolution)
 export(extract_samples)
 export(stancode)
diff --git a/R/coev_pred_series.R b/R/coev_pred_series.R
new file mode 100644
index 0000000..89f7c5d
--- /dev/null
+++ b/R/coev_pred_series.R
@@ -0,0 +1,128 @@
+#' Predict a co-evolutionary time series from a fitted \code{coevfit}
+#' object
+#'
+#' @param object An object of class \code{coevfit}
+#' @param eta_anc If \code{NULL} default, the starting values for the latent states \eqn{\eta} will be set to the estimated ancestral states at the root from the the \code{coevfit}. Otherwise, a numeric vector with length equal to the number of variable specifying the initial \eqn{\eta}  values.
+#' @param tmax A positive number indicating the total duration of time for which to predict. Set to 1 by default, corresponding to the entire time depth of the original phylogenetic tree(s).
+#' @param tmax A positive number indicating the total duration of time for which to predict. Set to 1 by default, corresponding to the entire time depth of the original phylogenetic tree(s).
+#' @param ntimes A positive integer indicating the total number of discrete time steps over which to make predictions. Each step corresponds to a time difference of dt = tmax/ntimes. Set to 30 by default.
+#' @param ndraws An integer indicating the number of draws to return. The
+#'   default and maximum number of draws is the size of the posterior sample.
+#' @param stochastic Logical (defaults to \code{FALSE}); indicator of whether predictions should include only the expected co-evolutionary change (deterministic selection), or also stochastic drift.
+#' @return An [ndraw, ntimes, nvariables] array of predicted \eqn{\eta}  values.
+#'
+#' @author Scott Claessens \email{scott.claessens@@gmail.com}, Erik Ringen
+#'   \email{erikjacob.ringen@@uzh.ch}
+#'
+#' @examples
+#' \dontrun{
+#' # fit dynamic coevolutionary model
+#' fit <- coev_fit(
+#'   data = authority$data,
+#'   variables = list(
+#'     political_authority = "ordered_logistic",
+#'     religious_authority = "ordered_logistic"
+#'   ),
+#'   id = "language",
+#'   tree = authority$phylogeny,
+#'   # additional arguments for cmdstanr::sample()
+#'   chains = 4,
+#'   parallel_chains = 4,
+#'   seed = 1
+#'   )
+#'
+#' # simulate trait co-evolution
+#' sims <- coev_pred_series(
+#'   object = fit,
+#'   stochastic = T
+#'   )
+#' 
+#' # expected trait co-evolution, no drift
+#' epreds <- coev_pred_series(
+#'   object = fit,
+#'   stochastic = F
+#'   )
+#' }
+#'
+#' @export
+coev_pred_series <- function(object, eta_anc = NULL, tmax = 1, ntimes = 30, ndraws = NULL, stochastic = FALSE){
+  if (!methods::is(object, "coevfit")) {
+    stop2(
+      paste0(
+        "Argument 'object' must be a fitted coevolutionary model ",
+        "of class coevfit."
+        )
+      )
+  }
+  # stop if eta_anc is neither NULL nor a vector of length equal to the number of variables
+  if (!is.null(eta_anc) & (!is.numeric(eta_anc) | length(eta_anc) != length(object$variables))) {
+    stop2("Argument 'eta_anc' must be numeric and equal in length to the number of variables.")
+  }
+  # stop if tmax is not numeric or not positive
+  if (!is.numeric(tmax)|length(tmax) != 1) {
+    stop2("Argument 'tmax' must be a single numeric value.")
+  }
+  else if (tmax <= 0) {
+    stop2("Argument 'tmax' must be positive.")
+  }
+ # stop if ndraws is not a single integer between 1 and the total num draws
+ if (!is.null(ndraws)) {
+  if (!as.integer(ndraws) == ndraws|ndraws <= 0|length(ndraws) != 1) {
+    stop2("Argument 'ndraws' must be a single integer.")
+  } else if (ndraws < 1 | ndraws > nrow(object$fit$draws())) {
+    stop2(
+      "Argument 'ndraws' must be between 1 and the total number of draws."
+      )
+  }
+}
+# stop if stochastic not logical
+  if (!is.logical(stochastic)) {
+    stop2("Argument 'stochastic' must be logical.")
+  }
+
+  post <- extract_samples(object)
+  J <- length(object$variables)
+
+  nsamps <- ifelse(is.null(ndraws), length(post$lp__), ndraws)
+
+  preds <- array(NA, dim = c(nsamps, ntimes+1, J), dimnames = list(samps = 1:nsamps, time = 1:(ntimes+1), response = names(object$variables)))
+
+  # User defined ancestral states
+  if (!is.null(eta_anc)) {
+    for (i in 1:nsamps) preds[i,1,] = eta_anc
+  }
+
+  # Model inferred ancestral states
+  else {
+    eta_anc_long <- post$eta_anc
+    ntrees <- dim(eta_anc_long)[2]
+    # shuffle dimenisions of tree dimension to get random draws
+    eta_anc_long <- eta_anc_long[,sample(1:ntrees, size = ntrees, replace = F),]
+
+    if (ntrees > 1) {
+    # make into 2d matrix by stacking trees
+    dim(eta_anc_long) <- c(dim(eta_anc_long)[1], prod(dim(eta_anc_long)[2:3]) )
+    }
+    for (i in 1:nsamps) preds[i,1,] = eta_anc_long[i,]
+    }
+
+  for (i in 1:nsamps) {
+    A <- post$A[i,,]
+    A_delta <- as.matrix(Matrix::expm(A * tmax/ntimes))
+    inv_A <- solve(A)
+    I <- diag(rep(1, J))
+    b <- post$b[i,]
+
+    Q_inf <- post$Q_inf[i,,]
+    VCV <- Q_inf - ((A_delta) %*% Q_inf %*% t(A_delta))
+    chol_VCV <- t(chol(VCV))
+
+    for (t in 1:ntimes) {
+      preds[i, t+1, ] = (A_delta %*% preds[i, t, ] + (inv_A %*% (A_delta - I) %*% b))[,1]
+      # add drift
+      if (stochastic == T) preds[i, t+1, ] = preds[i, t+1, ] + chol_VCV %*% rnorm(J, 0, 1)
+    }
+  }
+  # shuffle draws to de-correlate when doing simulations
+  return(preds[sample(1:nsamps, size = nsamps, replace = F),,])
+}
diff --git a/man/coev_pred_series.Rd b/man/coev_pred_series.Rd
new file mode 100644
index 0000000..f9466cc
--- /dev/null
+++ b/man/coev_pred_series.Rd
@@ -0,0 +1,72 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/coev_pred_series.R
+\name{coev_pred_series}
+\alias{coev_pred_series}
+\title{Predict a co-evolutionary time series from a fitted \code{coevfit}
+object}
+\usage{
+coev_pred_series(
+  object,
+  eta_anc = NULL,
+  tmax = 1,
+  ntimes = 30,
+  ndraws = NULL,
+  stochastic = FALSE
+)
+}
+\arguments{
+\item{object}{An object of class \code{coevfit}}
+
+\item{eta_anc}{If \code{NULL} default, the starting values for the latent states \eqn{\eta} will be set to the estimated ancestral states at the root from the the \code{coevfit}. Otherwise, a numeric vector with length equal to the number of variable specifying the initial \eqn{\eta}  values.}
+
+\item{tmax}{A positive number indicating the total duration of time for which to predict. Set to 1 by default, corresponding to the entire time depth of the original phylogenetic tree(s).}
+
+\item{ntimes}{A positive integer indicating the total number of discrete time steps over which to make predictions. Each step corresponds to a time difference of dt = tmax/ntimes. Set to 30 by default.}
+
+\item{ndraws}{An integer indicating the number of draws to return. The
+default and maximum number of draws is the size of the posterior sample.}
+
+\item{stochastic}{Logical (defaults to \code{FALSE}); indicator of whether predictions should include only the expected co-evolutionary change (deterministic selection), or also stochastic drift.}
+}
+\value{
+An \link{ndraw, ntimes, nvariables} array of predicted \eqn{\eta}  values.
+}
+\description{
+Predict a co-evolutionary time series from a fitted \code{coevfit}
+object
+}
+\examples{
+\dontrun{
+# fit dynamic coevolutionary model
+fit <- coev_fit(
+  data = authority$data,
+  variables = list(
+    political_authority = "ordered_logistic",
+    religious_authority = "ordered_logistic"
+  ),
+  id = "language",
+  tree = authority$phylogeny,
+  # additional arguments for cmdstanr::sample()
+  chains = 4,
+  parallel_chains = 4,
+  seed = 1
+  )
+
+# simulate trait co-evolution
+sims <- coev_pred_series(
+  object = fit,
+  stochastic = T
+  )
+
+# expected trait co-evolution, no drift
+epreds <- coev_pred_series(
+  object = fit,
+  stochastic = F
+  )
+}
+
+}
+\author{
+Scott Claessens \email{scott.claessens@gmail.com}, Erik Ringen
+\email{erikjacob.ringen@uzh.ch}
+}

From ae298e46e46f18b4f1bfa20c5ef03cd24a68fe94 Mon Sep 17 00:00:00 2001
From: erik-ringen <ringen.erik@gmail.com>
Date: Wed, 9 Oct 2024 10:21:18 +0200
Subject: [PATCH 2/5] plot function for coev_pred_series

---
 DESCRIPTION                  |   3 +-
 NAMESPACE                    |   1 +
 R/coev_plot_pred_series.R    | 107 +++++++++++++++++++++++++++++++++++
 man/coev_plot_pred_series.Rd |  59 +++++++++++++++++++
 4 files changed, 169 insertions(+), 1 deletion(-)
 create mode 100644 R/coev_plot_pred_series.R
 create mode 100644 man/coev_plot_pred_series.Rd

diff --git a/DESCRIPTION b/DESCRIPTION
index 5028ebf..9c1d390 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -41,7 +41,8 @@ Imports:
     rlang,
     stats,
     stringr,
-    tidyr
+    tidyr,
+    cubelyr
 VignetteBuilder: knitr
 Depends: 
     R (>= 3.5.0)
diff --git a/NAMESPACE b/NAMESPACE
index f21bd6b..396624f 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -13,6 +13,7 @@ export(coev_make_stancode)
 export(coev_make_standata)
 export(coev_plot_delta_theta)
 export(coev_plot_flowfield)
+export(coev_plot_pred_series)
 export(coev_plot_predictive_check)
 export(coev_plot_selection_gradient)
 export(coev_pred_series)
diff --git a/R/coev_plot_pred_series.R b/R/coev_plot_pred_series.R
new file mode 100644
index 0000000..a6c03df
--- /dev/null
+++ b/R/coev_plot_pred_series.R
@@ -0,0 +1,107 @@
+#' Plot a predicted co-evolutionary time series from a fitted \code{coevfit}
+#' object
+#'
+#' @param object An object of class \code{coevfit}
+#' @param prob A value between 0 and 1 indicating the desired probability
+#'   to be covered by the uncertainty intervals. The default is 0.95.
+#' @param ... Additional arguments passed to \code{\link{coev_pred_series}}
+#' @return A ggplot object.
+#'
+#' @author Scott Claessens \email{scott.claessens@@gmail.com}, Erik Ringen
+#'   \email{erikjacob.ringen@@uzh.ch}
+#'
+#' @examples
+#' \dontrun{
+#' # fit dynamic coevolutionary model
+#' fit <- coev_fit(
+#'   data = authority$data,
+#'   variables = list(
+#'     political_authority = "ordered_logistic",
+#'     religious_authority = "ordered_logistic"
+#'   ),
+#'   id = "language",
+#'   tree = authority$phylogeny,
+#'   # additional arguments for cmdstanr::sample()
+#'   chains = 4,
+#'   parallel_chains = 4,
+#'   seed = 1
+#'   )
+#'
+#' # simulated trait co-evolution
+#' coev_plot_pred_series(
+#'   object = fit,
+#'   stochastic = T
+#'   )
+#' 
+#' # expected trait co-evolution, no drift
+#' coev_plot_pred_series(
+#'   object = fit,
+#'   stochastic = F
+#'   )
+#' }
+#'
+#' @export
+coev_plot_pred_series <- function(object, prob = 0.95, ...){
+  if (!methods::is(object, "coevfit")) {
+    stop2(
+      paste0(
+        "Argument 'object' must be a fitted coevolutionary model ",
+        "of class coevfit."
+        )
+      )
+  }
+  # stop if prob is outside of range 0 - 1
+    if (prob <= 0 | prob >= 1) {
+      stop2("Argument 'prob' is not between 0 and 1.")
+  }
+  # user-supplied arguments
+  user_args <- list(...)
+  # default arguments from coev_pred_series
+  default_args <- as.list(formals(coev_pred_series))
+  combined_args <- modifyList(default_args, user_args)
+  combined_args$object <- object
+  preds <- do.call(coev_pred_series, combined_args)
+  
+  if (combined_args$stochastic == F) {
+      # for CI quantiles
+      probs <- c(((1 - prob) / 2), 1 - ((1 - prob) / 2))
+
+      epreds_long <- preds |>
+      cubelyr::as.tbl_cube(met_name = "est") |> 
+      as.data.frame()
+
+      epreds_summary <- epreds_long |> 
+      dplyr::group_by(response, time) |> 
+      dplyr::summarise(mean = mean(est),
+      lower_CI = quantile(est, probs[1]),
+      upper_CI = quantile(est, probs[2]),
+      )
+    
+      p <- ggplot2::ggplot(epreds_summary, aes(x = time, y = mean, fill = response, color = response, linetype = response)) +
+      geom_line(lwd = 1) + 
+        geom_ribbon(aes(ymin = lower_CI, ymax = upper_CI), alpha = 0.25, lwd = 0.25) +
+      theme_classic(base_size = 14) + 
+      scale_x_continuous(breaks = c(0, max(epreds_summary$time)), labels = c("LCA", "present")) +
+      ylab("trait value (latent scale)") + 
+      xlab("time") +
+      theme(legend.title = element_blank(),  strip.background = element_blank(),strip.text = element_blank()) + 
+      ggtitle("Expected trait coevolution")
+  }
+  else if (combined_args$stochastic == T) {
+    sims_long <- preds |>
+      cubelyr::as.tbl_cube(met_name = "est") |> 
+      as.data.frame() |> 
+      dplyr::mutate(sim = factor(paste("sim", samps), levels = paste("sim", sort(unique(samps)))))
+
+    p <- ggplot2::ggplot(sims_long |> dplyr::filter(samps <= 15), aes(x = time, y = est, color = response)) +
+    facet_wrap(~sim, ncol = 5) +
+    geom_line(lwd = 1) + 
+    theme_minimal(base_size = 14) + 
+    scale_x_continuous(breaks = c(0, max(sims_long$time)), labels = c("LCA", "present")) +
+    ylab("trait value (latent scale)") + 
+    xlab("time") +
+    theme(legend.title = element_blank(),  strip.background = element_blank(), panel.spacing.x = unit(1, "lines"), axis.text.x = element_blank(),  axis.ticks.x = element_blank()) + 
+    ggtitle("Predicted trait coevolution")
+  }
+  return(p)
+}
\ No newline at end of file
diff --git a/man/coev_plot_pred_series.Rd b/man/coev_plot_pred_series.Rd
new file mode 100644
index 0000000..9fff0da
--- /dev/null
+++ b/man/coev_plot_pred_series.Rd
@@ -0,0 +1,59 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/coev_plot_pred_series.R
+\name{coev_plot_pred_series}
+\alias{coev_plot_pred_series}
+\title{Plot a predicted co-evolutionary time series from a fitted \code{coevfit}
+object}
+\usage{
+coev_plot_pred_series(object, prob = 0.95, ...)
+}
+\arguments{
+\item{object}{An object of class \code{coevfit}}
+
+\item{prob}{A value between 0 and 1 indicating the desired probability
+to be covered by the uncertainty intervals. The default is 0.95.}
+
+\item{...}{Additional arguments passed to \code{\link{coev_pred_series}}}
+}
+\value{
+A ggplot object.
+}
+\description{
+Plot a predicted co-evolutionary time series from a fitted \code{coevfit}
+object
+}
+\examples{
+\dontrun{
+# fit dynamic coevolutionary model
+fit <- coev_fit(
+  data = authority$data,
+  variables = list(
+    political_authority = "ordered_logistic",
+    religious_authority = "ordered_logistic"
+  ),
+  id = "language",
+  tree = authority$phylogeny,
+  # additional arguments for cmdstanr::sample()
+  chains = 4,
+  parallel_chains = 4,
+  seed = 1
+  )
+
+# simulated trait co-evolution
+coev_plot_pred_series(
+  object = fit,
+  stochastic = T
+  )
+
+# expected trait co-evolution, no drift
+coev_plot_pred_series(
+  object = fit,
+  stochastic = F
+  )
+}
+
+}
+\author{
+Scott Claessens \email{scott.claessens@gmail.com}, Erik Ringen
+\email{erikjacob.ringen@uzh.ch}
+}

From 1ed1e02413f9cdc2ceb6e584fd1cad4b5374a205 Mon Sep 17 00:00:00 2001
From: erik-ringen <ringen.erik@gmail.com>
Date: Wed, 9 Oct 2024 10:58:37 +0200
Subject: [PATCH 3/5] unit tests for coev_pred_series()

---
 R/coev_pred_series.R                   | 24 +++----
 tests/testthat/test-coev_pred_series.R | 88 ++++++++++++++++++++++++++
 2 files changed, 100 insertions(+), 12 deletions(-)
 create mode 100644 tests/testthat/test-coev_pred_series.R

diff --git a/R/coev_pred_series.R b/R/coev_pred_series.R
index 89f7c5d..6568a84 100644
--- a/R/coev_pred_series.R
+++ b/R/coev_pred_series.R
@@ -49,8 +49,7 @@ coev_pred_series <- function(object, eta_anc = NULL, tmax = 1, ntimes = 30, ndra
   if (!methods::is(object, "coevfit")) {
     stop2(
       paste0(
-        "Argument 'object' must be a fitted coevolutionary model ",
-        "of class coevfit."
+        "Argument 'object' must be a fitted coevolutionary model of class coevfit."
         )
       )
   }
@@ -65,16 +64,16 @@ coev_pred_series <- function(object, eta_anc = NULL, tmax = 1, ntimes = 30, ndra
   else if (tmax <= 0) {
     stop2("Argument 'tmax' must be positive.")
   }
- # stop if ndraws is not a single integer between 1 and the total num draws
- if (!is.null(ndraws)) {
-  if (!as.integer(ndraws) == ndraws|ndraws <= 0|length(ndraws) != 1) {
-    stop2("Argument 'ndraws' must be a single integer.")
-  } else if (ndraws < 1 | ndraws > nrow(object$fit$draws())) {
-    stop2(
-      "Argument 'ndraws' must be between 1 and the total number of draws."
-      )
+  # stop if ndraws is not a single integer between 1 and the total num draws
+  if (!is.null(ndraws)) {
+    if (!is.integer(ndraws) | length(ndraws) != 1) {
+      stop2("Argument 'ndraws' must be a single integer.")
+    } else if (ndraws < 1 | ndraws > nrow(object$fit$draws())) {
+      stop2(
+        "Argument 'ndraws' must be between 1 and the total number of draws."
+        )
+    }
   }
-}
 # stop if stochastic not logical
   if (!is.logical(stochastic)) {
     stop2("Argument 'stochastic' must be logical.")
@@ -101,7 +100,8 @@ coev_pred_series <- function(object, eta_anc = NULL, tmax = 1, ntimes = 30, ndra
 
     if (ntrees > 1) {
     # make into 2d matrix by stacking trees
-    dim(eta_anc_long) <- c(dim(eta_anc_long)[1], prod(dim(eta_anc_long)[2:3]) )
+    eta_anc_long2 <- eta_anc_long[,1,]
+    for (t in 2:ntrees) eta_anc_long2 <- rbind(eta_anc_long2, eta_anc_long[,t,]) -> eta_anc_long
     }
     for (i in 1:nsamps) preds[i,1,] = eta_anc_long[i,]
     }
diff --git a/tests/testthat/test-coev_pred_series.R b/tests/testthat/test-coev_pred_series.R
new file mode 100644
index 0000000..6fc0a22
--- /dev/null
+++ b/tests/testthat/test-coev_pred_series.R
@@ -0,0 +1,88 @@
+test_that("coev_pred_series() produces expected errors and output", {
+  # load models
+  m1 <- readRDS(test_path("fixtures", "coevfit_example1.rds"))
+  m2 <- readRDS(test_path("fixtures", "coevfit_example2.rds"))
+  m3 <- readRDS(test_path("fixtures", "coevfit_example3.rds"))
+  m4 <- readRDS(test_path("fixtures", "coevfit_example4.rds"))
+  m5 <- readRDS(test_path("fixtures", "coevfit_example5.rds"))
+  m6 <- readRDS(test_path("fixtures", "coevfit_example6.rds"))
+  m7 <- readRDS(test_path("fixtures", "coevfit_example7.rds"))
+  m8 <- readRDS(test_path("fixtures", "coevfit_example8.rds"))
+  m9 <- readRDS(test_path("fixtures", "coevfit_example9.rds"))
+  m1 <- reload_fit(m1, filename = "coevfit_example1-1.csv")
+  m2 <- reload_fit(m2, filename = "coevfit_example2-1.csv")
+  m3 <- reload_fit(m3, filename = "coevfit_example3-1.csv")
+  m4 <- reload_fit(m4, filename = "coevfit_example4-1.csv")
+  m5 <- reload_fit(m5, filename = "coevfit_example5-1.csv")
+  m6 <- reload_fit(m6, filename = "coevfit_example6-1.csv")
+  m7 <- reload_fit(m7, filename = "coevfit_example7-1.csv")
+  m8 <- reload_fit(m8, filename = "coevfit_example8-1.csv")
+  m9 <- reload_fit(m9, filename = "coevfit_example9-1.csv")
+  # expect the following errors
+  expect_error(
+    coev_pred_series(object = "fail"),
+     "Argument 'object' must be a fitted coevolutionary model of class coevfit."
+  )
+  expect_error(
+    coev_pred_series(object = m1, eta_anc = c(1, "LCA")),
+    paste0(
+      "Argument 'eta_anc' must be numeric and equal in length to the number of variables."
+    )
+  )
+  expect_error(
+    coev_pred_series(object = m1, tmax = -1),
+    paste0(
+      "Argument 'tmax' must be positive."
+    )
+  )
+  expect_error(
+    coev_pred_series(object = m1, ndraws = "fail"),
+    "Argument 'ndraws' must be a single integer."
+  )
+  expect_error(
+    coev_pred_series(object = m1, ndraws = 0L),
+    "Argument 'ndraws' must be between 1 and the total number of draws."
+  )
+  expect_error(
+    coev_pred_series(
+      object = m1, ndraws = as.integer(nrow(m1$fit$draws()) + 1)
+      ),
+    "Argument 'ndraws' must be between 1 and the total number of draws."
+  )
+  expect_error(
+    coev_pred_series(object = m1, stochastic = "only"),
+    paste0(
+      "Argument 'stochastic' must be logical."
+    )
+  )
+  # suppress warnings
+  SW <- suppressWarnings
+  # should run without error and produce list of ggplot objects
+  expect_no_error(SW(coev_pred_series(m1)))
+  expect_no_error(SW(coev_pred_series(m2)))
+  expect_no_error(SW(coev_pred_series(m3)))
+  expect_no_error(SW(coev_pred_series(m4)))
+  expect_no_error(SW(coev_pred_series(m5)))
+  expect_no_error(SW(coev_pred_series(m6)))
+  expect_no_error(SW(coev_pred_series(m7)))
+  expect_no_error(SW(coev_pred_series(m8)))
+  expect_no_error(SW(coev_pred_series(m9)))
+  expect_no_error(SW(coev_pred_series(m1, stochastic = T)))
+  expect_no_error(SW(coev_pred_series(m2, stochastic = T)))
+  expect_no_error(SW(coev_pred_series(m3, stochastic = T)))
+  expect_no_error(SW(coev_pred_series(m4, stochastic = T)))
+  expect_no_error(SW(coev_pred_series(m5, stochastic = T)))
+  expect_no_error(SW(coev_pred_series(m6, stochastic = T)))
+  expect_no_error(SW(coev_pred_series(m7, stochastic = T)))
+  expect_no_error(SW(coev_pred_series(m8, stochastic = T)))
+  expect_no_error(SW(coev_pred_series(m9, stochastic = T)))
+  expect_no_error(SW(coev_pred_series(m1, ndraws = 1L)))
+  expect_no_error(SW(coev_pred_series(m2, ndraws = 1L)))
+  expect_no_error(SW(coev_pred_series(m3, ndraws = 1L)))
+  expect_no_error(SW(coev_pred_series(m4, ndraws = 1L)))
+  expect_no_error(SW(coev_pred_series(m5, ndraws = 1L)))
+  expect_no_error(SW(coev_pred_series(m6, ndraws = 1L)))
+  expect_no_error(SW(coev_pred_series(m7, ndraws = 1L)))
+  expect_no_error(SW(coev_pred_series(m8, ndraws = 1L)))
+  expect_no_error(SW(coev_pred_series(m9, ndraws = 1L)))
+})

From a48f53659c0009441b13c53254a006721a3ee2cd Mon Sep 17 00:00:00 2001
From: erik-ringen <ringen.erik@gmail.com>
Date: Wed, 9 Oct 2024 11:19:52 +0200
Subject: [PATCH 4/5] unit tests for coev_plot_pred_series()

---
 R/coev_plot_pred_series.R                   | 40 ++++++------
 tests/testthat/test-coev_plot_pred_series.R | 71 +++++++++++++++++++++
 2 files changed, 91 insertions(+), 20 deletions(-)
 create mode 100644 tests/testthat/test-coev_plot_pred_series.R

diff --git a/R/coev_plot_pred_series.R b/R/coev_plot_pred_series.R
index a6c03df..5797a4a 100644
--- a/R/coev_plot_pred_series.R
+++ b/R/coev_plot_pred_series.R
@@ -73,19 +73,19 @@ coev_plot_pred_series <- function(object, prob = 0.95, ...){
       epreds_summary <- epreds_long |> 
       dplyr::group_by(response, time) |> 
       dplyr::summarise(mean = mean(est),
-      lower_CI = quantile(est, probs[1]),
-      upper_CI = quantile(est, probs[2]),
+      lower_CI = stats::quantile(est, probs[1]),
+      upper_CI = stats::quantile(est, probs[2]),
       )
     
-      p <- ggplot2::ggplot(epreds_summary, aes(x = time, y = mean, fill = response, color = response, linetype = response)) +
-      geom_line(lwd = 1) + 
-        geom_ribbon(aes(ymin = lower_CI, ymax = upper_CI), alpha = 0.25, lwd = 0.25) +
-      theme_classic(base_size = 14) + 
-      scale_x_continuous(breaks = c(0, max(epreds_summary$time)), labels = c("LCA", "present")) +
-      ylab("trait value (latent scale)") + 
-      xlab("time") +
-      theme(legend.title = element_blank(),  strip.background = element_blank(),strip.text = element_blank()) + 
-      ggtitle("Expected trait coevolution")
+      p <- ggplot2::ggplot(epreds_summary, ggplot2::aes(x = time, y = mean, fill = response, color = response, linetype = response)) +
+      ggplot2::geom_line(lwd = 1) + 
+        ggplot2::geom_ribbon(ggplot2::aes(ymin = lower_CI, ymax = upper_CI), alpha = 0.25, lwd = 0.25) +
+      ggplot2::theme_classic(base_size = 14) + 
+      ggplot2::scale_x_continuous(breaks = c(0, max(epreds_summary$time)), labels = c("LCA", "present")) +
+      ggplot2::ylab("trait value (latent scale)") + 
+      ggplot2::xlab("time") +
+      ggplot2::theme(legend.title = ggplot2::element_blank(),  strip.background = ggplot2::element_blank(),strip.text = ggplot2::element_blank()) + 
+      ggplot2::ggtitle("Expected trait coevolution")
   }
   else if (combined_args$stochastic == T) {
     sims_long <- preds |>
@@ -93,15 +93,15 @@ coev_plot_pred_series <- function(object, prob = 0.95, ...){
       as.data.frame() |> 
       dplyr::mutate(sim = factor(paste("sim", samps), levels = paste("sim", sort(unique(samps)))))
 
-    p <- ggplot2::ggplot(sims_long |> dplyr::filter(samps <= 15), aes(x = time, y = est, color = response)) +
-    facet_wrap(~sim, ncol = 5) +
-    geom_line(lwd = 1) + 
-    theme_minimal(base_size = 14) + 
-    scale_x_continuous(breaks = c(0, max(sims_long$time)), labels = c("LCA", "present")) +
-    ylab("trait value (latent scale)") + 
-    xlab("time") +
-    theme(legend.title = element_blank(),  strip.background = element_blank(), panel.spacing.x = unit(1, "lines"), axis.text.x = element_blank(),  axis.ticks.x = element_blank()) + 
-    ggtitle("Predicted trait coevolution")
+    p <- ggplot2::ggplot(sims_long |> dplyr::filter(samps <= 15), ggplot2::aes(x = time, y = est, color = response)) +
+    ggplot2::facet_wrap(~sim, ncol = 5) +
+    ggplot2::geom_line(lwd = 1) + 
+    ggplot2::theme_minimal(base_size = 14) + 
+    ggplot2::scale_x_continuous(breaks = c(0, max(sims_long$time)), labels = c("LCA", "present")) +
+    ggplot2::ylab("trait value (latent scale)") + 
+    ggplot2::xlab("time") +
+    ggplot2::theme(legend.title = ggplot2::element_blank(),  strip.background = ggplot2::element_blank(), panel.spacing.x = ggplot2::unit(1, "lines"), axis.text.x = ggplot2::element_blank(),  axis.ticks.x = ggplot2::element_blank()) + 
+    ggplot2::ggtitle("Predicted trait coevolution")
   }
   return(p)
 }
\ No newline at end of file
diff --git a/tests/testthat/test-coev_plot_pred_series.R b/tests/testthat/test-coev_plot_pred_series.R
new file mode 100644
index 0000000..a05274c
--- /dev/null
+++ b/tests/testthat/test-coev_plot_pred_series.R
@@ -0,0 +1,71 @@
+test_that("coev_plot_pred_series() produces expected errors and output", {
+  # load models
+  m1 <- readRDS(test_path("fixtures", "coevfit_example1.rds"))
+  m2 <- readRDS(test_path("fixtures", "coevfit_example2.rds"))
+  m3 <- readRDS(test_path("fixtures", "coevfit_example3.rds"))
+  m4 <- readRDS(test_path("fixtures", "coevfit_example4.rds"))
+  m5 <- readRDS(test_path("fixtures", "coevfit_example5.rds"))
+  m6 <- readRDS(test_path("fixtures", "coevfit_example6.rds"))
+  m7 <- readRDS(test_path("fixtures", "coevfit_example7.rds"))
+  m8 <- readRDS(test_path("fixtures", "coevfit_example8.rds"))
+  m9 <- readRDS(test_path("fixtures", "coevfit_example9.rds"))
+  m1 <- reload_fit(m1, filename = "coevfit_example1-1.csv")
+  m2 <- reload_fit(m2, filename = "coevfit_example2-1.csv")
+  m3 <- reload_fit(m3, filename = "coevfit_example3-1.csv")
+  m4 <- reload_fit(m4, filename = "coevfit_example4-1.csv")
+  m5 <- reload_fit(m5, filename = "coevfit_example5-1.csv")
+  m6 <- reload_fit(m6, filename = "coevfit_example6-1.csv")
+  m7 <- reload_fit(m7, filename = "coevfit_example7-1.csv")
+  m8 <- reload_fit(m8, filename = "coevfit_example8-1.csv")
+  m9 <- reload_fit(m9, filename = "coevfit_example9-1.csv")
+  # expect the following errors
+  expect_error(
+    coev_plot_pred_series(object = "fail"),
+     "Argument 'object' must be a fitted coevolutionary model of class coevfit."
+  )
+  expect_error(
+    coev_plot_pred_series(object = m1, prob = 95),
+    paste0(
+      "Argument 'prob' is not between 0 and 1."
+    )
+  )
+  # suppress warnings
+  SW <- suppressWarnings
+  # should run without error and produce list of ggplot objects
+  expect_no_error(SW(coev_plot_pred_series(m1)))
+  expect_no_error(SW(coev_plot_pred_series(m2)))
+  expect_no_error(SW(coev_plot_pred_series(m3)))
+  expect_no_error(SW(coev_plot_pred_series(m4)))
+  expect_no_error(SW(coev_plot_pred_series(m5)))
+  expect_no_error(SW(coev_plot_pred_series(m6)))
+  expect_no_error(SW(coev_plot_pred_series(m7)))
+  expect_no_error(SW(coev_plot_pred_series(m8)))
+  expect_no_error(SW(coev_plot_pred_series(m9)))
+  expect_no_error(SW(coev_plot_pred_series(m1, stochastic = T)))
+  expect_no_error(SW(coev_plot_pred_series(m2, stochastic = T)))
+  expect_no_error(SW(coev_plot_pred_series(m3, stochastic = T)))
+  expect_no_error(SW(coev_plot_pred_series(m4, stochastic = T)))
+  expect_no_error(SW(coev_plot_pred_series(m5, stochastic = T)))
+  expect_no_error(SW(coev_plot_pred_series(m6, stochastic = T)))
+  expect_no_error(SW(coev_plot_pred_series(m7, stochastic = T)))
+  expect_no_error(SW(coev_plot_pred_series(m8, stochastic = T)))
+  expect_no_error(SW(coev_plot_pred_series(m9, stochastic = T)))
+  expect_no_error(SW(coev_plot_pred_series(m1, ndraws = 1L)))
+  expect_no_error(SW(coev_plot_pred_series(m2, ndraws = 1L)))
+  expect_no_error(SW(coev_plot_pred_series(m3, ndraws = 1L)))
+  expect_no_error(SW(coev_plot_pred_series(m4, ndraws = 1L)))
+  expect_no_error(SW(coev_plot_pred_series(m5, ndraws = 1L)))
+  expect_no_error(SW(coev_plot_pred_series(m6, ndraws = 1L)))
+  expect_no_error(SW(coev_plot_pred_series(m7, ndraws = 1L)))
+  expect_no_error(SW(coev_plot_pred_series(m8, ndraws = 1L)))
+  expect_no_error(SW(coev_plot_pred_series(m9, ndraws = 1L)))
+  expect_true(methods::is(SW(coev_plot_pred_series(m1)), "ggplot"))
+  expect_true(methods::is(SW(coev_plot_pred_series(m2)), "ggplot"))
+  expect_true(methods::is(SW(coev_plot_pred_series(m3)), "ggplot"))
+  expect_true(methods::is(SW(coev_plot_pred_series(m4)), "ggplot"))
+  expect_true(methods::is(SW(coev_plot_pred_series(m5)), "ggplot"))
+  expect_true(methods::is(SW(coev_plot_pred_series(m6)), "ggplot"))
+  expect_true(methods::is(SW(coev_plot_pred_series(m7)), "ggplot"))
+  expect_true(methods::is(SW(coev_plot_pred_series(m8)), "ggplot"))
+  expect_true(methods::is(SW(coev_plot_pred_series(m9)), "ggplot"))
+})

From 2939e5a7e612d61622bbfccd890ea4e0ad295452 Mon Sep 17 00:00:00 2001
From: Scott Claessens <scla896@UoA.auckland.ac.nz>
Date: Wed, 9 Oct 2024 14:39:31 +0100
Subject: [PATCH 5/5] Tidy and fix failing RMD check

---
 DESCRIPTION                                 |   5 +-
 R/coev_plot_pred_series.R                   | 153 ++++++++++++++------
 R/coev_pred_series.R                        | 148 +++++++++++++------
 man/coev_plot_pred_series.Rd                |  18 ++-
 man/coev_pred_series.Rd                     |  36 +++--
 tests/testthat/test-coev_plot_pred_series.R |  26 ++--
 tests/testthat/test-coev_pred_series.R      |  61 +++++---
 7 files changed, 305 insertions(+), 142 deletions(-)

diff --git a/DESCRIPTION b/DESCRIPTION
index 9c1d390..5949de9 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -29,8 +29,10 @@ Imports:
     bayesplot,
     cmdstanr,
     colorspace,
+    cubelyr,
     dplyr,
     ggplot2,
+    Matrix,
     methods,
     phangorn,
     phaseR,
@@ -41,8 +43,7 @@ Imports:
     rlang,
     stats,
     stringr,
-    tidyr,
-    cubelyr
+    tidyr
 VignetteBuilder: knitr
 Depends: 
     R (>= 3.5.0)
diff --git a/R/coev_plot_pred_series.R b/R/coev_plot_pred_series.R
index 5797a4a..62b302d 100644
--- a/R/coev_plot_pred_series.R
+++ b/R/coev_plot_pred_series.R
@@ -1,15 +1,26 @@
-#' Plot a predicted co-evolutionary time series from a fitted \code{coevfit}
+#' Plot a predicted coevolutionary time series from a fitted \code{coevfit}
 #' object
 #'
+#' This function plots a predicted coevolutionary time series using the
+#' estimated parameters from a fitted \code{coevfit} model. By default, the
+#' plot uses the posterior ancestral states estimated by the model as the
+#' starting values, but users can also set their own starting values for
+#' traits. Plots can be generated with or without stochastic drift. For more
+#' details on the underlying predictive function, see
+#' \code{help(coev_pred_series)}.
+#'
 #' @param object An object of class \code{coevfit}
 #' @param prob A value between 0 and 1 indicating the desired probability
 #'   to be covered by the uncertainty intervals. The default is 0.95.
 #' @param ... Additional arguments passed to \code{\link{coev_pred_series}}
+#'
 #' @return A ggplot object.
 #'
 #' @author Scott Claessens \email{scott.claessens@@gmail.com}, Erik Ringen
 #'   \email{erikjacob.ringen@@uzh.ch}
 #'
+#' @seealso \code{\link{coev_pred_series}}
+#'
 #' @examples
 #' \dontrun{
 #' # fit dynamic coevolutionary model
@@ -30,18 +41,19 @@
 #' # simulated trait co-evolution
 #' coev_plot_pred_series(
 #'   object = fit,
-#'   stochastic = T
+#'   stochastic = TRUE
 #'   )
-#' 
+#'
 #' # expected trait co-evolution, no drift
 #' coev_plot_pred_series(
 #'   object = fit,
-#'   stochastic = F
+#'   stochastic = FALSE
 #'   )
 #' }
 #'
 #' @export
-coev_plot_pred_series <- function(object, prob = 0.95, ...){
+coev_plot_pred_series <- function(object, prob = 0.95, ...) {
+  # stop if object not of class coevfit
   if (!methods::is(object, "coevfit")) {
     stop2(
       paste0(
@@ -51,57 +63,110 @@ coev_plot_pred_series <- function(object, prob = 0.95, ...){
       )
   }
   # stop if prob is outside of range 0 - 1
-    if (prob <= 0 | prob >= 1) {
-      stop2("Argument 'prob' is not between 0 and 1.")
+  if (prob <= 0 | prob >= 1) {
+    stop2("Argument 'prob' is not between 0 and 1.")
   }
   # user-supplied arguments
   user_args <- list(...)
   # default arguments from coev_pred_series
   default_args <- as.list(formals(coev_pred_series))
-  combined_args <- modifyList(default_args, user_args)
+  combined_args <- utils::modifyList(default_args, user_args)
   combined_args$object <- object
   preds <- do.call(coev_pred_series, combined_args)
-  
-  if (combined_args$stochastic == F) {
-      # for CI quantiles
+  # create plot
+  if (combined_args$stochastic == FALSE) {
+      # get CI quantiles
       probs <- c(((1 - prob) / 2), 1 - ((1 - prob) / 2))
-
-      epreds_long <- preds |>
-      cubelyr::as.tbl_cube(met_name = "est") |> 
-      as.data.frame()
-
-      epreds_summary <- epreds_long |> 
-      dplyr::group_by(response, time) |> 
-      dplyr::summarise(mean = mean(est),
-      lower_CI = stats::quantile(est, probs[1]),
-      upper_CI = stats::quantile(est, probs[2]),
+      # get expected values in cube format
+      epreds_long <-
+        preds |>
+        cubelyr::as.tbl_cube(
+          met_name = "est"
+          ) |>
+        as.data.frame()
+      # summarise expected values
+      epreds_summary <-
+        epreds_long |>
+        dplyr::group_by(.data$response, .data$time) |>
+        dplyr::summarise(
+          mean = mean(.data$est),
+          lower_CI = stats::quantile(.data$est, probs[1]),
+          upper_CI = stats::quantile(.data$est, probs[2])
       )
-    
-      p <- ggplot2::ggplot(epreds_summary, ggplot2::aes(x = time, y = mean, fill = response, color = response, linetype = response)) +
-      ggplot2::geom_line(lwd = 1) + 
-        ggplot2::geom_ribbon(ggplot2::aes(ymin = lower_CI, ymax = upper_CI), alpha = 0.25, lwd = 0.25) +
-      ggplot2::theme_classic(base_size = 14) + 
-      ggplot2::scale_x_continuous(breaks = c(0, max(epreds_summary$time)), labels = c("LCA", "present")) +
-      ggplot2::ylab("trait value (latent scale)") + 
+      # plot
+      p <-
+        ggplot2::ggplot(
+          epreds_summary,
+          ggplot2::aes(
+            x = .data$time,
+            y = .data$mean,
+            fill = .data$response,
+            color = .data$response,
+            linetype = .data$response
+            )
+          ) +
+      ggplot2::geom_line(lwd = 1) +
+      ggplot2::geom_ribbon(
+        ggplot2::aes(
+          ymin = .data$lower_CI,
+          ymax = .data$upper_CI
+          ),
+        alpha = 0.25,
+        lwd = 0.25
+        ) +
+      ggplot2::theme_classic(base_size = 14) +
+      ggplot2::scale_x_continuous(
+        breaks = c(0, max(epreds_summary$time)),
+        labels = c("LCA", "present")
+        ) +
+      ggplot2::ylab("trait value (latent scale)") +
       ggplot2::xlab("time") +
-      ggplot2::theme(legend.title = ggplot2::element_blank(),  strip.background = ggplot2::element_blank(),strip.text = ggplot2::element_blank()) + 
+      ggplot2::theme(
+        legend.title = ggplot2::element_blank(),
+        strip.background = ggplot2::element_blank(),
+        strip.text = ggplot2::element_blank()
+        ) +
       ggplot2::ggtitle("Expected trait coevolution")
-  }
-  else if (combined_args$stochastic == T) {
-    sims_long <- preds |>
-      cubelyr::as.tbl_cube(met_name = "est") |> 
-      as.data.frame() |> 
-      dplyr::mutate(sim = factor(paste("sim", samps), levels = paste("sim", sort(unique(samps)))))
-
-    p <- ggplot2::ggplot(sims_long |> dplyr::filter(samps <= 15), ggplot2::aes(x = time, y = est, color = response)) +
-    ggplot2::facet_wrap(~sim, ncol = 5) +
-    ggplot2::geom_line(lwd = 1) + 
-    ggplot2::theme_minimal(base_size = 14) + 
-    ggplot2::scale_x_continuous(breaks = c(0, max(sims_long$time)), labels = c("LCA", "present")) +
-    ggplot2::ylab("trait value (latent scale)") + 
+  } else if (combined_args$stochastic == TRUE) {
+    # get predictions
+    sims_long <-
+      preds |>
+      cubelyr::as.tbl_cube(met_name = "est") |>
+      as.data.frame() |>
+      dplyr::mutate(
+        sim = factor(
+          paste("sim", .data$samps),
+          levels = paste("sim", sort(unique(.data$samps)))
+          )
+        )
+    # plot
+    p <-
+      sims_long |>
+      dplyr::filter(.data$samps <= 15) |>
+      ggplot2::ggplot(
+        ggplot2::aes(
+          x = .data$time,
+          y = .data$est,
+          color = .data$response
+          )
+        ) +
+    ggplot2::facet_wrap( ~ .data$sim, ncol = 5) +
+    ggplot2::geom_line(lwd = 1) +
+    ggplot2::theme_minimal(base_size = 14) +
+    ggplot2::scale_x_continuous(
+      breaks = c(0, max(sims_long$time)),
+      labels = c("LCA", "present")
+      ) +
+    ggplot2::ylab("trait value (latent scale)") +
     ggplot2::xlab("time") +
-    ggplot2::theme(legend.title = ggplot2::element_blank(),  strip.background = ggplot2::element_blank(), panel.spacing.x = ggplot2::unit(1, "lines"), axis.text.x = ggplot2::element_blank(),  axis.ticks.x = ggplot2::element_blank()) + 
+    ggplot2::theme(
+      legend.title = ggplot2::element_blank(),
+      strip.background = ggplot2::element_blank(),
+      panel.spacing.x = ggplot2::unit(1, "lines"),
+      axis.text.x = ggplot2::element_blank(),
+      axis.ticks.x = ggplot2::element_blank()
+      ) +
     ggplot2::ggtitle("Predicted trait coevolution")
   }
   return(p)
-}
\ No newline at end of file
+}
diff --git a/R/coev_pred_series.R b/R/coev_pred_series.R
index 6568a84..ad3520b 100644
--- a/R/coev_pred_series.R
+++ b/R/coev_pred_series.R
@@ -1,19 +1,35 @@
-#' Predict a co-evolutionary time series from a fitted \code{coevfit}
-#' object
+#' Predict a co-evolutionary time series from a fitted \code{coevfit} object
+#'
+#' This function produces predicted values of traits over a time series, given
+#' the estimated parameters for a \code{coevfit} model. This function is used
+#' under the hood by the plotting function \code{\link{coev_pred_series}}.
 #'
 #' @param object An object of class \code{coevfit}
-#' @param eta_anc If \code{NULL} default, the starting values for the latent states \eqn{\eta} will be set to the estimated ancestral states at the root from the the \code{coevfit}. Otherwise, a numeric vector with length equal to the number of variable specifying the initial \eqn{\eta}  values.
-#' @param tmax A positive number indicating the total duration of time for which to predict. Set to 1 by default, corresponding to the entire time depth of the original phylogenetic tree(s).
-#' @param tmax A positive number indicating the total duration of time for which to predict. Set to 1 by default, corresponding to the entire time depth of the original phylogenetic tree(s).
-#' @param ntimes A positive integer indicating the total number of discrete time steps over which to make predictions. Each step corresponds to a time difference of dt = tmax/ntimes. Set to 30 by default.
+#' @param eta_anc If \code{NULL} (default), the starting values for the latent
+#'   states \eqn{\eta} will be set to the estimated ancestral states at the
+#'   root from the \code{coevfit} model. Otherwise, a named numeric vector with
+#'   length equal to the number of variables specifying the initial \eqn{\eta}
+#'   values. All variable names must be included in this vector.
+#' @param tmax A positive number indicating the total duration of time for which
+#'   to predict. Set to 1 by default, corresponding to the entire time depth of
+#'   the original phylogenetic tree(s).
+#' @param ntimes A positive integer indicating the total number of discrete time
+#'   steps over which to make predictions. Each step corresponds to a time
+#'   difference of dt = tmax/ntimes. Set to 30 by default.
 #' @param ndraws An integer indicating the number of draws to return. The
 #'   default and maximum number of draws is the size of the posterior sample.
-#' @param stochastic Logical (defaults to \code{FALSE}); indicator of whether predictions should include only the expected co-evolutionary change (deterministic selection), or also stochastic drift.
-#' @return An [ndraw, ntimes, nvariables] array of predicted \eqn{\eta}  values.
+#' @param stochastic Logical (defaults to \code{FALSE}); indicator of whether
+#'   predictions should include only the expected co-evolutionary change
+#'   due to deterministic selection (FALSE) or also stochastic drift (TRUE).
+#'
+#' @return An \[ndraw, ntimes, nvariables\] array of predicted \eqn{\eta}
+#'   values.
 #'
 #' @author Scott Claessens \email{scott.claessens@@gmail.com}, Erik Ringen
 #'   \email{erikjacob.ringen@@uzh.ch}
 #'
+#' @seealso \code{\link{coev_plot_pred_series}}
+#'
 #' @examples
 #' \dontrun{
 #' # fit dynamic coevolutionary model
@@ -34,34 +50,57 @@
 #' # simulate trait co-evolution
 #' sims <- coev_pred_series(
 #'   object = fit,
-#'   stochastic = T
+#'   stochastic = TRUE
 #'   )
-#' 
+#'
 #' # expected trait co-evolution, no drift
 #' epreds <- coev_pred_series(
 #'   object = fit,
-#'   stochastic = F
+#'   stochastic = FALSE
 #'   )
 #' }
 #'
 #' @export
-coev_pred_series <- function(object, eta_anc = NULL, tmax = 1, ntimes = 30, ndraws = NULL, stochastic = FALSE){
+coev_pred_series <- function(object, eta_anc = NULL, tmax = 1, ntimes = 30,
+                             ndraws = NULL, stochastic = FALSE) {
+  # stop if object is not of class coevfit
   if (!methods::is(object, "coevfit")) {
     stop2(
       paste0(
-        "Argument 'object' must be a fitted coevolutionary model of class coevfit."
+        "Argument 'object' must be a fitted coevolutionary model of class ",
+        "coevfit."
         )
       )
   }
-  # stop if eta_anc is neither NULL nor a vector of length equal to the number of variables
-  if (!is.null(eta_anc) & (!is.numeric(eta_anc) | length(eta_anc) != length(object$variables))) {
-    stop2("Argument 'eta_anc' must be numeric and equal in length to the number of variables.")
+  if (!is.null(eta_anc)) {
+    # stop if eta_anc is neither NULL nor
+    # a named vector of length equal to the number of variables
+    if (!is.numeric(eta_anc) | length(eta_anc) != length(object$variables)) {
+        stop2(
+          paste0(
+            "Argument 'eta_anc' must be numeric and equal in length to the ",
+            "number of variables."
+          )
+        )
+    }
+    # stop if eta_anc is not a named vector
+    if (is.null(names(eta_anc))) {
+      stop2("Argument 'eta_anc' is not a named vector.")
+    }
+    # stop if eta_anc does not have names equal to variables in the model
+    if (!identical(sort(names(eta_anc)), sort(names(object$variables)))) {
+      stop2(
+        paste0(
+          "Argument 'eta_anc' has names different to the variables included ",
+          "in the model."
+        )
+      )
+    }
   }
   # stop if tmax is not numeric or not positive
-  if (!is.numeric(tmax)|length(tmax) != 1) {
+  if (!is.numeric(tmax) | length(tmax) != 1) {
     stop2("Argument 'tmax' must be a single numeric value.")
-  }
-  else if (tmax <= 0) {
+  } else if (tmax <= 0) {
     stop2("Argument 'tmax' must be positive.")
   }
   # stop if ndraws is not a single integer between 1 and the total num draws
@@ -74,55 +113,74 @@ coev_pred_series <- function(object, eta_anc = NULL, tmax = 1, ntimes = 30, ndra
         )
     }
   }
-# stop if stochastic not logical
+  # stop if stochastic not logical
   if (!is.logical(stochastic)) {
     stop2("Argument 'stochastic' must be logical.")
   }
-
+  # get posterior samples and number of variables
   post <- extract_samples(object)
   J <- length(object$variables)
-
+  # get eta_anc in the correct order
+  eta_anc <- eta_anc[names(object$variables)]
+  # get number of samples
   nsamps <- ifelse(is.null(ndraws), length(post$lp__), ndraws)
-
-  preds <- array(NA, dim = c(nsamps, ntimes+1, J), dimnames = list(samps = 1:nsamps, time = 1:(ntimes+1), response = names(object$variables)))
-
-  # User defined ancestral states
+  # initialise empty array for predictions
+  preds <-
+    array(
+      NA,
+      dim = c(nsamps, ntimes + 1, J),
+      dimnames = list(
+        samps = 1:nsamps,
+        time = 1:(ntimes + 1),
+        response = names(object$variables)
+        )
+      )
+  # if user defined ancestral states
+  # otherwise, use model inferred ancestral states
   if (!is.null(eta_anc)) {
-    for (i in 1:nsamps) preds[i,1,] = eta_anc
-  }
-
-  # Model inferred ancestral states
-  else {
+    for (i in 1:nsamps) preds[i, 1, ] = eta_anc
+  } else {
+    # set posterior ancestral states
     eta_anc_long <- post$eta_anc
     ntrees <- dim(eta_anc_long)[2]
-    # shuffle dimenisions of tree dimension to get random draws
-    eta_anc_long <- eta_anc_long[,sample(1:ntrees, size = ntrees, replace = F),]
-
+    # shuffle dimensions of tree dimension to get random draws
+    eta_anc_long <-
+      eta_anc_long[, sample(1:ntrees, size = ntrees, replace = FALSE), ]
     if (ntrees > 1) {
-    # make into 2d matrix by stacking trees
-    eta_anc_long2 <- eta_anc_long[,1,]
-    for (t in 2:ntrees) eta_anc_long2 <- rbind(eta_anc_long2, eta_anc_long[,t,]) -> eta_anc_long
+      # make into 2d matrix by stacking trees
+      eta_anc_long2 <- eta_anc_long[, 1, ]
+      for (t in 2:ntrees) {
+        eta_anc_long2 <- rbind(eta_anc_long2, eta_anc_long[,t,])
+        eta_anc_long <- eta_anc_long2
+      }
     }
-    for (i in 1:nsamps) preds[i,1,] = eta_anc_long[i,]
+    for (i in 1:nsamps) {
+      preds[i, 1, ] = eta_anc_long[i, ]
     }
-
+  }
+  # calculate predictions
   for (i in 1:nsamps) {
+    # selection parameters
     A <- post$A[i,,]
-    A_delta <- as.matrix(Matrix::expm(A * tmax/ntimes))
+    A_delta <- as.matrix(Matrix::expm(A * tmax / ntimes))
     inv_A <- solve(A)
     I <- diag(rep(1, J))
     b <- post$b[i,]
-
+    # drift parameters
     Q_inf <- post$Q_inf[i,,]
     VCV <- Q_inf - ((A_delta) %*% Q_inf %*% t(A_delta))
     chol_VCV <- t(chol(VCV))
-
+    # calculate predictions over time
     for (t in 1:ntimes) {
-      preds[i, t+1, ] = (A_delta %*% preds[i, t, ] + (inv_A %*% (A_delta - I) %*% b))[,1]
+      preds[i, t + 1, ] <- (A_delta %*% preds[i, t, ] +
+                              (inv_A %*% (A_delta - I) %*% b))[,1]
       # add drift
-      if (stochastic == T) preds[i, t+1, ] = preds[i, t+1, ] + chol_VCV %*% rnorm(J, 0, 1)
+      if (stochastic == TRUE) {
+        preds[i, t + 1, ] <-
+          preds[i, t + 1, ] + (chol_VCV %*% stats::rnorm(J, 0, 1))
+      }
     }
   }
   # shuffle draws to de-correlate when doing simulations
-  return(preds[sample(1:nsamps, size = nsamps, replace = F),,])
+  preds[sample(1:nsamps, size = nsamps, replace = FALSE), , ]
 }
diff --git a/man/coev_plot_pred_series.Rd b/man/coev_plot_pred_series.Rd
index 9fff0da..c2d86c6 100644
--- a/man/coev_plot_pred_series.Rd
+++ b/man/coev_plot_pred_series.Rd
@@ -2,7 +2,7 @@
 % Please edit documentation in R/coev_plot_pred_series.R
 \name{coev_plot_pred_series}
 \alias{coev_plot_pred_series}
-\title{Plot a predicted co-evolutionary time series from a fitted \code{coevfit}
+\title{Plot a predicted coevolutionary time series from a fitted \code{coevfit}
 object}
 \usage{
 coev_plot_pred_series(object, prob = 0.95, ...)
@@ -19,8 +19,13 @@ to be covered by the uncertainty intervals. The default is 0.95.}
 A ggplot object.
 }
 \description{
-Plot a predicted co-evolutionary time series from a fitted \code{coevfit}
-object
+This function plots a predicted coevolutionary time series using the
+estimated parameters from a fitted \code{coevfit} model. By default, the
+plot uses the posterior ancestral states estimated by the model as the
+starting values, but users can also set their own starting values for
+traits. Plots can be generated with or without stochastic drift. For more
+details on the underlying predictive function, see
+\code{help(coev_pred_series)}.
 }
 \examples{
 \dontrun{
@@ -42,16 +47,19 @@ fit <- coev_fit(
 # simulated trait co-evolution
 coev_plot_pred_series(
   object = fit,
-  stochastic = T
+  stochastic = TRUE
   )
 
 # expected trait co-evolution, no drift
 coev_plot_pred_series(
   object = fit,
-  stochastic = F
+  stochastic = FALSE
   )
 }
 
+}
+\seealso{
+\code{\link{coev_pred_series}}
 }
 \author{
 Scott Claessens \email{scott.claessens@gmail.com}, Erik Ringen
diff --git a/man/coev_pred_series.Rd b/man/coev_pred_series.Rd
index f9466cc..a4b42ab 100644
--- a/man/coev_pred_series.Rd
+++ b/man/coev_pred_series.Rd
@@ -2,8 +2,7 @@
 % Please edit documentation in R/coev_pred_series.R
 \name{coev_pred_series}
 \alias{coev_pred_series}
-\title{Predict a co-evolutionary time series from a fitted \code{coevfit}
-object}
+\title{Predict a co-evolutionary time series from a fitted \code{coevfit} object}
 \usage{
 coev_pred_series(
   object,
@@ -17,23 +16,35 @@ coev_pred_series(
 \arguments{
 \item{object}{An object of class \code{coevfit}}
 
-\item{eta_anc}{If \code{NULL} default, the starting values for the latent states \eqn{\eta} will be set to the estimated ancestral states at the root from the the \code{coevfit}. Otherwise, a numeric vector with length equal to the number of variable specifying the initial \eqn{\eta}  values.}
+\item{eta_anc}{If \code{NULL} (default), the starting values for the latent
+states \eqn{\eta} will be set to the estimated ancestral states at the
+root from the \code{coevfit} model. Otherwise, a named numeric vector with
+length equal to the number of variables specifying the initial \eqn{\eta}
+values. All variable names must be included in this vector.}
 
-\item{tmax}{A positive number indicating the total duration of time for which to predict. Set to 1 by default, corresponding to the entire time depth of the original phylogenetic tree(s).}
+\item{tmax}{A positive number indicating the total duration of time for which
+to predict. Set to 1 by default, corresponding to the entire time depth of
+the original phylogenetic tree(s).}
 
-\item{ntimes}{A positive integer indicating the total number of discrete time steps over which to make predictions. Each step corresponds to a time difference of dt = tmax/ntimes. Set to 30 by default.}
+\item{ntimes}{A positive integer indicating the total number of discrete time
+steps over which to make predictions. Each step corresponds to a time
+difference of dt = tmax/ntimes. Set to 30 by default.}
 
 \item{ndraws}{An integer indicating the number of draws to return. The
 default and maximum number of draws is the size of the posterior sample.}
 
-\item{stochastic}{Logical (defaults to \code{FALSE}); indicator of whether predictions should include only the expected co-evolutionary change (deterministic selection), or also stochastic drift.}
+\item{stochastic}{Logical (defaults to \code{FALSE}); indicator of whether
+predictions should include only the expected co-evolutionary change
+due to deterministic selection (FALSE) or also stochastic drift (TRUE).}
 }
 \value{
-An \link{ndraw, ntimes, nvariables} array of predicted \eqn{\eta}  values.
+An [ndraw, ntimes, nvariables] array of predicted \eqn{\eta}
+values.
 }
 \description{
-Predict a co-evolutionary time series from a fitted \code{coevfit}
-object
+This function produces predicted values of traits over a time series, given
+the estimated parameters for a \code{coevfit} model. This function is used
+under the hood by the plotting function \code{\link{coev_pred_series}}.
 }
 \examples{
 \dontrun{
@@ -55,16 +66,19 @@ fit <- coev_fit(
 # simulate trait co-evolution
 sims <- coev_pred_series(
   object = fit,
-  stochastic = T
+  stochastic = TRUE
   )
 
 # expected trait co-evolution, no drift
 epreds <- coev_pred_series(
   object = fit,
-  stochastic = F
+  stochastic = FALSE
   )
 }
 
+}
+\seealso{
+\code{\link{coev_plot_pred_series}}
 }
 \author{
 Scott Claessens \email{scott.claessens@gmail.com}, Erik Ringen
diff --git a/tests/testthat/test-coev_plot_pred_series.R b/tests/testthat/test-coev_plot_pred_series.R
index a05274c..e13e31f 100644
--- a/tests/testthat/test-coev_plot_pred_series.R
+++ b/tests/testthat/test-coev_plot_pred_series.R
@@ -21,13 +21,13 @@ test_that("coev_plot_pred_series() produces expected errors and output", {
   # expect the following errors
   expect_error(
     coev_plot_pred_series(object = "fail"),
-     "Argument 'object' must be a fitted coevolutionary model of class coevfit."
+    "Argument 'object' must be a fitted coevolutionary model of class coevfit.",
+    fixed = TRUE
   )
   expect_error(
     coev_plot_pred_series(object = m1, prob = 95),
-    paste0(
-      "Argument 'prob' is not between 0 and 1."
-    )
+    "Argument 'prob' is not between 0 and 1.",
+    fixed = TRUE
   )
   # suppress warnings
   SW <- suppressWarnings
@@ -41,15 +41,15 @@ test_that("coev_plot_pred_series() produces expected errors and output", {
   expect_no_error(SW(coev_plot_pred_series(m7)))
   expect_no_error(SW(coev_plot_pred_series(m8)))
   expect_no_error(SW(coev_plot_pred_series(m9)))
-  expect_no_error(SW(coev_plot_pred_series(m1, stochastic = T)))
-  expect_no_error(SW(coev_plot_pred_series(m2, stochastic = T)))
-  expect_no_error(SW(coev_plot_pred_series(m3, stochastic = T)))
-  expect_no_error(SW(coev_plot_pred_series(m4, stochastic = T)))
-  expect_no_error(SW(coev_plot_pred_series(m5, stochastic = T)))
-  expect_no_error(SW(coev_plot_pred_series(m6, stochastic = T)))
-  expect_no_error(SW(coev_plot_pred_series(m7, stochastic = T)))
-  expect_no_error(SW(coev_plot_pred_series(m8, stochastic = T)))
-  expect_no_error(SW(coev_plot_pred_series(m9, stochastic = T)))
+  expect_no_error(SW(coev_plot_pred_series(m1, stochastic = TRUE)))
+  expect_no_error(SW(coev_plot_pred_series(m2, stochastic = TRUE)))
+  expect_no_error(SW(coev_plot_pred_series(m3, stochastic = TRUE)))
+  expect_no_error(SW(coev_plot_pred_series(m4, stochastic = TRUE)))
+  expect_no_error(SW(coev_plot_pred_series(m5, stochastic = TRUE)))
+  expect_no_error(SW(coev_plot_pred_series(m6, stochastic = TRUE)))
+  expect_no_error(SW(coev_plot_pred_series(m7, stochastic = TRUE)))
+  expect_no_error(SW(coev_plot_pred_series(m8, stochastic = TRUE)))
+  expect_no_error(SW(coev_plot_pred_series(m9, stochastic = TRUE)))
   expect_no_error(SW(coev_plot_pred_series(m1, ndraws = 1L)))
   expect_no_error(SW(coev_plot_pred_series(m2, ndraws = 1L)))
   expect_no_error(SW(coev_plot_pred_series(m3, ndraws = 1L)))
diff --git a/tests/testthat/test-coev_pred_series.R b/tests/testthat/test-coev_pred_series.R
index 6fc0a22..9102ce8 100644
--- a/tests/testthat/test-coev_pred_series.R
+++ b/tests/testthat/test-coev_pred_series.R
@@ -21,43 +21,60 @@ test_that("coev_pred_series() produces expected errors and output", {
   # expect the following errors
   expect_error(
     coev_pred_series(object = "fail"),
-     "Argument 'object' must be a fitted coevolutionary model of class coevfit."
+    "Argument 'object' must be a fitted coevolutionary model of class coevfit.",
+    fixed = TRUE
   )
   expect_error(
     coev_pred_series(object = m1, eta_anc = c(1, "LCA")),
     paste0(
-      "Argument 'eta_anc' must be numeric and equal in length to the number of variables."
-    )
+      "Argument 'eta_anc' must be numeric and equal in length to the number ",
+      "of variables."
+    ),
+    fixed = TRUE
   )
   expect_error(
-    coev_pred_series(object = m1, tmax = -1),
+    coev_pred_series(object = m2, eta_anc = c(1, 2)),
+    "Argument 'eta_anc' is not a named vector.",
+    fixed = TRUE
+  )
+  expect_error(
+    coev_pred_series(object = m2, eta_anc = c("a" = 1, "b" = 2)),
     paste0(
-      "Argument 'tmax' must be positive."
-    )
+      "Argument 'eta_anc' has names different to the variables included ",
+      "in the model."
+    ),
+    fixed = TRUE
+  )
+  expect_error(
+    coev_pred_series(object = m1, tmax = -1),
+    "Argument 'tmax' must be positive.",
+    fixed = TRUE
   )
   expect_error(
     coev_pred_series(object = m1, ndraws = "fail"),
-    "Argument 'ndraws' must be a single integer."
+    "Argument 'ndraws' must be a single integer.",
+    fixed = TRUE
   )
   expect_error(
     coev_pred_series(object = m1, ndraws = 0L),
-    "Argument 'ndraws' must be between 1 and the total number of draws."
+    "Argument 'ndraws' must be between 1 and the total number of draws.",
+    fixed = TRUE
   )
   expect_error(
     coev_pred_series(
       object = m1, ndraws = as.integer(nrow(m1$fit$draws()) + 1)
       ),
-    "Argument 'ndraws' must be between 1 and the total number of draws."
+    "Argument 'ndraws' must be between 1 and the total number of draws.",
+    fixed = TRUE
   )
   expect_error(
     coev_pred_series(object = m1, stochastic = "only"),
-    paste0(
-      "Argument 'stochastic' must be logical."
-    )
+    "Argument 'stochastic' must be logical.",
+    fixed = TRUE
   )
   # suppress warnings
   SW <- suppressWarnings
-  # should run without error and produce list of ggplot objects
+  # should run without error
   expect_no_error(SW(coev_pred_series(m1)))
   expect_no_error(SW(coev_pred_series(m2)))
   expect_no_error(SW(coev_pred_series(m3)))
@@ -67,15 +84,15 @@ test_that("coev_pred_series() produces expected errors and output", {
   expect_no_error(SW(coev_pred_series(m7)))
   expect_no_error(SW(coev_pred_series(m8)))
   expect_no_error(SW(coev_pred_series(m9)))
-  expect_no_error(SW(coev_pred_series(m1, stochastic = T)))
-  expect_no_error(SW(coev_pred_series(m2, stochastic = T)))
-  expect_no_error(SW(coev_pred_series(m3, stochastic = T)))
-  expect_no_error(SW(coev_pred_series(m4, stochastic = T)))
-  expect_no_error(SW(coev_pred_series(m5, stochastic = T)))
-  expect_no_error(SW(coev_pred_series(m6, stochastic = T)))
-  expect_no_error(SW(coev_pred_series(m7, stochastic = T)))
-  expect_no_error(SW(coev_pred_series(m8, stochastic = T)))
-  expect_no_error(SW(coev_pred_series(m9, stochastic = T)))
+  expect_no_error(SW(coev_pred_series(m1, stochastic = TRUE)))
+  expect_no_error(SW(coev_pred_series(m2, stochastic = TRUE)))
+  expect_no_error(SW(coev_pred_series(m3, stochastic = TRUE)))
+  expect_no_error(SW(coev_pred_series(m4, stochastic = TRUE)))
+  expect_no_error(SW(coev_pred_series(m5, stochastic = TRUE)))
+  expect_no_error(SW(coev_pred_series(m6, stochastic = TRUE)))
+  expect_no_error(SW(coev_pred_series(m7, stochastic = TRUE)))
+  expect_no_error(SW(coev_pred_series(m8, stochastic = TRUE)))
+  expect_no_error(SW(coev_pred_series(m9, stochastic = TRUE)))
   expect_no_error(SW(coev_pred_series(m1, ndraws = 1L)))
   expect_no_error(SW(coev_pred_series(m2, ndraws = 1L)))
   expect_no_error(SW(coev_pred_series(m3, ndraws = 1L)))