From 4f9ca66202ca3d992ea7417a977e58e1cc40c445 Mon Sep 17 00:00:00 2001 From: "Adam Ginsburg (keflavich)" Date: Sun, 16 Jun 2024 14:51:03 -0400 Subject: [PATCH 01/12] switch to tqdm's progressbar by default, but enhance astropy's so it can also provide a desc --- spectral_cube/analysis_utilities.py | 5 ++--- spectral_cube/spectral_cube.py | 16 ++++++++-------- spectral_cube/utils.py | 19 +++++++++++++++++++ spectral_cube/visualization-tools.py | 6 +++--- 4 files changed, 32 insertions(+), 14 deletions(-) diff --git a/spectral_cube/analysis_utilities.py b/spectral_cube/analysis_utilities.py index 0818e4983..f54167cec 100644 --- a/spectral_cube/analysis_utilities.py +++ b/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 diff --git a/spectral_cube/spectral_cube.py b/spectral_cube/spectral_cube.py index e5492c0e5..69b8d9d21 100644 --- a/spectral_cube/spectral_cube.py +++ b/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) diff --git a/spectral_cube/utils.py b/spectral_cube/utils.py index c2983b0d4..421ede7b5 100644 --- a/spectral_cube/utils.py +++ b/spectral_cube/utils.py @@ -6,6 +6,25 @@ bigdataurl = "https://spectral-cube.readthedocs.io/en/latest/big_data.html" +try: + from tqdm.auto import tqdm as ProgressBar +except ImportError: + from astropy.utils.console import ProgressBar as pb + class ProgressBar(pb): + def __init__(self, *args, desc=None, **kwargs): + super().__init__(*args, **kwargs) + self.desc = desc + + def _handle_resize(self, signum=None, frame=None): + terminal_width = get_terminal_size().columns + self._bar_length = terminal_width - 37 - (len(self.desc) if self.desc else 0) + + def __enter__(self): + if self.desc: + print(self.desc, end='') + return self + + def cached(func): """ Decorator to cache function calls diff --git a/spectral_cube/visualization-tools.py b/spectral_cube/visualization-tools.py index f4904db5f..1f95adacf 100644 --- a/spectral_cube/visualization-tools.py +++ b/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 From bfefc6333dd5332dbef0fec7022eec924c6cc864 Mon Sep 17 00:00:00 2001 From: "Adam Ginsburg (keflavich)" Date: Sun, 16 Jun 2024 15:00:09 -0400 Subject: [PATCH 02/12] astropy progressbar requires more work to support desc; instead, we just silently ignore it --- spectral_cube/utils.py | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/spectral_cube/utils.py b/spectral_cube/utils.py index 421ede7b5..83d9b98d7 100644 --- a/spectral_cube/utils.py +++ b/spectral_cube/utils.py @@ -10,19 +10,12 @@ from tqdm.auto import tqdm as ProgressBar except ImportError: from astropy.utils.console import ProgressBar as pb + from astropy.utils.console import get_terminal_size class ProgressBar(pb): def __init__(self, *args, desc=None, **kwargs): - super().__init__(*args, **kwargs) + # astropy progressbar won't show this self.desc = desc - - def _handle_resize(self, signum=None, frame=None): - terminal_width = get_terminal_size().columns - self._bar_length = terminal_width - 37 - (len(self.desc) if self.desc else 0) - - def __enter__(self): - if self.desc: - print(self.desc, end='') - return self + super().__init__(*args, **kwargs) def cached(func): From 11ef4a760a9b3e7be79a7978a533806fcc85d5c6 Mon Sep 17 00:00:00 2001 From: "Adam Ginsburg (keflavich)" Date: Sun, 16 Jun 2024 16:34:24 -0400 Subject: [PATCH 03/12] np.NaN -> np.nan From 13853078649c16ab68609c2dbcf34f5fdff0594b Mon Sep 17 00:00:00 2001 From: "Adam Ginsburg (keflavich)" Date: Sun, 16 Jun 2024 16:35:08 -0400 Subject: [PATCH 04/12] ModuleNotFoundError --- spectral_cube/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spectral_cube/utils.py b/spectral_cube/utils.py index 83d9b98d7..2fd8ab3ee 100644 --- a/spectral_cube/utils.py +++ b/spectral_cube/utils.py @@ -8,7 +8,7 @@ try: from tqdm.auto import tqdm as ProgressBar -except ImportError: +except (ImportError, ModuleNotFoundError): from astropy.utils.console import ProgressBar as pb from astropy.utils.console import get_terminal_size class ProgressBar(pb): From 50077ed730d6249678adb06f1f8af039534f3fb3 Mon Sep 17 00:00:00 2001 From: "Adam Ginsburg (keflavich)" Date: Fri, 11 Oct 2024 18:17:03 -0400 Subject: [PATCH 05/12] remove astropy progressbar --- spectral_cube/utils.py | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/spectral_cube/utils.py b/spectral_cube/utils.py index 2fd8ab3ee..220c40b30 100644 --- a/spectral_cube/utils.py +++ b/spectral_cube/utils.py @@ -6,16 +6,7 @@ bigdataurl = "https://spectral-cube.readthedocs.io/en/latest/big_data.html" -try: - from tqdm.auto import tqdm as ProgressBar -except (ImportError, ModuleNotFoundError): - from astropy.utils.console import ProgressBar as pb - from astropy.utils.console import get_terminal_size - class ProgressBar(pb): - def __init__(self, *args, desc=None, **kwargs): - # astropy progressbar won't show this - self.desc = desc - super().__init__(*args, **kwargs) +from tqdm.auto import tqdm as ProgressBar def cached(func): From 2ca3008c3bf59ebb95e65a0db1164a8eab9e2ae1 Mon Sep 17 00:00:00 2001 From: "Adam Ginsburg (keflavich)" Date: Fri, 11 Oct 2024 18:17:34 -0400 Subject: [PATCH 06/12] remove pbar from ytcube --- spectral_cube/ytcube.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spectral_cube/ytcube.py b/spectral_cube/ytcube.py index 2102d0afe..4bcd99e2d 100644 --- a/spectral_cube/ytcube.py +++ b/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 From 5f9f88555292b7bcb243cbaa53a466e6ac61eb42 Mon Sep 17 00:00:00 2001 From: "Adam Ginsburg (keflavich)" Date: Fri, 11 Oct 2024 18:36:12 -0400 Subject: [PATCH 07/12] add tqdm as dep --- setup.cfg | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.cfg b/setup.cfg index de940362c..a943e2b0d 100644 --- a/setup.cfg +++ b/setup.cfg @@ -22,6 +22,7 @@ install_requires = joblib casa-formats-io packaging + tqdm [options.extras_require] test = From ad3a45c819c7fbc2839b1f9aea607fbf09ff1885 Mon Sep 17 00:00:00 2001 From: "Adam Ginsburg (keflavich)" Date: Fri, 11 Oct 2024 18:52:00 -0400 Subject: [PATCH 08/12] progressbars were set up to take an integer number of iterations --- spectral_cube/utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/spectral_cube/utils.py b/spectral_cube/utils.py index 220c40b30..ad4200f28 100644 --- a/spectral_cube/utils.py +++ b/spectral_cube/utils.py @@ -6,7 +6,10 @@ bigdataurl = "https://spectral-cube.readthedocs.io/en/latest/big_data.html" -from tqdm.auto import tqdm as ProgressBar +from tqdm.auto import tqdm + +def ProgressBar(niter, **kwargs): + return tqdm(total=niter, **kwargs) def cached(func): From 61e6ed0e65bb786c5fdd649f526bda09b3e05716 Mon Sep 17 00:00:00 2001 From: "Adam Ginsburg (keflavich)" Date: Fri, 11 Oct 2024 17:54:05 -0400 Subject: [PATCH 09/12] add moment progressbar and allow kwargs in progressbar --- spectral_cube/_moments.py | 20 ++++++++++++++++++-- spectral_cube/spectral_cube.py | 24 ++++++++++++------------ 2 files changed, 30 insertions(+), 14 deletions(-) diff --git a/spectral_cube/_moments.py b/spectral_cube/_moments.py index b19bcbd16..aee69a815 100644 --- a/spectral_cube/_moments.py +++ b/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,6 +109,12 @@ 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)) @@ -115,11 +122,12 @@ def moment_slicewise(cube, order, axis): (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): diff --git a/spectral_cube/spectral_cube.py b/spectral_cube/spectral_cube.py index 69b8d9d21..ef24d1af4 100644 --- a/spectral_cube/spectral_cube.py +++ b/spectral_cube/spectral_cube.py @@ -1585,7 +1585,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 +1662,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 +1693,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 +1726,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): From db0743126d7e41e5c4bbe9a4d78a2bb644115aad Mon Sep 17 00:00:00 2001 From: "Adam Ginsburg (keflavich)" Date: Sat, 12 Oct 2024 15:57:39 -0400 Subject: [PATCH 10/12] pass progressbar kw for raywise --- spectral_cube/spectral_cube.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/spectral_cube/spectral_cube.py b/spectral_cube/spectral_cube.py index ef24d1af4..2a3160eb5 100644 --- a/spectral_cube/spectral_cube.py +++ b/spectral_cube/spectral_cube.py @@ -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) From 03cb090815c5215ba44d859f5ac3fd37467efb3b Mon Sep 17 00:00:00 2001 From: "Adam Ginsburg (keflavich)" Date: Sat, 12 Oct 2024 17:06:16 -0400 Subject: [PATCH 11/12] progressbar - make more consistent --- spectral_cube/spectral_cube.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/spectral_cube/spectral_cube.py b/spectral_cube/spectral_cube.py index 2a3160eb5..e34a4c02a 100644 --- a/spectral_cube/spectral_cube.py +++ b/spectral_cube/spectral_cube.py @@ -3479,9 +3479,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, desc: 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)] @@ -3492,7 +3493,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]), desc='Downsample: '): + for ii in range(newshape[axis]): view_fulldata = makeslice_local(ii*factor) view_newdata = makeslice_local(ii, nsteps=1) @@ -3501,6 +3502,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, From bb021f5356449e15b0ef6d07c709e2eba947a121 Mon Sep 17 00:00:00 2001 From: "Adam Ginsburg (keflavich)" Date: Sun, 13 Oct 2024 00:02:46 -0400 Subject: [PATCH 12/12] mad_std was missing a kwargs that prevents progressbar --- spectral_cube/spectral_cube.py | 1 + 1 file changed, 1 insertion(+) diff --git a/spectral_cube/spectral_cube.py b/spectral_cube/spectral_cube.py index e34a4c02a..96f463f9f 100644 --- a/spectral_cube/spectral_cube.py +++ b/spectral_cube/spectral_cube.py @@ -736,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,