Improve support (or documentation) for weighted draws

[avehtari]

?weight_draws says Add weights to draws objects, with one weight per draw, for use in subsequent weighting operations. Currently the only supported subsequent operation is resample_draws(), but the documentation doesn't make it clear. Ideally rvars support would be extended to support weighted draws, but if that is infeasible, the documentation should be clarified.

Example:

> x = rvar_rng(rnorm, 4, mean = 1:4, sd = 2)
> w=runif(4000)

weight_draws() is not defined for draws_rvars:

> xw = weight_draws(x, w, log=FALSE)
Error in UseMethod("weight_draws") : 
  no applicable method for 'weight_draws' applied to an object of class "c('rvar', 'vctrs_vctr', 'list')"

There is no error if weighted draws_df is converted to draws_rvars:

> xw = as_draws_rvars(weight_draws(as_draws_df(x), w, log=FALSE))

The default print method gives the unweighted result

> xw
# A draws_rvars: 4000 iterations, 1 chains, and 1 variables
$x: rvar<4000>[4] mean ± sd:
[1] 0.96 ± 2  1.95 ± 2  3.02 ± 2  4.05 ± 2 

and mean fails

> mean(xw)
[1] NA
Warning message:
In mean.default(xw) : argument is not numeric or logical: returning NA

[mjskay]

Thanks, I think there's a couple of different issues here.

weight_draws() is not defined for draws_rvars

It is, I think you were using it on an rvar not a draws_rvars object. Based on your example:

> x = rvar_rng(rnorm, 4, mean = 1:4, sd = 2)
> w = runif(4000)
> x
rvar<4000>[4] mean ± sd:
[1] 1 ± 2  2 ± 2  3 ± 2  4 ± 2 
> 
> d = draws_rvars(x = x)
> d
# A draws_rvars: 4000 iterations, 1 chains, and 1 variables
$x: rvar<4000>[4] mean ± sd:
[1] 1 ± 2  2 ± 2  3 ± 2  4 ± 2 

> 
> weight_draws(d, w)
# A draws_rvars: 4000 iterations, 1 chains, and 1 variables
$x: rvar<4000>[4] mean ± sd:
[1] 1 ± 2  2 ± 2  3 ± 2  4 ± 2 

# ... hidden reserved variables {'.log_weight'}

Perhaps we could have a more useful error message here directing people to put the rvar in a draws_rvars if they want to weight it. Currently rvars do not store / know anything about weights, but draws_rvars mimic the other format by storing weights internally as a separate variable:

> dw = weight_draws(d, w)
> dw$.log_weight
rvar<4000>[1] mean ± sd:
[1] -1 ± 1 

Quite possibly this is not the best approach. I think we could allow rvars to store weights, with some work. This would add a bit of overhead in that we'd want to check if weights are the same when people try to do operations on multiple rvars. We'd also want to decide on the semantics of those operations --- if weights are different do we just throw an error, forcing the user to think carefully about the correct solution (resampling before combining)? Are there situations where we don't want to throw an error --- e.g. if someone wants to (say) add two rvars x and y together and x has weights and y doesn't have any weights, is it correct to add them and give the results the weights of x?

There is no error if weighted draws_df is converted to draws_rvars:

Thanks for catching this, that is a bug --- the weights column should be carried over but is not.

The default print method fives the unweighted result

Yeah, the draws_rvars printing implementation just calls down to the rvar print, which doesn't know anything about the weights in the containing draws_rvars object. If we move weights into rvars this would be solved, or we could include a warning or make a smarter print function for a weighted draws_rvars object.

This raises another question, which is what to do with weighted draws objects in summarise_draws(). As far as I am aware (@paul-buerkner can correct me) the way summarise_draws() works it does not have a way to pass weights down to summary functions, so I don't think there's a straightforward way currently to get weighted means, medians, or quantiles with summarise_draws(). That seems like a bit of a potential headache to fix given its current architecture, but might be worth it? Since most summary functions use ... we could come up a standard argument for passing weights down. It complicates things slightly in that some natural functions, like mean(), don't support weights, so we'd need different implementations of those, perhaps similar to the approach with quantile2().

and mean fails

It that example you are asking for the mean of the entire draws_rvars object. You can pass and rvar to mean() but not a draws_rvars; for the latter I would probably use summarise_draws or lapply:

> summarise_draws(d, mean)
# A tibble: 4 x 2
  variable  mean
  <chr>    <dbl>
1 x[1]      1.02
2 x[2]      1.99
3 x[3]      2.97
4 x[4]      3.97
> lapply(d, mean)
$x
[1] 1.016734 1.986222 2.972354 3.970131

[avehtari]

Thanks for the clarifications. Support for weights would be cool and useful, but as it's also complicated, I'm fine that they are not supported. We could instead add explicit functions and examples of using weights. E.g. loo package has E_loo https://mc-stan.org/loo/reference/E_loo.html. With explicit functions it's possible to make it more clear what is expected to work with weights.

[mjskay]

I looked into this more for rvars and in principle I think it is possible without too much surgery. Since we already have a common function that is used to ensure the number of chains conform between rvars that are used together, we could use that as the same point to ensure weights conform. Beyond that, it's mostly a matter of implementing weighted versions of all summary functions.

For draws formats I think it's a bit more complicated, as I'm not sure if there's a great way to pass down weights in summarise_draws. I'd be curious what @paul-buerkner thinks about that.

[paul-buerkner]

Thank you both for discussing this issue. I have no good opinion on whether we shall support weights in rvars. I leave this @mjskay's discretion on whether he things that can be done and is useful enough for him to invest work in.

With regard to summarize_draws, I think the question relates more generally to https://github.com/stan-dev/posterior/issues/105 and I don't think we shall have this functionality in summarize_draws because the interface is not made for it. Rather, perhaps, we can come up with a second summarizing function that can handle things much like dplyr::summarize does.

[avehtari]

Another worry I have is that if posterior would work well with weights, that it could be a huge work to get, e.g., ggdist to work with weights, too. That is, there is a danger of endless work to extend support for weighted draws if one part works too well. I still think support for weighted draws would be cool, but I feel it would be easier to start with having functions that the user needs to explicitly use. The idea of including weights as .log_weights was to make it easier to delay the resampling step as that step loses useful information (including diagnostics on the distribution of the weights), but rvars is changing the expectations of the users a lot.

[mjskay]

That is, there is a danger of endless work to extend support for weighted draws if one part works too well. I still think support for weighted draws would be cool, but I feel it would be easier to start with having functions that the user needs to explicitly use.

Yeah, unless/until we have time, maybe the short term solution is explicit functions paired with some warnings or messages when people use summarise_draws or the rvar summary functions that summaries do not account for weights?

[n-kall]

I don't think there's a straightforward way currently to get weighted means, medians, or quantiles with summarise_draws()

As part of priorsense I've implemented weighted functions to match default_summary_measures(). Currently, not very user-friendly as they are mainly used internally, but work like this:

x <- example_draws()
x <- weight_draws(x, weights = rexp(ndraws(x)))
summarise_draws(
  x,
  priorsense::mean_weighted,
  priorsense::sd_weighted,
  priorsense::median_weighted,
  priorsense::quantile2_weighted,
  .args = list(weights = weights(x))
)

If it's decided to directly support summaries of weighted draws, the *_weighted functions could perhaps be moved to posterior (they are just wrappers around MatrixStats and Hmisc functions (see here).

[mjskay]

That's an interesting solution @n-kall. @paul-buerkner would a default approach for weighted draws be to pass through a weights argument? Do most of the diagnostic functions have ... in them anyway, which should eat up the weights parameter if it is unused. Then we could modify our diagnostic functions to support weights, and maybe supply weighted versions of base functions like mean/sd (like mean2() or mean_weighted()) but with output column names mapped onto the base names (like "mean"). If these were given in default_summary_measures() and the like, maybe that would work?

[paul-buerkner]

That could be an option, yes. Not sure yet though if it would be a good idea to add ... to all of our rhat and ess functions. We can definitely do it but first we should figure out if there may be another option that is more dplyr::summarize like, essentially a second summarize_draws (just named differently). That would require the user to write a bit more code but is much more flexibile for passing custom arguments.

[n-kall]

Since @mjskay has implemented weighted rvars (#331), I have been revisited the idea of passing rvar objects to the summary functions from summarise_draws (see discussion in #278). This should allow for weight support in summary functions that have rvar methods. I have an implementation in a fork which just adds an optional .use_rvars argument to summarise_draws, and if TRUE it will pass the draws as rvars to the individual summary functions (and if FALSE, passes arrays as currently).

Initial testing seems to indicate that this works without breaking anything, so I am hopeful that this could be a viable way to properly support weights (it would also allow for addressing #341).

[mjskay]

Cool! The core is this commit? Looks pretty straightforward.

I wonder about what we want the long term plan to be: do we want folks to migrate to the version that uses rvars? If so, how do we want to get there? Some ideas:

• We could provide a deprecation warning if .use_rvars is FALSE, then eventually make .use_rvars = TRUE the default.
• We could create a new function (maybe summarise_variables(), after mutate_variables()?) instead of an option on summarise_draws(), and deprecate summarise_draws().
• Or if we don't want to migrate folks to a new function, we could provide a warning (or error) if they try to do summaries on draws that have weights when .use_rvars = FALSE
• Or we could have a value like .use_rvars = "auto", which would pick the rvar approach when weights are present.

Personally, long-term I'd rather folks not have to manually set an option to get summaries that take weights into account, so some solution that gets us there would be ideal.

[n-kall]

Yea, that's the main change. I think it's fairly unlikely that users have been using weights and then unweighted summaries expecting weights to be ignored, so I think (at least) auto detecting should be fine. Although I didn't check what could happen if an rvar is passed to a custom summary function that expects an array

[n-kall]

Ok, so if one uses a summary function that does not support rvars, there can be an error, which we should somehow handle. Simple example is stats::sd. This is actually why I changed the default summaries to use sd rather than stats::sd in the same commit.

summarise_draws(x, stats::sd, .use_rvars = TRUE)
# Error in `as.double()`:
# ! Can't convert `x` <rvar<100,4>> to <double>.

[mjskay]

What about three values to .use_rvars:

• FALSE: doesn't use them
• TRUE: uses them, and if errors happen doesn't try to catch them (as the user has opted in)
• NA or "auto" or something like that (this would be default):
  • if no weights, doesn't use rvars
  • if weights, tries to use rvars
    • if that succeeds, good
    • if that fails, tries to use an array
      • if that succeeds, returns the result with a warning that weights are ignored and that the user should figure out a solution with rvars (either that or throw an error instead of a warning saying they should figure out how to get their summary using rvars since they have weights)
      • if that fails, just rethrow the error (or maybe the original error on rvar) since it's likely not a problem with rvars

[n-kall]

I like the idea. Having 3 options might be a little confusing for users but if the default is sensible most of the time, the TRUE and FALSE options would only be needed for special cases / backward compatibility.

We should be mindful of how this could scale with many variables and several summary functions, particularly catching errors and potential warnings/messages.

I can try to spin something up early next week and see how it feels.

[mjskay]

Cool thanks!

We should be mindful of how this could scale with many variables and several summary functions, particularly catching errors and potential warnings/messages.

Yeah good point. Ideally people should not be using code with the warnings in place, but actually fixing it so the warnings aren't there.... of course we all know that may be wishful thinking....