Improve floating point accuracy

R/E_loo.R

@@ -219,7 +219,10 @@ E_loo.matrix <-
   # sample size ESS is estimated with the generic target quantity invariant
   # estimate 1/sum(w^2), see e.g. "Monte Carlo theory, methods and examples"
   # by Owen (2013).
-  (sum(.wmean(x^2, w)) - sum(.wmean(x, w)^2)) / (1 - sum(w^2))
+  # The two-pass form avoids the cancellation in E[x^2] - E[x]^2 and is
+  # equivalent to it only because `w` sums to one.
+  weighted_mean <- .wmean(x, w)
+  sum(w * (x - weighted_mean)^2) / (1 - sum(w^2))
 }
 .wsd <- function(x, w, ...) {
   sqrt(.wvar(x, w))

R/helpers.R

@@ -26,6 +26,52 @@ colLogMeanExps <- function(x) {
   matrixStats::colLogSumExps(x) - logS
 }
 
+#' More stable version of `exp(x) - exp(y)`
+#'
+#' @noRd
+#' @param x A numeric vector.
+#' @param y A numeric scalar or vector recycled to the length of `x`.
+#'   Must satisfy `x >= y` elementwise.
+#' @return A numeric vector equal to `exp(x) - exp(y)`.
+#'
+exp_x_minus_exp_y <- function(x, y) {
+  out <- -exp(x) * expm1(y - x)
+  # which() drops the NA comparisons that NA or NaN inputs would produce
+  out[which(x == y)] <- 0
+  out
+}
+
+#' More stable version of `x^2 - y^2`
+#'
+#' @noRd
+#' @param x,y Numeric vectors of the same length.
+#' @return A numeric vector equal to `x^2 - y^2`.
+#'
+difference_of_squares <- function(x, y) {
+  (x - y) * (x + y)
+}
+
+#' More stable version of `(exp(a) - exp(b)) / exp(c)`
+#'
+#' @noRd
+#' @param a,b,c Numeric vectors of the same length.
+#' @return A numeric vector equal to `(exp(a) - exp(b)) / exp(c)`. Elements
+#'   with `a == b` are returned as an exact zero regardless of `c`; elsewhere
+#'   `NA` and `NaN` inputs propagate.
+#'
+exp_diff_over_exp <- function(a, b, c) {
+  # `a >= b` is NA if `a` or `b` is NA or NaN, and R silently ignores NA
+  # indices in `[<-`. Seed the result from the inputs and index with which()
+  # so that missing values propagate instead of leaving a zero behind.
+  out <- a + b + c
+  larger <- which(a >= b)
+  smaller <- which(a < b)
+  out[larger] <- exp(a[larger] - c[larger]) * -expm1(b[larger] - a[larger])
+  out[smaller] <- exp(b[smaller] - c[smaller]) * expm1(a[smaller] - b[smaller])
+  out[which(a == b)] <- 0
+  out
+}
+
 #' Compute point estimates and standard errors from pointwise vectors
 #'
 #' @noRd
@@ -58,11 +104,37 @@ validate_ll <- function(x) {
   } else if (anyNA(x)) {
     stop("NAs not allowed in input.")
   } else if (any(x == Inf)) {
-    stop("All input values must be finite or -Inf.")
+    # classed so that callers which negate a log-likelihood matrix can report
+    # the error in terms of the input the user actually supplied
+    stop(errorCondition(
+      "All input values must be finite or -Inf.",
+      class = "loo_positive_infinity_error"
+    ))
   }
   invisible(x)
 }
 
+#' Report `+Inf` log ratios in terms of the log likelihood that produced them
+#'
+#' `loo()` negates the log-likelihood before importance sampling, so a `-Inf`
+#' log-likelihood value reaches [validate_ll()] as `+Inf`. Wrap the importance
+#' sampling call so the user sees a message about their own input. The wrapped
+#' expression is only forced inside the handler, so there is no cost unless an
+#' error is raised.
+#'
+#' @noRd
+#' @param expr Expression that negates a log-likelihood and importance samples it.
+#' @return The value of `expr`.
+#'
+with_log_lik_error_message <- function(expr) {
+  withCallingHandlers(
+    expr,
+    loo_positive_infinity_error = function(cnd) {
+      stop("-Inf log-likelihood values are not allowed.", call. = FALSE)
+    }
+  )
+}
+
 #' Convert iter by chain by obs array to (iter * chain) by obs matrix
 #'
 #' @noRd

R/importance_sampling.R

@@ -12,6 +12,16 @@ importance_sampling <- function(log_ratios, method, ...) {
   UseMethod("importance_sampling")
 }
 
+validate_log_ratios <- function(x) {
+  validate_ll(x)
+  # validate_ll() has already ruled out NA and +Inf, so a column without a
+  # finite value is exactly a column whose maximum is -Inf
+  if (any(matrixStats::colMaxs(x) == -Inf)) {
+    stop("Each column of log ratios must contain at least one finite value.")
+  }
+  invisible(x)
+}
+
 
 #' @rdname importance_sampling
 #' @inheritParams psis
@@ -24,8 +34,8 @@ importance_sampling.array <-
     cores <- loo_cores(cores)
     stopifnot(length(dim(log_ratios)) == 3)
     assert_importance_sampling_method_is_implemented(method)
-    log_ratios <- validate_ll(log_ratios)
     log_ratios <- llarray_to_matrix(log_ratios)
+    log_ratios <- validate_log_ratios(log_ratios)
     r_eff <- prepare_psis_r_eff(r_eff, len = ncol(log_ratios))
     do_importance_sampling(log_ratios, r_eff = r_eff, cores = cores, method = method)
   }
@@ -40,7 +50,7 @@ importance_sampling.matrix <-
            cores = getOption("mc.cores", 1)) {
     cores <- loo_cores(cores)
     assert_importance_sampling_method_is_implemented(method)
-    log_ratios <- validate_ll(log_ratios)
+    log_ratios <- validate_log_ratios(log_ratios)
     r_eff <- prepare_psis_r_eff(r_eff, len = ncol(log_ratios))
     do_importance_sampling(log_ratios, r_eff = r_eff, cores = cores, method = method)
   }

R/loo.R

@@ -198,7 +198,9 @@ loo.array <-
            cores = getOption("mc.cores", 1),
            is_method = c("psis", "tis", "sis")) {
     is_method <- match.arg(is_method)
-    psis_out <- importance_sampling.array(log_ratios = -x, r_eff = r_eff, cores = cores, method = is_method)
+    psis_out <- with_log_lik_error_message(
+      importance_sampling.array(log_ratios = -x, r_eff = r_eff, cores = cores, method = is_method)
+    )
     ll <- llarray_to_matrix(x)
     pointwise <- pointwise_loo_calcs(ll, psis_out)
     importance_sampling_loo_object(
@@ -222,13 +224,14 @@ loo.matrix <-
            cores = getOption("mc.cores", 1),
            is_method = c("psis", "tis", "sis")) {
     is_method <- match.arg(is_method)
-    psis_out <-
+    psis_out <- with_log_lik_error_message(
       importance_sampling.matrix(
         log_ratios = -x,
         r_eff = r_eff,
         cores = cores,
         method = is_method
       )
+    )
     pointwise <- pointwise_loo_calcs(x, psis_out)
     importance_sampling_loo_object(
       pointwise = pointwise,
@@ -371,13 +374,14 @@ loo_i <-
     if (!is.matrix(ll_i)) {
       ll_i <- as.matrix(ll_i)
     }
-    psis_out <-
+    psis_out <- with_log_lik_error_message(
       importance_sampling.matrix(
         log_ratios = -ll_i,
         r_eff = r_eff,
         cores = 1,
         method = is_method
       )
+    )
     structure(
       list(
         pointwise = pointwise_loo_calcs(ll_i, psis_out),
@@ -488,29 +492,64 @@ importance_sampling_loo_object <- function(pointwise, diagnostics, dims,
 #' @return Vector of standard error estimates.
 #'
 mcse_elpd <- function(ll, lw, E_elpd, r_eff, n_samples = NULL) {
-  lik <- exp(ll)
-  w2 <- exp(lw)^2
-  E_epd <- exp(E_elpd)
-  if (length(r_eff) == 1 && !is.null(ncol(ll))) {
+  if (!is.matrix(ll)) {
+    ll <- as.matrix(ll)
+  }
+  if (!is.matrix(lw)) {
+    lw <- as.matrix(lw)
+  }
+  S <- nrow(ll)
+  if (length(r_eff) == 1) {
     r_eff <- rep(r_eff, ncol(ll))
   }
-  var_elpd <-
-    vapply(
-      seq_len(ncol(w2)),
-      FUN.VALUE = numeric(1),
-      FUN = function(i) {
-        # Variance in linear scale
-        # Equation (6) in Vehtari et al. (2024)
-        var_epd_i <- sum(w2[, i] * (lik[, i] - E_epd[i]) ^ 2) / r_eff[i]
-        # Compute variance in log scale by match the variance of a
-        # log-normal approximation
-        # https://en.wikipedia.org/wiki/Log-normal_distribution#Arithmetic_moments
-        log(1 + var_epd_i / E_epd[i]^2)
-      }
+  # Everything is computed relative to the loo predictive density, so that
+  # 1) exp() of the log likelihood never over- or underflows, and
+  # 2) expm1() avoids the cancellation in `exp(ll) - exp(E_elpd)`.
+  # `ll - E_elpd` is bounded above by `-lw`, so the product below cannot
+  # overflow for consistent (ll, lw, E_elpd); the fallback covers the rest.
+  #
+  # Variance in linear scale, relative to E_epd^2.
+  # Equation (6) in Vehtari et al. (2024)
+  var_epd_ratio <-
+    matrixStats::colSums2((exp(lw) * expm1(ll - rep(E_elpd, each = S)))^2) /
+      r_eff
+  # Variance in log scale by matching the variance of a log-normal
+  # https://en.wikipedia.org/wiki/Log-normal_distribution#Arithmetic_moments
+  var_elpd <- log1p(var_epd_ratio)
+  undefined <- is.infinite(E_elpd) & E_elpd < 0
+  overflow <- !is.finite(var_epd_ratio) & !undefined
+  if (any(overflow)) {
+    lvr <- log_var_epd_ratio(
+      ll[, overflow, drop = FALSE] - rep(E_elpd[overflow], each = S),
+      lw[, overflow, drop = FALSE],
+      r_eff[overflow]
     )
+    var_elpd[overflow] <-
+      ifelse(lvr > 0, lvr + log1p(exp(-lvr)), log1p(exp(lvr)))
+  }
+  var_elpd[undefined] <- NA_real_
   sqrt(var_elpd)
 }
 
+#' Log of the relative linear-scale ELPD variance, for the rare case where
+#' `exp(lw) * expm1(log_lik_ratio)` over- or underflows
+#'
+#' @noRd
+#' @param log_lik_ratio Matrix of `ll - E_elpd` values.
+#' @param lw Matrix of normalized log weights.
+#' @param r_eff Vector of relative effective sample sizes.
+#' @return Vector of `log(var_epd / E_epd^2)` values.
+#'
+log_var_epd_ratio <- function(log_lik_ratio, lw, r_eff) {
+  log_abs_diff <- log(abs(expm1(log_lik_ratio)))
+  big <- which(log_lik_ratio > 700)
+  if (length(big)) {
+    log_abs_diff[big] <-
+      log_lik_ratio[big] + log(-expm1(-log_lik_ratio[big]))
+  }
+  matrixStats::colLogSumExps(2 * (lw + log_abs_diff)) - log(r_eff)
+}
+
 
 #' Warning message if r_eff not specified
 #' @noRd

R/loo_model_weights.R

@@ -257,6 +257,9 @@ stacking_weights <-
     if (K < 2) {
       stop("At least two models are required for stacking weights.")
     }
+    if (any(rowSums(is.finite(lpd_point)) == 0)) {
+      stop("Each observation must have a finite predictive density for at least one model.")
+    }
 
     negative_log_score_loo <- function(w) {
       # objective function: log score
@@ -272,11 +275,15 @@ stacking_weights <-
       stopifnot(length(w) == K - 1)
       w_full <- c(w, 1 - sum(w))
       grad <- rep(0, K - 1)
-      # avoid over- and underflows using log weights, rowLogSumExps,
-      # and by subtracting the row maximum of lpd_point
-      mlpd <- matrixStats::rowMaxs(lpd_point)
+      mixture_lpd <- matrixStats::rowLogSumExps(
+        sweep(lpd_point, 2, log(w_full), "+")
+      )
       for (k in 1:(K - 1)) {
-        grad[k] <- sum((exp(lpd_point[, k] - mlpd) - exp(lpd_point[, K] - mlpd)) / exp(matrixStats::rowLogSumExps(sweep(lpd_point, 2, log(w_full), '+')) - mlpd))
+        grad[k] <- sum(exp_diff_over_exp(
+          lpd_point[, k],
+          lpd_point[, K],
+          mixture_lpd
+        ))
       }
       return(-grad)
     }
@@ -317,9 +324,12 @@ pseudobma_weights <-
     if (K < 2) {
       stop("At least two models are required for pseudo-BMA weights.")
     }
+    elpd <- colSums2(lpd_point)
+    if (!any(is.finite(elpd))) {
+      stop("At least one model must have a finite total predictive density.")
+    }
 
     if (!BB) {
-      elpd <- colSums2(lpd_point)
       uwts <- exp(elpd - max(elpd))
       wts <- structure(
         uwts / sum(uwts),

R/loo_moment_matching.R

@@ -408,7 +408,7 @@ loo_moment_match_i <- function(i,
   elpd_loo_i <- matrixStats::logSumExp(log_liki + lwi)
   mcse_elpd_loo <- mcse_elpd(
     ll = as.matrix(log_liki), lw = as.matrix(lwi),
-    E_elpd = exp(elpd_loo_i), r_eff = r_eff_i
+    E_elpd = elpd_loo_i, r_eff = r_eff_i
   )
 
   list(elpd_loo_i = elpd_loo_i,
@@ -531,11 +531,14 @@ shift_and_scale <- function(x, upars, lwi) {
   # compute moments using log weights
   S <- dim(upars)[1]
   mean_original <- colMeans(upars)
-  mean_weighted <- colSums(exp(lwi) * upars)
+  weights <- exp(lwi)
+  mean_weighted <- colSums(weights * upars)
   shift <- mean_weighted - mean_original
-  mii <- exp(lwi)* upars^2
-  mii <- colSums(mii) - mean_weighted^2
-  mii <- mii*S/(S-1)
+  # The two-pass form avoids the cancellation in E[x^2] - E[x]^2 and is
+  # equivalent to it only because `weights` sums to one.
+  centered <- sweep(upars, 2, mean_weighted)
+  mii <- colSums(weights * centered^2)
+  mii <- mii * S / (S - 1)
   scaling <- sqrt(mii / matrixStats::colVars(upars))
   # transform posterior draws
   upars_new <- sweep(upars, 2, mean_original, "-")

R/loo_subsample.R

@@ -1185,7 +1185,7 @@ srs_diff_est <- function(y_approx, y, y_idx) {
   t_pi_tilde <- sum(y_approx)
   t_pi2_tilde <- sum(y_approx^2)
   t_e <- N * mean(e_i)
-  t_hat_epsilon <- N * mean(y^2 - y_approx_m^2)
+  t_hat_epsilon <- N * mean(difference_of_squares(y, y_approx_m))
 
   est_list <- list(m = length(y), N = N)
   # eq (7)

R/psis.R

@@ -254,7 +254,7 @@ psis_smooth_tail <- function(x, cutoff) {
   exp_cutoff <- exp(cutoff)
 
   # save time not sorting since x already sorted
-  fit <- posterior::gpdfit(exp(x) - exp_cutoff, sort_x = FALSE)
+  fit <- posterior::gpdfit(exp_x_minus_exp_y(x, cutoff), sort_x = FALSE)
   k <- fit$k
   sigma <- fit$sigma
   if (is.na(k)) {

R/psis_approximate_posterior.R

@@ -34,10 +34,11 @@ psis_approximate_posterior <- function(log_p = NULL, log_g = NULL, log_liks = NU
   checkmate::assert_flag(save_psis)
 
   if (is.null(log_liks)) {
-    approx_correction <- log_p - log_g
-    # Handle underflow/overflow
-    approx_correction <- approx_correction - max(approx_correction)
+    approx_correction <- validate_approx_correction(log_p, log_g)
     log_ratios <- matrix(approx_correction, ncol = 1)
+    log_ratios <- validate_log_ratios(log_ratios)
+    # Handle underflow/overflow
+    log_ratios <- log_ratios - max(log_ratios)
   } else {
     log_ratios <- correct_log_ratios(log_ratios = -log_liks, log_p = log_p, log_g = log_g)
   }
@@ -65,9 +66,18 @@ psis_approximate_posterior <- function(log_p = NULL, log_g = NULL, log_liks = NU
 #' @inheritParams ap_psis
 #' @noRd
 #' @keywords internal
-correct_log_ratios <- function(log_ratios, log_p, log_g) {
+validate_approx_correction <- function(log_p, log_g) {
   approx_correction <- log_p - log_g
+  if (any(is.nan(approx_correction))) {
+    stop("The log density ratio is undefined for one or more draws.")
+  }
+  approx_correction
+}
+
+correct_log_ratios <- function(log_ratios, log_p, log_g) {
+  approx_correction <- validate_approx_correction(log_p, log_g)
   log_ratios <- log_ratios + approx_correction
+  log_ratios <- validate_log_ratios(log_ratios)
   # Handle underflow/overflow
   log_ratio_max <- apply(log_ratios, 2, max)
   log_ratios <- sweep(log_ratios, MARGIN = 2, STATS = log_ratio_max)