Pass kwargs through moments & add progressbar

setup.cfg

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

spectral_cube/_moments.py

@@ -4,6 +4,7 @@
 
 from .cube_utils import iterator_strategy
 from .np_compat import allbadtonan
+from .utils import ProgressBar
 
 """
 Functions to compute moment maps in a variety of ways
@@ -87,7 +88,7 @@ def _slice1(cube, axis):
     return result / weights
 
 
-def moment_slicewise(cube, order, axis):
+def moment_slicewise(cube, order, axis, progressbar=False):
     """
     Compute moments by accumulating the result 1 slice at a time
     """
@@ -108,18 +109,25 @@ def moment_slicewise(cube, order, axis):
     # possible for mom2, not sure about general case
     mom1 = _slice1(cube, axis)
 
+    if progressbar:
+        progressbar = ProgressBar(cube.shape[axis], desc='Moment raywise: ')
+        pbu = progressbar.update
+    else:
+        pbu = lambda: True
+
     for i in range(cube.shape[axis]):
         view[axis] = i
         plane = cube._get_filled_data(fill=0, view=tuple(view))
         result += (plane *
                    (pix_cen[tuple(view)] - mom1) ** order *
                    pix_size)
         weights += plane * pix_size
+        pbu()
 
     return (result / weights)
 
 
-def moment_raywise(cube, order, axis):
+def moment_raywise(cube, order, axis, progressbar=False):
     """
     Compute moments by accumulating the answer one ray at a time
     """
@@ -129,6 +137,12 @@ def moment_raywise(cube, order, axis):
     pix_cen = cube._pix_cen()[axis]
     pix_size = cube._pix_size_slice(axis)
 
+    if progressbar:
+        progressbar = ProgressBar(out.size, desc='Moment raywise: ')
+        pbu = progressbar.update
+    else:
+        pbu = lambda: True
+
     for x, y, slc in cube._iter_rays(axis):
         # the intensity, i.e. the weights
         include = cube._mask.include(data=cube._data, wcs=cube._wcs, view=slc,
@@ -151,6 +165,8 @@ def moment_raywise(cube, order, axis):
         ordern /= data.sum()
 
         out[x, y] = ordern
+
+        pbu()
     return out
 
 def moment_cubewise(cube, order, axis):

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
@@ -430,7 +430,8 @@ def apply_numpy_function(self, function, fill=np.nan,
                                          includemask=includemask,
                                          progressbar=progressbar, **kwargs)
         elif how == 'ray':
-            out = self.apply_function(function, **kwargs)
+            out = self.apply_function(function, progressbar=progressbar,
+                                      **kwargs)
         elif how not in ['auto', 'cube']:
             warnings.warn("Cannot use how=%s. Using how=cube" % how,
                           UnsupportedIterationStrategyWarning)
@@ -515,7 +516,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
@@ -735,6 +736,7 @@ def mad_std(self, axis=None, how='cube', **kwargs):
                                              unit=self.unit,
                                              projection=projection,
                                              ignore_nan=True,
+                                             **kwargs
                                             )
         elif how == 'slice' and hasattr(axis, '__len__') and len(axis) == 2:
             return self.apply_numpy_function(stats.mad_std,
@@ -1057,7 +1059,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
@@ -1585,7 +1587,7 @@ def _pix_size(self):
 
         return dspectral, dy, dx
 
-    def moment(self, order=0, axis=0, how='auto'):
+    def moment(self, order=0, axis=0, how='auto', **kwargs):
         """
         Compute moments along the spectral axis.
 
@@ -1662,7 +1664,7 @@ def moment(self, order=0, axis=0, how='auto'):
             return ValueError("Invalid how. Must be in %s" %
                               sorted(list(dispatch.keys())))
 
-        out = dispatch[how](self, order, axis)
+        out = dispatch[how](self, order, axis, **kwargs)
 
         # apply units
         if order == 0:
@@ -1693,31 +1695,31 @@ def moment(self, order=0, axis=0, how='auto'):
         return Projection(out, copy=False, wcs=new_wcs, meta=meta,
                           header=self._nowcs_header)
 
-    def moment0(self, axis=0, how='auto'):
+    def moment0(self, axis=0, how='auto', **kwargs):
         """
         Compute the zeroth moment along an axis.
 
         See :meth:`moment`.
         """
-        return self.moment(axis=axis, order=0, how=how)
+        return self.moment(axis=axis, order=0, how=how, **kwargs)
 
-    def moment1(self, axis=0, how='auto'):
+    def moment1(self, axis=0, how='auto', **kwargs):
         """
         Compute the 1st moment along an axis.
 
         For an explanation of the ``axis`` and ``how`` parameters, see :meth:`moment`.
         """
-        return self.moment(axis=axis, order=1, how=how)
+        return self.moment(axis=axis, order=1, how=how, **kwargs)
 
-    def moment2(self, axis=0, how='auto'):
+    def moment2(self, axis=0, how='auto', **kwargs):
         """
         Compute the 2nd moment along an axis.
 
         For an explanation of the ``axis`` and ``how`` parameters, see :meth:`moment`.
         """
-        return self.moment(axis=axis, order=2, how=how)
+        return self.moment(axis=axis, order=2, how=how, **kwargs)
 
-    def linewidth_sigma(self, how='auto'):
+    def linewidth_sigma(self, how='auto', **kwargs):
         """
         Compute a (sigma) linewidth map along the spectral axis.
 
@@ -1726,15 +1728,15 @@ def linewidth_sigma(self, how='auto'):
         with np.errstate(invalid='ignore'):
             with warnings.catch_warnings():
                 warnings.simplefilter("ignore", VarianceWarning)
-                return np.sqrt(self.moment2(how=how))
+                return np.sqrt(self.moment2(how=how, **kwargs))
 
-    def linewidth_fwhm(self, how='auto'):
+    def linewidth_fwhm(self, how='auto', **kwargs):
         """
         Compute a (FWHM) linewidth map along the spectral axis.
 
         For an explanation of the ``how`` parameter, see :meth:`moment`.
         """
-        return self.linewidth_sigma() * SIGMA2FWHM
+        return self.linewidth_sigma(**kwargs) * SIGMA2FWHM
 
     @property
     def spectral_axis(self):
@@ -2993,7 +2995,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 +3258,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)):
@@ -3478,9 +3480,10 @@ def makeslice_local(startpoint, axis=axis, nsteps=factor):
             newshape = tuple(newshape)
 
             if progressbar:
-                progressbar = ProgressBar
+                progressbar = ProgressBar(newshape[axis], desc='Downsample: ')
+                pbu = progressbar.update
             else:
-                progressbar = lambda x: x
+                pbu = lambda: True
 
             # 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 +3494,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 range(newshape[axis]):
                 view_fulldata = makeslice_local(ii*factor)
                 view_newdata = makeslice_local(ii, nsteps=1)
 
@@ -3500,6 +3503,7 @@ def makeslice_local(startpoint, axis=axis, nsteps=factor):
 
                 dsarr[view_newdata] = estimator(to_average, axis)[view_newaxis]
                 mask[view_newdata] = np.any(to_anyfy, axis).astype('bool')[view_newaxis]
+                pbu()
 
 
         # the slice should just start at zero; we had factor//2 here earlier,
@@ -4171,7 +4175,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