Replace fastprogress progress bars with rich

fonnesbeck commented 3 months ago

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

[x] 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

[x] New feature / enhancement
[ ] Bug fix
[ ] Documentation
[ ] Maintenance
[ ] Other (please specify):

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

fonnesbeck commented 3 months ago

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

fonnesbeck commented 3 months ago

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

fonnesbeck commented 3 months ago

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

codecov[bot] commented 3 months ago

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

[![Impacted file tree graph](https://app.codecov.io/gh/pymc-devs/pymc/pull/7233/graphs/tree.svg?width=650&height=150&src=pr&token=JFuXtOJ4Cb&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)](https://app.codecov.io/gh/pymc-devs/pymc/pull/7233?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) ```diff @@ 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](https://app.codecov.io/gh/pymc-devs/pymc/pull/7233?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | Coverage Δ | | |---|---|---| | [pymc/sampling/mcmc.py](https://app.codecov.io/gh/pymc-devs/pymc/pull/7233?src=pr&el=tree&filepath=pymc%2Fsampling%2Fmcmc.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zYW1wbGluZy9tY21jLnB5) | `87.71% <100.00%> (ø)` | | | [pymc/sampling/parallel.py](https://app.codecov.io/gh/pymc-devs/pymc/pull/7233?src=pr&el=tree&filepath=pymc%2Fsampling%2Fparallel.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zYW1wbGluZy9wYXJhbGxlbC5weQ==) | `88.38% <100.00%> (-0.05%)` | :arrow_down: | | [pymc/smc/sampling.py](https://app.codecov.io/gh/pymc-devs/pymc/pull/7233?src=pr&el=tree&filepath=pymc%2Fsmc%2Fsampling.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zbWMvc2FtcGxpbmcucHk=) | `99.20% <100.00%> (-0.10%)` | :arrow_down: | | [pymc/stats/log\_density.py](https://app.codecov.io/gh/pymc-devs/pymc/pull/7233?src=pr&el=tree&filepath=pymc%2Fstats%2Flog_density.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zdGF0cy9sb2dfZGVuc2l0eS5weQ==) | `98.33% <100.00%> (+0.02%)` | :arrow_up: | | [pymc/util.py](https://app.codecov.io/gh/pymc-devs/pymc/pull/7233?src=pr&el=tree&filepath=pymc%2Futil.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy91dGlsLnB5) | `81.09% <100.00%> (+0.16%)` | :arrow_up: | | [pymc/sampling/forward.py](https://app.codecov.io/gh/pymc-devs/pymc/pull/7233?src=pr&el=tree&filepath=pymc%2Fsampling%2Fforward.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zYW1wbGluZy9mb3J3YXJkLnB5) | `95.94% <94.11%> (+0.03%)` | :arrow_up: | | [pymc/sampling/population.py](https://app.codecov.io/gh/pymc-devs/pymc/pull/7233?src=pr&el=tree&filepath=pymc%2Fsampling%2Fpopulation.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zYW1wbGluZy9wb3B1bGF0aW9uLnB5) | `73.88% <94.73%> (+0.02%)` | :arrow_up: | | [pymc/tuning/starting.py](https://app.codecov.io/gh/pymc-devs/pymc/pull/7233?src=pr&el=tree&filepath=pymc%2Ftuning%2Fstarting.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy90dW5pbmcvc3RhcnRpbmcucHk=) | `91.30% <77.27%> (-0.44%)` | :arrow_down: | | [pymc/variational/inference.py](https://app.codecov.io/gh/pymc-devs/pymc/pull/7233?src=pr&el=tree&filepath=pymc%2Fvariational%2Finference.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy92YXJpYXRpb25hbC9pbmZlcmVuY2UucHk=) | `87.93% <91.42%> (+1.93%)` | :arrow_up: | ... and [7 files with indirect coverage changes](https://app.codecov.io/gh/pymc-devs/pymc/pull/7233/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)

ColCarroll commented 3 months ago

Can we haz screenshot plz?

fonnesbeck commented 3 months ago

pm.sample is quite simple:

fonnesbeck commented 3 months ago

pm.fit is similar:

fonnesbeck commented 3 months ago

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

zaxtax commented 3 months ago

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 https://github.com/pymc-devs/pymc/pull/7233#issuecomment-2030771060, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAACCUOSNG7AR2VGWAIJAPDY3HTCZAVCNFSM6AAAAABFPRMAVKVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDAMZQG43TCMBWGA . You are receiving this because you were mentioned.Message ID: @.***>

zaxtax commented 3 months ago

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 https://github.com/pymc-devs/pymc/pull/7233#issuecomment-2030771060, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAACCUOSNG7AR2VGWAIJAPDY3HTCZAVCNFSM6AAAAABFPRMAVKVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDAMZQG43TCMBWGA . You are receiving this because you were mentioned.Message ID: @.***>

fonnesbeck commented 3 months ago

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

zaxtax commented 3 months ago

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 commented 3 months ago

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

fonnesbeck commented 3 months ago

Running:

fonnesbeck commented 3 months ago

Complete:

fonnesbeck commented 3 months ago

Can now pass custom progressbar themes to samplers!

ricardoV94 commented 3 months ago

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

zaxtax commented 3 months ago

Bad at displaying more than one progress bar

https://github.com/fastai/fastprogress/issues/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 https://github.com/pymc-devs/pymc/pull/7233#issuecomment-2034569213, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAACCUO3EHEQE56EALF7Y6LY3P52JAVCNFSM6AAAAABFPRMAVKVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDAMZUGU3DSMRRGM . You are receiving this because you were mentioned.Message ID: @.***>

fonnesbeck commented 3 months ago

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 commented 3 months ago

mypy is not happy

Armavica commented 3 months ago

It looks great on terminal too!

fonnesbeck commented 3 months ago

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

Armavica commented 3 months ago

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 commented 3 months ago

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

jucor commented 4 days ago

Adding my thanks here!! 👏 👏 👏 🙌

fonnesbeck commented 4 days ago

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

jucor commented 4 days ago

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:

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 commented 2 days ago

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 https://github.com/rstudio/rstudio/issues/14942 , as well as an associated bug report https://github.com/rstudio/rstudio/issues/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?

pymc-devs / pymc