Python Nested Progress Bars for Spherical Harmonic Coefficients using Rich

56 Views Asked by At

I have a set of nested for loops for spherical harmonic coefficients, and I'd like to use Rich to show the progress of the nested components.

import numpy as np
from time import sleep

l = 32
for ll in np.arange(l):
    # I want a progress bar for ll here

    for mm in np.arange(-1*ll, ll+1):
        # I want a progress bar for mm here
        sleep(0.01)
        print(f"l: [{ll}], m: [{mm}]")

This github issue provides a partial solution https://github.com/Textualize/rich/discussions/2272, but assumes that each sub loop has the same number of entries:

from rich.progress import Progress
from time import sleep

with Progress() as pb:
    t1 = pb.add_task('inner', total=10)
    t2 = pb.add_task('outer', total=100)

    for i in range(100):
        for j in range(10):
            print(f"Verbose info! {i, j}")
            sleep(0.1)
            pb.update(task_id=t1, completed=j + 1)
        pb.update(task_id=t2, completed=i + 1)

While I love tqdm, and have achieved partial success using the leave=False keyword, I'm not satisfied with how I loose the last output once both loops have been completed

import numpy as np
from tqdm import tqdm
from time import sleep

l = 32
for ll in tqdm(
    np.arange(l),
    desc=r" ",
    disable=False,
):

    for mm in tqdm(
        np.arange(-1 * ll, ll + 1),
        leave=False,
        desc=r" ",
        disable=False,
    ):
        sleep(0.01)

Any help would be appreciated!

1

There are 1 best solutions below

0
aviso On

If I'm understanding correctly, you want the progress bars when it's running, but you don't care if they remain once the program is finished? But you do care about the print output?

Unfortunately, most progress bar libraries do not handle arbitrary printing to the screen when using progress bars well or at all. Enlighten was specifically written for this use case.

The example below explicitly calls the close and stop methods, but you can also use context managers to do this implicitly. You can also adjust the value of leave to suit your needs. Many more options, such as color and custom formatting, are described in the documentation.

from time import sleep

import enlighten
import numpy as np

manager = enlighten.get_manager()

l = 32
l_pbar = manager.counter(total=l, desc='l', leave=False)
for ll in np.arange(l):
    l_pbar.update()
    m_pbar = manager.counter(total=ll * 2 + 1, desc='m', leave=False)

    for mm in np.arange(-1*ll, ll+1):
        m_pbar.update()
        sleep(0.01)
        print(f"l: [{ll}], m: [{mm}]")

    m_pbar.close()
l_pbar.close()
manager.stop()