Fix progressbar with nested compound step samplers #7776

ricardoV94 · 2025-05-05T09:54:43Z

Alternative to #7730 that fixes #7721 and fixes #7724

aloctavodia · 2025-06-21T15:28:55Z

Currently, PyMC-BART works without this. However, a recent PR, which allows multiple BART RVs, fails without this. I have not tried with the alternative #7730.

Copilot

Pull Request Overview

This PR refactors and centralizes progress bar functionality, updates sampling step methods to use the new API, and adds a regression test for nested compound samplers.

Move CustomProgress, RecolorOnFailureBarColumn, ProgressBarManager, and related logic from pymc/util.py into a new pymc/progress_bar.py module.
Update all imports to point to pymc.progress_bar and adjust references to the new _make_progressbar_update_functions API.
Add an integration test (test_progressbar_nested_compound) to guard against regressions in nested CompoundStep progress bars.

Reviewed Changes

Copilot reviewed 18 out of 18 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/test_progress_bar.py	Add regression test for nested compound-step sampling progress.
pymc/variational/inference.py	Update imports to use new `CustomProgress` location.
pymc/util.py	Remove old progress-bar classes and imports.
pymc/tuning/starting.py	Update progress-bar import to new module.
pymc/step_methods/slicer.py	Rename update-stats API to `_make_progressbar_update_functions`.
pymc/step_methods/metropolis.py	Rename and adjust update-stats API.
pymc/step_methods/hmc/*.py	Rename and adjust update-stats API in HMC and NUTS.
pymc/step_methods/compound.py	Rename and aggregate update-stat functions.
pymc/smc/sampling.py	Update progress-bar import.
pymc/sampling/*.py	Update progress-bar imports in various sampling pathways.
pymc/progress_bar.py	New module consolidating progress-bar logic and classes.
.github/workflows/tests.yml	Include the new test file in CI.

Comments suppressed due to low confidence (1)

pymc/progress_bar.py:37

[nitpick] Consider adding unit tests for CustomProgress (e.g., toggling disable and include_headers) to verify that tasks are added, updated, or suppressed as expected.

class CustomProgress(Progress):

pymc/step_methods/slicer.py

codecov · 2025-07-09T15:39:36Z

Codecov Report

Attention: Patch coverage is 85.92233% with 29 lines in your changes missing coverage. Please review.

Project coverage is 92.88%. Comparing base (ae43026) to head (6748186).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
pymc/progress_bar.py	83.55%	25 Missing ⚠️
pymc/step_methods/hmc/nuts.py	20.00%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #7776      +/-   ##
==========================================
- Coverage   92.92%   92.88%   -0.04%     
==========================================
  Files         107      108       +1     
  Lines       18299    18313      +14     
==========================================
+ Hits        17004    17010       +6     
- Misses       1295     1303       +8

Files with missing lines	Coverage Δ
pymc/backends/arviz.py	`95.78% <100.00%> (+0.01%)`	⬆️
pymc/sampling/forward.py	`96.75% <100.00%> (+0.01%)`	⬆️
pymc/sampling/mcmc.py	`91.37% <100.00%> (+0.01%)`	⬆️
pymc/sampling/parallel.py	`87.12% <100.00%> (+0.04%)`	⬆️
pymc/sampling/population.py	`70.83% <100.00%> (ø)`
pymc/smc/sampling.py	`99.30% <100.00%> (+<0.01%)`	⬆️
pymc/step_methods/compound.py	`97.87% <100.00%> (-0.02%)`	⬇️
pymc/step_methods/hmc/base_hmc.py	`92.25% <100.00%> (+0.34%)`	⬆️
pymc/step_methods/hmc/hmc.py	`94.59% <100.00%> (+1.04%)`	⬆️
pymc/step_methods/metropolis.py	`93.13% <100.00%> (-0.10%)`	⬇️
... and 6 more

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

jessegrabowski

@ricardoV94 can you give the changes i pushed a final once over then we can merge this

ricardoV94 · 2025-07-09T16:03:08Z

@aloctavodia can you check if this also fixes the problem on your end?

ricardoV94 · 2025-07-09T16:04:28Z

pymc/step_methods/metropolis.py

+        # Tune is always False, because this step does not have any tunable parameters.
+        stats = {"tune": False}
+        return q, [stats]


Why is it a dict inside a list, and why not just an empty dict?

All step samplers return a dict inside a list, idk

Because the CategoricalGibbsMetropolis has tune declared in its stats_dtypes_shapes, which I interpreted as a promise to return tune from its step function. It's incorrect that it wasn't doing this before.

Okay, I would just remove it but fine. Weird thing for the dict inside a list xD, open an issue to stop doing that?

technically It's your PR. I think removing the stats_dtypes_shapes is also a good way to go.

We still need to return an empty dictionary though, otherwise the zip(update_stats_fns, stats) fails (because there is always an update function, but sometimes no stats)

The shape is used by mcbackend I thing

empty dict is fine?

No need to do in this PR though

…ctions`

Every step sampler can now decide whether sampling is failing or not by setting "failing" in the returned update dict

ricardoV94 · 2025-07-15T10:30:48Z

@jessegrabowski I cleanup a bit the PR, if tests pass I think we're good. Wanna take a last look?

jessegrabowski

One docstring out of date, and a small issue I want to bring to your attention.

pymc/progress_bar.py

pymc/step_methods/hmc/hmc.py

ricardoV94 · 2025-07-15T14:51:07Z

Oops my PR got merged stale and a test is failing

ricardoV94 · 2025-07-15T14:51:38Z

The other failure is ours

jessegrabowski · 2025-07-15T15:55:22Z

There was a test with a hardcoded list of expected stats. It should pass now (not sure about the pre-commit problem)

ricardoV94 · 2025-07-15T16:40:53Z

pre-commit seems legit

ricardoV94 requested review from jessegrabowski and Copilot May 5, 2025 09:54

github-actions bot added the bug label May 5, 2025

This comment was marked as outdated.

Sign in to view

aloctavodia mentioned this pull request Jun 16, 2025

Add support for multiple BART random variables per model. pymc-devs/pymc-bart#231

Open

ricardoV94 force-pushed the fix_progress_bar branch from 46226a8 to 4b4c28d Compare July 9, 2025 12:08

jessegrabowski requested a review from Copilot July 9, 2025 15:23

Copilot AI reviewed Jul 9, 2025

View reviewed changes

pymc/step_methods/slicer.py Show resolved Hide resolved

jessegrabowski approved these changes Jul 9, 2025

View reviewed changes

ricardoV94 commented Jul 9, 2025

View reviewed changes

ricardoV94 marked this pull request as ready for review July 9, 2025 16:04

jessegrabowski mentioned this pull request Jul 10, 2025

Fix progress bar error when nested CompoundStep samplers are assigned #7730

Closed

11 tasks

ricardoV94 and others added 6 commits July 15, 2025 11:59

Fix progressbar with nested compound step samplers

c3be085

Move progressbar code to its own module

96be267

Rename make_update_stats_functions to `_make_progressbar_update_fun…

09be7da

…ctions`

Cleanup progress_bar

bfd9189

Abstract special behavior of NUTS divergences in ProgressBar

129d988

Every step sampler can now decide whether sampling is failing or not by setting "failing" in the returned update dict

Correct stats_dtypes_shapes is CategoricalGibbsMetropolis

6748186

ricardoV94 force-pushed the fix_progress_bar branch from c341491 to 6748186 Compare July 15, 2025 10:29

jessegrabowski approved these changes Jul 15, 2025

View reviewed changes

pymc/progress_bar.py Outdated Show resolved Hide resolved

pymc/step_methods/hmc/hmc.py Outdated Show resolved Hide resolved

jessegrabowski added 2 commits July 15, 2025 22:23

Track divergences in HMC stats

3297370

Track divergences in HMC stats

a27c9a4

docstring update

61c2473

Uh oh!

Fix progressbar with nested compound step samplers #7776

Are you sure you want to change the base?

Fix progressbar with nested compound step samplers #7776

Conversation

ricardoV94 commented May 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

aloctavodia commented Jun 21, 2025

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

codecov bot commented Jul 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented Jul 9, 2025

Uh oh!

ricardoV94 Jul 9, 2025

Choose a reason for hiding this comment

Uh oh!

jessegrabowski Jul 9, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jul 9, 2025

Choose a reason for hiding this comment

Uh oh!

jessegrabowski Jul 9, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jul 9, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jul 9, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jul 9, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented Jul 15, 2025

Uh oh!

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ricardoV94 commented Jul 15, 2025

Uh oh!

ricardoV94 commented Jul 15, 2025

Uh oh!

jessegrabowski commented Jul 15, 2025

Uh oh!

ricardoV94 commented Jul 15, 2025

Uh oh!

Uh oh!

ricardoV94 commented May 5, 2025 •

edited

Loading

codecov bot commented Jul 9, 2025 •

edited

Loading