teareduce 0.4.9__py3-none-any.whl → 0.5.0__py3-none-any.whl

teareduce/cleanest/reviewcosmicray.py

@@ -10,13 +10,20 @@
 """Define the ReviewCosmicRay class."""
 
 import tkinter as tk
+from tkinter import messagebox
 from tkinter import simpledialog
 
 import matplotlib.pyplot as plt
 from matplotlib.backend_bases import key_press_handler
 from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk
 import numpy as np
-from scipy import ndimage
+from rich import print
+
+from .definitions import MAX_PIXEL_DISTANCE_TO_CR
+from .imagedisplay import ImageDisplay
+from .interpolation_a import interpolation_a
+from .interpolation_x import interpolation_x
+from .interpolation_y import interpolation_y
 
 from ..imshow import imshow
 from ..sliceregion import SliceRegion2D
@@ -26,74 +33,182 @@ import matplotlib
 matplotlib.use("TkAgg")
 
 
-class ReviewCosmicRay():
+class ReviewCosmicRay(ImageDisplay):
     """Class to review suspected cosmic ray pixels."""
 
-    def __init__(self, root, data, mask_fixed, mask_crfound):
+    def __init__(self, root, data, auxdata, cleandata_lacosmic, cr_labels, num_features,
+                 first_cr_index=1, single_cr=False,
+                 last_dilation=None, last_npoints=None, last_degree=None):
         """Initialize the review window.
 
         Parameters
         ----------
-        root : tk.Tk
-            The main Tkinter window.
+        root : tk.Toplevel
+            The parent Tkinter root window.
         data : 2D numpy array
             The original image data.
+        auxdata : 2D numpy array or None
+            The auxiliary image data.
+        cleandata_lacosmic: 2D numpy array or None
+            The cleaned image data from L.A.Cosmic.
+        cr_labels : 2D numpy array
+            Labels of connected cosmic ray pixel groups.
+        num_features : int
+            Number of connected cosmic ray pixel groups.
+        first_cr_index : int, optional
+            The index of the first cosmic ray to review (default is 1).
+        single_cr : bool, optional
+            Whether to review a single cosmic ray (default is False).
+            If True, the review window will close after reviewing the
+            selected first cosmic ray.
+        last_dilation : int or None, optional
+            The last used dilation parameter employed after L.A.Cosmic
+            detection. If > 0, the replacement by the L.A.Cosmic cleaned
+            data will not be allowed.
+        last_npoints : int or None, optional
+            The last used number of points parameter for interpolation.
+        last_degree : int or None, optional
+            The last used degree parameter for interpolation.
+
+        Methods
+        -------
+        create_widgets()
+            Create the GUI widgets for the review window.
+        update_display(cleaned=False)
+            Update the display to show the current cosmic ray.
+        set_ndeg()
+            Set the Npoints and Degree parameters for interpolation.
+        interp_x()
+            Perform X-interpolation for the current cosmic ray.
+        interp_y()
+            Perform Y-interpolation for the current cosmic ray.
+        interp_a(method)
+            Perform interpolation using the specified method for the current cosmic ray.
+        use_lacosmic()
+            Replace cosmic ray pixels with L.A.Cosmic cleaned data.
+        use_auxdata()
+            Replace cosmic ray pixels with auxiliary data.
+        remove_crosses()
+            Remove all pixels of the current cosmic ray from the review.
+        restore_cr()
+            Restore all pixels of the current cosmic ray to their original values.
+        continue_cr()
+            Move to the next cosmic ray for review.
+        exit_review()
+            Close the review window.
+        on_key(event)
+            Handle key press events for shortcuts.
+        on_click(event)
+            Handle mouse click events to mark/unmark pixels as cosmic rays.
+
+        Attributes
+        ----------
+        root : tk.Toplevel
+            The parent Tkinter root window.
+        data : 2D numpy array
+            The original image data.
+        auxdata : 2D numpy array or None
+            The auxiliary image data.
+        cleandata_lacosmic: 2D numpy array or None
+            The cleaned image data from L.A.Cosmic.
+        cr_labels : 2D numpy array
+            Labels of connected cosmic ray pixel groups.
+        num_features : int
+            Number of connected cosmic ray pixel groups.
+        num_cr_cleaned : int
+            Number of cosmic rays cleaned during the review.
         mask_fixed : 2D numpy array
-            Mask of previously corrected pixels.
-        mask_crfound : 2D numpy array
-            Mask of new pixels identified as cosmic rays.
+            Mask of pixels fixed during the review.
+        first_plot : bool
+            Flag to indicate if it's the first plot.
+        degree : int
+            Degree parameter for interpolation.
+        npoints : int
+            Number of points parameter for interpolation.
+        last_dilation : int or None
+            The last used dilation parameter employed after L.A.Cosmic
+            detection.
         """
         self.root = root
+        self.root.title("Review Cosmic Rays")
+        self.auxdata = auxdata
+        if self.auxdata is not None:
+            self.root.geometry("1200x700+100+100")
+        else:
+            self.root.geometry("800x700+100+100")
         self.data = data
+        self.cleandata_lacosmic = cleandata_lacosmic
         self.data_original = data.copy()
-        self.mask_fixed = mask_fixed
-        self.mask_crfound = mask_crfound
+        self.cr_labels = cr_labels
+        self.num_features = num_features
+        self.num_cr_cleaned = 0
+        self.mask_fixed = np.zeros(self.data.shape, dtype=bool)  # Mask of pixels fixed during review
         self.first_plot = True
-        self.degree = 1    # Degree of polynomial for interpolation
-        self.npoints = 2   # Number of points at each side of the CR pixel for interpolation
-        # Label connected components in the mask; note that by default,
-        # structure is a cross [0,1,0;1,1,1;0,1,0], but we want to consider
-        # diagonal connections too, so we define a 3x3 square.
-        structure = [[1, 1, 1], [1, 1, 1], [1, 1, 1]]
-        self.cr_labels, self.num_features = ndimage.label(self.mask_crfound, structure=structure)
+        self.degree = last_degree if last_degree is not None else 1
+        self.npoints = last_npoints if last_npoints is not None else 2
+        self.last_dilation = last_dilation
         # Make a copy of the original labels to allow pixel re-marking
         self.cr_labels_original = self.cr_labels.copy()
-        print(f"Number of cosmic ray pixels detected: {np.sum(self.mask_crfound)}")
-        print(f"Number of cosmic rays detected: {self.num_features}")
+        print(f"Number of cosmic ray pixels detected..: {np.sum(self.cr_labels > 0)}")
+        print(f"Number of cosmic rays (grouped pixels): {self.num_features}")
         if self.num_features == 0:
             print('No CR hits found!')
         else:
-            self.cr_index = 1
+            self.cr_index = first_cr_index
+            self.single_cr = single_cr
             self.create_widgets()
 
     def create_widgets(self):
-        self.review_window = tk.Toplevel(self.root)
-        self.review_window.title("Review Cosmic Rays")
-        self.review_window.geometry("800x700+100+100")
-
-        self.button_frame1 = tk.Frame(self.review_window)
+        """Create the GUI widgets for the review window."""
+        # Row 1 of buttons
+        self.button_frame1 = tk.Frame(self.root)
         self.button_frame1.pack(pady=5)
-        self.remove_crosses_button = tk.Button(self.button_frame1, text="remove all X's", command=self.remove_crosses)
+        self.ndeg_label = tk.Button(self.button_frame1, text=f"Npoints={self.npoints}, Degree={self.degree}",
+                                    command=self.set_ndeg)
+        self.ndeg_label.pack(side=tk.LEFT, padx=5)
+        self.remove_crosses_button = tk.Button(self.button_frame1, text="remove all x's", command=self.remove_crosses)
         self.remove_crosses_button.pack(side=tk.LEFT, padx=5)
-        self.restore_cr_button = tk.Button(self.button_frame1, text="[r]estore CR", command=self.restore_cr)
+        self.restore_cr_button = tk.Button(self.button_frame1, text="[r]estore CR data", command=self.restore_cr)
         self.restore_cr_button.pack(side=tk.LEFT, padx=5)
         self.restore_cr_button.config(state=tk.DISABLED)
         self.next_button = tk.Button(self.button_frame1, text="[c]ontinue", command=self.continue_cr)
         self.next_button.pack(side=tk.LEFT, padx=5)
+        self.exit_button = tk.Button(self.button_frame1, text="[e]xit review", command=self.exit_review)
+        self.exit_button.pack(side=tk.LEFT, padx=5)
 
-        self.button_frame2 = tk.Frame(self.review_window)
+        # Row 2 of buttons
+        self.button_frame2 = tk.Frame(self.root)
         self.button_frame2.pack(pady=5)
-        self.ndeg_label = tk.Button(self.button_frame2, text=f"deg={self.degree}, n={self.npoints}",
-                                    command=self.set_ndeg)
-        self.ndeg_label.pack(side=tk.LEFT, padx=5)
         self.interp_x_button = tk.Button(self.button_frame2, text="[x] interp.", command=self.interp_x)
         self.interp_x_button.pack(side=tk.LEFT, padx=5)
         self.interp_y_button = tk.Button(self.button_frame2, text="[y] interp.", command=self.interp_y)
         self.interp_y_button.pack(side=tk.LEFT, padx=5)
-        self.interp_s_button = tk.Button(self.button_frame2, text="[s] interp.", command=self.interp_s)
+        # it is important to use lambda here to pass the method argument correctly
+        # (avoiding the execution of the function at button creation time, which would happen
+        # if we didn't use lambda; in that case, the function would be called immediately and
+        # its return value (None) would be assigned to the command parameter; furthermore,
+        # the function is trying to deactivate the buttons before they are created, which
+        # would lead to an error; in addition, since I have two buttons calling the same function
+        # with different arguments, using lambda allows to differentiate them)
+        self.interp_s_button = tk.Button(self.button_frame2, text="[s]urface interp.",
+                                         command=lambda: self.interp_a('surface'))
         self.interp_s_button.pack(side=tk.LEFT, padx=5)
-
-        self.button_frame3 = tk.Frame(self.review_window)
+        self.interp_m_button = tk.Button(self.button_frame2, text="[m]edian",
+                                         command=lambda: self.interp_a('median'))
+        self.interp_m_button.pack(side=tk.LEFT, padx=5)
+        self.interp_l_button = tk.Button(self.button_frame2, text="[l]acosmic", command=self.use_lacosmic)
+        self.interp_l_button.pack(side=tk.LEFT, padx=5)
+        if self.last_dilation is not None and self.last_dilation > 0:
+            self.interp_l_button.config(state=tk.DISABLED)
+        if self.cleandata_lacosmic is None:
+            self.interp_l_button.config(state=tk.DISABLED)
+        self.interp_aux_button = tk.Button(self.button_frame2, text="[a]ux. data", command=self.use_auxdata)
+        self.interp_aux_button.pack(side=tk.LEFT, padx=5)
+        if self.auxdata is None:
+            self.interp_aux_button.config(state=tk.DISABLED)
+
+        # Row 3 of buttons
+        self.button_frame3 = tk.Frame(self.root)
         self.button_frame3.pack(pady=5)
         vmin, vmax = zscale(self.data)
         self.vmin_button = tk.Button(self.button_frame3, text=f"vmin: {vmin:.2f}", command=self.set_vmin)
@@ -104,11 +219,15 @@ class ReviewCosmicRay():
         self.set_minmax_button.pack(side=tk.LEFT, padx=5)
         self.set_zscale_button = tk.Button(self.button_frame3, text="zscale [/]", command=self.set_zscale)
         self.set_zscale_button.pack(side=tk.LEFT, padx=5)
-        self.exit_button = tk.Button(self.button_frame3, text="[e]xit review", command=self.exit_review)
-        self.exit_button.pack(side=tk.LEFT, padx=5)
 
-        self.fig, self.ax = plt.subplots(figsize=(8, 5))
-        self.canvas = FigureCanvasTkAgg(self.fig, master=self.review_window)
+        # Figure
+        if self.auxdata is not None:
+            self.fig, (self.ax_aux, self.ax) = plt.subplots(
+                ncols=2, figsize=(10, 5), constrained_layout=True)
+        else:
+            self.fig, self.ax = plt.subplots(figsize=(8, 5), constrained_layout=True)
+        self.canvas = FigureCanvasTkAgg(self.fig, master=self.root)
+        self.canvas.get_tk_widget().pack(padx=5, pady=5)
         # The next two instructions prevent a segmentation fault when pressing "q"
         self.canvas.mpl_disconnect(self.canvas.mpl_connect("key_press_event", key_press_handler))
         self.canvas.mpl_connect("key_press_event", self.on_key)
@@ -117,286 +236,206 @@ class ReviewCosmicRay():
         self.canvas_widget.pack(fill=tk.BOTH, expand=True)
 
         # Matplotlib toolbar
-        self.toolbar_frame = tk.Frame(self.review_window)
+        self.toolbar_frame = tk.Frame(self.root)
         self.toolbar_frame.pack(fill=tk.X, expand=False, pady=5)
         self.toolbar = NavigationToolbar2Tk(self.canvas, self.toolbar_frame)
         self.toolbar.update()
 
         self.update_display()
 
-        self.root.wait_window(self.review_window)
+    def update_display(self, cleaned=False):
+        """Update the display to show the current cosmic ray.
 
-    def update_display(self):
+        Parameters
+        ----------
+        cleaned : bool, optional
+            Whether the cosmic ray has been cleaned (default is False).
+            If True, the cosmic ray pixels will be marked differently.
+        """
         ycr_list, xcr_list = np.where(self.cr_labels == self.cr_index)
         ycr_list_original, xcr_list_original = np.where(self.cr_labels_original == self.cr_index)
         if self.first_plot:
             print(f"Cosmic ray {self.cr_index}: "
                   f"Number of pixels = {len(xcr_list)}, "
-                  f"Centroid = ({np.mean(xcr_list):.2f}, {np.mean(ycr_list):.2f})")
+                  f"Centroid = ({np.mean(xcr_list)+1:.2f}, {np.mean(ycr_list)+1:.2f})")
         # Use original positions to define the region to display in order
         # to avoid image shifts when some pixels are unmarked or new ones are marked
         i0 = int(np.mean(ycr_list_original) + 0.5)
         j0 = int(np.mean(xcr_list_original) + 0.5)
-        jmin = j0 - 15 if j0 - 15 >= 0 else 0
-        jmax = j0 + 15 if j0 + 15 < self.data.shape[1] else self.data.shape[1] - 1
-        imin = i0 - 15 if i0 - 15 >= 0 else 0
-        imax = i0 + 15 if i0 + 15 < self.data.shape[0] else self.data.shape[0] - 1
+        semiwidth = MAX_PIXEL_DISTANCE_TO_CR
+        jmin = j0 - semiwidth if j0 - semiwidth >= 0 else 0
+        jmax = j0 + semiwidth if j0 + semiwidth < self.data.shape[1] else self.data.shape[1] - 1
+        imin = i0 - semiwidth if i0 - semiwidth >= 0 else 0
+        imax = i0 + semiwidth if i0 + semiwidth < self.data.shape[0] else self.data.shape[0] - 1
+        # Force the region to be of size (2*semiwidth + 1) x (2*semiwidth + 1)
+        if jmin == 0:
+            jmax = 2 * semiwidth
+        elif jmax == self.data.shape[1] - 1:
+            jmin = self.data.shape[1] - 1 - 2 * semiwidth
+        if imin == 0:
+            imax = 2 * semiwidth
+        elif imax == self.data.shape[0] - 1:
+            imin = self.data.shape[0] - 1 - 2 * semiwidth
         self.region = SliceRegion2D(f'[{jmin+1}:{jmax+1}, {imin+1}:{imax+1}]', mode='fits').python
         self.ax.clear()
         vmin = self.get_vmin()
         vmax = self.get_vmax()
         xlabel = 'X pixel (from 1 to NAXIS1)'
         ylabel = 'Y pixel (from 1 to NAXIS2)'
-        self.image_review, _, _ = imshow(self.fig, self.ax, self.data[self.region], colorbar=False,
-                                         xlabel=xlabel, ylabel=ylabel,
-                                         vmin=vmin, vmax=vmax)
-        self.image_review.set_extent([jmin + 0.5, jmax + 1.5, imin + 0.5, imax + 1.5])
+        self.image, _, _ = imshow(self.fig, self.ax, self.data[self.region], colorbar=False,
+                                  xlabel=xlabel, ylabel=ylabel,
+                                  vmin=vmin, vmax=vmax)
+        self.image.set_extent([jmin + 0.5, jmax + 1.5, imin + 0.5, imax + 1.5])
+        if self.auxdata is not None:
+            self.ax_aux.clear()
+            self.image_aux, _, _ = imshow(self.fig, self.ax_aux, self.auxdata[self.region],
+                                          colorbar=False,
+                                          xlabel=xlabel, ylabel=ylabel,
+                                          vmin=vmin, vmax=vmax)
+            self.image_aux.set_extent([jmin + 0.5, jmax + 1.5, imin + 0.5, imax + 1.5])
+            self.ax_aux.set_title("Auxiliary data")
+        # Overplot cosmic ray pixels
         xlim = self.ax.get_xlim()
         ylim = self.ax.get_ylim()
         for xcr, ycr in zip(xcr_list, ycr_list):
             xcr += 1  # from index to pixel
             ycr += 1  # from index to pixel
-            self.ax.plot([xcr - 0.5, xcr + 0.5], [ycr + 0.5, ycr - 0.5], 'r-')
-            self.ax.plot([xcr - 0.5, xcr + 0.5], [ycr - 0.5, ycr + 0.5], 'r-')
+            if cleaned:
+                self.ax.plot(xcr, ycr, 'C1o', markersize=4)
+            else:
+                self.ax.plot([xcr - 0.5, xcr + 0.5], [ycr + 0.5, ycr - 0.5], 'r-')
+                self.ax.plot([xcr - 0.5, xcr + 0.5], [ycr - 0.5, ycr + 0.5], 'r-')
         self.ax.set_xlim(xlim)
         self.ax.set_ylim(ylim)
         self.ax.set_title(f"Cosmic ray #{self.cr_index}/{self.num_features}")
         if self.first_plot:
             self.first_plot = False
-            self.fig.tight_layout()
-        self.canvas.draw()
-
-    def set_vmin(self):
-        old_vmin = self.get_vmin()
-        new_vmin = simpledialog.askfloat("Set vmin", "Enter new vmin:", initialvalue=old_vmin)
-        if new_vmin is None:
-            return
-        self.vmin_button.config(text=f"vmin: {new_vmin:.2f}")
-        self.image_review.set_clim(vmin=new_vmin)
-        self.canvas.draw()
-
-    def set_vmax(self):
-        old_vmax = self.get_vmax()
-        new_vmax = simpledialog.askfloat("Set vmax", "Enter new vmax:", initialvalue=old_vmax)
-        if new_vmax is None:
-            return
-        self.vmax_button.config(text=f"vmax: {new_vmax:.2f}")
-        self.image_review.set_clim(vmax=new_vmax)
-        self.canvas.draw()
-
-    def get_vmin(self):
-        return float(self.vmin_button.cget("text").split(":")[1])
-
-    def get_vmax(self):
-        return float(self.vmax_button.cget("text").split(":")[1])
-
-    def set_minmax(self):
-        vmin_new = np.min(self.data[self.region])
-        vmax_new = np.max(self.data[self.region])
-        self.vmin_button.config(text=f"vmin: {vmin_new:.2f}")
-        self.vmax_button.config(text=f"vmax: {vmax_new:.2f}")
-        self.image_review.set_clim(vmin=vmin_new)
-        self.image_review.set_clim(vmax=vmax_new)
-        self.canvas.draw()
-
-    def set_zscale(self):
-        vmin_new, vmax_new = zscale(self.data[self.region])
-        self.vmin_button.config(text=f"vmin: {vmin_new:.2f}")
-        self.vmax_button.config(text=f"vmax: {vmax_new:.2f}")
-        self.image_review.set_clim(vmin=vmin_new)
-        self.image_review.set_clim(vmax=vmax_new)
-        self.canvas.draw()
+        self.canvas.draw_idle()
 
     def set_ndeg(self):
-        new_degree = simpledialog.askinteger("Set degree", "Enter new degree (min=0):",
+        """Set the number of points and degree for interpolation."""
+        new_npoints = simpledialog.askinteger("Set Npoints", "Enter Npoints:",
+                                              initialvalue=self.npoints,  minvalue=1)
+        if new_npoints is None:
+            return
+        new_degree = simpledialog.askinteger("Set degree", "Enter Degree (min=0):",
                                              initialvalue=self.degree, minvalue=0)
         if new_degree is None:
             return
-        new_npoints = simpledialog.askinteger("Set n", f"Enter new n (min={2*new_degree}):",
-                                              initialvalue=self.npoints, minvalue=2*new_degree)
-        if new_npoints is None:
-            return
         self.degree = new_degree
         self.npoints = new_npoints
-        self.ndeg_label.config(text=f"deg={self.degree}, n={self.npoints}")
+        self.ndeg_label.config(text=f"Npoints={self.npoints}, Degree={self.degree}")
 
-    def interp_x(self):
-        print(f"X-interpolation of cosmic ray {self.cr_index}")
-        ycr_list, xcr_list = np.where(self.cr_labels == self.cr_index)
-        ycr_min = np.min(ycr_list)
-        ycr_max = np.max(ycr_list)
-        xfit_all = []
-        yfit_all = []
-        for ycr in range(ycr_min, ycr_max + 1):
-            xmarked = xcr_list[np.where(ycr_list == ycr)]
-            if len(xmarked) > 0:
-                jmin = np.min(xmarked)
-                jmax = np.max(xmarked)
-                # mark intermediate pixels too
-                for ix in range(jmin, jmax + 1):
-                    self.cr_labels[ycr, ix] = self.cr_index
-                xmarked = xcr_list[np.where(ycr_list == ycr)]
-                xfit = []
-                zfit = []
-                for i in range(jmin - self.npoints, jmin):
-                    if 0 <= i < self.data.shape[1]:
-                        xfit.append(i)
-                        xfit_all.append(i)
-                        yfit_all.append(ycr)
-                        zfit.append(self.data[ycr, i])
-                for i in range(jmax + 1, jmax + 1 + self.npoints):
-                    if 0 <= i < self.data.shape[1]:
-                        xfit.append(i)
-                        xfit_all.append(i)
-                        yfit_all.append(ycr)
-                        zfit.append(self.data[ycr, i])
-                if len(xfit) > self.degree:
-                    p = np.polyfit(xfit, zfit, self.degree)
-                    for i in range(jmin, jmax + 1):
-                        if 0 <= i < self.data.shape[1]:
-                            self.data[ycr, i] = np.polyval(p, i)
-                            self.mask_fixed[ycr, i] = True
-                else:
-                    print(f"Not enough points to fit at y={ycr+1}")
-                    self.update_display()
-                    return
+    def set_buttons_after_cleaning_cr(self):
+        """Set the state of buttons after cleaning a cosmic ray."""
         self.restore_cr_button.config(state=tk.NORMAL)
         self.remove_crosses_button.config(state=tk.DISABLED)
         self.interp_x_button.config(state=tk.DISABLED)
         self.interp_y_button.config(state=tk.DISABLED)
         self.interp_s_button.config(state=tk.DISABLED)
-        self.update_display()
+        self.interp_m_button.config(state=tk.DISABLED)
+        self.interp_l_button.config(state=tk.DISABLED)
+        self.interp_aux_button.config(state=tk.DISABLED)
+
+    def interp_x(self):
+        """Perform x-direction interpolation to clean a cosmic ray."""
+        if 2 * self.npoints <= self.degree:
+            messagebox.showerror("Input Error", "2*Npoints must be greater than Degree for x interpolation.")
+            return
+        print(f"X-interpolation of cosmic ray {self.cr_index}")
+        interpolation_performed, xfit_all, yfit_all = interpolation_x(
+            data=self.data,
+            mask_fixed=self.mask_fixed,
+            cr_labels=self.cr_labels,
+            cr_index=self.cr_index,
+            npoints=self.npoints,
+            degree=self.degree
+        )
+        if interpolation_performed:
+            self.num_cr_cleaned += 1
+        self.set_buttons_after_cleaning_cr()
+        self.update_display(cleaned=interpolation_performed)
         if len(xfit_all) > 0:
             self.ax.plot(np.array(xfit_all) + 1, np.array(yfit_all) + 1, 'mo', markersize=4)  # +1: from index to pixel
-            self.canvas.draw()
+            self.canvas.draw_idle()
 
     def interp_y(self):
+        """Perform y-direction interpolation to clean a cosmic ray."""
+        if 2 * self.npoints <= self.degree:
+            messagebox.showerror("Input Error", "2*Npoints must be greater than Degree for y interpolation.")
+            return
         print(f"Y-interpolation of cosmic ray {self.cr_index}")
-        ycr_list, xcr_list = np.where(self.cr_labels == self.cr_index)
-        xcr_min = np.min(xcr_list)
-        xcr_max = np.max(xcr_list)
-        xfit_all = []
-        yfit_all = []
-        for xcr in range(xcr_min, xcr_max + 1):
-            ymarked = ycr_list[np.where(xcr_list == xcr)]
-            if len(ymarked) > 0:
-                imin = np.min(ymarked)
-                imax = np.max(ymarked)
-                # mark intermediate pixels too
-                for iy in range(imin, imax + 1):
-                    self.cr_labels[iy, xcr] = self.cr_index
-                ymarked = ycr_list[np.where(xcr_list == xcr)]
-                yfit = []
-                zfit = []
-                for i in range(imin - self.npoints, imin):
-                    if 0 <= i < self.data.shape[0]:
-                        yfit.append(i)
-                        yfit_all.append(i)
-                        xfit_all.append(xcr)
-                        zfit.append(self.data[i, xcr])
-                for i in range(imax + 1, imax + 1 + self.npoints):
-                    if 0 <= i < self.data.shape[0]:
-                        yfit.append(i)
-                        yfit_all.append(i)
-                        xfit_all.append(xcr)
-                        zfit.append(self.data[i, xcr])
-                if len(yfit) > self.degree:
-                    p = np.polyfit(yfit, zfit, self.degree)
-                    for i in range(imin, imax + 1):
-                        if 0 <= i < self.data.shape[1]:
-                            self.data[i, xcr] = np.polyval(p, i)
-                            self.mask_fixed[i, xcr] = True
-                else:
-                    print(f"Not enough points to fit at x={xcr+1}")
-                    self.update_display()
-                    return
-        self.restore_cr_button.config(state=tk.NORMAL)
-        self.remove_crosses_button.config(state=tk.DISABLED)
-        self.interp_x_button.config(state=tk.DISABLED)
-        self.interp_y_button.config(state=tk.DISABLED)
-        self.interp_s_button.config(state=tk.DISABLED)
-        self.update_display()
+        interpolation_performed, xfit_all, yfit_all = interpolation_y(
+            data=self.data,
+            mask_fixed=self.mask_fixed,
+            cr_labels=self.cr_labels,
+            cr_index=self.cr_index,
+            npoints=self.npoints,
+            degree=self.degree
+        )
+        if interpolation_performed:
+            self.num_cr_cleaned += 1
+        self.set_buttons_after_cleaning_cr()
+        self.update_display(cleaned=interpolation_performed)
         if len(xfit_all) > 0:
             self.ax.plot(np.array(xfit_all) + 1, np.array(yfit_all) + 1, 'mo', markersize=4)  # +1: from index to pixel
-            self.canvas.draw()
+            self.canvas.draw_idle()
 
-    def interp_s(self):
-        print(f"S-interpolation of cosmic ray {self.cr_index}")
-        ycr_list, xcr_list = np.where(self.cr_labels == self.cr_index)
-        ycr_min = np.min(ycr_list)
-        ycr_max = np.max(ycr_list)
-        xfit_all = []
-        yfit_all = []
-        zfit_all = []
-        # First do horizontal lines
-        for ycr in range(ycr_min, ycr_max + 1):
-            xmarked = xcr_list[np.where(ycr_list == ycr)]
-            if len(xmarked) > 0:
-                jmin = np.min(xmarked)
-                jmax = np.max(xmarked)
-                # mark intermediate pixels too
-                for ix in range(jmin, jmax + 1):
-                    self.cr_labels[ycr, ix] = self.cr_index
-                xmarked = xcr_list[np.where(ycr_list == ycr)]
-                for i in range(jmin - self.npoints, jmin):
-                    if 0 <= i < self.data.shape[1]:
-                        xfit_all.append(i)
-                        yfit_all.append(ycr)
-                        zfit_all.append(self.data[ycr, i])
-                for i in range(jmax + 1, jmax + 1 + self.npoints):
-                    if 0 <= i < self.data.shape[1]:
-                        xfit_all.append(i)
-                        yfit_all.append(ycr)
-                        zfit_all.append(self.data[ycr, i])
-                    xcr_min = np.min(xcr_list)
-        # Now do vertical lines
-        xcr_max = np.max(xcr_list)
-        for xcr in range(xcr_min, xcr_max + 1):
-            ymarked = ycr_list[np.where(xcr_list == xcr)]
-            if len(ymarked) > 0:
-                imin = np.min(ymarked)
-                imax = np.max(ymarked)
-                # mark intermediate pixels too
-                for iy in range(imin, imax + 1):
-                    self.cr_labels[iy, xcr] = self.cr_index
-                ymarked = ycr_list[np.where(xcr_list == xcr)]
-                for i in range(imin - self.npoints, imin):
-                    if 0 <= i < self.data.shape[0]:
-                        yfit_all.append(i)
-                        xfit_all.append(xcr)
-                        zfit_all.append(self.data[i, xcr])
-                for i in range(imax + 1, imax + 1 + self.npoints):
-                    if 0 <= i < self.data.shape[0]:
-                        yfit_all.append(i)
-                        xfit_all.append(xcr)
-                        zfit_all.append(self.data[i, xcr])
-        if len(xfit_all) > 3:
-            # Construct the design matrix for a 2D polynomial fit to a plane,
-            # where each row corresponds to a point (x, y, z) and the model
-            # is z = C[0]*x + C[1]*y + C[2]
-            A = np.c_[xfit_all, yfit_all, np.ones(len(xfit_all))]
-            # Least squares polynomial fit
-            C, _, _, _ = np.linalg.lstsq(A, zfit_all, rcond=None)
-            # recompute all CR pixels to take into account "holes" between marked pixels
-            ycr_list, xcr_list = np.where(self.cr_labels == self.cr_index)
-            for iy, ix in zip(ycr_list, xcr_list):
-                self.data[iy, ix] = C[0] * ix + C[1] * iy + C[2]
-                self.mask_fixed[iy, ix] = True
-        else:
-            print("Not enough points to fit a plane")
-            self.update_display()
-            return
-        self.restore_cr_button.config(state=tk.NORMAL)
-        self.remove_crosses_button.config(state=tk.DISABLED)
-        self.interp_x_button.config(state=tk.DISABLED)
-        self.interp_y_button.config(state=tk.DISABLED)
-        self.interp_s_button.config(state=tk.DISABLED)
-        self.update_display()
+    def interp_a(self, method):
+        """Perform interpolation using the specified method to clean a cosmic ray.
+
+        Parameters
+        ----------
+        method : str
+            The interpolation method to use ('surface' or 'median').
+        """
+        print(f"{method} interpolation of cosmic ray {self.cr_index}")
+        interpolation_performed, xfit_all, yfit_all = interpolation_a(
+            data=self.data,
+            mask_fixed=self.mask_fixed,
+            cr_labels=self.cr_labels,
+            cr_index=self.cr_index,
+            npoints=self.npoints,
+            method=method
+        )
+        if interpolation_performed:
+            self.num_cr_cleaned += 1
+        self.set_buttons_after_cleaning_cr()
+        self.update_display(cleaned=interpolation_performed)
         if len(xfit_all) > 0:
             self.ax.plot(np.array(xfit_all) + 1, np.array(yfit_all) + 1, 'mo', markersize=4)  # +1: from index to pixel
-            self.canvas.draw()
+            self.canvas.draw_idle()
+
+    def use_lacosmic(self):
+        """Use L.A.Cosmic cleaned data to clean a cosmic ray."""
+        if self.cleandata_lacosmic is None:
+            print("L.A.Cosmic cleaned data not available.")
+            return
+        print(f"L.A.Cosmic interpolation of cosmic ray {self.cr_index}")
+        ycr_list, xcr_list = np.where(self.cr_labels == self.cr_index)
+        for iy, ix in zip(ycr_list, xcr_list):
+            self.data[iy, ix] = self.cleandata_lacosmic[iy, ix]
+            self.mask_fixed[iy, ix] = True
+        self.num_cr_cleaned += 1
+        self.set_buttons_after_cleaning_cr()
+        self.update_display(cleaned=True)
+
+    def use_auxdata(self):
+        """Use auxiliary data to clean a cosmic ray."""
+        if self.auxdata is None:
+            print("Auxiliary data not available.")
+            return
+        print(f"Auxiliary data interpolation of cosmic ray {self.cr_index}")
+        ycr_list, xcr_list = np.where(self.cr_labels == self.cr_index)
+        for iy, ix in zip(ycr_list, xcr_list):
+            self.data[iy, ix] = self.auxdata[iy, ix]
+            self.mask_fixed[iy, ix] = True
+        self.num_cr_cleaned += 1
+        self.set_buttons_after_cleaning_cr()
+        self.update_display(cleaned=True)
 
     def remove_crosses(self):
+        """Remove all pixels of the current cosmic ray from the review."""
         ycr_list, xcr_list = np.where(self.cr_labels == self.cr_index)
         for iy, ix in zip(ycr_list, xcr_list):
             self.cr_labels[iy, ix] = 0
@@ -405,35 +444,61 @@ class ReviewCosmicRay():
         self.interp_x_button.config(state=tk.DISABLED)
         self.interp_y_button.config(state=tk.DISABLED)
         self.interp_s_button.config(state=tk.DISABLED)
+        self.interp_m_button.config(state=tk.DISABLED)
+        self.interp_l_button.config(state=tk.DISABLED)
+        self.interp_aux_button.config(state=tk.DISABLED)
         self.update_display()
 
     def restore_cr(self):
+        """Restore all pixels of the current cosmic ray to their original values."""
         ycr_list, xcr_list = np.where(self.cr_labels == self.cr_index)
         for iy, ix in zip(ycr_list, xcr_list):
             self.data[iy, ix] = self.data_original[iy, ix]
+            self.mask_fixed[iy, ix] = False
             self.interp_x_button.config(state=tk.NORMAL)
             self.interp_y_button.config(state=tk.NORMAL)
             self.interp_s_button.config(state=tk.NORMAL)
+            self.interp_m_button.config(state=tk.NORMAL)
+            if self.cleandata_lacosmic is not None:
+                if self.last_dilation is None or self.last_dilation == 0:
+                    self.interp_l_button.config(state=tk.NORMAL)
+            if self.auxdata is not None:
+                self.interp_aux_button.config(state=tk.NORMAL)
         print(f"Restored all pixels of cosmic ray {self.cr_index}")
+        self.num_cr_cleaned -= 1
         self.remove_crosses_button.config(state=tk.NORMAL)
         self.restore_cr_button.config(state=tk.DISABLED)
         self.update_display()
 
     def continue_cr(self):
+        """Move to the next cosmic ray for review."""
+        if self.single_cr:
+            self.exit_review()
+            return  # important: do not remove (to avoid errors)
         self.cr_index += 1
         if self.cr_index > self.num_features:
-            self.cr_index = 1
+            self.exit_review()
+            return  # important: do not remove (to avoid errors)
         self.first_plot = True
         self.restore_cr_button.config(state=tk.DISABLED)
         self.interp_x_button.config(state=tk.NORMAL)
         self.interp_y_button.config(state=tk.NORMAL)
         self.interp_s_button.config(state=tk.NORMAL)
+        self.interp_m_button.config(state=tk.NORMAL)
+        if self.cleandata_lacosmic is not None:
+            if self.last_dilation is None or self.last_dilation == 0:
+                self.interp_l_button.config(state=tk.NORMAL)
+        if self.auxdata is not None:
+            self.interp_aux_button.config(state=tk.NORMAL)
+        self.remove_crosses_button.config(state=tk.NORMAL)
         self.update_display()
 
     def exit_review(self):
-        self.review_window.destroy()
+        """Close the review window."""
+        self.root.destroy()
 
     def on_key(self, event):
+        """Handle key press events."""
         if event.key == 'q':
             pass  # Ignore the "q" key to prevent closing the window
         elif event.key == 'r':
@@ -447,7 +512,16 @@ class ReviewCosmicRay():
                 self.interp_y()
         elif event.key == 's':
             if self.interp_s_button.cget("state") != "disabled":
-                self.interp_s()
+                self.interp_a('surface')
+        elif event.key == 'm':
+            if self.interp_m_button.cget("state") != "disabled":
+                self.interp_a('median')
+        elif event.key == 'l':
+            if self.interp_l_button.cget("state") != "disabled":
+                self.use_lacosmic()
+        elif event.key == 'a':
+            if self.interp_aux_button.cget("state") != "disabled":
+                self.use_auxdata()
         elif event.key == 'right' or event.key == 'c':
             self.continue_cr()
         elif event.key == ',':
@@ -456,13 +530,14 @@ class ReviewCosmicRay():
             self.set_zscale()
         elif event.key == 'e':
             self.exit_review()
+            return  # important: do not remove (to avoid errors)
         else:
             print(f"Key pressed: {event.key}")
 
     def on_click(self, event):
-        if event.inaxes:
+        """Handle mouse click events on the image."""
+        if event.inaxes == self.ax:
             x, y = event.xdata, event.ydata
-            print(f"Clicked at image coordinates: ({x:.2f}, {y:.2f})")
             ix = int(x+0.5) - 1  # from pixel to index
             iy = int(y+0.5) - 1  # from pixel to index
             if int(self.cr_labels[iy, ix]) == self.cr_index:
@@ -470,17 +545,26 @@ class ReviewCosmicRay():
                 print(f"Pixel ({ix+1}, {iy+1}) unmarked as cosmic ray.")
             else:
                 self.cr_labels[iy, ix] = self.cr_index
-                print(f"Pixel ({ix+1}, {iy+1}) marked as cosmic ray.")
+                print(f"Pixel ({ix+1}, {iy+1}), with signal {self.data[iy, ix]}, marked as cosmic ray.")
             xcr_list, ycr_list = np.where(self.cr_labels == self.cr_index)
             if len(xcr_list) == 0:
                 self.interp_x_button.config(state=tk.DISABLED)
                 self.interp_y_button.config(state=tk.DISABLED)
                 self.interp_s_button.config(state=tk.DISABLED)
+                self.interp_m_button.config(state=tk.DISABLED)
+                self.interp_l_button.config(state=tk.DISABLED)
+                self.interp_aux_button.config(state=tk.DISABLED)
                 self.remove_crosses_button.config(state=tk.DISABLED)
             else:
                 self.interp_x_button.config(state=tk.NORMAL)
                 self.interp_y_button.config(state=tk.NORMAL)
                 self.interp_s_button.config(state=tk.NORMAL)
+                self.interp_m_button.config(state=tk.NORMAL)
+                if self.cleandata_lacosmic is not None:
+                    if self.last_dilation is None or self.last_dilation == 0:
+                        self.interp_l_button.config(state=tk.NORMAL)
+                if self.auxdata is not None:
+                    self.interp_aux_button.config(state=tk.NORMAL)
                 self.remove_crosses_button.config(state=tk.NORMAL)
             # Update the display to reflect the change
             self.update_display()