R/temporal_core.R

In tulpa: Templated Unified Library for Posterior Approximation in Bayesian Hierarchical Models

#' RW1 temporal structure (first-order random walk)
#'
#' @description
#' Specify a first-order random walk temporal random effect. RW1 penalizes
#' first differences, so adjacent time points are smoothed toward each other:
#' `phi[t] - phi[t-1] ~ N(0, sigma^2)`.
#'
#' @param time_var A formula (`~ time`) or single character string naming the
#'   time variable in the data.
#' @param group_var Optional formula (`~ g`) or character string naming a
#'   grouping variable. When supplied, a separate random walk is fit per group;
#'   `NULL` (default) fits a single walk shared across all observations.
#' @param cyclic Logical. If `TRUE`, the random walk wraps around so the last
#'   time point is a neighbour of the first (cyclic boundary, e.g. month of
#'   year). Default `FALSE`.
#' @param shared Whether the temporal effect is shared across processes in a
#'   multi-process model. `NULL` (default) shares the effect; `FALSE` fits
#'   process-specific effects and emits a warning about unshared confounding.
#'
#' @return A `tulpa_temporal` object.
#'
#' @details
#' The precision matrix is rank `T - 1` (one constraint needed). RW1 is the
#' least smooth of the random-walk priors; for smoother trends see
#' [temporal_rw2()], and for a stationary alternative see [temporal_ar1()].
#'
#' @examples
#' # Create temporal RW1 specification
#' temporal_rw1("year")
#' temporal_rw1("month", cyclic = TRUE)
#'
#' @seealso [temporal_rw2()], [temporal_ar1()] for other temporal priors.
#'
#' @export
temporal_rw1 <- function(time_var, group_var = NULL, cyclic = FALSE,
                         shared = NULL) {

  time_var <- .coerce_var_arg(time_var, "time_var", "~ time")
  if (is.null(time_var)) {
    stop("`time_var` must be a formula (~ time) or single character string",
         call. = FALSE)
  }
  group_var <- .coerce_var_arg(group_var, "group_var")

  if (isFALSE(shared)) .warn_nonshared("temporal effects")

 structure(
   list(
     type = "rw1",
     time_var = time_var,
     group_var = group_var,
     cyclic = cyclic,
     shared = shared,
     # These will be filled in when validated against data
     n_times = NULL,
     n_groups = NULL,
     time_index = NULL,
     group_index = NULL
   ),
   class = c("tulpa_temporal", "list")
 )
}


#' RW2 temporal structure (second-order random walk)
#'
#' @description
#' Specify a second-order random walk temporal random effect.
#' RW2 penalizes deviations from linearity, resulting in smoother trends
#' than RW1: `phi[t] - 2*phi[t-1] + phi[t-2] ~ N(0, sigma^2)`.
#'
#' @inheritParams temporal_rw1
#'
#' @return A `tulpa_temporal` object
#'
#' @details
#' RW2 produces smoother curves than RW1 because it penalizes the second
#' derivative (curvature) rather than the first derivative (slope).
#' It requires at least 3 time points.
#'
#' The precision matrix is rank T-2 (two constraints needed).
#'
#' @examples
#' # Create temporal RW2 specification
#' temporal_rw2("year")
#'
#' \donttest{
#' # Smooth temporal trend
#' set.seed(126)
#' df <- data.frame(
#'   year = rep(1:20, each = 4),
#'   x = rnorm(80)
#' )
#' trend <- sin(seq(0, 2, length.out = 20))
#' df$count <- rpois(80, exp(1 + 0.3 * df$x + trend[df$year]))
#'
#' fit <- tulpa(
#'   count ~ x,
#'   data = df,
#'   family = "poisson",
#'   temporal = temporal_rw2("year"),
#'   mode = "auto"
#' )
#' summary(fit)
#' }
#'
#' @export
temporal_rw2 <- function(time_var, group_var = NULL, cyclic = FALSE,
                         shared = NULL) {

  time_var <- .coerce_var_arg(time_var, "time_var", "~ time")
  if (is.null(time_var)) {
    stop("`time_var` must be a formula (~ time) or single character string",
         call. = FALSE)
  }
  group_var <- .coerce_var_arg(group_var, "group_var")

  if (isFALSE(shared)) .warn_nonshared("temporal effects")

 structure(
   list(
     type = "rw2",
     time_var = time_var,
     group_var = group_var,
     cyclic = cyclic,
     shared = shared,
     n_times = NULL,
     n_groups = NULL,
     time_index = NULL,
     group_index = NULL
   ),
   class = c("tulpa_temporal", "list")
 )
}


#' AR1 temporal structure (first-order autoregressive)
#'
#' @description
#' Specify a first-order autoregressive temporal random effect.
#' AR1 models temporal correlation where each time point depends on
#' the previous one: `phi[t] = rho * phi[t-1] + epsilon[t]`.
#'
#' Unlike RW1/RW2, AR1 is a proper (stationary) model with an estimated
#' autocorrelation parameter rho.
#'
#' @inheritParams temporal_rw1
#' @param rho_prior Prior for the AR(1) autocorrelation. Default `NULL` is a
#'   Uniform(-1, 1) prior. Supply a `prior_beta(alpha, beta)` to place a
#'   Beta(alpha, beta) prior on `u = (rho + 1)/2` for a more informative prior
#'   (e.g. `prior_beta(5, 2)` favours positive autocorrelation). Honoured on
#'   both the exact-sampler and nested-Laplace paths.
#'
#' @return A `tulpa_temporal` object
#'
#' @details
#' The AR1 process has marginal variance sigma^2 / (1 - rho^2) and
#' correlation between time points t and s of rho^|t-s|.
#'
#' The precision matrix is tridiagonal and full rank, so no constraints
#' are needed.
#'
#' @examples
#' # Create temporal AR1 specification
#' temporal_ar1("year")
#' temporal_ar1("year", group = "site")
#'
#' \donttest{
#' # AR1 temporal correlation
#' set.seed(127)
#' df <- data.frame(
#'   year = rep(1:20, each = 3),
#'   x = rnorm(60)
#' )
#' f <- as.numeric(arima.sim(list(ar = 0.7), 20, sd = 0.4))
#' df$count <- rpois(60, exp(1 + 0.3 * df$x + f[df$year]))
#'
#' fit <- tulpa(
#'   count ~ x,
#'   data = df,
#'   family = "poisson",
#'   temporal = temporal_ar1("year"),
#'   mode = "auto"
#' )
#' summary(fit)
#' }
#'
#' @export
temporal_ar1 <- function(time_var, group_var = NULL, shared = NULL,
                         rho_prior = NULL) {

  time_var <- .coerce_var_arg(time_var, "time_var", "~ time")
  if (is.null(time_var)) {
    stop("`time_var` must be a formula (~ time) or single character string",
         call. = FALSE)
  }
  group_var <- .coerce_var_arg(group_var, "group_var")
  if (isFALSE(shared)) .warn_nonshared("temporal effects")
  if (!is.null(rho_prior) && !inherits(rho_prior, "tulpa_prior_beta")) {
    stop("`rho_prior` must be a Beta prior (prior_beta(alpha, beta)) on ",
         "u = (rho + 1)/2, or NULL for the default Uniform(-1, 1).",
         call. = FALSE)
  }

  structure(
   list(
     type = "ar1",
     time_var = time_var,
     group_var = group_var,
     cyclic = FALSE,  # AR1 is not cyclic
     shared = shared,
     rho_prior = rho_prior,
     n_times = NULL,
     n_groups = NULL,
     time_index = NULL,
     group_index = NULL
   ),
   class = c("tulpa_temporal", "list")
 )
}

# Beta(a, b) shape parameters for the AR1 rho prior on u = (rho + 1)/2. NULL (the
# default) is the Uniform(-1, 1) prior = Beta(1, 1). Single source for the
# sampler (ModelData ar1_rho_prior_a/b) and the nested-Laplace outer reweight.
.ar1_rho_beta_ab <- function(rho_prior) {
  if (is.null(rho_prior)) return(c(1, 1))
  c(rho_prior$alpha, rho_prior$beta)
}


#' Print method for tulpa_temporal
#'
#' @param x A tulpa_temporal object
#' @param ... Ignored
#'
#' @return The input `x`, returned invisibly. Called for the side effect of
#'   printing the temporal specification to the console.
#'
#' @noRd
#' @export
print.tulpa_temporal <- function(x, ...) {
 cat("tulpa temporal specification\n")
 cat("============================\n\n")

 type_name <- switch(x$type,
   rw1 = "RW1 (First-order Random Walk)",
   rw2 = "RW2 (Second-order Random Walk)",
   ar1 = "AR1 (First-order Autoregressive)",
   x$type
 )

 cat("Type:", type_name, "\n")
 cat("Time variable:", x$time_var, "\n")

 if (!is.null(x$group_var)) {
   cat("Group variable:", x$group_var, "\n")
 }

 if (x$type %in% c("rw1", "rw2") && x$cyclic) {
   cat("Cyclic: Yes\n")
 }

 cat("Shared:", if (!isFALSE(x$shared)) "Yes (enters both processes)" else "No", "\n")

 if (!is.null(x$n_times)) {
   cat("Time points:", x$n_times, "\n")
 }
 if (!is.null(x$n_groups) && x$n_groups > 1) {
   cat("Groups:", x$n_groups, "\n")
 }

 invisible(x)
}


# The time column of a temporal field as a factor whose level order IS the
# chain order, so `as.integer()` of it is each observation's node. The one
# resolver behind every temporal door (validate_temporal(),
# validate_temporal_multiscale(), validate_tvc(), the inline temporal() field):
#   * factor            -> its level order, which the caller stated;
#   * numeric / Date    -> sorted by value;
#   * character labels that all parse as numbers
#                       -> sorted by the number they spell, so "10" follows
#                          "9" rather than "1";
#   * other character   -> refused: a character vector states no order, and
#                          the sort order of its labels is not a time order
#                          ("Apr" < "Feb" < "Jan"). It built the walk out of
#                          order silently (gcol33/tulpa#900).
.resolve_time_index <- function(time_vals, time_var = "time") {
  if (anyNA(time_vals)) {
    stop("Temporal variable '", time_var, "' has missing values.",
         call. = FALSE)
  }
  if (is.factor(time_vals)) return(time_vals)
  if (is.character(time_vals)) {
    num <- suppressWarnings(as.numeric(time_vals))
    if (anyNA(num)) {
      stop("Temporal variable '", time_var, "' holds character labels that ",
           "are not numbers, so they carry no time order (their sort order ",
           "is alphabetical). Pass a numeric or Date column, or a factor ",
           "whose levels are in time order.", call. = FALSE)
    }
    lev <- unique(time_vals[order(num)])
    return(factor(time_vals, levels = lev))
  }
  factor(time_vals, levels = sort(unique(time_vals)))
}


#' Validate temporal specification against data
#'
#' @param temporal tulpa_temporal object
#' @param data Data frame
#'
#' @return Updated tulpa_temporal object with indices computed
#' @noRd
validate_temporal <- function(temporal, data) {
 if (is.null(temporal)) return(NULL)

 # Check time variable exists
 if (!(temporal$time_var %in% names(data))) {
   stop(sprintf("Temporal variable '%s' not found in data",
                temporal$time_var), call. = FALSE)
 }

 # Check group variable exists if specified
 if (!is.null(temporal$group_var)) {
   if (!(temporal$group_var %in% names(data))) {
     stop(sprintf("Temporal group variable '%s' not found in data",
                  temporal$group_var), call. = FALSE)
   }
 }

 .check_time_complete(data[[temporal$time_var]], temporal$time_var)
 time_factor <- .resolve_time_index(data[[temporal$time_var]],
                                    temporal$time_var)

 temporal$n_times <- nlevels(time_factor)
 temporal$time_index <- as.integer(time_factor)
 temporal$time_levels <- levels(time_factor)

 # Check minimum time points for RW2
 if (temporal$type == "rw2" && temporal$n_times < 3) {
   stop("RW2 requires at least 3 time points", call. = FALSE)
 }
 # A chain over a single time point has no increment to penalise (RW1) and no
 # lag to correlate (AR1); it used to surface as an internal "unvalidated
 # spec" error from the nested kernel (gcol33/tulpa#909).
 if (temporal$type %in% c("rw1", "ar1") && temporal$n_times < 2) {
   stop(sprintf(paste0("%s requires at least 2 distinct time points; '%s' has ",
                       "%d."), toupper(temporal$type), temporal$time_var,
                temporal$n_times), call. = FALSE)
 }

 # Handle grouping
 if (!is.null(temporal$group_var)) {
   group_vals <- data[[temporal$group_var]]
   group_factor <- as.factor(group_vals)
   temporal$n_groups <- nlevels(group_factor)
   temporal$group_index <- as.integer(group_factor)
   temporal$group_levels <- levels(group_factor)

   # Total temporal parameters = n_times * n_groups
   temporal$n_temporal_params <- temporal$n_times * temporal$n_groups
 } else {
   temporal$n_groups <- 1L
   temporal$group_index <- rep(1L, nrow(data))
   temporal$n_temporal_params <- temporal$n_times
 }

 temporal
}

# Every observation of a temporal field needs a time: a missing one has no
# position on the chain, and it used to reach a kernel as an NA index (an RW
# field) or as "TVC `time_index[3]` is -2147483648" (a TVC; gcol33/tulpa#909).
# Shared by validate_temporal() and validate_tvc().
.check_time_complete <- function(time_vals, time_var) {
  if (anyNA(time_vals)) {
    stop(sprintf(paste0("Time variable '%s' has %d missing value(s) (first at ",
                        "row %d); every observation needs a time."),
                 time_var, sum(is.na(time_vals)),
                 which(is.na(time_vals))[1L]), call. = FALSE)
  }
  invisible(TRUE)
}