Working now

tests/test_script.py

@@ -1,6 +1,15 @@
-from winternlc.config import cor_dir, DEFAULT_CUTOFF, DEFAULT_IMG_PATH, save_dir
-from winternlc.example import apply_mask_mef, apply_nlc_mef
+from winternlc.apply import apply_mask_mef, apply_nlc_mef
+from winternlc.config import (
+    DEFAULT_CORRECTION_DIR,
+    DEFAULT_CUTOFF,
+    DEFAULT_IMG_PATH,
+    DEFAULT_SAVE_DIR,
+)
 
 if __name__ == "__main__":
-    apply_nlc_mef(DEFAULT_IMG_PATH, cor_dir, save_dir, DEFAULT_CUTOFF)
-    apply_mask_mef(DEFAULT_IMG_PATH, cor_dir, save_dir, DEFAULT_CUTOFF)
+    apply_nlc_mef(
+        DEFAULT_IMG_PATH, DEFAULT_CORRECTION_DIR, DEFAULT_SAVE_DIR, DEFAULT_CUTOFF
+    )
+    apply_mask_mef(
+        DEFAULT_IMG_PATH, DEFAULT_CORRECTION_DIR, DEFAULT_SAVE_DIR, DEFAULT_CUTOFF
+    )

winternlc/__init__.py

@@ -1 +1 @@
-from winternlc.example import apply_mask_mef, apply_nlc_mef
+from winternlc.apply import apply_mask_mef, apply_nlc_mef

winternlc/apply.py

@@ -4,9 +4,14 @@
 
 from astropy.io import fits
 
-from winternlc.config import cor_dir, DEFAULT_CUTOFF, DEFAULT_IMG_PATH, save_dir
-from winternlc.non_linear_correction import nlc_single
+from winternlc.config import (
+    DEFAULT_CORRECTION_DIR,
+    DEFAULT_CUTOFF,
+    DEFAULT_IMG_PATH,
+    DEFAULT_SAVE_DIR,
+)
 from winternlc.mask import mask_single
+from winternlc.non_linear_correction import nlc_single
 
 
 def apply_nlc_mef(
@@ -31,9 +36,7 @@ def apply_nlc_mef(
             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
-                )
+                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
@@ -47,9 +50,7 @@ def apply_nlc_mef(
         print(f"Corrected FITS file saved to {corrected_fits_file}")
 
 
-def apply_mask_mef(
-    fits_file: str | Path, cor_dir: str | Path, save_dir: str | Path
-):
+def apply_mask_mef(fits_file: str | Path, cor_dir: str | Path, save_dir: str | Path):
     """
     Process a multi-extension FITS file, applying nonlinearity correction to each
     extension, and write the corrected FITS file to disk.
@@ -83,5 +84,7 @@ def apply_mask_mef(
 
 
 if __name__ == "__main__":
-    apply_nlc_mef(DEFAULT_IMG_PATH, cor_dir, save_dir, DEFAULT_CUTOFF)
-    apply_mask_mef(DEFAULT_IMG_PATH, cor_dir, save_dir)
+    apply_nlc_mef(
+        DEFAULT_IMG_PATH, DEFAULT_CORRECTION_DIR, DEFAULT_SAVE_DIR, DEFAULT_CUTOFF
+    )
+    apply_mask_mef(DEFAULT_IMG_PATH, DEFAULT_CORRECTION_DIR, DEFAULT_SAVE_DIR)

winternlc/config.py

@@ -2,12 +2,12 @@
 
 code_dir = Path(__file__).parent
 
-data_dir = code_dir / "data"
+data_dir = code_dir.parent / "data"
 
 # paths
 DEFAULT_IMG_PATH = data_dir / "example_data/example_science_image_mef.fits"
-save_dir = data_dir / "example_data"
-cor_dir = "/home/winter/GIT/winter_linearity/data/linearity_corrections"
+DEFAULT_SAVE_DIR = data_dir / "example_data"
+DEFAULT_CORRECTION_DIR = "/home/winter/GIT/winter_linearity/data/linearity_corrections"
 test_directory = "/data/flats_iwr/20240610"
 output_directory = data_dir / "linearity_corrections"
 

winternlc/make_generic_rational_corrections.py

@@ -17,9 +17,9 @@
 from scipy.optimize import curve_fit
 
 from winternlc.config import DEFAULT_CUTOFF, output_directory, test_directory
-from winternlc.utils import extract_pixel_values, find_median_files, get_exposure_time
-from winternlc.non_linear_correction import get_coeffs_path
 from winternlc.mask import get_mask_path
+from winternlc.non_linear_correction import get_coeffs_path
+from winternlc.utils import extract_pixel_values, find_median_files, get_exposure_time
 
 
 def create_rational_func(num_params: int) -> Callable[[float], float]:
@@ -385,5 +385,9 @@ def fitted_func(x):
 
         # Load and plot the fitted rational functions for the central pixel
         load_and_plot_rational(
-            median_files, output_directory, DEFAULT_CUTOFF, num_params, test_pixel=test_pixel
+            median_files,
+            output_directory,
+            DEFAULT_CUTOFF,
+            num_params,
+            test_pixel=test_pixel,
         )

winternlc/make_polynomial_corrections.py

@@ -308,5 +308,9 @@ def load_and_plot_polynomials(
 
         # Load and plot the fitted polynomials for the central pixel
         load_and_plot_polynomials(
-            median_files, output_directory, DEFAULT_CUTOFF, poly_order, test_pixel=test_pixel
+            median_files,
+            output_directory,
+            DEFAULT_CUTOFF,
+            poly_order,
+            test_pixel=test_pixel,
         )

winternlc/make_rationial8_multithread_corrections.py

@@ -7,10 +7,10 @@
 from config import DEFAULT_CUTOFF, output_directory, test_directory
 from scipy.optimize import curve_fit
 
-from winternlc.utils import extract_pixel_values, find_median_files, get_exposure_time
-from winternlc.rational import rational_func
-from winternlc.non_linear_correction import get_coeffs_path
 from winternlc.mask import get_mask_path
+from winternlc.non_linear_correction import get_coeffs_path
+from winternlc.rational import rational_func
+from winternlc.utils import extract_pixel_values, find_median_files, get_exposure_time
 
 
 def fit_rational_to_pixel(args):
@@ -310,7 +310,11 @@ def fitted_func(x):
 
         # Load and plot the fitted rational functions for the test pixel
         load_and_plot_rational(
-            median_files, output_directory, DEFAULT_CUTOFF, test_pixel=test_pixel, test=test
+            median_files,
+            output_directory,
+            DEFAULT_CUTOFF,
+            test_pixel=test_pixel,
+            test=test,
         )
     else:
         # Plot central pixel signal for all pixels

winternlc/mask.py

@@ -1,6 +1,7 @@
-import numpy as np
 from pathlib import Path
 
+import numpy as np
+
 
 def get_mask_path(cor_dir: Path, board_id: int) -> Path:
     """
@@ -25,16 +26,13 @@ def load_mask(cor_dir: Path | str, board_id: int) -> np.ndarray:
 
     if not mask_path.exists():
         raise FileNotFoundError(
-            f"Bad pixel mask file not found at {mask_path} "
-            f"for board_id {board_id}"
+            f"Bad pixel mask file not found at {mask_path} " f"for board_id {board_id}"
         )
 
     return np.load(str(mask_path))
 
 
-def apply_mask(
-    image: np.ndarray, mask: np.ndarray
-) -> np.ndarray:
+def apply_mask(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
     """
     Applies a bad pixel mask to an image.
 
@@ -46,9 +44,7 @@ def apply_mask(
     return image
 
 
-def mask_single(
-    image: np.ndarray, board_id: int, cor_dir: str
-) -> np.ndarray:
+def mask_single(image: np.ndarray, board_id: int, cor_dir: str) -> np.ndarray:
     """
     Applies a bad pixel mask to an image.
 

winternlc/non_linear_correction.py

@@ -3,10 +3,10 @@
 """
 
 import os
+from pathlib import Path
 
 import numpy as np
 
-from pathlib import Path
 from winternlc.config import DEFAULT_CUTOFF
 from winternlc.rational import rational_func
 
@@ -42,7 +42,7 @@ def load_rational_coeffs(cor_dir: Path | str, board_id: int) -> np.ndarray:
 
 
 def apply_nonlinearity_correction(
-        image: np.ndarray, coeffs: np.ndarray, cutoff: float = DEFAULT_CUTOFF
+    image: np.ndarray, coeffs: np.ndarray, cutoff: float = DEFAULT_CUTOFF
 ) -> np.ndarray:
     """
     Applies nonlinearity correction to an image using precomputed rational coefficients.