diff --git a/NAMESPACE b/NAMESPACE index 216cf1ce..f72712e6 100755 --- a/NAMESPACE +++ b/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) diff --git a/NEWS.md b/NEWS.md index bcca5839..d6df0696 100644 --- a/NEWS.md +++ b/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 diff --git a/R/convergence.R b/R/convergence.R index 894c6bef..26569a99 100644 --- a/R/convergence.R +++ b/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) } diff --git a/R/pareto_smooth.R b/R/pareto_smooth.R index 1b283eb5..55c8f128 100644 --- a/R/pareto_smooth.R +++ b/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 diff --git a/man/ps_tail_sorted_log_ratios.Rd b/man/ps_tail_sorted_log_ratios.Rd new file mode 100644 index 00000000..3f6cf756 --- /dev/null +++ b/man/ps_tail_sorted_log_ratios.Rd @@ -0,0 +1,33 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/pareto_smooth.R +\name{ps_tail_sorted_log_ratios} +\alias{ps_tail_sorted_log_ratios} +\title{Pareto smooth sorted tail of log ratios} +\usage{ +ps_tail_sorted_log_ratios(x, cutoff, smooth_draws = TRUE, ...) +} +\arguments{ +\item{x}{Vector of log ratios in the tail, sorted in ascending order and +shifted so that \code{exp(x)} does not overflow, e.g. by subtracting the +maximum of all log ratios.} + +\item{cutoff}{The largest log ratio not in the tail, on the same shifted +scale as \code{x}. Tail values may be equal to it.} + +\item{smooth_draws}{(logical) Should the tail be smoothed? If \code{FALSE}, +only \code{k} is estimated and \code{tail} is \code{x} unchanged.} + +\item{...}{Arguments passed to \code{\link[=gpdfit]{gpdfit()}}.} +} +\value{ +A list with elements \code{tail}, the (smoothed) tail of log ratios, and +\code{k}, the estimated shape parameter. Unlike \code{\link[=ps_tail]{ps_tail()}}, the full vector +of draws is not returned. +} +\description{ +Function to Pareto smooth a sorted tail of log ratios. Exported +for usage in other packages, not by users. +} +\seealso{ +\code{\link{ps_tail}} for the general tail smoothing function. +} diff --git a/tests/testthat/test-convergence.R b/tests/testthat/test-convergence.R index 3046a7b5..404bfa90 100644 --- a/tests/testthat/test-convergence.R +++ b/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))) +}) diff --git a/tests/testthat/test-pareto_smooth.R b/tests/testthat/test-pareto_smooth.R index a79e0bd8..78cb3500 100644 --- a/tests/testthat/test-pareto_smooth.R +++ b/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) + ) + } + } +})