[avehtari]

Summary

Related to #260.

#290 will replace this code with posterior::ps_tail(), and stan-dev/posterior#452 adds the same partial sort there. loo has to work with the posterior version on CRAN until a posterior release contains #452, so this PR makes the same change in loo's own code now. When #290 is merged, this code goes away.

do_psis_i() sorted all S log ratios of each observation to find the tail of tail_len values and the cutoff below it. It now finds the cutoff with a partial sort (sort.int(partial = S - tail_len)) and sorts only the values at or above the cutoff. The results are identical: order() and sort.int(index.return = TRUE) sort doubles stably, and the candidate set contains all values that tie with the cutoff, so the selected tail and its order do not change.

Speed

Median time of psis() on one core (Intel i5-1335U, R 4.6.1), log ratios rnorm() + rexp(rate = 2), r_eff = 1:

S x N	master	this PR	speed-up
1000 x 1000	0.163 s	0.147 s	1.11x
4000 x 1000	0.479 s	0.379 s	1.26x
16 000 x 1000	1.526 s	1.067 s	1.43x
4000 x 10 000	4.528 s	3.613 s	1.25x

Tests

psis() returns identical() objects on master and with this PR for 8 inputs: example_loglik_matrix() and example_loglik_array() (which have tied draws), normal and t(1) log ratios, rounded log ratios with many ties, a column with -Inf values, S = 30, and a constant tail.

New test in tests/testthat/test_psis.R: do_psis_i() gives an identical result to the previous full-sort version for two vectors with ties; in the first one, 9 tail values equal the cutoff.

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

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 92.49%. Comparing base (fd910fe) to head (328af19).
⚠️ Report is 5 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #413      +/-   ##
==========================================
+ Coverage   92.48%   92.49%   +0.01%     
==========================================
  Files          31       31              
  Lines        3139     3146       +7     
==========================================
+ Hits         2903     2910       +7     
  Misses        236      236              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

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

• 🚀loo_function: 1.91s -> 1.65s [-13.83%, -13.24%]
• 🚀loo_matrix: 1.87s -> 1.64s [-13.07%, -12.29%]
  Further explanation regarding interpretation and methodology can be found in the documentation.

[jgabry]

Two tiny suggestions, otherwise this looks good.

[avehtari]

Thanks, applied both suggestions. Results are still identical to master on the same test inputs

[jgabry]

Looks good. Ok to merge when CI passes