Sort only the tail in do_psis_i()

R/psis.R

@@ -215,20 +215,25 @@ do_psis_i <- function(log_ratios_i, tail_len_i, ...) {
   khat <- Inf
 
   if (enough_tail_samples(tail_len_i)) {
-    ord <- sort.int(lw_i, index.return = TRUE)
-    tail_ids <- seq(S - tail_len_i + 1, S)
-    lw_tail <- ord$x[tail_ids]
+    # only the tail and the cutoff need sorting: a partial sort finds the
+    # cutoff value, and the values at or above it are usually few
+    cut_pos <- S - tail_len_i
+    cut_val <- sort.int(lw_i, partial = cut_pos)[cut_pos]
+    top <- which(lw_i >= cut_val)
+    top <- top[order(lw_i[top])]
+    tail_ids <- top[seq(length(top) - tail_len_i + 1, length(top))]
+    lw_tail <- lw_i[tail_ids]
     if (abs(max(lw_tail) - min(lw_tail)) < .Machine$double.eps / 100) {
       warning(
         "Can't fit generalized Pareto distribution ",
         "because all tail values are the same.",
         call. = FALSE
       )
     } else {
-      cutoff <- ord$x[min(tail_ids) - 1] # largest value smaller than tail values
+      cutoff <- cut_val # largest value not in the tail
       smoothed <- psis_smooth_tail(lw_tail, cutoff)
       khat <- smoothed$k
-      lw_i[ord$ix[tail_ids]] <- smoothed$tail
+      lw_i[tail_ids] <- smoothed$tail
     }
   }
 

tests/testthat/test_psis.R

@@ -178,3 +178,23 @@ test_that("psis_smooth_tail returns original tail values if k is infinite", {
   expect_equal(val$k, Inf)
 })
 
+test_that("do_psis_i selects the same tail as a full sort, also with ties", {
+  # reference: the tail from a full sort, as in loo 2.10.1
+  do_psis_i_full_sort <- function(log_ratios_i, tail_len_i) {
+    S <- length(log_ratios_i)
+    lw_i <- log_ratios_i - max(log_ratios_i)
+    ord <- sort.int(lw_i, index.return = TRUE)
+    tail_ids <- seq(S - tail_len_i + 1, S)
+    smoothed <- psis_smooth_tail(ord$x[tail_ids], ord$x[min(tail_ids) - 1])
+    lw_i[ord$ix[tail_ids]] <- smoothed$tail
+    lw_i[lw_i > 0] <- 0
+    list(log_weights = lw_i + max(log_ratios_i), pareto_k = smoothed$k)
+  }
+  set.seed(260)
+  # rounding gives many ties, also at the cutoff
+  for (x in list(round(rnorm(4000), 1), round(rt(1000, df = 2), 2))) {
+    tail_len <- n_pareto(1, length(x))
+    expect_identical(do_psis_i(x, tail_len), do_psis_i_full_sort(x, tail_len))
+  }
+})
+