From 4d68bd823496e2cdc804fc4cb066637d337d0bcb Mon Sep 17 00:00:00 2001 From: Max West Date: Fri, 23 Aug 2024 15:58:55 -0700 Subject: [PATCH] add PROGRESS_BAR global variable --- src/kbmod/__init__.py | 2 ++ src/kbmod/reprojection.py | 33 +++++++++++++++++---------------- src/kbmod/work_unit.py | 11 ++++++----- 3 files changed, 25 insertions(+), 21 deletions(-) diff --git a/src/kbmod/__init__.py b/src/kbmod/__init__.py index 4f9b7d1e3..9fdfcbcea 100644 --- a/src/kbmod/__init__.py +++ b/src/kbmod/__init__.py @@ -13,6 +13,8 @@ # Import the rest of the package from kbmod.search import Logging +PROGRESS_BAR = bool(int(os.environ.get("KB_PROGRESS_BARS", 1))) + # there are ways for this to go to a file, but is it worth it? # Then we have to roll a whole logging.config_from_shared_config thing _SHARED_LOGGING_CONFIG = { diff --git a/src/kbmod/reprojection.py b/src/kbmod/reprojection.py index ead045344..6dcc6a0e1 100644 --- a/src/kbmod/reprojection.py +++ b/src/kbmod/reprojection.py @@ -9,6 +9,7 @@ from kbmod.work_unit import WorkUnit from kbmod.tqdm_utils import TQDMUtils from kbmod.wcs_utils import append_wcs_to_hdu_header +from kbmod import PROGRESS_BAR from astropy.io import fits import os from copy import copy @@ -80,7 +81,7 @@ def reproject_work_unit( write_output=False, directory=None, filename=None, - disable_progress=False, + progress=PROGRESS_BAR, ): """Given a WorkUnit and a WCS, reproject all of the images in the ImageStack into a common WCS. @@ -109,8 +110,8 @@ def reproject_work_unit( The directory where output will be written if `write_output` is set to True. filename : `str` The base filename where output will be written if `write_output` is set to True. - disable_progress : `bool` - Whether or not to disable the `tqdm` progress bar. + progress : `bool` + Whether or not to enable the `tqdm` progress bar. Returns ---------- @@ -127,7 +128,7 @@ def reproject_work_unit( max_parallel_processes=max_parallel_processes, directory=directory, filename=filename, - disable_progress=disable_progress, + progress=progress, ) if parallelize: return _reproject_work_unit_in_parallel( @@ -138,7 +139,7 @@ def reproject_work_unit( write_output=write_output, directory=directory, filename=filename, - disable_progress=disable_progress, + progress=progress, ) else: return _reproject_work_unit( @@ -148,7 +149,7 @@ def reproject_work_unit( write_output=write_output, directory=directory, filename=filename, - disable_progress=disable_progress, + progress=progress, ) @@ -159,7 +160,7 @@ def _reproject_work_unit( write_output=False, directory=None, filename=None, - disable_progress=False, + progress=PROGRESS_BAR, ): """Given a WorkUnit and a WCS, reproject all of the images in the ImageStack into a common WCS. @@ -196,7 +197,7 @@ def _reproject_work_unit( enumerate(zip(unique_obstimes, unique_obstime_indices)), bar_format=TQDMUtils.DEFAULT_TQDM_BAR_FORMAT, desc="Reprojecting", - disable=disable_progress, + disable=not progress, ): time, indices = o_i science_add = np.zeros(common_wcs.array_shape, dtype=np.float32) @@ -310,7 +311,7 @@ def _reproject_work_unit_in_parallel( write_output=False, directory=None, filename=None, - disable_progress=False, + progress=PROGRESS_BAR, ): """Given a WorkUnit and a WCS, reproject all of the images in the ImageStack into a common WCS. This function uses multiprocessing to reproject the images @@ -337,8 +338,8 @@ def _reproject_work_unit_in_parallel( The directory where output will be written if `write_output` is set to True. filename : `str` The base filename where output will be written if `write_output` is set to True. - disable_progress : `bool` - Whether or not to disable the `tqdm` progress bar. + progress : `bool` + Whether or not to enable the `tqdm` progress bar. Returns ---------- @@ -405,7 +406,7 @@ def _reproject_work_unit_in_parallel( total=len(future_reprojections), bar_format=TQDMUtils.DEFAULT_TQDM_BAR_FORMAT, desc="Reprojecting", - disable=disable_progress, + disable=not progress, ) ) @@ -466,7 +467,7 @@ def reproject_lazy_work_unit( filename, frame="original", max_parallel_processes=MAX_PROCESSES, - disable_progress=False, + progress=PROGRESS_BAR, ): """Given a WorkUnit and a WCS, reproject all of the images in the ImageStack into a common WCS. This function is used with lazily evaluated `WorkUnit`s and @@ -495,8 +496,8 @@ def reproject_lazy_work_unit( The maximum number of parallel processes to use when reprojecting. Default is 8. For more see `concurrent.futures.ProcessPoolExecutor` in the Python docs. - disable_progress : `bool` - Whether or not to disable the `tqdm` progress bar. + progress : `bool` + Whether or not to enable the `tqdm` progress bar. """ if not work_unit.lazy: raise ValueError("WorkUnit must be lazily loaded.") @@ -535,7 +536,7 @@ def reproject_lazy_work_unit( total=len(future_reprojections), bar_format=TQDMUtils.DEFAULT_TQDM_BAR_FORMAT, desc="Reprojecting", - disable=disable_progress, + disable=not progress, ) ) diff --git a/src/kbmod/work_unit.py b/src/kbmod/work_unit.py index cbcc54a17..5a980edcc 100644 --- a/src/kbmod/work_unit.py +++ b/src/kbmod/work_unit.py @@ -25,6 +25,7 @@ ) from kbmod.reprojection_utils import invert_correct_parallax from kbmod.tqdm_utils import TQDMUtils +from kbmod import PROGRESS_BAR logger = Logging.getLogger(__name__) @@ -225,7 +226,7 @@ def get_num_images(self): return len(self._per_image_indices) @classmethod - def from_fits(cls, filename, disable_progress=False): + def from_fits(cls, filename, progress=PROGRESS_BAR): """Create a WorkUnit from a single FITS file. The FITS file will have at least the following extensions: @@ -241,8 +242,8 @@ def from_fits(cls, filename, disable_progress=False): ---------- filename : `str` The file to load. - disable_progress : `bool` - Whether or not to disable the `tqdm` progress bar. + progress : `bool` + Whether or not to enable the `tqdm` progress bar. Returns ------- @@ -292,7 +293,7 @@ def from_fits(cls, filename, disable_progress=False): range(num_images), bar_format=TQDMUtils.DEFAULT_TQDM_BAR_FORMAT, desc="Loading images", - disable=disable_progress, + disable=not progress, ): sci_hdu = hdul[f"SCI_{i}"] @@ -321,7 +322,7 @@ def from_fits(cls, filename, disable_progress=False): range(n_constituents), bar_format=TQDMUtils.DEFAULT_TQDM_BAR_FORMAT, desc="Loading WCS", - disable=disable_progress, + disable=not progress, ): # Extract the per-image WCS if one exists. per_image_wcs.append(extract_wcs_from_hdu_header(hdul[f"WCS_{i}"].header))