Progressbar: switch to tqdm, add description

setup.cfg

@@ -22,6 +22,7 @@ install_requires =
     joblib
     casa-formats-io
     packaging
+    tqdm
 
 [options.extras_require]
 test =

spectral_cube/analysis_utilities.py

@@ -3,11 +3,10 @@
 from astropy import units as u
 from six.moves import zip, range
 from astropy.wcs import WCS
-from astropy.utils.console import ProgressBar
 from astropy import log
 import warnings
 
-from .utils import BadVelocitiesWarning
+from .utils import BadVelocitiesWarning, ProgressBar
 from .cube_utils import _map_context
 from .lower_dimensional_structures import VaryingResolutionOneDSpectrum, OneDSpectrum
 from .spectral_cube import VaryingResolutionSpectralCube, SpectralCube
@@ -284,7 +283,7 @@ def stack_spectra(cube, velocity_surface, v0=None,
     # Create chunks of spectra for read-out.
     chunks = get_chunks(len(xy_posns[0]), chunk_size)
     if progressbar:
-        iterat = ProgressBar(chunks)
+        iterat = ProgressBar(chunks, desc='Stack Spectra: ')
     else:
         iterat = chunks
 

spectral_cube/spectral_cube.py

@@ -20,7 +20,6 @@
 import astropy.wcs
 from astropy import units as u
 from astropy.io.fits import PrimaryHDU, BinTableHDU, Header, Card, HDUList
-from astropy.utils.console import ProgressBar
 from astropy import log
 from astropy import wcs
 from astropy import convolution
@@ -35,6 +34,7 @@
 from . import cube_utils
 from . import wcs_utils
 from . import spectral_axis
+from .utils import ProgressBar
 from .masks import (LazyMask, LazyComparisonMask, BooleanArrayMask, MaskBase,
                     is_broadcastable_and_smaller)
 from .ytcube import ytCube
@@ -515,7 +515,7 @@ def _reduce_slicewise(self, function, fill, check_endian,
         result = next(planes)
 
         if progressbar:
-            progressbar = ProgressBar(self.shape[iterax])
+            progressbar = ProgressBar(self.shape[iterax], desc='Slicewise: ')
             pbu = progressbar.update
         else:
             pbu = lambda: True
@@ -1057,7 +1057,7 @@ def apply_function(self, function, axis=None, weights=None, unit=None,
         out = np.empty([nz, nx, ny]) * np.nan
 
         if progressbar:
-            progressbar = ProgressBar(nx*ny)
+            progressbar = ProgressBar(nx*ny, desc='Apply: ')
             pbu = progressbar.update
         elif update_function is not None:
             pbu = update_function
@@ -2993,7 +2993,7 @@ def apply_async(self, func, callback=None):
             if update_function is not None:
                 pbu = update_function
             elif verbose > 0:
-                progressbar = ProgressBar(self.shape[1]*self.shape[2])
+                progressbar = ProgressBar(self.shape[1]*self.shape[2], desc='Apply parallel: ')
                 pbu = progressbar.update
             else:
                 pbu = object
@@ -3256,7 +3256,7 @@ def spectral_interpolate(self, spectral_grid,
 
         yy,xx = np.indices(self.shape[1:])
         if update_function is None:
-            pb = ProgressBar(xx.size)
+            pb = ProgressBar(xx.size, desc='Spectral Interpolate: ')
             update_function = pb.update
 
         for ix, iy in (zip(xx.flat, yy.flat)):
@@ -3480,7 +3480,7 @@ def makeslice_local(startpoint, axis=axis, nsteps=factor):
             if progressbar:
                 progressbar = ProgressBar
             else:
-                progressbar = lambda x: x
+                progressbar = lambda x, desc: x
 
             # Create a view that will add a blank newaxis at the right spot
             view_newaxis = [slice(None) for ii in range(self.ndim)]
@@ -3491,7 +3491,7 @@ def makeslice_local(startpoint, axis=axis, nsteps=factor):
             dsarr = np.memmap(ntf, mode='w+', shape=newshape, dtype=float)
             ntf2 = tempfile.NamedTemporaryFile()
             mask = np.memmap(ntf2, mode='w+', shape=newshape, dtype=bool)
-            for ii in progressbar(range(newshape[axis])):
+            for ii in progressbar(range(newshape[axis]), desc='Downsample: '):
                 view_fulldata = makeslice_local(ii*factor)
                 view_newdata = makeslice_local(ii, nsteps=1)
 
@@ -4171,7 +4171,7 @@ def convolve_to(self, beam, allow_smaller=False,
             beam_ratio_factors = [1.] * len(convolution_kernels)
 
         if update_function is None:
-            pb = ProgressBar(self.shape[0])
+            pb = ProgressBar(self.shape[0], desc='Convolve: ')
             update_function = pb.update
 
         newdata = np.empty(self.shape)

spectral_cube/utils.py

@@ -6,6 +6,12 @@
 
 bigdataurl = "https://spectral-cube.readthedocs.io/en/latest/big_data.html"
 
+from tqdm.auto import tqdm
+
+def ProgressBar(niter, **kwargs):
+    return tqdm(total=niter, **kwargs)
+
+
 def cached(func):
     """
     Decorator to cache function calls

spectral_cube/visualization-tools.py

@@ -1,6 +1,6 @@
 import os
 import numpy as np
-from astropy.utils.console import ProgressBar
+from .utils import ProgressBar
 
 def check_ffmpeg(ffmpeg_cmd):
     returncode = os.system(f'{ffmpeg_cmd} > /dev/null 2&> /dev/null')
@@ -46,7 +46,7 @@ def make_rgb_movie(cube, prefix, v1, v2, vmin, vmax, ffmpeg_cmd='ffmpeg'):
     p1 = cube.closest_spectral_channel(v1)
     p2 = cube.closest_spectral_channel(v2)
 
-    for jj, ii in enumerate(ProgressBar(range(p1, p2-1))):
+    for jj, ii in enumerate(ProgressBar(range(p1, p2-1), desc='RGB: ')):
         rgb = np.array([cube[ii+2], cube[ii+1], cube[ii]]).T.swapaxes(0, 1)
 
         # in case you manually set min/max
@@ -114,7 +114,7 @@ def make_multispecies_rgb(cube_r, cube_g, cube_b, prefix, v1, v2, vmin, vmax,
     p1 = cube_r.closest_spectral_channel(v1)
     p2 = cube_r.closest_spectral_channel(v2)
 
-    for jj, ii in enumerate(ProgressBar(range(p1, p2+1))):
+    for jj, ii in enumerate(ProgressBar(range(p1, p2+1), desc='RGB multi: ')):
         rgb = np.array([cube_r[ii].value,
                         cube_g[ii].value,
                         cube_b[ii].value

spectral_cube/ytcube.py

@@ -5,7 +5,7 @@
 import subprocess
 import numpy as np
 import time
-from astropy.utils.console import ProgressBar
+from .utils import ProgressBar
 from astropy import log
 import warnings