Add function for taking expectations wrt smoothed weights

[sethaxen]

This PR adds the API function expectation, which allows one to take some expectation wrt to PSIS-smoothed weights. This is the main function needed to compute expectations wrt LOO posteriors, though it is not limited to this application.

Some notes:

• Currently the expectand must have shape (draws, [chains, [params...]]), i.e. it might share the data dimensions. This allows the function to be used for data-dependent scalar expectations e.g. loo_pit. However, one might have data-independent dimensions with or without data dimensions, e.g. for computing the LOO posterior means of all marginals. Supporting this could require a keyword that specifies which (trailing) dims in x are unrelated to the data or draws.
• Currently we do not compute function-specific diagnostics. This requires something like Make smoothing optional #23, so that we can compute diagnostics without then smoothing.
• Maybe a better way to compute expectations like mean is on the log-scale. However, this would require us to limit which expectations we can support.

Summary by CodeRabbit

• New Features

  • Expanded statistical calculation capabilities with the addition of an expectation function for computing the expectation of values using smoothed importance weights.

• Documentation

  • Updated API documentation to reflect new statistical functions and usage.

• Refactor

  • Integrated the StatsBase module to enhance statistical operations within the module.

[codecov]

Codecov Report

Attention: 24 lines in your changes are missing coverage. Please review.

Comparison is base (3e6e5ec) 97.44% compared to head (7dea818) 89.59%.

Files	Patch %	Lines
src/expectation.jl	0.00%	24 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #61      +/-   ##
==========================================
- Coverage   97.44%   89.59%   -7.85%     
==========================================
  Files           6        7       +1     
  Lines         274      298      +24     
==========================================
  Hits          267      267              
- Misses          7       31      +24     

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

[sethaxen]

@coderabbitai review

[coderabbitai]

Walkthrough

The Julia package has been expanded with new statistical capabilities. Specifically, an "expectation.jl" file was introduced, adding an expectation function for calculating the expectation of values using smoothed importance weights. This enhancement integrates with the "PSIS.jl" file, which now includes the "StatsBase" module, allowing for a broader range of statistical operations.

Changes

File(s)	Change Summary
docs/src/api.md	The unidiff snippet provided was incomplete, lacking the context needed for a summary.
src/PSIS.jl, src/.../expectation.jl	Added "expectation.jl" and "StatsBase" module to "PSIS.jl". Introduced expectation function for computing weighted expectations.

🐇✨
To code anew, we hop with glee,
With stats and weights, we set data free.
In Julia's realm, we craft and weave,
Expectations formed, in numbers we believe. 🌟

Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media?

Share

• X
• Mastodon
• Reddit
• LinkedIn

---

Tips

Chat with CodeRabbit Bot (@coderabbitai)

• You can directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>.
  • Generate unit-tests for this file.
• You can tag CodeRabbit on specific lines of code or entire files in the PR by tagging @coderabbitai in a comment. Examples:
  • @coderabbitai generate unit tests for this file
  • @coderabbitai modularize this function
• You can tag @coderabbitai in a PR comment and ask questions about the PR and the codebase. Examples:
  • @coderabbitai gather interesting statistics about this repository and render them in a table.
  • @coderabbitai show all the console.log statements in this repository.
  • @coderabbitai generate unit tests for the src/utils.ts file.

CodeRabbit Commands (invoked as PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger a review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai help to get help.

Additionally, you can add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.

CodeRabbit Configration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• The JSON schema for the configuration file is available here.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/coderabbit-overrides.v2.json

CodeRabbit Discord Community

Join our Discord Community to get help, request features, and share feedback.

[coderabbitai]

Review Status

Actionable comments generated: 0

Configuration used: CodeRabbit UI

Commits

Files that changed from the base of the PR and between 3e6e5ec and 7dea818.

Files ignored due to filter (1)

• Project.toml

Files selected for processing (3)

• docs/src/api.md (1 hunks)
• src/PSIS.jl (2 hunks)
• src/expectation.jl (1 hunks)

Additional comments: 9

docs/src/api.md (1)

• 9-9: Ensure that the documentation for PSIS.expectation is complete and provides accurate usage examples and descriptions.

src/PSIS.jl (2)

• 6-6: Ensure that all necessary functions from StatsBase are being used and that there are no unnecessary imports.

• 16-16: Confirm that expectation.jl is properly included and that all its dependencies are resolved.

src/expectation.jl (6)

• 1-31: The documentation for the expectation function is clear and informative.

• 32-32: The function signature for expectation is appropriate for the described functionality.

• 36-53: Good implementation of dimension and axis checks to ensure consistency between x and psis_result.weights.

• 58-62: Proper use of StatsBase.AnalyticWeights for computing weighted expectations.

• 64-66: Correct handling of scalar return values when values is a zero-dimensional array.

• 69-73: Effective use of multiple dispatch to handle different types of expectation functions, including quantiles.