From 0dd0a8597a98ff219ab1c99643d41b31595aeaf4 Mon Sep 17 00:00:00 2001
From: Marshall Perrin <mperrin@stsci.edu>
Date: Thu, 7 Mar 2024 11:43:02 -0500
Subject: [PATCH] generalize progressbar to work on calc_psf as well as
 calc_datacube

---
 poppy/instrument.py | 10 +---------
 poppy/poppy_core.py |  9 ++++++++-
 poppy/utils.py      | 26 ++++++++++++++++++++++++++
 3 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/poppy/instrument.py b/poppy/instrument.py
index 18e771bc..aacb7b1c 100644
--- a/poppy/instrument.py
+++ b/poppy/instrument.py
@@ -335,15 +335,7 @@ def calc_datacube(self, wavelengths, progressbar=False, *args, **kwargs):
             cube[ext].data[0] = psf[ext].data
             cube[ext].header[label_wl(0)] = wavelengths[0]
 
-        if progressbar:
-            # set up an optional progressbar wrapper
-            from tqdm import tqdm
-            import functools
-            iterate_wrapper = functools.partial(tqdm, ncols=80)
-        else:
-            # null wrapper that does nothing, for no progress bar
-            iterate_wrapper = lambda x: x
-
+        iterate_wrapper = utils.get_progressbar_wrapper(progressbar, nwaves=nwavelengths)
         # iterate rest of wavelengths
         for i in iterate_wrapper(range(1, nwavelengths)):
             wl = wavelengths[i]
diff --git a/poppy/poppy_core.py b/poppy/poppy_core.py
index 41a59a56..847ba72d 100644
--- a/poppy/poppy_core.py
+++ b/poppy/poppy_core.py
@@ -1579,6 +1579,7 @@ def calc_psf(self, wavelength=1e-6,
                  source=None,
                  normalize='first',
                  display_intermediates=False,
+                 progressbar=False,
                  inwave=None):
         """Calculate a PSF, either multi-wavelength or monochromatic.
 
@@ -1612,6 +1613,11 @@ def calc_psf(self, wavelength=1e-6,
         display_intermediates: bool, optional
             Display intermediate optical planes? Default is False. This option is incompatible with
             parallel calculations using `multiprocessing`. (If calculating in parallel, it will have no effect.)
+        progressbar : bool
+            Optionally display a progress bar indicator for status
+            while iterating over wavelengths. Note, this requires the
+            optional dependency package 'tqdm', which is not included as
+            a requirement.
 
         Returns
         -------
@@ -1725,7 +1731,8 @@ def calc_psf(self, wavelength=1e-6,
         else:  # ######### single-threaded computations (may still use multi cores if FFTW enabled ######
             if display:
                 plt.clf()
-            for wlen, wave_weight in zip(wavelength, normwts):
+            iterate_wrapper = utils.get_progressbar_wrapper(progressbar, nwaves=len(wavelength))
+            for wlen, wave_weight in iterate_wrapper(zip(wavelength, normwts)):
                 mono_psf, mono_intermediate_wfs = self.propagate_mono(
                     wlen,
                     retain_intermediates=retain_intermediates,
diff --git a/poppy/utils.py b/poppy/utils.py
index d1e355bb..f44cba27 100644
--- a/poppy/utils.py
+++ b/poppy/utils.py
@@ -1835,3 +1835,29 @@ def fftw_load_wisdom(filename=None):
             "optimization measurements (automatically). ")
 
     _loaded_fftw_wisdom = True
+
+# ##################################################################
+#   Progress bar (optional convenience)
+#
+
+def get_progressbar_wrapper(progressbar=True, nwaves=None):
+    """ Utility function to return an iterator that MAY display a progress bar,
+    or may not, depending
+    """
+    if progressbar:
+        # this relies on an optional dependency, tqdm
+        # if it's not present, just don't try to display a progressbar
+        try:
+            from tqdm import tqdm
+        except ImportError:
+            progressbar = False
+
+    if progressbar:
+        import functools
+        # set up an optional progressbar wrapper
+        iterate_wrapper = functools.partial(tqdm, ncols=80, total=nwaves)
+    else:
+        # null wrapper that does nothing, for no progress bar
+        iterate_wrapper = lambda x: x
+
+    return iterate_wrapper