[VisruthSK]

This PR was made with assistance from Codex and Claude.

Summary

• ps_tail() no longer sorts the whole vector to find the tail. The internal .ps_tail_selection() finds the cutoff with a partial sort (sort.int(partial = ...)), takes the values at or above it, including all values that tie with the cutoff, and sorts only those. The radix order is stable, so the selected tail and its order are the same as with a full sort. This affects all functions that call ps_tail(): pareto_smooth(), pareto_khat() and pareto_diags().
• The new exported ps_tail_sorted_log_ratios() smooths a tail of log ratios that is already sorted. Rely on posterior for pareto smooth tails loo#290 will use it, so that loo selects the tail once and does not sort the log ratios twice (see this comment). If the tail is not smoothed (smooth_draws = FALSE or a non-finite k), it returns x unchanged, as loo's psis_smooth_tail() did. It gives an error if x is not sorted; the check costs less than 1 % of the call. ps_tail() and ps_tail_sorted_log_ratios() share the internal .ps_tail_sorted().
• mcse_quantile() sorted all draws to read two order statistics. It now uses sort.int(partial = c(i1, i2)), which gives the same two values.

sort.int() does not allow index.return = TRUE together with partial, which is why the tail selection gets the indices with which() and order() of the few values at or above the cutoff.

Speed

Speed-up relative to master, 500 vectors, median of 5 runs, one core (Intel i5-1335U, R 4.6.1):

function	data	S = 4000	S = 16 000
ps_tail(are_log_weights = TRUE)	continuous	1.28x	1.64x
	ties at the cutoff (values rounded to 0.1)	1.09x	1.76x
	-Inf log weights	1.60x	1.78x
mcse_quantile(probs = c(0.05, 0.5, 0.95))	continuous	1.10x	1.15x

mcse_quantile() gains less because ess_quantile() takes most of its time.

For loo, the tail selection of stan-dev/loo#413 followed by ps_tail_sorted_log_ratios() is 1.12x (S = 4000) and 1.08x (S = 16 000) faster than ps_tail() on the whole vector, for 1000 observations, median of 7 runs. The log weights and k are identical.

Results

ps_tail() (right and left tail, values and log weights), pareto_smooth(tail = "both"), pareto_khat(), pareto_pit() and mcse_quantile() return results identical to master in 63 of 63 checks: 7 vectors (normal, t(1), rounded values with many ties, -Inf log weights, ties at the cutoff, 40 draws, and a constant tail). ps_tail(are_log_weights = TRUE) is also identical for the 3000 vectors of the speed test, and mcse_quantile() for 200 vectors at 3 probabilities.

New tests:

• ps_tail_sorted_log_ratios() agrees with ps_tail(), also for a constant tail.
• ps_tail_sorted_log_ratios(smooth_draws = FALSE) returns x unchanged when exp(x) underflows.
• ps_tail_sorted_log_ratios() gives an error for a tail that is not sorted.
• .ps_tail_selection() gives the same indices and cutoff as a full sort for 4 vectors with ties and -Inf values and 3 tail lengths.
• mcse_quantile() gives the same values as the computation with sort() for 5 probabilities, and NA for draws with NA.

devtools::test() with NOT_CRAN=true passes all 2290 tests and gives no warnings (R 4.6.1, testthat 3.3.2, Linux x86_64).

Copyright and Licensing

By submitting this pull request, the copyright holder is agreeing to license the submitted work under the following licenses:

• Code: BSD 3-clause (https://opensource.org/licenses/BSD-3-Clause)
• Documentation: CC-BY 4.0 (https://creativecommons.org/licenses/by/4.0/)

[github-actions]

This is how benchmark results would change (along with a 95% confidence interval in relative change) if 16e594c is merged into master:

• ✔️as_draws_array: 164ms -> 165ms [-0.18%, +1.13%]
• ✔️as_draws_df: 81.1ms -> 80.8ms [-1.4%, +0.74%]
• ✔️as_draws_list: 166ms -> 167ms [-0.62%, +1.59%]
• ✔️as_draws_matrix: 30.1ms -> 30.1ms [-0.64%, +0.73%]
• ❗🐌as_draws_rvars: 124ms -> 130ms [+4%, +5.31%]
• ✔️summarise_draws_100_variables: 744ms -> 745ms [-0.39%, +0.67%]
• ✔️summarise_draws_10_variables: 81.5ms -> 82.2ms [-1.64%, +3.15%]
  Further explanation regarding interpretation and methodology can be found in the documentation.

[avehtari]

exported function might need NEWS.md item

[avehtari]

Looks good

[jgabry]

@VisruthSK is this waiting on anything else? Also does this include the performance benefit from that loo PR that we closed?

[VisruthSK]

I think I need to implement that after which it's gtg

[avehtari]

I pushed two commits and changed the title and description, because the PR changes ps_tail() itself and so all functions that call it.

The first commit changes the tie handling in .ps_tail_selection(). When a tail value ties with the cutoff, which(x > cutoff) finds fewer than ndraws_tail values, and the function fell back to a full order(x) after it had already done the partial sort. Ties at the cutoff are common for log ratios, for example duplicated draws and log ratios that underflow to the same value, and in these cases ps_tail() was slower than on master (0.79x and 0.86x for 4000 and 16 000 draws). The commit takes which(x >= cutoff), which includes all values that tie with the cutoff, and sorts only those. The radix order is stable, so the last ndraws_tail of them are the tail of a full sort. With ties at the cutoff, ps_tail() is now 1.08x and 1.67x faster than master, and the other cases did not change beyond measurement noise.

The second commit uses the same idea in mcse_quantile(), which sorted all draws to read two order statistics.

Before and after these commits, the results are identical to master. The description lists the checks.

A C routine for the tail selection (#453) gave about the same speed-up in my measurements as the R version (1.43x and 1.50x for continuous data), so the R version seems enough.

The previous description said that the warning for fewer than 5 tail draws was removed, but ps_tail() still gives it, and the test still expects it. I removed that sentence from the description.

[avehtari]

@VisruthSK @jgabry, this needs a new review since I added two commits

[jgabry]

The results look identical to master for ps_tail() and mcse_quantile() in my checks too. I made some review comments. A couple of other things:

1. Do we need to export ps_tail_sorted_log_ratios()? The description says loo#290 uses it, but loo#290 currently calls posterior::ps_tail(), and loo#413 says #290 will use ps_tail() too. Now that ps_tail() only sorts the tail, the remaining extra sort in loo#290 is the full sort.int(lw_i) it does just to warn about a constant tail. ps_tail() already returns k = NA in that case, so loo could drop its own sort and check is.na(smoothed$k). But if loo is going to use it, then can you update loo#290 to call it and I guess check if it's better than calling ps_tail() directly?

2. NEWS still needs an entry. Maybe something like: "ps_tail(), and therefore pareto_smooth(), pareto_khat() and pareto_diags(), and mcse_quantile(), no longer sort all draws (#452)".

[avehtari]

Do we need to export ps_tail_sorted_log_ratios()?

The idea is that this is a fast version not doing any checks and branches that ps_tail(), and loo doesn't call this yet as it is not yet available

[jgabry]

Do we need to export ps_tail_sorted_log_ratios()?

The idea is that this is a fast version not doing any checks and branches that ps_tail(), and loo doesn't call this yet as it is not yet available

Makes sense, thanks. I asked Codex to check and selecting the tail in loo and then calling ps_tail_sorted_log_ratios() is about 1.12x faster than calling ps_tail() for S = 4000 and 1.17x faster for S = 16000 (1000 observations, identical results). So I'm fine keeping the export.

In that case my other comments on it still apply (the log(exp(x)) fix and the docs), and the NEWS entry should mention the new function too. Maybe also change the description to say loo#290 will use it.

Regarding the is.unsorted() check I suggeasted, it only checks the tail, so it probably wouldn't cost anything noticeable, but if you prefer to keep the function without checks that's OK.

[avehtari]

Thanks for the review. I pushed four commits:

• 16b007f: ps_tail_sorted_log_ratios() returns x unchanged when the tail is not smoothed (your suggestion). The new test with expect_identical() uses values where exp(x) underflows and where k is finite, so that it tests the smooth_draws = FALSE branch. Before the fix, the test fails with -Inf for 3 values.
• b39cd62: the docs. cutoff in ps_tail_sorted_log_ratios() and .ps_tail_sorted() now says that tail values can be equal to it. ps_tail_sorted_log_ratios() documents smooth_draws and ... directly instead of @inheritParams ps_tail. .ps_tail_selection() has a @noRd header that says what it returns.
• 017e7c1: removed x <- x[!is.na(x)] in .mcse_quantile(). You are right that it never has an effect.
• be0f57d: the is.unsorted() check you suggested, with a test for a reversed tail. The check costs about 1 us per call: 0.85 us of 114 us (0.75 %) for S = 4000 (tail of 190 draws), and 1.05 us of 203 us (0.52 %) for S = 16 000 (tail of 380 draws).

I repeated your measurement with the tail selection of loo#413 followed by ps_tail_sorted_log_ratios(), against ps_tail() on the whole vector (1000 observations, median of 7 runs, one core). I got 1.12x for S = 4000 and 1.08x for S = 16 000, with identical log weights and k. I changed the description to say that loo#290 will use it.

Results are still identical to master in all checks, and devtools::test() with NOT_CRAN=true passes all 2290 tests without warnings.

[jgabry]

The updates look good, thanks! From my perspective this is ready to merge when CI passes.

[jgabry]

@VisruthSK since this was your PR originally did you want to take a look again before merging?

[VisruthSK]

Looks good to me!