Vignette masking and cosmic ray masking tasks

pyproject.toml

@@ -31,7 +31,8 @@ dependencies = [
     "pydantic-pint >=0.4",
     "typer>=0.12",
     "seaborn>=0.13",
-    "h5netcdf[h5py]>=1.3"
+    "h5netcdf[h5py]>=1.3",
+    "lacosmic>=1.4"
 ]
 
 [project.scripts]

src/eregion/core/image_operations.py

@@ -3,6 +3,8 @@
 import numpy as np
 from astropy.stats import sigma_clip
 from eregion.core.image_stats import do_statistics
+from scipy import ndimage
+from lacosmic import remove_cosmics
 
 
 def median_combine(images: Sequence[np.ndarray | np.ma.MaskedArray]) -> dict[str, np.ndarray]:
@@ -72,6 +74,157 @@ def sigma_clip_image(image: np.ndarray | np.ma.MaskedArray,
     masked = sigma_clip(image, sigma=sigma, axis=axis, **kwargs)
     return masked
 
+def _disk_structure(radius: int) -> np.ndarray:
+    """
+    Boolean disk-shaped structuring element for binary morphology.
+    :param radius: disk radius in pixels
+    :return: 2D boolean array of shape (2*radius+1, 2*radius+1)
+    """
+    yy, xx = np.ogrid[-radius:radius + 1, -radius:radius + 1]
+    return (yy ** 2 + xx ** 2) <= radius ** 2
+
+
+def find_vignetted_region(
+        image: np.ndarray,
+        threshold: float = 0.8,
+        smooth_size: int = 15,
+        reference_percentile: float = 90.0,
+        morph_radius: int = 5,
+        min_area: int = 100,
+        edge_connected: bool = True,
+) -> np.ndarray:
+    """
+    Find vignetted (under-illuminated) regions in an illuminated, zero-level corrected (bias/overscan subtracted) image.
+    Pixels whose median-smoothed value is below threshold * reference level are flagged, where the reference level is
+    a high percentile of the image (the fully illuminated level). Mask borders are smoothed with binary opening, closing
+    and hole filling using a disk structuring element.
+
+    :param image: 2D numpy array of the illuminated image region (no prescan/overscan).
+    :param threshold: fraction of the reference level below which a pixel is vignetted. Must be in (0, 1).
+    :param smooth_size: median filter size used to suppress noise, cosmics and bad pixels before thresholding.
+    :param reference_percentile: percentile of finite pixels used as the fully illuminated reference level.
+    :param morph_radius: radius in pixels of the disk used for opening/closing. 0 disables opening/closing.
+    :param min_area: connected regions with fewer pixels than this are dropped.
+    :param edge_connected: if True, keep only regions touching the image border (vignetting comes in from the edges).
+    :return: 2D boolean array, True where vignetted.
+    """
+    image = np.asarray(image)
+    if image.ndim != 2:
+        raise ValueError(f"image must be 2D, got {image.ndim}D.")
+    if not 0 < threshold < 1:
+        raise ValueError(f"threshold must be in (0, 1), got {threshold}.")
+    if not 0 < reference_percentile <= 100:
+        raise ValueError(f"reference_percentile must be in (0, 100], got {reference_percentile}.")
+    if smooth_size < 1:
+        raise ValueError(f"smooth_size must be >= 1, got {smooth_size}.")
+    if morph_radius < 0:
+        raise ValueError(f"morph_radius must be >= 0, got {morph_radius}.")
+    if min_area < 0:
+        raise ValueError(f"min_area must be >= 0, got {min_area}.")
+
+    finite = np.isfinite(image)
+    if not finite.any():
+        raise ValueError("image has no finite pixels.")
+    reference = np.percentile(image[finite], reference_percentile)
+    if reference <= 0:
+        raise ValueError(f"Reference illumination level is {reference} (<= 0). "
+                         f"Input must be an illuminated, bias/overscan subtracted image.")
+
+    # non-finite pixels are set to the reference level so they are never flagged
+    filled = np.where(finite, image, reference)
+    smoothed = ndimage.median_filter(filled, size=smooth_size, mode="nearest")
+    mask = smoothed < threshold * reference
+
+    if morph_radius > 0:
+        # pad by edge replication so regions touching the border are not eroded by the array boundary
+        structure = _disk_structure(morph_radius)
+        padded = np.pad(mask, morph_radius, mode="edge")
+        padded = ndimage.binary_opening(padded, structure=structure)
+        padded = ndimage.binary_closing(padded, structure=structure)
+        mask = padded[morph_radius:-morph_radius, morph_radius:-morph_radius]
+    mask = ndimage.binary_fill_holes(mask)
+
+    labels, nlabels = ndimage.label(mask)
+    if nlabels == 0:
+        return mask
+    areas = np.bincount(labels.ravel(), minlength=nlabels + 1)
+    keep = areas >= min_area
+    if edge_connected:
+        on_edge = np.zeros(nlabels + 1, dtype=bool)
+        on_edge[np.unique(np.concatenate([labels[0], labels[-1], labels[:, 0], labels[:, -1]]))] = True
+        keep &= on_edge
+    keep[0] = False  # background label
+    return keep[labels]
+
+
+def find_cosmic_rays(
+        image: np.ndarray,
+        contrast: float = 2.0,
+        cr_threshold: float = 5.0,
+        neighbor_threshold: float = 3.0,
+        effective_gain: float | None = None,
+        readnoise: float | None = None,
+        error: np.ndarray | None = None,
+        mask: np.ndarray | None = None,
+        maxiter: int = 4,
+        border_mode: str = "mirror",
+) -> np.ndarray:
+    """
+    Find cosmic rays in an image with the L.A.Cosmic algorithm (van Dokkum 2001) via lacosmic.remove_cosmics.
+    Non-finite pixels, and pixels with non-finite or non-positive error, are excluded from detection and replaced by
+    the finite median so they do not spread through the median filters.
+
+    :param image: 2D numpy array, ideally zero-level corrected (bias/overscan subtracted).
+    :param contrast: contrast threshold between the Laplacian and fine-structure images (~2 for critically sampled data).
+    :param cr_threshold: Laplacian signal-to-noise threshold for cosmic ray detection.
+    :param neighbor_threshold: Laplacian signal-to-noise threshold for pixels neighboring detected cosmic rays.
+    :param effective_gain: ratio of electrons to image units (e-/ADU). Required with readnoise if error is not given.
+    :param readnoise: read noise in electrons. Required with effective_gain if error is not given.
+    :param error: optional 2D 1-sigma error array (e.g. a 'std' plane), same shape as image. Overrides effective_gain
+        and readnoise.
+    :param mask: optional 2D boolean array, same shape as image. True pixels are ignored and never flagged.
+    :param maxiter: maximum number of detection iterations.
+    :param border_mode: scipy.ndimage border mode for convolution and median filtering.
+    :return: 2D boolean array, True where a cosmic ray was found.
+    """
+    image = np.asarray(image, dtype=float)
+    if image.ndim != 2:
+        raise ValueError(f"image must be 2D, got {image.ndim}D.")
+    if error is None and (effective_gain is None or readnoise is None):
+        raise ValueError("effective_gain and readnoise must both be given when error is not given.")
+    if error is not None:
+        error = np.asarray(error, dtype=float)
+        if error.shape != image.shape:
+            raise ValueError(f"error shape {error.shape} does not match image shape {image.shape}.")
+    if mask is not None:
+        mask = np.asarray(mask, dtype=bool)
+        if mask.shape != image.shape:
+            raise ValueError(f"mask shape {mask.shape} does not match image shape {image.shape}.")
+    for pname, pval in (("contrast", contrast), ("cr_threshold", cr_threshold),
+                        ("neighbor_threshold", neighbor_threshold)):
+        if pval <= 0:
+            raise ValueError(f"{pname} must be > 0, got {pval}.")
+    if maxiter < 1:
+        raise ValueError(f"maxiter must be >= 1, got {maxiter}.")
+
+    finite = np.isfinite(image)
+    if error is not None:
+        good_error = np.isfinite(error) & (error > 0)
+        finite &= good_error
+        if good_error.any():
+            error = np.where(good_error, error, np.median(error[good_error]))
+    if not finite.any():
+        raise ValueError("image has no finite pixels with finite, positive error.")
+    if not finite.all():
+        image = np.where(finite, image, np.median(image[finite]))
+        mask = ~finite if mask is None else (mask | ~finite)
+
+    _, cr_mask = remove_cosmics(image, contrast, cr_threshold, neighbor_threshold, error=error, mask=mask,
+                                effective_gain=effective_gain, readnoise=readnoise, maxiter=maxiter,
+                                border_mode=border_mode)
+    return cr_mask
+
+
 def flip_and_rotate(image: np.ndarray, angle: float, flip_x: bool=False, flip_y: bool=False) -> np.ndarray:
     """
     Flip and rotate an image. Rotation angle is assumed to be in degrees and positive for counter-clockwise direction,

src/eregion/datamodels/image.py

@@ -15,9 +15,7 @@
 from eregion.utils import (
     ensure_dataarray,
     slice_data,
-    ensure_numpy,
     configure_logger,
-    decrease_slicer_stop_index,
     set_slice_in_data,
     XRDATA
 )
@@ -101,7 +99,8 @@ class Output(Mappable):
     )
     masks: Optional[xr.Dataset] = Field(
         default=None,
-        description="Optional xr.Dataset containing masks for this output.",
+        description="Optional xr.Dataset containing masks for this output. "
+        "Write masks via set_mask() so the parent DetImage's full mask is rebuilt.",
         exclude=True,
     )
 
@@ -137,6 +136,22 @@ def get_output_mask(self, mask_key: str | Literal['combined', 'all'] = 'combined
         else:
             return None
 
+    def set_mask(self, mask_key: str, mask: xr.DataArray, combine: bool = True):
+        """
+        Set a mask for this output and clear the parent's full mask so it is rebuilt on next access.
+        :param mask_key: str, name of the mask data_var in self.masks.
+        :param mask: xr.DataArray, boolean mask with the same coords as this output.
+        :param combine: bool, if True OR with an existing mask under mask_key, else replace it.
+        """
+        if self.masks is None:
+            self.masks = mask.to_dataset(name=mask_key)
+        elif combine and mask_key in self.masks:
+            self.masks[mask_key] = self.masks[mask_key] | mask
+        else:
+            self.masks[mask_key] = mask
+        if self.parent is not None:
+            self.parent.masks = None
+
     def get_masked_output_data(self, data_var: str = 'data', mask_key: str = 'combined', fill_value=np.nan):
         mask = self.get_output_mask(mask_key)
         datarr = self.get_output_data(data_var)
@@ -256,7 +271,8 @@ def _select_header_and_data(
 
         return header, data
 
-    def show(self, ax=None, save=None, with_mask=True, mask_key='sigma_clip_mask', data_var='data', **imshow_kwargs):
+    def show(self, ax=None, save=None, with_mask=True, mask_key='sigma_clip_mask', data_var='data',
+             **imshow_kwargs):
         """
         Plot this output's image.
         :param ax: Matplotlib Axes object to plot on. If None, a new figure and axes are created.
@@ -398,8 +414,8 @@ def get_overscan(
         """
         return self.get_scan(axis=axis, kind='overscan', corner=corner, data_var=data_var)
 
-    def show(self, ax=None, shade_regions=False, save=None, with_mask=True, mask_key='sigma_clip_mask', data_var='data',
-             **imshow_kwargs):
+    def show(self, ax=None, save=None, with_mask=True, mask_key='sigma_clip_mask', data_var='data',
+             shade_regions=False, **imshow_kwargs):
         ax = super().show(ax=ax, save=None, with_mask=with_mask, mask_key=mask_key, data_var=data_var, **imshow_kwargs)
 
         if shade_regions:
@@ -697,8 +713,6 @@ def _build_full_mask(self) -> bool:
         Build a full mask from the individual output masks.
         :return: True if a full mask was built, False otherwise.
         """
-        if self.masks is not None:
-            return True
         self.masks = xr.Dataset(coords=self.data.coords)
         for out_id, output in self.outputs.items():
             if output.masks is not None:
@@ -722,7 +736,8 @@ def get_mask(self, mask_key: str | Literal['all', 'combined'] = 'combined') -> X
         :param mask_key: str or 'all' or 'combined'. Returns dataset if 'all'. Returns DataArray if key is specified. Returns OR combined mask from all keys if 'combined'.
         :return: xr.DataArray or xr.Dataset
         """
-        if self._build_full_mask():
+        built = self._build_full_mask() if self.masks is None else True
+        if built:
             match mask_key:
                 case 'all':
                     return self.masks
@@ -739,7 +754,8 @@ def get_mask(self, mask_key: str | Literal['all', 'combined'] = 'combined') -> X
 
     @property
     def combined_mask(self):
-        if self._build_full_mask():
+        built = self._build_full_mask() if self.masks is None else True
+        if built:
             combmask = np.any(self.masks.to_array().to_numpy(), axis=0)
             return xr.DataArray(combmask, dims=self.masks.dims, coords=self.masks.coords)
         else:

src/eregion/plotting/ptc.py

@@ -45,7 +45,7 @@ def plot(self, ax=None, **kwargs):
             ax = self.new_axes()
 
         table = self.result.ptc_table
-        variance = table["std_diff"] ** 2
+        variance = table["std"] ** 2
 
         # Relabel the hue column through the descriptor up front
         hue_data = table[self.hue].map(self.descriptor.resolve_legend_label) if self.hue in table.columns else None

src/eregion/tasks/linbin.py

@@ -164,7 +164,10 @@ def _bin_image(self, img: DetImage):
             # bin mask if it exists
             if immask is not None:
                 binned_mask = do_digital_binning(immask.astype(int), binsizes=self.bins, binaxis=binaxint).astype(bool)
-                output.masks[mask_key] = set_slice_in_data(ensure_dataarray(binned_mask), output.masks[mask_key], imslc)
+                output.set_mask(mask_key,
+                                set_slice_in_data(ensure_dataarray(binned_mask), output.masks[mask_key], imslc),
+                                combine=False
+                                )
 
         return img