Updating tqdm progress bar inside vectorized function

Question

I have a function which has the following call signature:

import numpy as np
@np.vectorize
def evolve_system(a0, e0, beta, m1, m2, lt, pbar):
    ...
    pbar.update(1)
    ...

    return

and is called like this:

from tqdm import tqdm

with tqdm(total=len(df)) as pbar:
    n, m, ef, Pf, c = evolve_system(df['a0'].values,
                                    df['e0'].values,
                                    df['beta'].values,
                                    df['m1'].values,
                                    df['m2'].values,
                                    df['lifetime'].values,
                                    pbar
                                   )

where df is a pandas DataFrame. Upon running the code, I get an error with the following traceback:

Traceback (most recent call last):
  File "/home/sean/anaconda3/lib/python3.7/site-packages/julia/pseudo_python_cli.py", line 308, in main
    python(**vars(ns))
  File "/home/sean/anaconda3/lib/python3.7/site-packages/julia/pseudo_python_cli.py", line 59, in python
    scope = runpy.run_path(script, run_name="__main__")
  File "/home/sean/anaconda3/lib/python3.7/runpy.py", line 263, in run_path
    pkg_name=pkg_name, script_name=fname)
  File "/home/sean/anaconda3/lib/python3.7/runpy.py", line 96, in _run_module_code
    mod_name, mod_spec, pkg_name, script_name)
  File "/home/sean/anaconda3/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "PhaseSpace.py", line 15, in <module>
    df_B = takahe.evolve.period_eccentricity(dataframes_Bray[Z].sample(1000))
  File "/home/sean/Documents/takahe/takahe/evolve.py", line 209, in period_eccentricity
    pbar
  File "/home/sean/.local/lib/python3.7/site-packages/numpy/lib/function_base.py", line 2108, in __call__
    return self._vectorize_call(func=func, args=vargs)
  File "/home/sean/.local/lib/python3.7/site-packages/numpy/lib/function_base.py", line 2198, in _vectorize_call
    for x, t in zip(outputs, otypes)])
  File "/home/sean/.local/lib/python3.7/site-packages/numpy/lib/function_base.py", line 2198, in <listcomp>
    for x, t in zip(outputs, otypes)])
ValueError: setting an array element with a sequence.

As far as I can tell, this is due to the pbar argument -- as omitting it from the definition and call causes the code to run.

Is there a clean way around this? Can I call pbar.update() within a vectorized function?

Answer 1

Thanks to the advice of @hpaulj - the simplest solution appears to be to put pbar in the global scope and use it from there, viz

from tqdm import tqdm
import numpy as np

@np.vectorize
def evolve_system(a0, e0, beta, m1, m2, lt):
    global pbar
    ...
    pbar.update(1)
    ...

    return

def main():
    global pbar
    ...
    with tqdm(total=len(df)) as pbar:
        n, m, ef, Pf, c = evolve_system(df['a0'].values,
                                        df['e0'].values,
                                        df['beta'].values,
                                        df['m1'].values,
                                        df['m2'].values,
                                        df['lifetime'].values,
                                       )