Add n_samples=0 support and HDI_Prob in plot_hdi

[pkaf]

feat(plot): Add n_samples=0 support, list hdi_prob in plot_hdi, and include_mean option #2225

Description

1. Allow n_samples=0 in plot_curve()

Previously passing n_samples=0 to plot_curve() would break the function. Now users are allow to put any positive integer for n_samples, 0 or None. For 0 or None users should be now able to plot just the HDI bands without any individual sample lines.

Implementation notes:

When n_samples=0 or None, plot_curve() skips calling plot_samples() entirely.

2. Allow hdi_prob to accept a list in plot_hdi()

Previously plot_hdi() only accepted a single float value for hdi_prob. Not it can accept a list e.g. [0.5, 0.8, 0.95].

Loop from plot_curve is now redundant as plot_hdi can directly handle the loop.

• Closes #

• Related to feat(plot): Add n_samples=0 support, list hdi_prob in plot_hdi, and include_mean option #2225

• Checked that the pre-commit linting/style checks pass. Feel free to comment pre-commit.ci autofix to auto-fix.

• Included tests that prove the fix is effective or that the new feature works

• Added necessary documentation (docstrings and/or example notebooks) using numpydoc format.

• If you are a pro: each commit corresponds to a relevant logical change

---

📚 Documentation preview 📚: https://pymc-marketing--2257.org.readthedocs.build/en/2257/

[cursor]

PR Summary

Medium Risk
Behavior changes in core plotting helpers (axes/subplot reuse and repeated HDI rendering) could alter output or legend/alpha stacking for existing users, but changes are localized to visualization code.

Overview
Updates plotting APIs to support HDI-only curves and multiple HDI bands.

plot_curve now allows n_samples=0 to skip plot_samples and render only HDIs (while still wiring subplot_kwargs/axes appropriately). plot_hdi now accepts a sequence of hdi_prob values (including handling an empty list) and internally loops to draw multiple HDI intervals on the same axes, removing the need for callers to loop externally.

^{Written by Cursor Bugbot for commit 47a8854. This will update automatically on new commits. Configure here.}

[codecov]

Codecov Report

❌ Patch coverage is 88.23529% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 93.14%. Comparing base (c061466) to head (47a8854).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
pymc_marketing/plot.py	88.23%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2257      +/-   ##
==========================================
- Coverage   93.17%   93.14%   -0.03%     
==========================================
  Files          78       78              
  Lines       12460    12472      +12     
==========================================
+ Hits        11609    11617       +8     
- Misses        851      855       +4     

☔ View full report in Codecov by Sentry.
📢 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.

[pkaf]

Up until the second last commit (62aa9e6) all comments are addressed. Somehow I meshed up and made the last commit which rebased. Need help from this point. Think if 0cb0117 can be stashed, we are good.

[juanitorduz]

Thanks @pkaf . I will take a look into the pre-commit hooks error later 🙏

[juanitorduz]

Hi @pkaf the pre-commit errors are

pymc_marketing/plot.py:481: error: List item 0 has incompatible type "None"; expected "float"  [list-item]
pymc_marketing/plot.py:483: error: List item 0 has incompatible type "float | None"; expected "float"  [list-item]
pymc_marketing/plot.py:517: error: Incompatible return value type (got "tuple[Any, ndarray[tuple[Any, ...], dtype[Any]] | None]", expected "tuple[Any, ndarray[tuple[Any, ...], dtype[Any]]]")  [return-value]
pymc_marketing/mmm/budget_optimizer.py:665: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs  [annotation-unchecked]
pymc_marketing/mmm/additive_effect.py:461: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs  [annotation-unchecked]
Found 3 errors in 1 file (checked 80 source files)

You must ensure consistency of the types. You can test locally by running pre-commit run --all-files.

Please let me know if you get stuck so that I can try to fix it.

[pkaf]

Hi @pkaf the pre-commit errors are

pymc_marketing/plot.py:481: error: List item 0 has incompatible type "None"; expected "float"  [list-item]
pymc_marketing/plot.py:483: error: List item 0 has incompatible type "float | None"; expected "float"  [list-item]
pymc_marketing/plot.py:517: error: Incompatible return value type (got "tuple[Any, ndarray[tuple[Any, ...], dtype[Any]] | None]", expected "tuple[Any, ndarray[tuple[Any, ...], dtype[Any]]]")  [return-value]
pymc_marketing/mmm/budget_optimizer.py:665: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs  [annotation-unchecked]
pymc_marketing/mmm/additive_effect.py:461: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs  [annotation-unchecked]
Found 3 errors in 1 file (checked 80 source files)

You must ensure consistency of the types. You can test locally by running pre-commit run --all-files.

Please let me know if you get stuck so that I can try to fix it.

I understood the issue regarding type setting, but I think am stuck. Was further confused as the other piece of code in a similar method was passing the test without a fuss. Can you please help me with this?

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

[cursor]

Dict merge order reversal silently overrides user-provided hdi_prob

Medium Severity

The dict merge order at current_hdi_kwargs = {**hdi_kwargs, **dict(hdi_prob=ihdi_prob)} was reversed compared to the old code ({**dict(hdi_prob=hdi_prob), **hdi_kwargs}). Now ihdi_prob always overrides any hdi_prob the user supplied in hdi_kwargs. When hdi_prob defaults to None, this silently discards a user-provided hdi_kwargs["hdi_prob"], replacing it with None (arviz default). Downstream callers like base.py:plot_curve_hdi, which only expose hdi_kwargs and not hdi_prob directly, lose the ability to control the HDI probability.

 

[juanitorduz]

@pkaf could you please look into this and add a test in https://github.com/pymc-labs/pymc-marketing/blob/main/tests/test_plot.py ?

[juanitorduz]

Thanks @pkaf (and sorry for the late reply). Would you please add some basic tests to ensure the expected behaviour of the changes in https://github.com/pymc-labs/pymc-marketing/blob/main/tests/test_plot.py 🙏 ?