Pareto k diagnostics and Pareto smoothing

[n-kall]

Summary

This PR adds functionality for Pareto smoothing, following on from work by @ozanstats and @avehtari to address #237. It adds two primary user-facing functions: pareto_smooth and pareto_khat. pareto_smooth performs smoothing on the tail(s) of the input and returns the smoothed draws and diagnostic values, whereas pareto_khat calls pareto_smooth but only returns the diagnostic values.

pareto_khat is designed for use in similar ways to the convergence functions rhat and ess_bulk. It is made to work in summarise_draws. On the other hand, pareto_smooth is probably most useful if one wants to do Pareto smoothing but not use the loo package for this (for example to implement moment-matching + psis in the iwmm package).

I believe the only functionality it does not implement that is mentioned in #237 is allowing r_eff to be a vector of different values (one for each variable). If r_eff is specified, it must be just a single value, as I'm not sure how to allow different r_eff values for each variable when using summarise_draws. However, when it is automatically calculated (default), then it works as intended and is calculated separately for each variable.

Current status

The functions should be working as intended, and the documentation is more or less complete (ready for comments). There are no tests yet, so any suggestions for recommended tests would be helpful.

Example functionality

mu <- extract_variable_matrix(example_draws(), "mu")
pareto_khat(mu)
# > 0.270711

ex <- example_draws()
summarise_draws(ex, pareto_khat, .args = list(extra_diags = TRUE))
# # A tibble: 10 × 5
#   variable      k min_ss khat_threshold convergence_rate
#   <chr>     <dbl>  <dbl>          <dbl>            <dbl>
# 1 mu       0.271    23.5          0.616            0.971
# 2 tau      0.0151   10.4          0.616            1.00
# 3 theta[1] 0.123    13.8          0.616            0.994
# 4 theta[2] 0.0681   11.8          0.616            0.998
# 5 theta[3] 0.362    37.0          0.616            0.937
# 6 theta[4] 0.211    18.5          0.616            0.984
# 7 theta[5] 0.145    14.8          0.616            0.992
# 8 theta[6] 0.181    16.6          0.616            0.988
# 9 theta[7] 0.265    22.9          0.616            0.973
#10 theta[8] 0.0888   12.5          0.616            0.997

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/)

[paul-buerkner]

Please let me know once this PR is ready for review.

[n-kall]

I've now added tests for the pareto_khat and pareto_smooth functions. So I think now it is ready for review @paul-buerkner @avehtari.

I suppose tests for gpd functions could be imported from loo or the PR from @ozanstats here

[paul-buerkner]

This PR looks quite good already. I hope my comments help to further improve it. Thank you for working on this!

[paul-buerkner]

Would it be sensible to make all the functions of the GPD vectorized across mu, sigma, and k?

Als the names are pretty short. How about (d|p...)generalized_pareto?

[n-kall]

After thinking about these functions, I realised that only dgpd is used internally for the calculation of the marginal posterior in gpdpost. extraDistr already provides GPD (using a C++ implementation), so I'm not sure posterior needs to have reimplementations of them. @avehtari can you see a clear use for having the GPD functions (q*, p*, r*) directly available via posterior or could we just have dgpd for internal use and only export gpdpost?

[n-kall]

For now, I have removed most of these functions from the PR, leaving only the necessary ones for point estimate diagnostics and pareto smoothing. I have also changed the remaining to internal functions, not exported.

[paul-buerkner]

@mjskay could you have a look and check if this pattern for rvars is the most sensible approach (same for pareto_smooth below)?

[mjskay]

I took a look, and I'm not sure the resulting output format is quite what I'd expect for rvars given how we've tended to define other diagnostics and transformations on rvars, which tend to return rvars (the above returns the draws and the diagnostics of the smoothed rvar as a flattened list of arrays).

If you do something like this for pareto_smooth:

pareto_smooth.rvar <- function(x, ...) {
  draws_diags <- summarise_rvar_by_element_with_chains(x, pareto_smooth.default, ...)
  dim(draws_diags) <- dim(draws_diags) %||% length(draws_diags)
  margins <- seq_along(dim(draws_diags))
  list(
    x = rvar(apply(draws_diags, margins, function(x) x[[1]]$x), nchains = nchains(x)),
    khat = apply(draws_diags, margins, function(x) x[[1]]$diagnostics$khat)
  )
}

Then the smoothed sample is returned as an rvar with the same structure as the input, and the khat diagnostic is returned as an array with the same dimensions as the input; e.g.:

set.seed(1234)
x = rvar_rng(rnorm, 4, 1:4)
dim(x) = c(2,2)
pareto_smooth(x)
# $x
# rvar<4000>[2,2] mean ± sd:
#      [,1]      [,2]     
# [1,] 1 ± 1.00  3 ± 0.99 
# [2,] 2 ± 0.99  4 ± 1.00 
# 
# $khat
#             [,1]        [,2]
# [1,] -0.12174198 -0.07504085
# [2,] -0.05649535 -0.01960201

I think something along those lines would make the most sense to me, unless I've misunderstood something about the intent here.

[n-kall]

Thanks for this suggestion. I have modified the code accordingly, and also made it handle the extra diagnostics. The output is now like this for rvars:

pareto_smooth(x, extra_diags = TRUE)
$x
rvar<4000>[2,2] mean ± sd:
     [,1]      [,2]
[1,] 1 ± 1.00  3 ± 0.99
[2,] 2 ± 0.99  4 ± 1.00

$diagnostics
$diagnostics$khat
            [,1]        [,2]
[1,] -0.12174198 -0.07504085
[2,] -0.05649535 -0.01960201

$diagnostics$min_ss
     [,1] [,2]
[1,]   10   10
[2,]   10   10

$diagnostics$khat_threshold
          [,1]      [,2]
[1,] 0.7223811 0.7223811
[2,] 0.7223811 0.7223811

$diagnostics$convergence_rage
     [,1] [,2]
[1,]    1    1
[2,]    1    1

[paul-buerkner]

@mjskay can you check if this is what you had in mind?

[mjskay]

yeah, looks good to me!

[n-kall]

Thanks for the review and comments @paul-buerkner and @mjskay! I'll have an updated version in the coming days with these suggestions addressed.

[n-kall]

After thinking about this PR more, I am going to split it into two separate PRs:

1. This PR: the basic Pareto smoothing functions and (point estimate) diagnostics (pretty much complete)
2. Another PR: the functions for the marginal posterior of the Pareto k (needs more work).

As I still need to think a bit more about 2, I will remove the parts pertaining to it from this PR, so they do not hold it up.

[n-kall]

Note, a fix is needed for the extra_diags:

• if khat > 1, min_ss should be Inf

[paul-buerkner]

Please let me know, once the PR is ready to merge from your side.

[n-kall]

There may be a slight delay with finishing this PR as the repository of the fork seems to be corrupted on GitHub and is not able to accept new commits or even be cloned. I've filed a ticket with GitHub, hoping for a quick response.

[n-kall]

After discussion with @avehtari we decided to make some further adjustments:

• separate the extra_diags into another function pareto_diags
• allow pareto_smooth to only return the smoothed draws, and not diagnostics (for use in mutate_variables for example)
• add interpretation details into the documentation

[github-actions]

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

•   :ballot_box_with_check:as_draws_array: 158ms -> 158ms [-1.11%, +0.64%]
•   :ballot_box_with_check:as_draws_df: 128ms -> 127ms [-3.07%, +1.89%]
•   :ballot_box_with_check:as_draws_list: 272ms -> 275ms [-0.71%, +2.47%]
•   :ballot_box_with_check:as_draws_matrix: 49.1ms -> 48.9ms [-1.94%, +1.28%]
•   :ballot_box_with_check:as_draws_rvars: 247ms -> 250ms [-0.27%, +2.74%]
•   :ballot_box_with_check:summarise_draws_100_variables: 1.08s -> 1.08s [-0.76%, +0.81%]
•   :rocket:summarise_draws_10_variables: 193ms -> 121ms [-37.75%, -36.63%]
  Further explanation regarding interpretation and methodology can be found in the documentation.

[github-actions]

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

•   :ballot_box_with_check:as_draws_array: 131ms -> 130ms [-4.8%, +3.55%]
•   :ballot_box_with_check:as_draws_df: 111ms -> 112ms [-3.44%, +4.67%]
•   :ballot_box_with_check:as_draws_list: 230ms -> 236ms [-0.31%, +5.25%]
•   :ballot_box_with_check:as_draws_matrix: 40.2ms -> 40.8ms [-2.79%, +5.94%]
•   :ballot_box_with_check:as_draws_rvars: 211ms -> 212ms [-4.11%, +4.5%]
•   :ballot_box_with_check:summarise_draws_100_variables: 871ms -> 888ms [-2.98%, +6.69%]
•   :rocket:summarise_draws_10_variables: 157ms -> 98.3ms [-39.8%, -34.82%]
  Further explanation regarding interpretation and methodology can be found in the documentation.

[github-actions]

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

•   :ballot_box_with_check:as_draws_array: 165ms -> 166ms [-0.76%, +1.2%]
•   :rocket:as_draws_df: 142ms -> 57.6ms [-60.72%, -58.26%]
•   :ballot_box_with_check:as_draws_list: 301ms -> 303ms [-0.68%, +1.7%]
•   :ballot_box_with_check:as_draws_matrix: 51.4ms -> 51.1ms [-2.54%, +1.62%]
•   :ballot_box_with_check:as_draws_rvars: 276ms -> 277ms [-1.05%, +1.79%]
•   :rocket:summarise_draws_100_variables: 1.16s -> 1.15s [-2.02%, -0.64%]
•   :ballot_box_with_check:summarise_draws_10_variables: 128ms -> 126ms [-3.74%, +1.17%]
  Further explanation regarding interpretation and methodology can be found in the documentation.

[avehtari]

Went through the doc parts. Found things to improve and some typos

[avehtari]

Ah, now that I rehink this, ess_tail would be more appropriate when r_eff is used to determine tail_length

[github-actions]

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

•   :ballot_box_with_check:as_draws_array: 163ms -> 162ms [-0.77%, +0.09%]
•   :ballot_box_with_check:as_draws_df: 55.6ms -> 55.6ms [-1.03%, +0.98%]
•   :ballot_box_with_check:as_draws_list: 288ms -> 288ms [-1.01%, +0.41%]
•   :ballot_box_with_check:as_draws_matrix: 50.2ms -> 50ms [-2.41%, +1.5%]
•   :ballot_box_with_check:as_draws_rvars: 264ms -> 266ms [-0.64%, +2.06%]
•   :rocket:summarise_draws_100_variables: 1.12s -> 1.11s [-1.55%, -0.38%]
•   :ballot_box_with_check:summarise_draws_10_variables: 124ms -> 124ms [-0.32%, +0.64%]
  Further explanation regarding interpretation and methodology can be found in the documentation.

[avehtari]

docs ok now

[paul-buerkner]

Great! @avehtari can I merge the PR now?

[avehtari]

I approve merging

[paul-buerkner]

Perfect, thank you both so much for adding this functionality to posterior!

[jgabry]

Perfect, thank you both so much for adding this functionality to posterior!

Indeed, thank you! I've been preoccupied with a bunch of other Stan R package issues, so I hadn't looked through this yet. Very cool!