From 221b74a403be2e43fa53977585f6ab339206f043 Mon Sep 17 00:00:00 2001
From: "Adam Ginsburg (keflavich)" <keflavich@gmail.com>
Date: Fri, 11 Oct 2024 17:54:05 -0400
Subject: [PATCH] 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):