diff --git a/R/psis.R b/R/psis.R index 303da755..85f8883c 100644 --- a/R/psis.R +++ b/R/psis.R @@ -215,9 +215,14 @@ 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 ", @@ -225,10 +230,10 @@ do_psis_i <- function(log_ratios_i, tail_len_i, ...) { 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 } } diff --git a/tests/testthat/test_psis.R b/tests/testthat/test_psis.R index 9e033715..cd63699d 100644 --- a/tests/testthat/test_psis.R +++ b/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)) + } +}) +