Added sorted log ratio path 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)

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,120 @@ 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)
+    smoothed <- ps_tail_sorted_log_ratios(
+      x = draws_tail,
+      cutoff = cutoff,
+      smooth_draws = smooth_draws,
+      ...
+    )
   } else {
-    tail_excesses <- draws_tail - cutoff
+    smoothed <- .ps_tail_sorted(
+      x = draws_tail,
+      excess = draws_tail - cutoff,
+      cutoff = cutoff,
+      smooth_draws = smooth_draws,
+      ...
+    )
   }
-  fit <- gpdfit(tail_excesses, sort_x = FALSE, ...)
-  k <- fit$k
-  sigma <- fit$sigma
+
+  k <- smoothed$k
+
   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)
+    x[ix_tail] <- smoothed$tail
+  }
+
+  if (tail == "left") {
+    x <- -x
+  }
+
+  list(x = x, k = k)
+}
+
+.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]
+    ix <- which(x > cutoff)
+    if (length(ix) == ndraws_tail) {
+      return(list(ix = ix[order(x[ix], method = "radix")], cutoff = cutoff))
     }
-  } else {
-    smoothed <- NULL
   }
 
-  # truncate at max of raw draws
-  if (!is.null(smoothed)) {
-    smoothed[smoothed > max_tail] <- max_tail
-    x[ord$ix[tail_ids]] <- smoothed
+  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
+#'   already shifted to account for overflow.
+#' @param cutoff Largest value smaller than the tail values.
+#' @inheritParams ps_tail
+#' @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_constant(x)) {
+    return(list(tail = x, k = NA))
   }
 
-  if (tail == "left") {
-    x <- -x
+  smoothed <- .ps_tail_sorted(
+    x = exp(x),
+    excess = exp_x_minus_exp_y(x, cutoff),
+    cutoff = exp(cutoff),
+    smooth_draws = smooth_draws,
+    ...
+  )
+
+  smoothed$tail <- log(smoothed$tail)
+
+  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 Largest value smaller than the tail values.
+#' @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
   }
 
-  out <- list(x = x, k = k)
+  # truncate at max of raw draws
+  tail[tail > max(x)] <- max(x)
 
-  return(out)
+  list(tail = tail, k = k)
 }
 
 #' Extra Pareto-k diagnostics

tests/testthat/test-pareto_smooth.R

@@ -332,19 +332,38 @@ 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("check ps_tail behavior for constant draws_tail", {
   x <- log(replicate(10, 0.3))