Avoid full sorts in ps_tail() and mcse_quantile(), and add ps_tail_sorted_log_ratios() for loo

NAMESPACE

@@ -493,6 +493,7 @@ export(ps_khat_threshold)
 export(ps_min_ss)
 export(ps_tail)
 export(ps_tail_length)
+export(ps_tail_sorted_log_ratios)
 export(qgeneralized_pareto)
 export(quantile2)
 export(r_scale)

NEWS.md

@@ -2,6 +2,8 @@
 
 ### Enhancements
 
+* Improve performance by avoiding full sorts in `ps_tail()` and `mcse_quantile()`.
+* Add `ps_tail_sorted_log_ratios()` for loo.
 * `bind_draws()` for `draws_array` objects is faster and holds one copy of the
 result instead of about three.
 * `subset_draws()` returns its input without copying it when every variable is

R/convergence.R

@@ -425,10 +425,17 @@ mcse_median <- function(x, ...) {
   ess <- ess_quantile(x, prob)
   p <- c(0.1586553, 0.8413447)
   a <- qbeta(p, ess * prob + 1, ess * (1 - prob) + 1)
-  ssims <- sort(x)
-  S <- length(ssims)
-  th1 <- ssims[max(floor(a[1] * S), 1)]
-  th2 <- ssims[min(ceiling(a[2] * S), S)]
+  x <- as.vector(x)
+  S <- length(x)
+  i1 <- max(floor(a[1] * S), 1)
+  i2 <- min(ceiling(a[2] * S), S)
+  if (anyNA(c(i1, i2))) {
+    return(NA_real_)
+  }
+  # only the two order statistics are needed
+  ssims <- sort.int(x, partial = unique(c(i1, i2)))
+  th1 <- ssims[i1]
+  th2 <- ssims[i2]
   as.vector((th2 - th1) / 2)
 }
 

R/pareto_smooth.R

@@ -503,11 +503,9 @@ ps_tail <- function(x,
     x <- -x
   }
 
-  ndraws <- length(x)
-  tail_ids <- seq(ndraws - ndraws_tail + 1, ndraws)
-
-  ord <- sort.int(x, index.return = TRUE)
-  draws_tail <- ord$x[tail_ids]
+  selection <- .ps_tail_selection(x, ndraws_tail)
+  ix_tail <- selection$ix
+  draws_tail <- x[ix_tail]
 
   if (is_constant(draws_tail)) {
     if (tail == "left") {
@@ -518,42 +516,146 @@ ps_tail <- function(x,
 
   # largest value not in the tail; tail values equal to it give zero
   # excesses, which are in the support of the generalized Pareto distribution
-  cutoff <- ord$x[min(tail_ids) - 1]
-
-  max_tail <- max(draws_tail)
+  cutoff <- selection$cutoff
 
   if (are_log_weights) {
-    tail_excesses <- exp_x_minus_exp_y(draws_tail, cutoff)
-    cutoff <- exp(cutoff)
-  } else {
-    tail_excesses <- draws_tail - cutoff
-  }
-  fit <- gpdfit(tail_excesses, sort_x = FALSE, ...)
-  k <- fit$k
-  sigma <- fit$sigma
-  if (is.finite(k) && smooth_draws) {
-    p <- (seq_len(ndraws_tail) - 0.5) / ndraws_tail
-    smoothed <- qgeneralized_pareto(p = p, mu = cutoff, k = k, sigma = sigma)
-    if (are_log_weights) {
-      smoothed <- log(smoothed)
-    }
+    smoothed <- ps_tail_sorted_log_ratios(
+      x = draws_tail,
+      cutoff = cutoff,
+      smooth_draws = smooth_draws,
+      ...
+    )
   } else {
-    smoothed <- NULL
+    smoothed <- .ps_tail_sorted(
+      x = draws_tail,
+      excess = draws_tail - cutoff,
+      cutoff = cutoff,
+      smooth_draws = smooth_draws,
+      ...
+    )
   }
 
-  # truncate at max of raw draws
-  if (!is.null(smoothed)) {
-    smoothed[smoothed > max_tail] <- max_tail
-    x[ord$ix[tail_ids]] <- smoothed
+  k <- smoothed$k
+
+  if (is.finite(k) && smooth_draws) {
+    x[ix_tail] <- smoothed$tail
   }
 
   if (tail == "left") {
     x <- -x
   }
 
-  out <- list(x = x, k = k)
+  list(x = x, k = k)
+}
 
-  return(out)
+#' Select the right tail without a full sort
+#'
+#' @param x Vector of draws.
+#' @param ndraws_tail Number of draws in the tail.
+#' @return A list with elements `ix`, the indices of the `ndraws_tail` largest
+#'   values of `x` in ascending order of the values, and `cutoff`, the largest
+#'   value not in the tail. The result is the same as with
+#'   `order(x, method = "radix")`, also with ties.
+#'
+#' @noRd
+.ps_tail_selection <- function(x, ndraws_tail) {
+  ndraws <- length(x)
+  start <- ndraws - ndraws_tail + 1
+
+  if (start > 1 && !anyNA(x)) {
+    cutoff <- sort.int(x, partial = start - 1)[start - 1]
+    # The candidates include all values that tie with the cutoff. The radix
+    # order is stable, so the last ndraws_tail candidates are the tail of a
+    # full sort, also with ties.
+    ix <- which(x >= cutoff)
+    ix <- ix[order(x[ix], method = "radix")]
+    n_ix <- length(ix)
+    return(list(ix = ix[seq(n_ix - ndraws_tail + 1, n_ix)], cutoff = cutoff))
+  }
+
+  ix <- order(x, method = "radix")
+  list(ix = ix[seq(start, ndraws)], cutoff = x[ix[start - 1]])
+}
+
+#' Pareto smooth sorted tail of log ratios
+#'
+#' Function to Pareto smooth a sorted tail of log ratios. Exported
+#' for usage in other packages, not by users.
+#'
+#' @param x Vector of log ratios in the tail, sorted in ascending order and
+#'   shifted so that `exp(x)` does not overflow, e.g. by subtracting the
+#'   maximum of all log ratios.
+#' @param cutoff The largest log ratio not in the tail, on the same shifted
+#'   scale as `x`. Tail values may be equal to it.
+#' @param smooth_draws (logical) Should the tail be smoothed? If `FALSE`,
+#'   only `k` is estimated and `tail` is `x` unchanged.
+#' @param ... Arguments passed to [gpdfit()].
+#' @return A list with elements `tail`, the (smoothed) tail of log ratios, and
+#'   `k`, the estimated shape parameter. Unlike [ps_tail()], the full vector
+#'   of draws is not returned.
+#' @seealso [`ps_tail`] for the general tail smoothing function.
+#'
+#' @export
+ps_tail_sorted_log_ratios <- function(x, cutoff, smooth_draws = TRUE, ...) {
+  if (is.unsorted(x)) {
+    stop_no_call("'x' must be sorted in ascending order.")
+  }
+  if (is_constant(x)) {
+    return(list(tail = x, k = NA))
+  }
+
+  smoothed <- .ps_tail_sorted(
+    x = exp(x),
+    excess = exp_x_minus_exp_y(x, cutoff),
+    cutoff = exp(cutoff),
+    smooth_draws = smooth_draws,
+    ...
+  )
+
+  # log(exp(x)) is not x: it rounds, and it gives -Inf where exp(x) underflows
+  if (is.finite(smoothed$k) && smooth_draws) {
+    smoothed$tail <- log(smoothed$tail)
+  } else {
+    smoothed$tail <- x
+  }
+
+  smoothed
+}
+
+#' Pareto smooth sorted tail
+#'
+#' Internal function to Pareto smooth a sorted tail.
+#'
+#' @param x Vector of tail elements already sorted in ascending order.
+#' @param excess Tail elements minus the cutoff, computed without overflow.
+#' @param cutoff The largest value not in the tail. Tail values may be equal
+#'   to it.
+#' @param smooth_draws Should the tail values be smoothed?
+#' @param ... Passed to [gpdfit()].
+#'
+#' @noRd
+.ps_tail_sorted <- function(x, excess, cutoff, smooth_draws = TRUE, ...) {
+  len <- length(x)
+
+  fit <- gpdfit(excess, sort_x = FALSE, ...)
+  k <- fit$k
+  sigma <- fit$sigma
+
+  if (is.finite(k) && smooth_draws) {
+    tail <- qgeneralized_pareto(
+      p = (seq_len(len) - 0.5) / len,
+      mu = cutoff,
+      k = k,
+      sigma = sigma
+    )
+  } else {
+    tail <- x
+  }
+
+  # truncate at max of raw draws
+  tail[tail > max(x)] <- max(x)
+
+  list(tail = tail, k = k)
 }
 
 #' Extra Pareto-k diagnostics

tests/testthat/test-convergence.R

@@ -186,3 +186,22 @@ test_that("autocovariance returns correct results", {
   ac2 <- acf(x, type = "covariance", lag.max = length(x), plot = FALSE)$acf[, 1, 1]
   expect_equal(ac1, ac2)
 })
+
+test_that("mcse_quantile gives the order statistics of a full sort", {
+  # reference: the computation with sort() of posterior 1.7.1
+  mcse_quantile_full_sort <- function(x, prob) {
+    ess <- ess_quantile(x, prob)
+    a <- qbeta(c(0.1586553, 0.8413447), ess * prob + 1, ess * (1 - prob) + 1)
+    ssims <- sort(x)
+    S <- length(ssims)
+    (ssims[min(ceiling(a[2] * S), S)] - ssims[max(floor(a[1] * S), 1)]) / 2
+  }
+  tau <- extract_variable_matrix(example_draws(), "tau")
+  for (prob in c(0.001, 0.05, 0.5, 0.95, 0.999)) {
+    expect_identical(
+      unname(mcse_quantile(tau, probs = prob)),
+      as.vector(mcse_quantile_full_sort(tau, prob))
+    )
+  }
+  expect_true(is.na(mcse_quantile(c(rnorm(100), NA), probs = 0.5)))
+})

tests/testthat/test-pareto_smooth.R

@@ -332,19 +332,57 @@ test_that("pareto_smooth does not return negative weights for ties at 0", {
 
 
 test_that("check ps_tail behavior for ndraws_tail less than 5", {
-  w <- c(1:25, 1e3, 1e3, 1e3)
-  lw <- log(w)
+  lw <- log(c(1:25, 1e3, 1e3, 1e3))
 
-  # prints correct warning
   expect_warning(
     tail <- ps_tail(lw, ndraws_tail = 4, tail = "right"),
     "Can't fit generalized Pareto distribution because ndraws_tail is less than 5."
   )
-  # output has expected shape and k = NA
   expect_equal(names(tail), c("x", "k"))
   expect_true(is.na(tail$k))
 })
 
+test_that("ps_tail_sorted_log_ratios matches ps_tail", {
+  lw <- sort(log(c(1:25, 1e3, 1e2, 1e3)))
+  lw <- lw - max(lw)
+  tail_ids <- 19:28
+
+  expected <- ps_tail(lw, ndraws_tail = 10, are_log_weights = TRUE)
+  actual <- ps_tail_sorted_log_ratios(lw[tail_ids], cutoff = lw[min(tail_ids) - 1])
+
+  expect_equal(names(actual), c("tail", "k"))
+  expect_equal(actual$k, expected$k)
+  expect_equal(actual$tail, expected$x[tail_ids])
+})
+
+test_that("ps_tail_sorted_log_ratios handles constant tails", {
+  x <- log(rep(0.3, 10))
+
+  out <- ps_tail_sorted_log_ratios(x, cutoff = log(0.2))
+
+  expect_equal(out$tail, x)
+  expect_true(is.na(out$k))
+})
+
+test_that("ps_tail_sorted_log_ratios returns x unchanged if not smoothed", {
+  # exp() underflows to 0 for the smallest values
+  x <- c(-850, -800, -760, log(seq(0.05, 1, length.out = 20)))
+
+  out <- ps_tail_sorted_log_ratios(x, cutoff = -900, smooth_draws = FALSE)
+
+  expect_identical(out$tail, x)
+  expect_true(is.finite(out$k))
+})
+
+test_that("ps_tail_sorted_log_ratios requires a sorted tail", {
+  x <- log(seq(0.05, 1, length.out = 20))
+
+  expect_error(
+    ps_tail_sorted_log_ratios(rev(x), cutoff = log(0.01)),
+    "must be sorted in ascending order"
+  )
+})
+
 test_that("check ps_tail behavior for constant draws_tail", {
   x <- log(replicate(10, 0.3))
 
@@ -404,3 +442,22 @@ test_that("pareto diagnostics and smoothing do not depend on the scale of the dr
     expect_equal(ps_tail(s * x, ndraws_tail = 50)$k, ps_tail(x, ndraws_tail = 50)$k)
   }
 })
+
+test_that(".ps_tail_selection selects the tail of a full sort, also with ties", {
+  full_sort_selection <- function(x, ndraws_tail) {
+    ix <- order(x, method = "radix")
+    start <- length(x) - ndraws_tail + 1
+    list(ix = ix[seq(start, length(x))], cutoff = x[ix[start - 1]])
+  }
+  set.seed(452)
+  # rounding and repeated values give ties at the cutoff
+  for (x in list(rnorm(1000), round(rnorm(4000), 1), c(rep(0, 300), runif(100)),
+                 c(rep(-Inf, 50), rnorm(200)))) {
+    for (ndraws_tail in c(5, 19, 60)) {
+      expect_identical(
+        .ps_tail_selection(x, ndraws_tail),
+        full_sort_selection(x, ndraws_tail)
+      )
+    }
+  }
+})