Add response_type parameter for expectation vs prediction intervals#643
Open
drbenvincent wants to merge 10 commits intomainfrom
Open
Add response_type parameter for expectation vs prediction intervals#643drbenvincent wants to merge 10 commits intomainfrom
drbenvincent wants to merge 10 commits intomainfrom
Conversation
Fixes #295 This adds an hdi_type parameter to _bayesian_plot methods across all experiment classes, allowing users to choose between: - 'expectation': HDI of model expectation (μ), excluding observation noise (default, current behavior) - 'prediction': HDI of posterior predictive (ŷ), including observation noise Changes: - Add HdiType type alias and add_hdi_annotation() helper to plot_utils.py - Update _bayesian_plot() methods in: - InterruptedTimeSeries - SyntheticControl - DifferenceInDifferences - PrePostNEGD - RegressionDiscontinuity - RegressionKink - Add text annotation to figures explaining what the HDI represents - Update docstrings with comprehensive parameter documentation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #643 +/- ##
==========================================
+ Coverage 94.35% 94.52% +0.16%
==========================================
Files 44 46 +2
Lines 7517 7853 +336
Branches 456 478 +22
==========================================
+ Hits 7093 7423 +330
- Misses 262 265 +3
- Partials 162 165 +3 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
- Move annotation text lower to avoid overlap with axis labels - Add subplots_adjust to make room for annotation below axes - Add show_hdi_annotation parameter (default True) to all _bayesian_plot methods - Users can set show_hdi_annotation=False to hide the text annotation
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
drbenvincent
changed the title
Introduces the hdi_type parameter to allow users to choose between 'expectation' (model mean, excluding observation noise) and 'prediction' (posterior predictive, including observation noise) for HDI intervals in plots, effect summaries, and data extraction methods. Updates plotting and reporting utilities to support this option, adds logging and annotation for HDI type, and refactors experiment classes to calculate impacts and counterfactuals on demand based on hdi_type. Documentation and notebook examples are updated to reflect these changes.
drbenvincent
changed the title
hdi_type parameter for expectation vs prediction HDI intervals
- Add comprehensive tests in test_hdi_type.py for hdi_type='prediction' and show_hdi_annotation=True across all experiment classes (ITS, SC, DiD, RD, RKink, PrePostNEGD) - Add tests for effect_summary and get_plot_data_bayesian with prediction - Add unit tests for add_hdi_annotation and logging functions in test_plot_utils.py - Fix bug in _extract_window where it assumed result.post_y exists (ITS only); now correctly handles result.datapost_treated for SyntheticControl as well - These tests specifically cover the partial coverage branches reported by Codecov
OriolAbril
reviewed
Jan 17, 2026
Contributor
OriolAbril
left a comment
OriolAbril
left a comment
There was a problem hiding this comment.
I would use a different name for the argument so it doesn't include hdi in it. Some options could be response, response_type, effect_type or target (bambi uses a mix of string arg target and boolean pps argument for a very similar purpose; I am not a very big fan of this split though).
The main reason is avoiding otherwise unnecessary future changes. I see quite probable at some point the kind credible interval will be a user choice between HDI or ETI, so having hdi in the name would be a bit confusing when using quantile based intervals instead. It might also be interesting to add different ways to visualize the uncertainty in addition to the shaded band generated from the credible interval such as "spaghetti"/ensemble plots or projected histogram view
Align plotting and summary APIs with response_type naming, update logging helpers, and adjust tests for the new parameter name.
drbenvincent
changed the title
hdi_type parameter for expectation vs prediction HDI intervals
Cover datapre/datapost properties for ITS and synthetic control, and assert RegressionKink gradient change uses the instance epsilon.
Cover prediction response_type in three-period ITS analysis and refresh the its_covid notebook outputs.
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Allow causal impact plots and effect summaries to show posterior predictive outcomes, not just model expectation.
This PR adds a
response_typeparameter across experiment classes, allowing users to choose whether intervals represent:\"expectation\"(default): model expectation (μ), excludes observation noise. Narrower intervals for parameter uncertainty.\"prediction\": posterior predictive (ŷ), includes observation noise (σ). Wider intervals for full predictive uncertainty.Fixes #295
Changes
Core infrastructure (
causalpy/plot_utils.py)ResponseTypetype alias forLiteral[\"expectation\", \"prediction\"]add_hdi_annotation()helper to annotate plots_log_response_type_info_once()and_log_response_type_effect_summary_once()for one-time guidanceModel changes (
causalpy/pymc_models.py)response_typeparameter tocalculate_impact()response_typeto internal variable names (mu/y_hat)Experiment classes updated
All
_bayesian_plot()methods:InterruptedTimeSeries,SyntheticControl,DifferenceInDifferences,PrePostNEGD,RegressionDiscontinuity,RegressionKinkChanges to each:
response_typeparameter with default\"expectation\"show_hdi_annotationparameter with defaultFalsemuory_hatfromposterior_predictivebased onresponse_typeAdditional methods with
response_typesupport (ITS & SC only):effect_summary()get_plot_data_bayesian()analyze_persistence()(ITS only)Reporting changes (
causalpy/reporting.py)response_typeparameter to_extract_window()and_extract_counterfactual()y_hat-based impacts whenresponse_type=\"prediction\"Base class (
causalpy/experiments/base.py)response_typeparameter to abstracteffect_summary()signatureTests
causalpy/tests/test_experiment_refactor.pyanalyze_persistence(response_type=\"prediction\")in ITSNotebook
docs/source/notebooks/its_covid.ipynbto update outputsResponse type applicability
User communication strategy
Instead of cluttering plots with annotations by default:
show_hdi_annotation=Falseshow_hdi_annotation=TrueUsage
Design decisions
\"expectation\"/\"prediction\"rather than internal variable names\"mu\"/\"y_hat\"muby default;y_hat-based impacts calculated on requestTesting
pre-commit run --all-filesChecklist
ResponseTypealias used throughout)its_covid.ipynb)response_typeparameter📚 Documentation preview 📚: https://causalpy--643.org.readthedocs.build/en/643/