diff --git a/README.md b/README.md index 1bef637..a7a3744 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,38 @@ pre-commit install Set the directory ## Get Started + +You can use winternlc directly from the command line. + +```bash +winternlc-apply /path/to/data.fits ``` -winter_corrections/example.py + +This will apply the nonlinearity correction to the data and save the corrected data to a new file. + +You can also run the correction on multiple files at once. + +```bash +winternlc-apply /path/to/data1.fits /path/to/data2.fits +``` + +Alternatively, you can specify a directory and all the files in the directory will be corrected. + +```bash +winternlc-apply /path/to/directory ``` + +In all cases, you can also specify the output directory. + +```bash +winternlc-apply /path/to/data.fits --output-dir /path/to/output +``` + +If you do not specify an output directory, the corrected files will be saved in the same directory as the input files. + +See the help message for more information. + +```bash +winternlc --help +``` + diff --git a/pyproject.toml b/pyproject.toml index 38d89da..633953f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,6 +46,9 @@ dev = [ [project.urls] Homepage = "https://github.com/winter-telescope/winternlc" +[project.scripts] +winternlc-apply = "winternlc.apply:nlc_cli" + [tool.setuptools] packages = ["winternlc"] diff --git a/winternlc/apply.py b/winternlc/apply.py index 645704d..d796478 100644 --- a/winternlc/apply.py +++ b/winternlc/apply.py @@ -3,28 +3,27 @@ from pathlib import Path from astropy.io import fits - -from winternlc.config import ( - corrections_dir, - DEFAULT_CUTOFF, - EXAMPLE_IMG_PATH, - example_data_dir, -) +import logging from winternlc.mask import mask_single from winternlc.non_linear_correction import nlc_single +from winternlc.config import DEFAULT_CUTOFF, corrections_dir +import argparse + +logger = logging.getLogger(__name__) def apply_nlc_mef( - fits_file: str | Path, cor_dir: str | Path, save_dir: str | Path, cutoff: float + fits_file: Path, save_dir: Path | str | None = None, cor_dir: Path = corrections_dir, cutoff: float = DEFAULT_CUTOFF, output_suffix: str | None = "corrected_" ): """ Process a multi-extension FITS file, applying nonlinearity correction to each extension, and write the corrected FITS file to disk. :param fits_file: Path to the FITS file - :param cor_dir: Directory containing the correction files :param save_dir: Directory to save the corrected FITS file - :param cutoff: Cutoff value for the image + :param cor_dir: Directory containing the correction files + :param cutoff: Cutoff value for the correction + :param output_suffix: Suffix to append to the output file name :return: None """ @@ -33,19 +32,13 @@ def apply_nlc_mef( header = hdul[ext].header image = hdul[ext].data board_id = header.get("BOARD_ID", None) - if board_id is not None: - print(f"Processing extension {ext} with BOARD_ID {board_id}") - start = time.time() - corrected_image = nlc_single(image, board_id, cor_dir, cutoff) - end = time.time() - print(f"took {end-start} s to execute") - hdul[ext].data = corrected_image - else: - print(f"Skipping extension {ext} as it does not have a BOARD_ID") + corrected_image = nlc_single(image, board_id, cor_dir, cutoff) + hdul[ext].data = corrected_image - corrected_fits_file = os.path.join( - save_dir, "corrected_" + os.path.basename(fits_file) - ) + if save_dir is None: + save_dir = Path(fits_file).parent + + corrected_fits_file = save_dir / f"{output_suffix}{fits_file.name}" hdul.writeto(corrected_fits_file, overwrite=True) print(f"Corrected FITS file saved to {corrected_fits_file}") @@ -66,25 +59,53 @@ def apply_mask_mef(fits_file: str | Path, cor_dir: str | Path, save_dir: str | P header = hdul[ext].header image = hdul[ext].data board_id = header.get("BOARD_ID", None) - if board_id is not None: - print(f"Masking extension {ext} with BOARD_ID {board_id}") - start = time.time() - corrected_image = mask_single(image, board_id, cor_dir) - end = time.time() - print(f"took {end-start} s to execute") - hdul[ext].data = corrected_image - else: - print(f"Skipping extension {ext} as it does not have a BOARD_ID") + corrected_image = mask_single(image, board_id, cor_dir) + hdul[ext].data = corrected_image corrected_fits_file = os.path.join( save_dir, "masked_" + os.path.basename(fits_file) ) hdul.writeto(corrected_fits_file, overwrite=True) - print(f"Corrected FITS file saved to {corrected_fits_file}") + logger.info(f"Corrected FITS file saved to {corrected_fits_file}") -if __name__ == "__main__": - apply_nlc_mef( - EXAMPLE_IMG_PATH, corrections_dir, example_data_dir, DEFAULT_CUTOFF +def nlc_cli(): + """ + Command-line interface for applying non-linearity + correction to multi-extension FITS file(s) + """ + parser = argparse.ArgumentParser( + description="Apply non-linearity correction to multi-extension FITS file" ) - apply_mask_mef(EXAMPLE_IMG_PATH, corrections_dir, example_data_dir) + parser.add_argument( + "-o", "--output_dir", default=None, + type=str, + help="Directory to save the corrected FITS file(s) " + "(default: same as input file)", + ) + parser.add_argument("files", nargs="+", help="FITS file(s) to correct") + + args = parser.parse_args() + + logger.info("Applying non-linearity correction to multi-extension FITS file") + + file_paths = [] + + for f_name in args.files: + path = Path(f_name) + if not path.exists(): + raise FileNotFoundError(f"File not found at {path}") + elif path.is_dir(): + file_paths.extend(list(path.glob("*.fits"))) + else: + file_paths.append(path) + + for path in file_paths: + apply_nlc_mef( + fits_file=path, + save_dir=args.output_dir, + ) + + +if __name__ == "__main__": + nlc_cli() diff --git a/winternlc/config.py b/winternlc/config.py index fb131bd..b2bd54d 100644 --- a/winternlc/config.py +++ b/winternlc/config.py @@ -17,7 +17,7 @@ _corrections_dir = os.getenv("WINTERNLC_DIR") if _corrections_dir is None: corrections_dir = Path.home() / "Data/winternlc/" - logger.warning(f"No data directory set, using {corrections_dir}") + logger.warning(f"No correction data directory set, using {corrections_dir}") else: corrections_dir = Path(_corrections_dir)