R/temporal_gp.R

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

#' Multi-scale temporal structure
#'
#' @description
#' Specify a temporal random effect that decomposes variation into separate
#' scales: a smooth trend, an optional seasonal cycle, and a short-term
#' component. Each scale uses its own prior, letting slow and fast dynamics be
#' modelled jointly.
#'
#' @param time_var Single character string naming the time variable in the data.
#' @param trend Prior for the smooth long-term trend. One of `"rw2"`, `"rw1"`,
#'   or `"none"`.
#' @param seasonal Optional integer period (`>= 2`) of a seasonal cycle, e.g.
#'   `12` for monthly data with an annual cycle. `NULL` (default) omits the
#'   seasonal component.
#' @param short_term Prior for the short-term component. One of `"ar1"`,
#'   `"iid"`, or `"none"`.
#' @param group_var Optional character string naming a grouping variable for
#'   group-specific temporal effects.
#' @param shared Whether the effect is shared across processes in a
#'   multi-process model. `NULL` (default) shares it; `FALSE` fits
#'   process-specific effects and emits a warning.
#'
#' @return A `tulpa_temporal_multiscale` object.
#'
#' @details
#' At least one of `trend`, `seasonal`, or `short_term` must be active.
#'
#' @seealso [temporal_rw1()], [temporal_rw2()], [temporal_ar1()] for
#'   single-scale temporal priors.
#'
#' @examples
#' # Trend + annual seasonal cycle + AR1 short-term component on monthly data
#' temporal_multiscale("month", trend = "rw2", seasonal = 12, short_term = "ar1")
#'
#' @export
temporal_multiscale <- function(time_var,
                                trend = c("rw2", "rw1", "none"),
                                seasonal = NULL,
                                short_term = c("ar1", "iid", "none"),
                                group_var = NULL,
                                shared = NULL) {

  trend <- match.arg(trend)
  short_term <- match.arg(short_term)

  if (!is.character(time_var) || length(time_var) != 1) {
    stop("`time_var` must be a single character string", call. = FALSE)
  }

  if (!is.null(group_var)) {
    if (!is.character(group_var) || length(group_var) != 1) {
      stop("`group_var` must be a single character string", call. = FALSE)
    }
  }

  # Validate seasonal
  if (!is.null(seasonal)) {
    if (!is.numeric(seasonal) || length(seasonal) != 1 || seasonal < 2) {
      stop("`seasonal` must be an integer >= 2 (the period)", call. = FALSE)
    }
    seasonal <- as.integer(seasonal)
  }

  # Check at least one component is active
  if (trend == "none" && is.null(seasonal) && short_term == "none") {
    stop("At least one temporal component must be active", call. = FALSE)
  }

  # Determine which components are active
  components <- character(0)
  if (trend != "none") components <- c(components, "trend")
  if (!is.null(seasonal) && seasonal >= 2) components <- c(components, "seasonal")
  if (short_term != "none") components <- c(components, "short_term")

  structure(
    list(
      type = "multiscale",
      time_var = time_var,
      group_var = group_var,
      shared = shared,
      # Component specifications
      trend = trend,
      seasonal = seasonal,
      short_term = short_term,
      components = components,
      # Filled in during validation
      n_times = NULL,
      n_groups = NULL,
      time_index = NULL,
      group_index = NULL,
      precision_structures = NULL
    ),
    class = c("tulpa_temporal_multiscale", "tulpa_temporal", "list")
  )
}


#' @rdname temporal_multiscale
#' @param x A `tulpa_temporal_multiscale` specification.
#' @param ... Ignored.
#' @export
print.tulpa_temporal_multiscale <- function(x, ...) {
  cat("tulpa Multi-Scale temporal specification\n")
  cat("=========================================\n\n")

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

  cat("\nComponents:\n")

  if (x$trend != "none") {
    trend_name <- switch(x$trend,
      rw1 = "RW1 (first-order random walk)",
      rw2 = "RW2 (second-order random walk, smooth)"
    )
    cat("  Trend:", trend_name, "\n")
  }

  if (!is.null(x$seasonal) && x$seasonal >= 2) {
    cat("  Seasonal: Period =", x$seasonal, "(cyclic RW1)\n")
  }

  if (x$short_term != "none") {
    short_name <- switch(x$short_term,
      ar1 = "AR(1) (autocorrelated residuals)",
      iid = "IID (independent residuals)"
    )
    cat("  Short-term:", short_name, "\n")
  }

  cat("\nShared:", 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)
}


#' Validate a multi-scale temporal specification against data
#'
#' Fills in the time index and point count that [tulpa_gibbs()] needs from a
#' `temporal =` spec built with [temporal_multiscale()], since that
#' design-matrix door takes no `data` argument to validate one itself.
#' [tulpa()] calls this internally, so it only needs to be called directly
#' when fitting through a design-matrix door.
#'
#' @param temporal A [temporal_multiscale()] object (or a single-component
#'   temporal spec, dispatched to `validate_temporal()`).
#' @param data Data frame containing the time column named in `temporal`.
#'
#' @return Updated object with indices computed.
#' @seealso [validate_gp()], [temporal_multiscale()]
#' @examples
#' tm <- temporal_multiscale("t", trend = "rw1", seasonal = 12)
#' tm <- validate_temporal_multiscale(tm, data.frame(t = 1:36))
#' class(tm)
#' @export
validate_temporal_multiscale <- function(temporal, data) {
  if (is.null(temporal)) return(NULL)
  if (!inherits(temporal, "tulpa_temporal_multiscale")) {
    # Use regular validation for single-component temporal
    return(validate_temporal(temporal, data))
  }

  assert_columns_exist(temporal$time_var, data, role = "Temporal")
  if (!is.null(temporal$group_var)) {
    assert_columns_exist(temporal$group_var, data, role = "Temporal group")
  }

  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)

  # Validate minimum time points
  if (temporal$trend == "rw2" && temporal$n_times < 3) {
    stop("RW2 trend requires at least 3 time points", call. = FALSE)
  }

  if (!is.null(temporal$seasonal) && temporal$n_times < temporal$seasonal) {
    warning(sprintf(
      "Number of time points (%d) is less than seasonal period (%d).\n",
      temporal$n_times, temporal$seasonal
    ), "Seasonal component may not be well-identified.", 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)
  } else {
    temporal$n_groups <- 1L
    temporal$group_index <- rep(1L, nrow(data))
  }

  # Build precision structures for each component
  temporal$precision_structures <- list()

  if (temporal$trend != "none") {
    temporal$precision_structures$trend <- list(
      type = temporal$trend,
      n_times = temporal$n_times,
      cyclic = FALSE,
      rank_deficiency = if (temporal$trend == "rw1") 1 else 2
    )
  }

  if (!is.null(temporal$seasonal) && temporal$seasonal >= 2) {
    temporal$precision_structures$seasonal <- list(
      type = "rw1",
      n_times = temporal$seasonal,
      cyclic = TRUE,
      rank_deficiency = 1
    )
  }

  if (temporal$short_term == "ar1") {
    temporal$precision_structures$short_term <- list(
      type = "ar1",
      n_times = temporal$n_times,
      rank_deficiency = 0
    )
  } else if (temporal$short_term == "iid") {
    temporal$precision_structures$short_term <- list(
      type = "iid",
      n_times = temporal$n_times,
      rank_deficiency = 0
    )
  }

  # Calculate total number of temporal parameters
  n_params <- 0
  if (temporal$trend != "none") {
    n_params <- n_params + temporal$n_times
  }
  if (!is.null(temporal$seasonal) && temporal$seasonal >= 2) {
    n_params <- n_params + temporal$seasonal
  }
  if (temporal$short_term != "none") {
    n_params <- n_params + temporal$n_times
  }

  temporal$n_temporal_params <- n_params * temporal$n_groups

  temporal
}


# =============================================================================
# Gaussian Process Temporal Structure
# =============================================================================

# The kernel choices a continuous-time temporal GP has a closed form for,
# checked at construction.
#
# Matern is evaluated from its closed form, which exists at nu in {1/2, 3/2,
# 5/2}; general nu needs a Bessel function of the second kind and there is no
# autodiff-friendly form of it here. Rejecting is the point: an unsupported nu
# used to be accepted and then run as exponential (gcol33/tulpa#288).
#
# One check, asked by `temporal_gp()` and by `temporal_tvc(structure = "gp")`,
# which reach the same `src/temporal_gp_kernel.h`; the C++ spec reader
# (`read_temporal_gp_kernel`) is the same check at the other boundary, for a
# consumer package building the spec by hand. `who` names the door.
#' @keywords internal
.check_temporal_gp_kernel <- function(cov, nu, period, who) {
  if (identical(cov, "matern")) {
    if (!is.numeric(nu) || length(nu) != 1 || nu <= 0) {
      stop("`", who, "`: `nu` must be a positive number for Matern covariance.",
           call. = FALSE)
    }
    if (!any(abs(nu - c(0.5, 1.5, 2.5)) < 1e-12)) {
      stop("`", who, "`: `nu` must be 0.5, 1.5 or 2.5 for Matern covariance; ",
           "smoothnesses between them have no closed form here. Got ", nu, ".",
           call. = FALSE)
    }
  }
  if (identical(cov, "periodic")) {
    if (is.null(period) || !is.numeric(period) || length(period) != 1 ||
        period <= 0) {
      stop("`", who, "`: `period` must be a positive number when using ",
           "periodic covariance.", call. = FALSE)
    }
  }
  invisible(NULL)
}

#' Gaussian process temporal structure
#'
#' @description
#' Specify a Gaussian Process (GP) temporal random effect for irregularly-spaced
#' or continuous time points. Unlike RW1/RW2/AR1 which assume equally-spaced
#' observations, GP temporal effects model correlation as a function of time
#' distance.
#'
#' This is particularly useful for:
#' - Irregularly-spaced time series
#' - Continuous time (e.g., exact timestamps)
#' - Smooth temporal trends with uncertainty
#'
#' @param time_var Name of the time variable in data. Can be a formula
#'   (e.g., `~ year`) or a character string (e.g., `"year"`). Should be
#'   numeric (continuous time) or convertible to numeric.
#' @param cov Covariance function: `"exponential"` (default, rough),
#'   `"matern"` (tunable smoothness), `"gaussian"` (very smooth), or
#'   `"periodic"` (for seasonal patterns).
#' @param nu Smoothness parameter for Matern covariance, one of:
#'   - 0.5: equivalent to exponential (rough)
#'   - 1.5: once differentiable (moderate smoothness)
#'   - 2.5: twice differentiable (smooth)
#'   These are the smoothnesses with a closed form; anything between them needs
#'   a Bessel function of the second kind and is rejected. Ignored for
#'   non-Matern covariance functions.
#' @param period Period for periodic covariance (e.g., 12 for monthly, 365 for
#'   daily data with annual cycle). Only used when `cov = "periodic"`.
#' @param group_var Optional name of grouping variable for panel data.
#'   If provided, separate GPs are estimated for each group.
#' @param shared Logical; if TRUE (default), temporal effect enters both
#'   all processes.
#' @param scale_coords Logical; if TRUE (default), time values are scaled to
#'   unit variance before computing distances. Either way the lengthscale's
#'   support is placed relative to the spread of the times, and
#'   [temporal_corr()] reports the lengthscale in the original time units.
#' @param parameterization Parameterization for GP effects:
#'   `"noncentered"` (default) stores z ~ N(0,1) and scales by covariance
#'   (better for weakly-informed effects);
#'   `"centered"` stores effects directly (better for strongly-informed effects).
#'
#' @return A `tulpa_temporal_gp` object
#'
#' @details
#' The GP temporal model adds a time-correlated random effect:
#'
#' \deqn{\eta(t) = X\beta + f(t)}
#'
#' where \eqn{f(t)} follows a Gaussian process:
#' \deqn{f(t) \sim GP(0, \sigma^2 C(|t - t'|; \phi))}
#'
#' The correlation function \eqn{C(d; \phi)} depends on time distance \eqn{d}:
#'
#' - **Exponential**: \eqn{C(d) = \exp(-d/\phi)} - continuous but not differentiable
#' - **Matern**: Smooth with tunable roughness via \eqn{\nu}
#' - **Gaussian**: \eqn{C(d) = \exp(-(d/\phi)^2)} - infinitely differentiable
#' - **Periodic**: \eqn{C(d) = \exp(-2\sin^2(\pi d/p)/\phi^2)} - for seasonal data
#'
#' **Implementation**: the exponential kernel (equivalently Matern with
#' `nu = 0.5`) is an Ornstein-Uhlenbeck process, so its joint density
#' factorizes into a first-order Markov chain and evaluates in O(T) with no
#' matrix. The other kernels have no finite-dimensional state-space form and
#' are evaluated from a dense T x T Cholesky, where T is the number of distinct
#' time points.
#'
#' The field is fit on the sampler path: its two hyperparameters
#' (`log_sigma2_temporal_gp`, `logit_phi_temporal_gp`) are sampled jointly with
#' the field and the fixed effects, so pass a sampler `mode` such as
#' `mode = "hmc"`. There is no nested-Laplace kernel for it.
#'
#' @examples
#' # Create GP temporal specification
#' temporal_gp("timestamp")
#' temporal_gp("day", cov = "matern", nu = 1.5)
#' temporal_gp("month", cov = "periodic", period = 12)
#'
#' \donttest{
#' # Irregularly-spaced time series
#' set.seed(140)
#' times <- sort(runif(40, 0, 10))
#' df <- data.frame(
#'   time = times,
#'   x = rnorm(40),
#'   y = rbinom(40, 1, 0.5)
#' )
#'
#' fit <- tulpa(
#'   y ~ x,
#'   data = df,
#'   family = "binomial",
#'   temporal = temporal_gp("time"),
#'   mode = "hmc",
#'   control = list(n_iter = 200, warmup = 100, n_chains = 1)
#' )
#' }
#'
#' @seealso [temporal_rw1()], [temporal_ar1()] for equally-spaced temporal effects,
#'   [spatial_gp()] for spatial GP effects
#'
#' @export
temporal_gp <- function(time_var,
                        cov = c("exponential", "matern", "gaussian", "periodic"),
                        nu = 1.5,
                        period = NULL,
                        group_var = NULL,
                        shared = NULL,
                        scale_coords = TRUE,
                        parameterization = c("noncentered", "centered")) {

  cov <- match.arg(cov)
  parameterization <- match.arg(parameterization)

  if (!is.character(time_var) || length(time_var) != 1) {
    stop("`time_var` must be a single character string", call. = FALSE)
  }

  if (!is.null(group_var)) {
    if (!is.character(group_var) || length(group_var) != 1) {
      stop("`group_var` must be a single character string", call. = FALSE)
    }
  }

  .check_temporal_gp_kernel(cov, nu, period, "temporal_gp()")

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

  structure(
    list(
      type = "gp",
      time_var = time_var,
      group_var = group_var,
      cov = cov,
      nu = if (cov == "matern") nu else NULL,
      period = if (cov == "periodic") period else NULL,
      shared = shared,
      scale_coords = scale_coords,
      parameterization = parameterization,
      cyclic = FALSE,
      # Filled in during validation
      n_times = NULL,
      n_groups = NULL,
      n_obs = NULL,
      time_values = NULL,
      time_index = NULL,
      group_index = NULL,
      time_distances = NULL
    ),
    class = c("tulpa_temporal_gp", "tulpa_temporal", "list")
  )
}


#' @rdname temporal_gp
#' @param x A `tulpa_temporal_gp` specification.
#' @param ... Ignored.
#' @export
print.tulpa_temporal_gp <- function(x, ...) {
  cat("tulpa Gaussian Process temporal specification\n")
  cat("==============================================\n\n")

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

  cov_str <- x$cov
  if (x$cov == "matern" && !is.null(x$nu)) {
    cov_str <- sprintf("%s (nu = %.1f)", x$cov, x$nu)
  } else if (x$cov == "periodic" && !is.null(x$period)) {
    cov_str <- sprintf("%s (period = %.1f)", x$cov, x$period)
  }
  cat("Covariance:", cov_str, "\n")
  cat("Shared:", if (!isFALSE(x$shared)) "Yes (enters both processes)" else "No", "\n")

  if (!is.null(x$n_obs)) {
    cat("\nObservations:", x$n_obs, "\n")
    cat("Unique time points:", x$n_times, "\n")
  }

  invisible(x)
}


#' Validate temporal GP specification against data
#'
#' @param temporal tulpa_temporal_gp object
#' @param data Data frame
#'
#' @return Updated object with computed time structure
#' @noRd
validate_temporal_gp <- function(temporal, data) {
  if (is.null(temporal)) return(NULL)
  if (!inherits(temporal, "tulpa_temporal_gp")) return(temporal)

  N <- nrow(data)

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

  # Extract time values
  time_vals <- data[[temporal$time_var]]
  # The distinct instants in the user's own units (and class), which is what a
  # read-back labels the field by; the kernel sees the scaled numeric values
  # below, in the same order (gcol33/tulpa#905).
  temporal$time_levels <- sort(unique(time_vals))

  # Convert to numeric
  if (!is.numeric(time_vals)) {
    if (inherits(time_vals, c("Date", "POSIXt"))) {
      time_vals <- as.numeric(time_vals)
    } else {
      stop("Time variable must be numeric or a date/time type", call. = FALSE)
    }
  }

  # Check for missing values
  if (any(is.na(time_vals))) {
    stop("Time variable contains missing values", call. = FALSE)
  }

  # Scale time values if requested. `period` is a LAG, stated by the user in
  # the raw time units, and the kernel evaluates it against these values -- so
  # it has to make the same trip. Centring cancels in a lag; the divisor does
  # not, and leaving it untransformed is what collapsed a monthly period = 12
  # against a scaled lag range of ~3 (gcol33/tulpa#687). `period` keeps the
  # declared number so it prints and validates as the user wrote it;
  # `period_scaled` is that period in the units time_values now carries, and is
  # what every fit path reads.
  time_scale <- 1
  if (temporal$scale_coords) {
    scaled <- scale(time_vals)
    time_scale <- as.numeric(attr(scaled, "scaled:scale"))
    if (!is.finite(time_scale) || time_scale <= 0) time_scale <- 1
    time_vals <- as.vector(scaled)
  }
  temporal$time_scale <- time_scale
  # The lengthscale support, laid out on the spread of the times the kernel
  # sees: 1 once scaled, sd(time) in raw units (gcol33/tulpa#907).
  phi_b <- .gp_phi_bounds(stats::sd(time_vals))
  temporal$phi_prior_lower <- phi_b[["lower"]]
  temporal$phi_prior_upper <- phi_b[["upper"]]
  temporal$period_scaled <-
    if (is.null(temporal$period)) NULL else as.numeric(temporal$period) / time_scale

  temporal$n_obs <- N

  # Extract unique time values (sorted) and create time index
  # GP uses one effect per unique time, not per observation
  unique_times <- sort(unique(time_vals))
  temporal$unique_time_values <- unique_times  # Unique times for GP cov matrix
  temporal$n_times <- length(unique_times)

  # Create time_index: maps each obs to its unique time point (1-based)
  temporal$time_index <- match(time_vals, unique_times)

  # Keep original time_values for backward compatibility
  temporal$time_values <- unique_times  # Now unique times, not per-obs

  # Handle grouping
  if (!is.null(temporal$group_var)) {
    if (!(temporal$group_var %in% names(data))) {
      stop(sprintf("Group variable '%s' not found in data", temporal$group_var),
           call. = FALSE)
    }
    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)
  } else {
    temporal$n_groups <- 1L
    temporal$group_index <- rep(1L, N)
  }

  # Compute time distances for GP on unique times (not N x N)
  n_unique <- temporal$n_times
  if (n_unique <= 500) {
    temporal$time_distances <- outer(unique_times, unique_times,
                                     FUN = function(x, y) abs(x - y))
  } else {
    # For large data, we'll compute distances on-the-fly in C++
    temporal$time_distances <- NULL
  }

  # Total temporal parameters = n_times (one effect per unique time), not N
  temporal$n_temporal_params <- temporal$n_times * temporal$n_groups

  temporal
}