R/priors.R

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

#' Prior specification for tulpa models
#'
#' @description
#' Specify priors for model parameters. Supports both PC (penalized complexity)
#' priors for variance components and standard distributions for other parameters.
#'
#' A `tulpa_priors` object is what the simulation tools draw parameters from:
#' [prior_predict()] and [tulpa_simulate()] with `theta = NULL`. The model
#' fitters do not read one. [tulpa()] carries its own defaults, set per backend
#' through its `beta_prior`, `re_prior` and `hyperprior` arguments:
#'
#' - Fixed effects: `prior_normal(0, 2.5)` on every backend -- the same default
#'   as `beta` here.
#' - Random-effect SDs: a PC prior with `P(sigma > 3) = 0.01` (plus an LKJ(2)
#'   prior on a correlated block's correlations) on the routes that integrate or
#'   maximize over them (nested Laplace, `re_cov_nested`, `mode = "eb"`);
#'   half-Cauchy(2.5) on the samplers (`re_prior$sigma_re_scale`, and
#'   `re_prior$prior_sigma_scale` for `mode = "gibbs"`); an inverse-Wishart
#'   under `control$re_cov = "gibbs"`.
#'
#' The `sigma`, `phi`, `rho_temporal` and `rho_spatial` defaults below are
#' therefore the simulation defaults, not the priors a fit ran under.
#'
#' @param beta Prior for fixed effects. Default: `prior_normal(0, 2.5)`.
#' @param sigma Prior for random effect SDs. Default: PC prior with
#'   P(sigma > 1) = 0.01.
#' @param phi Prior for overdispersion parameter. Default: PC prior with
#'   P(phi > 10) = 0.01.
#' @param rho_temporal Prior for temporal autocorrelation. Default:
#'   `prior_beta(2, 2)` centered at 0.5.
#' @param rho_spatial Prior for spatial proportion (BYM2). Default:
#'   `prior_beta(1, 1)` (uniform).
#'
#' @return A `tulpa_priors` object
#'
#' @details
#' PC priors (Simpson et al., 2017) provide principled regularization that:
#' - Favors simpler models (smaller variance components)
#' - Has interpretable parameters (tail probabilities)
#' - Prevents overfitting with sparse data
#'
#' For other parameters, standard distributions are available via helper
#' functions: `prior_normal()`, `prior_half_normal()`, `prior_half_cauchy()`,
#' `prior_gamma()`, `prior_beta()`, `prior_exponential()`.
#'
#' @references
#' Simpson, D., Rue, H., Riebler, A., Martins, T. G., & Sorbye, S. H. (2017).
#' Penalising model component complexity: A principled, practical approach to
#' constructing priors. Statistical Science, 32(1), 1-28.
#'
#' @examples
#' # Default priors
#' tulpa_priors()
#'
#' # Custom fixed effect prior
#' tulpa_priors(beta = prior_normal(0, 1))
#'
#' # Tighter random effect prior
#' tulpa_priors(sigma = prior_pc(U = 0.5, alpha = 0.01))
#'
#' # Half-Cauchy for random effect SD
#' tulpa_priors(sigma = prior_half_cauchy(2.5))
#'
#' # Informative prior for temporal correlation
#' tulpa_priors(rho_temporal = prior_beta(5, 2))  # Prior mode at ~0.8
#'
#' @export
tulpa_priors <- function(
    beta = NULL,
    sigma = NULL,
    phi = NULL,
    rho_temporal = NULL,
    rho_spatial = NULL
) {

  # Set defaults
  if (is.null(beta)) beta <- prior_normal(0, 2.5)
  if (is.null(sigma)) sigma <- prior_pc(U = 1.0, alpha = 0.01)
  if (is.null(phi)) phi <- prior_pc(U = 10.0, alpha = 0.01)
  if (is.null(rho_temporal)) rho_temporal <- prior_beta(2, 2)
  if (is.null(rho_spatial)) rho_spatial <- prior_beta(1, 1)

  # Validate each prior
  validate_prior(beta, "beta")
  validate_prior(sigma, "sigma")
  validate_prior(phi, "phi")
  validate_prior(rho_temporal, "rho_temporal")
  validate_prior(rho_spatial, "rho_spatial")

  structure(
    list(
      beta = beta,
      sigma = sigma,
      phi = phi,
      rho_temporal = rho_temporal,
      rho_spatial = rho_spatial
    ),
    class = "tulpa_priors"
  )
}


#' Validate a prior specification
#'
#' @param prior A prior object
#' @param name Parameter name for error messages
#' @noRd
validate_prior <- function(prior, name) {
  if (!inherits(prior, "tulpa_prior")) {
    stop(sprintf("'%s' must be a prior object (use prior_*() functions)", name),
         call. = FALSE)
  }
}

# Constructor for the elementary tulpa_prior objects: validate every parameter
# as one finite number (the `positive` ones also > 0) through the shared
# `.check_scalar()`, so every prior_*() constructor checks the same way and
# the error names the argument; then box the distribution name and its
# parameters with the standard two-level class. Single source of truth for the
# simple prior_*() constructors; priors carrying derived fields (e.g.
# prior_pc()'s rate) build their object directly. Checking the positive ones
# alone let an NA or character location through, and turned an NA scale into
# "missing value where TRUE/FALSE needed" (#896).
.make_prior <- function(distribution, params, positive = character()) {
  for (nm in names(params)) {
    .check_scalar(params[[nm]], nm, positive = nm %in% positive)
  }
  structure(
    c(list(distribution = distribution), params),
    class = c(paste0("tulpa_prior_", distribution), "tulpa_prior")
  )
}


#' Normal prior
#'
#' @description
#' Specify a normal prior for a parameter.
#'
#' @param mean Prior mean. Default 0.
#' @param sd Prior standard deviation. Must be positive.
#'
#' @return A `tulpa_prior` object
#'
#' @examples
#' prior_normal(0, 2.5)
#' prior_normal(0, 1)
#'
#' @export
prior_normal <- function(mean = 0, sd = 2.5) {
  .make_prior("normal", list(mean = mean, sd = sd), positive = "sd")
}


#' Half-normal prior
#'
#' @description
#' Specify a half-normal (truncated at 0) prior for a positive parameter.
#'
#' @param sd Scale parameter. Must be positive.
#'
#' @return A `tulpa_prior` object
#'
#' @examples
#' prior_half_normal(1)
#'
#' @export
prior_half_normal <- function(sd = 1) {
  .make_prior("half_normal", list(sd = sd), positive = "sd")
}


#' Half-Cauchy prior
#'
#' @description
#' Specify a half-Cauchy prior for a positive parameter.
#' Has heavier tails than half-normal, often used for variance parameters.
#'
#' @param scale Scale parameter. Must be positive.
#'
#' @return A `tulpa_prior` object
#'
#' @examples
#' prior_half_cauchy(2.5)
#'
#' @export
prior_half_cauchy <- function(scale = 2.5) {
  .make_prior("half_cauchy", list(scale = scale), positive = "scale")
}


#' Gamma prior
#'
#' @description
#' Specify a gamma prior for a positive parameter.
#'
#' @param shape Shape parameter (alpha). Must be positive.
#' @param rate Rate parameter (beta). Must be positive.
#'
#' @return A `tulpa_prior` object
#'
#' @details
#' Mean = shape/rate, Variance = shape/rate^2.
#'
#' @examples
#' prior_gamma(2, 0.1)  # Mean = 20, weakly informative
#' prior_gamma(1, 1)    # Exponential(1)
#'
#' @export
prior_gamma <- function(shape = 2, rate = 0.1) {
  .make_prior("gamma", list(shape = shape, rate = rate),
              positive = c("shape", "rate"))
}


#' Exponential prior
#'
#' @description
#' Specify an exponential prior for a positive parameter.
#'
#' @param rate Rate parameter (lambda). Must be positive.
#'
#' @return A `tulpa_prior` object
#'
#' @examples
#' prior_exponential(1)
#'
#' @export
prior_exponential <- function(rate = 1) {
  .make_prior("exponential", list(rate = rate), positive = "rate")
}


#' Beta prior
#'
#' @description
#' Specify a beta prior for a parameter bounded in (0, 1).
#'
#' @param alpha First shape parameter. Must be positive.
#' @param beta Second shape parameter. Must be positive.
#'
#' @return A `tulpa_prior` object
#'
#' @details
#' - alpha = beta = 1: Uniform(0, 1)
#' - alpha = beta = 2: Symmetric, peaked at 0.5
#' - alpha > beta: Skewed toward 1
#' - alpha < beta: Skewed toward 0
#'
#' @examples
#' prior_beta(1, 1)   # Uniform
#' prior_beta(2, 2)   # Symmetric, centered at 0.5
#' prior_beta(5, 2)   # Skewed toward 1 (for high autocorrelation)
#'
#' @export
prior_beta <- function(alpha = 1, beta = 1) {
  .make_prior("beta", list(alpha = alpha, beta = beta),
              positive = c("alpha", "beta"))
}


#' Penalized complexity (PC) prior
#'
#' @description
#' Specify a PC prior for a positive parameter (typically a standard deviation).
#' PC priors shrink toward simpler models by penalizing deviation from a base model.
#'
#' @param U Upper bound. P(x > U) = alpha.
#' @param alpha Tail probability. Default 0.01.
#'
#' @return A `tulpa_prior` object
#'
#' @details
#' The PC prior is specified via: P(sigma > U) = alpha
#' - U is the upper bound you consider "large"
#' - alpha is the probability of exceeding it
#'
#' This implies an exponential prior with rate = -log(alpha) / U.
#'
#' @references
#' Simpson, D., Rue, H., Riebler, A., Martins, T. G., & Sorbye, S. H. (2017).
#' Penalising model component complexity: A principled, practical approach to
#' constructing priors. Statistical Science, 32(1), 1-28.
#'
#' @examples
#' prior_pc(U = 1, alpha = 0.01)    # P(sigma > 1) = 0.01
#' prior_pc(U = 0.5, alpha = 0.05)  # Tighter, P(sigma > 0.5) = 0.05
#'
#' @export
prior_pc <- function(U = 1, alpha = 0.01) {
  .check_pc_anchors(U, alpha, "U", "alpha", "prior_pc()")

  rate <- -log(alpha) / U

  structure(
    list(
      distribution = "pc",
      U = U,
      alpha = alpha,
      rate = rate
    ),
    class = c("tulpa_prior_pc", "tulpa_prior")
  )
}


#' @rdname tulpa_priors
#' @param x A `tulpa_priors` specification, or a single prior built by a
#'   `prior_*()` constructor.
#' @param ... Ignored.
#' @export
print.tulpa_priors <- function(x, ...) {
  cat("tulpa prior specification\n")
  cat("=========================\n\n")

  cat("Fixed effects (beta):\n")
  print_prior(x$beta, indent = "  ")
  cat("\n")

  cat("Random effect SD (sigma):\n")
  print_prior(x$sigma, indent = "  ")
  cat("\n")

  cat("Overdispersion (phi):\n")
  print_prior(x$phi, indent = "  ")
  cat("\n")

  cat("Temporal autocorrelation (rho_temporal):\n")
  print_prior(x$rho_temporal, indent = "  ")
  cat("\n")

  cat("Spatial proportion (rho_spatial):\n")
  print_prior(x$rho_spatial, indent = "  ")

  invisible(x)
}


#' @rdname tulpa_priors
#' @export
print.tulpa_prior <- function(x, ...) {
  print_prior(x, indent = "")
  invisible(x)
}


#' Format a single prior for printing
#'
#' @description
#' Thin wrapper around `format()` for `tulpa_prior` objects. Each
#' `prior_*()` constructor attaches a subclass (e.g. `tulpa_prior_normal`)
#' so adding a new distribution is one new `format.tulpa_prior_<dist>()`
#' method -- no central if/else to extend.
#'
#' @param prior A tulpa_prior object
#' @param indent Indentation string
#' @noRd
print_prior <- function(prior, indent = "") {
  cat(format(prior, indent = indent), "\n", sep = "")
}

#' @rdname tulpa_priors
#' @param indent String prefixed to each formatted line.
#' @return `format()` returns the prior as a character string.
#' @export
format.tulpa_prior_normal <- function(x, indent = "", ...) {
  sprintf("%sNormal(%.2f, %.2f)", indent, x$mean, x$sd)
}

#' @rdname tulpa_priors
#' @export
format.tulpa_prior_half_normal <- function(x, indent = "", ...) {
  sprintf("%sHalf-Normal(%.2f)", indent, x$sd)
}

#' @rdname tulpa_priors
#' @export
format.tulpa_prior_half_cauchy <- function(x, indent = "", ...) {
  sprintf("%sHalf-Cauchy(%.2f)", indent, x$scale)
}

#' @rdname tulpa_priors
#' @export
format.tulpa_prior_gamma <- function(x, indent = "", ...) {
  sprintf("%sGamma(%.2f, %.2f)  [mean = %.2f]",
          indent, x$shape, x$rate, x$shape / x$rate)
}

#' @rdname tulpa_priors
#' @export
format.tulpa_prior_exponential <- function(x, indent = "", ...) {
  sprintf("%sExponential(%.2f)  [mean = %.2f]",
          indent, x$rate, 1 / x$rate)
}

#' @rdname tulpa_priors
#' @export
format.tulpa_prior_beta <- function(x, indent = "", ...) {
  sprintf("%sBeta(%.2f, %.2f)  [mean = %.2f]",
          indent, x$alpha, x$beta, x$alpha / (x$alpha + x$beta))
}

#' @rdname tulpa_priors
#' @export
format.tulpa_prior_pc <- function(x, indent = "", ...) {
  paste0(
    sprintf("%sPC prior: P(x > %.2f) = %.3f", indent, x$U, x$alpha),
    "\n",
    sprintf("%s  => Exponential(%.3f)", indent, x$rate)
  )
}

#' @rdname tulpa_priors
#' @export
format.tulpa_prior <- function(x, indent = "", ...) {
  sprintf("%s%s", indent, x$distribution %||% "<unknown>")
}


#' Show default priors for a tulpa family
#'
#' @description
#' Display the default [tulpa_priors()] specification: the priors the
#' simulation tools ([prior_predict()], [tulpa_simulate()]) draw from, and a
#' starting point for customizing them. [tulpa()] does not read a
#' `tulpa_priors` object; the no-family printout closes with the defaults a
#' `tulpa()` fit does use, which differ for the random-effect SDs (see
#' [tulpa_priors()]).
#'
#' @param family `NULL` (default) for the general printout, or a
#'   `tulpa_family` object -- a model package's family constructor, such as a
#'   tulpaRatio ratio family, or one built with [tulpa_family()] -- for the
#'   parameters that family carries. A family NAME (`"poisson"`) is not a
#'   `tulpa_family` and is refused; the built-in [family_names()] share the
#'   general printout.
#' @param spatial Logical; if TRUE, include spatial priors. Default FALSE.
#' @param temporal Logical; if TRUE, include temporal priors. Default FALSE.
#'
#' @return Invisibly returns a `tulpa_priors` object with the defaults.
#'   Primarily called for its side effect of printing.
#'
#' @details
#' The family-specific overdispersion lines name the ratio families a model
#' package registers (`negbin_negbin`, `negbin_gamma`, `poisson_gamma`,
#' `gamma_gamma`, `beta_binomial`); they are not among tulpa's own
#' [family_names()].
#'
#' Default priors in [tulpa_priors()] follow these principles:
#'
#' - **Fixed effects (beta)**: Normal(0, 2.5) - weakly informative, allows
#'   coefficients roughly in \[-5, 5\] on the link scale.
#'
#' - **Random effect SD (sigma)**: PC prior with P(sigma > 1) = 0.01 - favors
#'   simpler models with smaller variance components.
#'
#' - **Overdispersion (phi)**: PC prior with P(phi > 10) = 0.01 on the NB2 size
#'   `phi` (larger `phi` is less overdispersion; `phi -> Inf` is the Poisson
#'   limit). The default keeps `phi` finite, allowing overdispersion while
#'   penalising extreme values.
#'
#' - **Temporal correlation (rho)**: Beta(2, 2) - symmetric prior centered
#'   at 0.5, appropriate for AR(1) correlation.
#'
#' - **Spatial mixing (rho_spatial)**: Beta(1, 1) = Uniform(0, 1) - no
#'   prior preference for structured vs. unstructured spatial variation.
#'
#' @examples
#' # Defaults for all families (no family argument)
#' priors_default()
#'
#' # Family-specific defaults take a tulpa_family object. Model packages
#' # (e.g. tulpaRatio) register rich families; a minimal one is enough here.
#' fam <- tulpa_family(
#'   name = "poisson_gamma",
#'   simulate_fn = function(eta, params, n_obs, ...) rpois(n_obs, exp(eta[[1]]))
#' )
#' priors_default(fam)
#'
#' # Including spatial parameters
#' priors_default(fam, spatial = TRUE)
#'
#' # Use as a starting point for customization
#' my_priors <- priors_default(fam)
#' my_priors$beta <- prior_normal(0, 1)  # Tighter prior on fixed effects
#'
#' @seealso [tulpa_priors()] for creating custom priors
#'
#' @export
priors_default <- function(family = NULL, spatial = FALSE, temporal = FALSE) {

  # Get default priors
  defaults <- tulpa_priors()

  if (is.null(family)) {
    # Show defaults for all families
    cat("Default priors for tulpa models\n")
    cat("================================\n\n")

    cat("These are tulpa_priors()'s defaults: what prior_predict() and\n")
    cat("tulpa_simulate() draw from. tulpa() fits do not read them; their own\n")
    cat("defaults are listed at the end.\n\n")

    cat("Fixed effects (beta):\n")
    cat("  Normal(0, 2.5)\n")
    cat("  Interpretation: Coefficients roughly in [-5, 5] on link scale\n")
    cat("  Customization: prior_normal(mean, sd)\n\n")

    cat("Random effect SD (sigma):\n")
    cat("  PC prior: P(sigma > 1) = 0.01\n")
    cat("  Interpretation: Favors smaller variance components\n")
    cat("  Customization: prior_pc(U, alpha) or prior_half_normal(sd)\n\n")

    cat("Overdispersion (phi) [negbin/poisson_gamma only]:\n")
    cat("  PC prior: P(phi > 10) = 0.01\n")
    cat("  Interpretation: NB2 size; keeps phi finite (allows overdispersion),\n")
    cat("                  phi -> Inf is the Poisson limit\n")
    cat("  Customization: prior_pc(U, alpha) or prior_gamma(shape, rate)\n\n")

    if (temporal) {
      cat("Temporal autocorrelation (rho_temporal) [AR models only]:\n")
      cat("  Beta(2, 2)\n")
      cat("  Interpretation: Symmetric, centered at 0.5\n")
      cat("  Customization: prior_beta(alpha, beta)\n\n")
    }

    if (spatial) {
      cat("Spatial mixing (rho_spatial) [BYM2 only]:\n")
      cat("  Beta(1, 1) = Uniform(0, 1)\n")
      cat("  Interpretation: No preference for structured vs. unstructured\n")
      cat("  Customization: prior_beta(alpha, beta)\n\n")
    }

    cat("Family-specific notes (model-package ratio families):\n")
    cat("  negbin_negbin: Uses phi for both processes\n")
    cat("  binomial: No overdispersion parameter (unless beta_binomial)\n")
    cat("  poisson_gamma: Uses phi for gamma shape parameter\n\n")

    cat("What a tulpa() fit uses instead (set via beta_prior / re_prior):\n")
    cat("  Fixed effects: Normal(0, 2.5) on every backend\n")
    cat("  Random-effect SD: PC prior P(sigma > 3) = 0.01 (+ LKJ(2) on\n")
    cat("    correlations) on nested Laplace / re_cov_nested / eb;\n")
    cat("    half-Cauchy(2.5) on the samplers and mode = \"gibbs\"\n")

  } else {
    # Show defaults for specific family
    if (!inherits(family, "tulpa_family")) {
      stop("`family` must be NULL or a tulpa_family object (a model package's ",
           "family constructor, or one built with tulpa_family()); got ",
           .arg_repr(family), ". A built-in family name shares the general ",
           "printout: call priors_default() with no family.", call. = FALSE)
    }

    family_name <- family$name

    cat(sprintf("Default priors for %s family\n", family_name))
    cat(paste(rep("=", nchar(family_name) + 24), collapse = ""), "\n\n")

    cat("Fixed effects (beta):\n")
    print_prior(defaults$beta, indent = "  ")
    cat("  Used for: All regression coefficients in all processes\n\n")

    cat("Random effect SD (sigma):\n")
    print_prior(defaults$sigma, indent = "  ")
    cat("  Used for: Standard deviation of group-level effects\n\n")

    # Family-specific parameters
    if (family_name %in% c("negbin_negbin", "negbin_gamma", "poisson_gamma", "gamma_gamma")) {
      cat("Overdispersion (phi):\n")
      print_prior(defaults$phi, indent = "  ")
      if (family_name %in% c("negbin_negbin", "negbin_gamma")) {
        cat("  Used for: NB2 size parameter (larger = less overdispersion)\n\n")
      } else if (family_name == "poisson_gamma") {
        cat("  Used for: Gamma shape parameter for effort/exposure\n\n")
      } else {
        cat("  Used for: Gamma shape parameters for both processes\n\n")
      }
    }

    if (family_name %in% c("beta_binomial_fixed", "beta_binomial")) {
      cat("Overdispersion (phi):\n")
      print_prior(defaults$phi, indent = "  ")
      cat("  Used for: Beta-binomial concentration parameter\n\n")
    }

    if (temporal) {
      cat("Temporal autocorrelation (rho_temporal):\n")
      print_prior(defaults$rho_temporal, indent = "  ")
      cat("  Used for: AR(1) correlation coefficient\n\n")
    }

    if (spatial) {
      cat("Spatial mixing (rho_spatial):\n")
      print_prior(defaults$rho_spatial, indent = "  ")
      cat("  Used for: BYM2 mixing proportion (structured vs. unstructured)\n\n")
    }

    cat("To customize, use tulpa_priors():\n")
    cat("  priors <- tulpa_priors(\n")
    cat("    beta = prior_normal(0, 1),\n")
    cat("    sigma = prior_pc(U = 0.5, alpha = 0.01)\n")
    cat("  )\n")
  }

  invisible(defaults)
}