Replace fastprogress progress bars with rich

[fonnesbeck]

Description

This PR aims to improve the robustness of PyMC's progress bars by implementing them with the rich library in place of fastprogress

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

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

[fonnesbeck]

Still some testing to do, but they appear to play nicely with Jupyter.

[fonnesbeck]

I'm sure someone like @zaxtax could be fancier with the progress bar styling. Feel free to embellish!

[fonnesbeck]

Also, making the progress bars work involved changing how parallel processing works in SMC sampling. Please have a close look.

[codecov]

Codecov Report

Attention: Patch coverage is 93.56436% with 13 lines in your changes are missing coverage. Please review.

Project coverage is 92.29%. Comparing base (4f6831e) to head (3f107e9).
Report is 5 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #7233      +/-   ##
==========================================
- Coverage   92.30%   92.29%   -0.02%     
==========================================
  Files         100      101       +1     
  Lines       16895    16891       -4     
==========================================
- Hits        15595    15589       -6     
- Misses       1300     1302       +2     

Files	Coverage Δ
pymc/sampling/mcmc.py	87.71% <100.00%> (ø)
pymc/sampling/parallel.py	88.38% <100.00%> (-0.05%)	⬇️
pymc/smc/sampling.py	99.20% <100.00%> (-0.10%)	⬇️
pymc/stats/log_density.py	98.33% <100.00%> (+0.02%)	⬆️
pymc/util.py	81.09% <100.00%> (+0.16%)	⬆️
pymc/sampling/forward.py	95.94% <94.11%> (+0.03%)	⬆️
pymc/sampling/population.py	73.88% <94.73%> (+0.02%)	⬆️
pymc/tuning/starting.py	91.30% <77.27%> (-0.44%)	⬇️
pymc/variational/inference.py	87.93% <91.42%> (+1.93%)	⬆️

... and 7 files with indirect coverage changes

[ColCarroll]

Can we haz screenshot plz?

[fonnesbeck]

pm.sample is quite simple:

[fonnesbeck]

pm.fit is similar:

[fonnesbeck]

For SMC, since the number of iterations is not fixed I use a spinner for each chain:

[zaxtax]

I love so much of this! Looking forward to reviewing

…

On Tue, 2 Apr 2024, 01:16 Chris Fonnesbeck, ***@***.***> wrote: For SMC, since the number of iterations is not fixed I use a spinner for each chain: image.png (view on web) <https://github.com/pymc-devs/pymc/assets/81476/8a7e49fa-5784-4bcb-b624-475df12d4a9a> — Reply to this email directly, view it on GitHub <#7233 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUOSNG7AR2VGWAIJAPDY3HTCZAVCNFSM6AAAAABFPRMAVKVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDAMZQG43TCMBWGA> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[zaxtax]

Would having a way for people to theme their progress bars be in scope? If so I can offer some suggestions.

…

On Tue, 2 Apr 2024, 01:18 Rob Zinkov, ***@***.***> wrote: I love so much of this! Looking forward to reviewing On Tue, 2 Apr 2024, 01:16 Chris Fonnesbeck, ***@***.***> wrote: > For SMC, since the number of iterations is not fixed I use a spinner for > each chain: > > image.png (view on web) > <https://github.com/pymc-devs/pymc/assets/81476/8a7e49fa-5784-4bcb-b624-475df12d4a9a> > > — > Reply to this email directly, view it on GitHub > <#7233 (comment)>, > or unsubscribe > <https://github.com/notifications/unsubscribe-auth/AAACCUOSNG7AR2VGWAIJAPDY3HTCZAVCNFSM6AAAAABFPRMAVKVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDAMZQG43TCMBWGA> > . > You are receiving this because you were mentioned.Message ID: > ***@***.***> >

[fonnesbeck]

Yes! Feel free to point me toward relevant examples/resources.

[zaxtax]

Yes! Feel free to point me toward relevant examples/resources.

So there are two ways to style. For simple stuff, you can just pass in a styles to track:

import time
from rich.progress import track

for i in track(range(20), description="Processing...", complete_style="rgb(83,145,210)"):
    time.sleep(1) 

But for theming it's probably easiest to take a theme or console object and pass it to Progress

from rich.theme import Theme
custom_theme = Theme({
    "bar.complete": "rgb(83,145,210)",
})

from rich.progress import Progress
from rich.console import Console
import time

with Progress(console=Console(theme=custom_theme)) as progress:
    task1 = progress.add_task("[red]Downloading...", total=100)
    task2 = progress.add_task("[green]Processing...", total=100)
    task3 = progress.add_task("[cyan]Cooking...", total=100)

    while not progress.finished:
        progress.update(task1, advance=0.5)
        progress.update(task2, advance=0.3)
        progress.update(task3, advance=0.9)
        time.sleep(0.02)

The style page (https://rich.readthedocs.io/en/latest/style.html) is the best resource for styling.

[zaxtax]

Stylistically, I associate red with error. I prefer blue to indicate normal operation, green for completion, and red for some error state.

[fonnesbeck]

Running:

[fonnesbeck]

Complete:

[fonnesbeck]

Can now pass custom progressbar themes to samplers!

[zaxtax]

I don't feel confident in reviewing the sampling.parallel code, but otherwise looks good!

[ricardoV94]

PR looks good to me. Was there any open / known issue with the old library?

[zaxtax]

Bad at displaying more than one progress bar fastai/fastprogress#39

…

On Wed, 3 Apr 2024, 15:12 Ricardo Vieira, ***@***.***> wrote: PR looks good to me. Was there any open / known issue with the old library? — Reply to this email directly, view it on GitHub <#7233 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUO3EHEQE56EALF7Y6LY3P52JAVCNFSM6AAAAABFPRMAVKVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDAMZUGU3DSMRRGM> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[fonnesbeck]

These should be more robust progress bars by virtue of being rich text-based. Also the rich project itself appears to be much better-maintained right now.

[ricardoV94]

mypy is not happy

[Armavica]

It looks great on terminal too!

[fonnesbeck]

mypy seems to be complaining about things that have nothing to do with this PR

[Armavica]

mypy seems to be complaining about things that have nothing to do with this PR

It looks like mcmc.py is the file that causes the error with

pymc/sampling/mcmc.py:1195: error: Argument "progressbar_theme" to "ParallelSampler" has incompatible type "Optional[Theme]"; expected "Theme"

[juanitorduz]

I just say this is a fantastic addition 🙌 ! Thanks!

[jucor]

Adding my thanks here!! 👏 👏 👏 🙌

[fonnesbeck]

Thanks! I'm sure they can be improved, so feel free to suggest additional changes.

[jucor]

I've spotted a bug with the progressbar when PyMC is called from R via Reticulate, and have opened a bug report at both Reticulate and Rich. Mentioned them below in case you want to follow what resolution there might be:

• Reticulate does not play nice with Rich progressbars rstudio/reticulate#1632
• [BUG] Rich not playing nice when called from Reticulate Textualize/rich#3416

Thanks for the extra option progressbar=False that you have in PyMC, which allows to override, but those progress bars are so useful due to long computations, that I really still need them :)

[jucor]

I have found the root cause for this problem with progress bars! Rstudio Console does not support cursor motion ANSI control sequences. I have opened a feature request upstream in rstudio/rstudio#14942 , as well as an associated bug report rstudio/rstudio#14941 .

Meanwhile, is there a way to have some basic SMC progress bars working without cursor motion control sequences, even if in a much simpler capacity? Basically just printing out the current version of Beta and a timestamp?