Add optional progress bar display during datacube calculations

poppy/instrument.py

@@ -298,14 +298,20 @@ def calc_psf(self, outfile=None, source=None, nlambda=None, monochromatic=None,
         else:
             return result
 
-    def calc_datacube(self, wavelengths, *args, **kwargs):
+    def calc_datacube(self, wavelengths, progressbar=False, *args, **kwargs):
         """Calculate a spectral datacube of PSFs
 
         Parameters
         -----------
         wavelengths : iterable of floats
             List or ndarray or tuple of floating point wavelengths in meters, such as
             you would supply in a call to calc_psf via the "monochromatic" option
+        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.
+
         """
 
         # Allow up to 10,000 wavelength slices. The number matters because FITS
@@ -329,8 +335,17 @@ def calc_datacube(self, wavelengths, *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 rest of wavelengths
-        for i in range(1, nwavelengths):
+        for i in iterate_wrapper(range(1, nwavelengths)):
             wl = wavelengths[i]
             psf = self.calc_psf(*args, monochromatic=wl, **kwargs)
             for ext in range(len(psf)):