Reorganise

tests/test_script.py

@@ -1,6 +1,6 @@
-from winternlc.config import cor_dir, cutoff, fits_file, save_dir
-from winternlc.example import test_mask, test_nonlinearity
+from winternlc.config import cor_dir, DEFAULT_CUTOFF, DEFAULT_IMG_PATH, save_dir
+from winternlc.example import apply_mask_mef, apply_nlc_mef
 
 if __name__ == "__main__":
-    test_nonlinearity(fits_file, cor_dir, save_dir, cutoff)
-    test_mask(fits_file, cor_dir, save_dir, cutoff)
+    apply_nlc_mef(DEFAULT_IMG_PATH, cor_dir, save_dir, DEFAULT_CUTOFF)
+    apply_mask_mef(DEFAULT_IMG_PATH, cor_dir, save_dir, DEFAULT_CUTOFF)

winternlc/__init__.py

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

winternlc/example.py → winternlc/apply.py

@@ -4,11 +4,12 @@
 
 from astropy.io import fits
 
-from winternlc.config import cor_dir, cutoff, fits_file, save_dir
-from winternlc.corrections import mask_bad_pixels, nonlinearity_correction
+from winternlc.config import cor_dir, DEFAULT_CUTOFF, DEFAULT_IMG_PATH, save_dir
+from winternlc.non_linear_correction import nlc_single
+from winternlc.mask import mask_single
 
 
-def test_nonlinearity(
+def apply_nlc_mef(
     fits_file: str | Path, cor_dir: str | Path, save_dir: str | Path, cutoff: float
 ):
     """
@@ -30,8 +31,8 @@ def test_nonlinearity(
             if board_id is not None:
                 print(f"Processing extension {ext} with BOARD_ID {board_id}")
                 start = time.time()
-                corrected_image = nonlinearity_correction(
-                    image, board_id, ext, cor_dir, cutoff
+                corrected_image = nlc_single(
+                    image, board_id, cor_dir, cutoff
                 )
                 end = time.time()
                 print(f"took {end-start} s to execute")
@@ -46,8 +47,8 @@ def test_nonlinearity(
         print(f"Corrected FITS file saved to {corrected_fits_file}")
 
 
-def test_mask(
-    fits_file: str | Path, cor_dir: str | Path, save_dir: str | Path, cutoff: float
+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
@@ -56,7 +57,6 @@ def test_mask(
     :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
 
     :return: None
     """
@@ -68,7 +68,7 @@ def test_mask(
             if board_id is not None:
                 print(f"Masking extension {ext} with BOARD_ID {board_id}")
                 start = time.time()
-                corrected_image = mask_bad_pixels(image, board_id, ext, cor_dir)
+                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
@@ -83,5 +83,5 @@ def test_mask(
 
 
 if __name__ == "__main__":
-    test_nonlinearity(fits_file, cor_dir, save_dir, cutoff)
-    test_mask(fits_file, cor_dir, save_dir, cutoff)
+    apply_nlc_mef(DEFAULT_IMG_PATH, cor_dir, save_dir, DEFAULT_CUTOFF)
+    apply_mask_mef(DEFAULT_IMG_PATH, cor_dir, save_dir)

winternlc/config.py

@@ -5,11 +5,11 @@
 data_dir = code_dir / "data"
 
 # paths
-fits_file = data_dir / "example_data/example_science_image_mef.fits"
+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"
 test_directory = "/data/flats_iwr/20240610"
 output_directory = data_dir / "linearity_corrections"
 
 # variables
-cutoff = 56000
+DEFAULT_CUTOFF = 56000

winternlc/make_generic_rational_corrections.py

@@ -16,8 +16,10 @@
 from astropy.io import fits
 from scipy.optimize import curve_fit
 
-from winternlc.config import cutoff, output_directory, test_directory
+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
 
 
 def create_rational_func(num_params: int) -> Callable[[float], float]:
@@ -186,11 +188,11 @@ def save_rational_coefficients(
             )
         else:
             np.save(
-                os.path.join(output_dir, f"rat_coeffs_board_{board_id}_ext_{ext}.npy"),
+                str(get_coeffs_path(output_dir, board_id)),
                 rat_coeffs,
             )
             np.save(
-                os.path.join(output_dir, f"bad_pix_board_{board_id}_ext_{ext}.npy"),
+                str(get_mask_path(output_dir, board_id)),
                 bad_pix,
             )
 
@@ -347,14 +349,14 @@ def fitted_func(x):
 
     if test:
         # Plot pixel signal for the test pixel
-        plot_pixel_signal(median_files, cutoff=cutoff, test_pixel=test_pixel)
+        plot_pixel_signal(median_files, cutoff=DEFAULT_CUTOFF, test_pixel=test_pixel)
 
         # Save rational coefficients for the test pixel
         save_rational_coefficients(
             median_files,
             num_params,
             output_dir=output_directory,
-            cutoff=cutoff,
+            cutoff=DEFAULT_CUTOFF,
             test=test,
             test_pixel=test_pixel,
         )
@@ -363,25 +365,25 @@ def fitted_func(x):
         load_and_plot_rational(
             median_files,
             output_directory,
-            cutoff,
+            DEFAULT_CUTOFF,
             num_params,
             test_pixel=test_pixel,
             test=test,
         )
     else:
         # Plot central pixel signal for all pixels
-        plot_pixel_signal(median_files, cutoff=cutoff, test_pixel=test_pixel)
+        plot_pixel_signal(median_files, cutoff=DEFAULT_CUTOFF, test_pixel=test_pixel)
 
         # Save rational coefficients for all pixels
         save_rational_coefficients(
             median_files,
             num_params,
             output_dir=output_directory,
-            cutoff=cutoff,
+            cutoff=DEFAULT_CUTOFF,
             test=test,
         )
 
         # Load and plot the fitted rational functions for the central pixel
         load_and_plot_rational(
-            median_files, output_directory, cutoff, num_params, test_pixel=test_pixel
+            median_files, output_directory, DEFAULT_CUTOFF, num_params, test_pixel=test_pixel
         )

winternlc/make_polynomial_corrections.py

@@ -10,7 +10,7 @@
 from collections import defaultdict
 
 import numpy as np
-from config import cutoff, output_directory, test_directory
+from config import DEFAULT_CUTOFF, output_directory, test_directory
 
 from winternlc.utils import extract_pixel_values, find_median_files, get_exposure_time
 
@@ -270,15 +270,15 @@ def load_and_plot_polynomials(
     if test:
         # Plot pixel signal for the test pixel
         print("plot pix sig")
-        plot_pixel_signal(median_files, cutoff=cutoff, test_pixel=test_pixel)
+        plot_pixel_signal(median_files, cutoff=DEFAULT_CUTOFF, test_pixel=test_pixel)
 
         print("save poly coeff")
         # Save polynomial coefficients for the test pixel
         save_polynomial_coefficients(
             median_files,
             poly_order,
             output_dir=output_directory,
-            cutoff=cutoff,
+            cutoff=DEFAULT_CUTOFF,
             test=test,
             test_pixel=test_pixel,
         )
@@ -288,25 +288,25 @@ def load_and_plot_polynomials(
         load_and_plot_polynomials(
             median_files,
             output_directory,
-            cutoff,
+            DEFAULT_CUTOFF,
             poly_order,
             test_pixel=test_pixel,
             test=test,
         )
     else:
         # Plot central pixel signal for all pixels
-        plot_pixel_signal(median_files, cutoff=cutoff, test_pixel=test_pixel)
+        plot_pixel_signal(median_files, cutoff=DEFAULT_CUTOFF, test_pixel=test_pixel)
 
         # Save polynomial coefficients for all pixels
         save_polynomial_coefficients(
             median_files,
             poly_order,
             output_dir=output_directory,
-            cutoff=cutoff,
+            cutoff=DEFAULT_CUTOFF,
             test=test,
         )
 
         # Load and plot the fitted polynomials for the central pixel
         load_and_plot_polynomials(
-            median_files, output_directory, 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

@@ -4,16 +4,13 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
-from config import cutoff, output_directory, test_directory
+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
-
-
-def rational_func(x, a0, a1, a2, a3, b0, b1, b2, b3):
-    return (a0 + a1 * x + a2 * x**2 + a3 * x**3) / (
-        1 + b0 * x + b1 * x**2 + b2 * x**3 + b3 * x**4
-    )
+from winternlc.rational import rational_func
+from winternlc.non_linear_correction import get_coeffs_path
+from winternlc.mask import get_mask_path
 
 
 def fit_rational_to_pixel(args):
@@ -159,13 +156,11 @@ def save_rational_coefficients(
                 )
         else:
             np.save(
-                os.path.join(output_dir, f"rat_coeffs_board_{board_id}_ext_{ext}.npy"),
+                str(get_coeffs_path(output_dir, board_id)),
                 rat_coeffs,
             )
             np.save(
-                os.path.join(
-                    output_dir, f"bad_pixel_mask_board_{board_id}_ext_{ext}.npy"
-                ),
+                str(get_mask_path(output_dir, board_id)),
                 bad_pixel_mask,
             )
 
@@ -315,15 +310,15 @@ def fitted_func(x):
 
         # Load and plot the fitted rational functions for the test pixel
         load_and_plot_rational(
-            median_files, output_directory, 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
         # plot_pixel_signal(median_files, cutoff=cutoff, test_pixel=test_pixel)
 
         # Save rational coefficients for all pixels
         save_rational_coefficients(
-            median_files, output_dir=output_directory, cutoff=cutoff, test=test
+            median_files, output_dir=output_directory, cutoff=DEFAULT_CUTOFF, test=test
         )
 
         # Load and plot the fitted rational functions for the central pixel