Add cr_method dial (per-frame deepCR route) and PSF star-pass decoupling (#61, #62)

autoreduce/drizzle/combine.py

@@ -13,9 +13,16 @@
 from pathlib import Path
 from typing import Dict, List, Tuple
 
+import numpy as np
+
 from ..instruments import InstrumentAdapter
 from ..target import TargetSpec
 
+# AstroDrizzle's cosmic-ray DQ bit. The cr_method="deepcr" route writes its
+# per-frame masks into this same bit, so every downstream consumer of CR
+# flags (final drizzle bit masking, frame products) reads one convention.
+CR_DQ_BIT = 4096
+
 
 def drizzle_kwargs_for(spec: TargetSpec, adapter: InstrumentAdapter, n_exposures: int) -> Dict:
     """Assemble the AstroDrizzle keyword set; pure function, unit-testable."""
@@ -36,9 +43,87 @@ def drizzle_kwargs_for(spec: TargetSpec, adapter: InstrumentAdapter, n_exposures
         median=multi,
         blot=multi,
     )
+    if multi and spec.cr_method == "deepcr":
+        # Per-frame route (issue #61): CR masks are already in the DQ arrays
+        # (apply_per_frame_cr_masks), so the stack rejection is off and the
+        # combine is a plain weighted mean. resetbits=0 is load-bearing —
+        # AstroDrizzle's default resetbits (4096) clears exactly the DQ bit
+        # the masks were written into, silently producing an unmasked mosaic.
+        kwargs.update(driz_cr=False, median=False, blot=False, resetbits=0)
     return kwargs
 
 
+def star_pass_kwargs_for(
+    spec: TargetSpec, adapter: InstrumentAdapter, n_exposures: int
+) -> Dict:
+    """
+    Keyword set for the dedicated PSF-star drizzle pass
+    (``psf_star_pass="no_cr"``, issue #62): the science kwargs with every
+    stack CR-rejection step off and previously written CR DQ flags treated
+    as good (``final_bits`` gains the CR bit). ``resetbits=0`` so the
+    science pass's DQ flags survive in the input exposures — frame products
+    and re-runs still see them. Pure function, unit-testable.
+    """
+    kwargs = drizzle_kwargs_for(spec, adapter, n_exposures)
+    kwargs.update(
+        driz_cr=False,
+        median=False,
+        blot=False,
+        resetbits=0,
+        final_bits=int(kwargs.get("final_bits", 0)) | CR_DQ_BIT,
+    )
+    return kwargs
+
+
+def dq_with_cr_flags(dq: np.ndarray, cr_mask: np.ndarray, cr_bit: int = CR_DQ_BIT) -> np.ndarray:
+    """
+    One chip's DQ array with the CR bit rewritten from ``cr_mask``: the bit
+    is cleared everywhere first (so re-runs are idempotent, mirroring
+    AstroDrizzle's own resetbits behaviour) then set where the mask flags a
+    cosmic ray. Pure function; preserves the DQ dtype.
+    """
+    dq = np.asarray(dq)
+    cr_mask = np.asarray(cr_mask, dtype=bool)
+    if dq.shape != cr_mask.shape:
+        raise ValueError(f"DQ/mask shape mismatch: {dq.shape} vs {cr_mask.shape}")
+    cleared = (dq & ~cr_bit).astype(dq.dtype)
+    return np.where(cr_mask, cleared | cr_bit, cleared).astype(dq.dtype)
+
+
+def apply_per_frame_cr_masks(exposures: List[Path], adapter: InstrumentAdapter) -> Dict:
+    """
+    The ``cr_method="deepcr"`` masking step (issue #61): per-frame deepCR
+    cosmic-ray masks written into each exposure's DQ arrays as the
+    AstroDrizzle CR bit, replacing driz_cr's blotted-median stack rejection
+    (which reads systematically low on steep gradients and flags genuine
+    core flux). Mutates the exposures' DQ in place — exactly as driz_cr
+    itself does — idempotently (the CR bit is cleared before rewriting).
+    Returns the provenance fragment.
+    """
+    from astropy.io import fits
+
+    from ..package import cosmic_rays as cr_mod
+
+    masker = cr_mod.masker_for(adapter.key)
+    n_cr_pixels = {}
+    for path in exposures:
+        with fits.open(path, mode="update") as hdul:
+            total = 0
+            for hdu in hdul:
+                if hdu.name != "SCI":
+                    continue
+                mask = masker(hdu.data)
+                dq_hdu = hdul["DQ", hdu.ver]
+                dq_hdu.data = dq_with_cr_flags(dq_hdu.data, mask)
+                total += int(mask.sum())
+            n_cr_pixels[Path(path).name] = total
+    return {
+        **cr_mod.cr_method_record(adapter.key),
+        "cr_bit": CR_DQ_BIT,
+        "n_cr_pixels": n_cr_pixels,
+    }
+
+
 def combine(
     exposures: List[Path],
     spec: TargetSpec,
@@ -75,6 +160,9 @@ def combine(
     # the scratch dir with a relative, already-lowercase output root. This
     # also keeps AstroDrizzle's cwd scratch files contained.
     output_name = f"{spec.name}_{spec.filter_name}".lower()
+    per_frame_cr = None
+    if len(exposures) > 1 and spec.cr_method == "deepcr":
+        per_frame_cr = apply_per_frame_cr_masks(exposures, adapter)
     kwargs = drizzle_kwargs_for(spec, adapter, len(exposures))
     with chdir_scratch(output_dir) as output_dir:
         astrodrizzle.AstroDrizzle(
@@ -95,12 +183,55 @@ def _one(suffix: str) -> Path:
     sci = _one("_sci.fits")
     wht = _one("_wht.fits")
 
+    tail = {"cr_method": spec.cr_method}
+    if per_frame_cr is not None:
+        tail["per_frame_cr"] = per_frame_cr
     provenance = combine_provenance(
         spec,
         adapter,
         exposures,
         fits.getdata(wht),
         kwargs_key="drizzle_kwargs",
         kwargs={k: kwargs[k] for k in sorted(kwargs)},
+        tail=tail,
     )
     return sci, wht, provenance
+
+
+def combine_star_pass(
+    exposures: List[Path],
+    spec: TargetSpec,
+    adapter: InstrumentAdapter,
+    output_dir: Path,
+) -> Tuple[Path, Dict]:
+    """
+    Drizzle the dedicated CR-flag-ignoring PSF-star mosaic
+    (``psf_star_pass="no_cr"``, issue #62) onto the same grid/geometry as
+    the science combine; returns (sci_path, kwargs). AstroDrizzle backend
+    only — the caller gates on ``adapter.combine_backend``.
+    """
+    if adapter.combine_backend != "astrodrizzle":
+        raise ValueError(
+            "combine_star_pass supports the astrodrizzle backend only "
+            f"(instrument {adapter.key!r})"
+        )
+
+    from drizzlepac import astrodrizzle
+
+    from ._common import chdir_scratch
+
+    output_name = f"{spec.name}_{spec.filter_name}_starpass".lower()
+    kwargs = star_pass_kwargs_for(spec, adapter, len(exposures))
+    with chdir_scratch(output_dir) as output_dir:
+        astrodrizzle.AstroDrizzle(
+            input=[str(p) for p in exposures],
+            output=output_name,
+            **kwargs,
+        )
+    hits = sorted(glob.glob(f"{output_dir / output_name}*_sci.fits"))
+    if len(hits) != 1:
+        raise FileNotFoundError(
+            f"expected exactly one star-pass _sci.fits for "
+            f"{output_dir / output_name}, got {hits}"
+        )
+    return Path(hits[0]), kwargs

autoreduce/pipeline.py

@@ -25,6 +25,7 @@
 from .drizzle.diagnostics import check_weight_uniformity
 from .instruments import InstrumentAdapter
 from .noise import rms as rms_mod
+from .package import cosmic_rays as cr_mod
 from .package import cutout as cutout_mod
 from .package import frames as frames_mod
 from .package import provenance as provenance_mod
@@ -481,6 +482,57 @@ def _noise(ctx: _StageContext, sci, wht, exptime: float) -> np.ndarray:
     return noise
 
 
+def _star_pass_image(ctx: _StageContext, sci, header):
+    """
+    The mosaic that PSF star finding + stamp extraction run on (issue #62),
+    decoupled from the shipped science mosaic: driz_cr's blotted-median
+    rejection reads systematically low on star cores, holing PSF stars
+    before DAOStarFinder sees them. Returns (image, header, provenance) —
+    provenance always records which pass fed the stars, so the coupling
+    cannot silently regress.
+
+    ``psf_star_pass="auto"`` (default) never adds a drizzle: it uses the
+    science mosaic, which is already the least-CR-rejected pass at zero
+    cost when no stack rejection ran (single-exposure branch, per-frame
+    cr_method, non-astrodrizzle backends) — and otherwise is simply the
+    only pass available without doubling combine time. The explicit
+    ``"no_cr"`` opt-in buys the dedicated CR-flag-ignoring pass with a
+    second full AstroDrizzle run.
+    """
+    spec, adapter = ctx.spec, ctx.adapter
+    single = bool(ctx.record["drizzle"].get("single_exposure_branch"))
+    if spec.psf_star_pass == "no_cr" and not single:
+        from astropy.io import fits
+
+        sci_path, kwargs = combine_mod.combine_star_pass(
+            ctx.exposures, spec, adapter, ctx.work_dir
+        )
+        with fits.open(sci_path) as hdul:
+            star_sci = hdul[0].data.astype(float)
+            star_header = hdul[0].header.copy()
+        return star_sci, star_header, {
+            "star_source_pass": "no_cr_drizzle",
+            "star_pass_kwargs": {k: kwargs[k] for k in sorted(kwargs)},
+        }
+    if single:
+        reason = "single-exposure branch: no CR rejection ran"
+    elif getattr(adapter, "combine_backend", None) != "astrodrizzle":
+        reason = (
+            f"combine backend {adapter.combine_backend!r}: the dedicated "
+            "star pass applies to the astrodrizzle path only"
+        )
+    elif spec.cr_method != "driz_cr":
+        reason = f"cr_method={spec.cr_method!r}: per-frame CR masks, no stack rejection"
+    elif spec.psf_star_pass == "science":
+        reason = "psf_star_pass='science': pinned to the shipped mosaic"
+    else:
+        reason = (
+            "psf_star_pass='auto': no cheaper less-CR-rejected pass exists; "
+            "opt into 'no_cr' for a dedicated star drizzle"
+        )
+    return sci, header, {"star_source_pass": "science", "star_source_reason": reason}
+
+
 def _psf(ctx: _StageContext, sci, header, noise=None):
     from astropy.io import fits
     from astropy.wcs import WCS
@@ -521,7 +573,8 @@ def _psf(ctx: _StageContext, sci, header, noise=None):
         )
         ctx.record["psf"] = diag
         return psf, psf_full
-    target_xy = WCS(header).world_to_pixel_values(spec.ra, spec.dec)
+    star_sci, star_header, star_pass_prov = _star_pass_image(ctx, sci, header)
+    target_xy = WCS(star_header).world_to_pixel_values(spec.ra, spec.dec)
     selection = stars_mod.StarSelection()
     max_single_exptime = adapter.max_single_exposure_seconds(
         [fits.getheader(p) for p in ctx.exposures]
@@ -532,7 +585,7 @@ def _psf(ctx: _StageContext, sci, header, noise=None):
         else None
     )
     stars = stars_mod.find_stars(
-        sci,
+        star_sci,
         selection,
         target_xy=(float(target_xy[0]), float(target_xy[1])),
         peak_max=peak_max,
@@ -548,18 +601,21 @@ def _psf(ctx: _StageContext, sci, header, noise=None):
                 "psf_backend='starred' needs the noise map (STARRED weights "
                 "stars by per-pixel noise); pipeline did not pass it"
             )
+        # Stamps come from the star pass; the science-pass noise map is a
+        # faithful weight there (same exposures/geometry — the passes differ
+        # only in which pixels CR rejection removed).
         psf, psf_full, psf_diag = starred_mod.build_starred_epsf(
-            sci, noise, stars, spec.psf_shape, spec.psf_full_shape
+            star_sci, noise, stars, spec.psf_shape, spec.psf_full_shape
         )
     else:
         psf, psf_full, psf_diag = epsf_mod.build_epsf(
-            sci, stars, spec.psf_shape, spec.psf_full_shape
+            star_sci, stars, spec.psf_shape, spec.psf_full_shape
         )
     if adapter.observatory == "keck":
         # Tier B: in-field ePSF. Usable, but an AO PSF from field stars at a
         # different anisoplanatic angle is still provisional by contract.
         psf_diag = {"psf_provisional": True, **psf_diag}
-    ctx.record["psf"] = psf_diag
+    ctx.record["psf"] = {**star_pass_prov, **psf_diag}
     return psf, psf_full
 
 
@@ -705,6 +761,28 @@ def reduce_target(
             "psf_from_frames supports HST and JWST only "
             f"(instrument {spec.instrument!r})"
         )
+    backend = getattr(adapter, "combine_backend", None)
+    if spec.cr_method != "driz_cr" and backend != "astrodrizzle":
+        raise ValueError(
+            f"cr_method={spec.cr_method!r} supports the HST astrodrizzle "
+            f"path only (instrument {spec.instrument!r})"
+        )
+    if spec.cr_method == "deepcr" and spec.instrument not in cr_mod.DEEPCR_MODELS:
+        raise ValueError(
+            f"cr_method='deepcr' needs a registered deepCR model; none for "
+            f"{spec.instrument!r} (known: {sorted(cr_mod.DEEPCR_MODELS)}; "
+            "wfc3_ir cosmic rays are already per-frame flagged by calwf3 "
+            "ramp fitting)"
+        )
+    if spec.psf_star_pass == "no_cr" and (
+        backend != "astrodrizzle" or spec.psf_from_frames
+    ):
+        raise ValueError(
+            "psf_star_pass='no_cr' builds a second AstroDrizzle pass for "
+            "mosaic star finding — HST astrodrizzle path without "
+            f"psf_from_frames only (instrument {spec.instrument!r}, "
+            f"psf_from_frames={spec.psf_from_frames})"
+        )
     if spec.inject_image and adapter.domain != "visibility" and (
         observatory, getattr(adapter, "combine_backend", None)
     ) not in (

autoreduce/target.py

@@ -42,6 +42,20 @@ class TargetSpec:
     final_pixfrac: float = 0.8
     final_kernel: str = "square"
 
+    # Cosmic-ray rejection route for multi-exposure AstroDrizzle combines
+    # (issue #61). "driz_cr" (STScI default): blotted-median reference +
+    # driz_cr — on steep gradients (galaxy cores, stars) the sub-pixel-
+    # dithered median reads systematically low, so genuine core flux can be
+    # flagged as CR. "deepcr": per-frame deepCR masks written into each
+    # exposure's DQ arrays (the AstroDrizzle CR bit, 4096), then a plain
+    # weighted-mean drizzle with median/blot/driz_cr off and resetbits=0
+    # (the default resetbits, 4096, would clear exactly the bit the masks
+    # were written into). #61's reporter used LACosmic; the pipeline's
+    # established per-frame CR machinery is deepCR (package/cosmic_rays.py),
+    # reused here instead of adding a dependency. Default deliberately
+    # unchanged pending SLACS validation (see hst_acs_pipeline.md stage 3).
+    cr_method: str = "driz_cr"
+
     # PSF products (design doc stage 5).
     psf_shape: Tuple[int, int] = (21, 21)
     psf_full_shape: Tuple[int, int] = (61, 61)
@@ -64,6 +78,18 @@ class TargetSpec:
     # when psf_from_frames or the Keck tier-A path own the PSF stage.
     psf_backend: str = "epsf"
 
+    # Which drizzle pass feeds mosaic PSF star finding + stamp extraction
+    # (issue #62): driz_cr's core rejection can hole PSF-star cores before
+    # DAOStarFinder sees them, so the star pass is decoupled from the
+    # shipped science mosaic. "auto" (default): the least-CR-rejected pass
+    # that costs no extra drizzle — today that is always the science mosaic
+    # (recorded with the reason); "science": pin star finding to the shipped
+    # mosaic; "no_cr": build a dedicated CR-flag-ignoring drizzle for the
+    # stars — an explicit opt-in, because it is a second full AstroDrizzle
+    # pass (HST astrodrizzle path only). Provenance records which pass fed
+    # the stars either way.
+    psf_star_pass: str = "auto"
+
     # Alignment: residual (pixels) above which TweakReg refinement triggers.
     alignment_tolerance_pix: float = 0.1
 
@@ -135,6 +161,15 @@ def __post_init__(self):
             raise ValueError(f"dec out of range: {self.dec}")
         if not 0.0 < self.final_pixfrac <= 1.0:
             raise ValueError(f"final_pixfrac must be in (0, 1]: {self.final_pixfrac}")
+        if self.cr_method not in ("driz_cr", "deepcr"):
+            raise ValueError(
+                f"cr_method must be 'driz_cr' or 'deepcr': {self.cr_method!r}"
+            )
+        if self.psf_star_pass not in ("auto", "science", "no_cr"):
+            raise ValueError(
+                "psf_star_pass must be 'auto', 'science' or 'no_cr': "
+                f"{self.psf_star_pass!r}"
+            )
         if self.alma_width < 0:
             raise ValueError(
                 f"alma_width must be >= 0 (0 = collapse the spw): {self.alma_width}"