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 95.12195% with 10 lines in your changes missing coverage. Please review.

Project coverage is 92.99%. Comparing base (dea3aef) to head (2b50d67).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
pymc/progress_bar.py	93.37%	10 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #7776      +/-   ##
==========================================
+ Coverage   89.19%   92.99%   +3.80%     
==========================================
  Files         107      108       +1     
  Lines       18313    18327      +14     
==========================================
+ Hits        16334    17043     +709     
+ Misses       1979     1284     -695

Files with missing lines	Coverage Δ
pymc/backends/arviz.py	`95.78% <100.00%> (+0.01%)`	⬆️
pymc/logprob/abstract.py	`97.05% <100.00%> (+7.50%)`	⬆️
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%> (+2.51%)`	⬆️
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%)`	⬆️
... and 7 more

... and 18 files with indirect coverage changes

🚀 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?

pymc/step_methods/metropolis.py

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

aloctavodia · 2025-07-16T07:33:57Z

I can confirm that this fixes the problem in PyMC-BART

…ctions`

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

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

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

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 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

Fix missing import

9f3a119

ricardoV94 force-pushed the fix_progress_bar branch from 61c2473 to 6ff478a Compare July 16, 2025 09:57

ricardoV94 requested a review from jessegrabowski July 16, 2025 09:57

ricardoV94 added 4 commits July 16, 2025 12:03

Fix progressbar with nested compound step samplers

7164bc6

Move progressbar code to its own module

fd230d3

Rename make_update_stats_functions to `_make_progressbar_update_fun…

8704886

…ctions`

Cleanup progress_bar

84f92f6

ricardoV94 and others added 2 commits July 16, 2025 12:03

Abstract special behavior of NUTS divergences in ProgressBar

5482faf

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

2b50d67

ricardoV94 force-pushed the fix_progress_bar branch from 6ff478a to 2b50d67 Compare July 16, 2025 10:04

jessegrabowski approved these changes Jul 16, 2025

View reviewed changes

jessegrabowski merged commit 8a436d8 into pymc-devs:main Jul 16, 2025
25 checks passed

ricardoV94 deleted the fix_progress_bar branch July 16, 2025 14:06

Uh oh!

Fix progressbar with nested compound step samplers #7776

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!

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!

aloctavodia commented Jul 16, 2025

Uh oh!

Uh oh!

Uh oh!

ricardoV94 commented May 5, 2025 •

edited

Loading

codecov bot commented Jul 9, 2025 •

edited

Loading