add drizzle implementation for coaddition. Add cell-based metadetection on the resulting coadd

import numpy as np

import pandas as pd

from astropy.io import fits

from astropy.io.fits import Header

from astropy.wcs import WCS

from drizzle.resample import Drizzle

from drizzle.utils import calc_pixmap

def read_fits_image_wcs(path, ext=1):

    print(path)

    with fits.open(path, memmap=True) as hdul:

        data = hdul[ext].data.astype(np.float32, copy=False)

        wcs = WCS(hdul[ext].header)

    ny, nx = data.shape

    return data, wcs, nx, ny

def load_tile_from_parquet(parquet_path, tile_id):

    df = pd.read_parquet(

        parquet_path,

        columns=[

            "tile_id",

            "wcs_header",

            "nx_total",

            "ny_total",

            "ra_bbox_min",

            "ra_bbox_max",

            "dec_bbox_min",

            "dec_bbox_max",

            "ra_center",

            "dec_center",

        ],

    )

    row = df.loc[df.tile_id == tile_id].iloc[0]

    hdr = Header.fromstring(row["wcs_header"], sep="\n")

    tile_wcs = WCS(hdr)

    shape_out = (int(row["ny_total"]), int(row["nx_total"]))  # (ny, nx)

    tile_bbox = (

        float(row["ra_bbox_min"]),

        float(row["ra_bbox_max"]),

        float(row["dec_bbox_min"]),

        float(row["dec_bbox_max"]),

    )

    tile_center = (float(row["ra_center"]), float(row["dec_center"]))

    return tile_wcs, shape_out, tile_bbox, tile_center

# ---------------------------------------------------------------------

# Optional helpers for weights / flags / headers

# ---------------------------------------------------------------------

def read_fits_image(path, ext=1, dtype=np.float32):

    with fits.open(path, memmap=True) as hdul:

        return hdul[ext].data.astype(dtype, copy=False)

def read_header_value(path, key, ext=1, default=None):

    with fits.open(path, memmap=True) as hdul:

        return hdul[ext].header.get(key, default)

def build_input_weight(image, weight_path=None, flag_path=None,

                       weight_ext=1, flag_ext=1, bad_flag_bits=None):

    """

    Build drizzle inwht from optional weight and flag maps.

    Parameters

    ----------

    image : 2D ndarray

    weight_path : str or None

        Optional weight map FITS path.

    flag_path : str or None

        Optional flag / DQ FITS path.

    bad_flag_bits : int or None

        If None: any nonzero flag is treated as bad.

        If int: pixels with (flag & bad_flag_bits) != 0 are treated as bad.

    """

    inwht = np.ones(image.shape, dtype=np.float32)

    if weight_path is not None:

        w = read_fits_image(weight_path, ext=weight_ext, dtype=np.float32)

        if w.shape != image.shape:

            raise ValueError(f"weight map shape {w.shape} != image shape {image.shape}")

        inwht *= np.where(np.isfinite(w), w, 0.0)

    if flag_path is not None:

        flg = read_fits_image(flag_path, ext=flag_ext, dtype=np.int32)

        if flg.shape != image.shape:

            raise ValueError(f"flag map shape {flg.shape} != image shape {image.shape}")

        if bad_flag_bits is None:

            bad = (flg != 0)

        else:

            bad = ((flg & bad_flag_bits) != 0)

        inwht[bad] = 0.0

    inwht[~np.isfinite(image)] = 0.0

    return inwht

# ---------------------------------------------------------------------

# Simple overlap prefilter

# ---------------------------------------------------------------------

def get_image_corners_sky(wcs, nx, ny):

    x = np.array([0, nx - 1, nx - 1, 0], dtype=float)

    y = np.array([0, 0, ny - 1, ny - 1], dtype=float)

    ra, dec = wcs.all_pix2world(x, y, 0)

    return np.asarray(ra), np.asarray(dec)

def normalize_ra_deg(ra):

    return np.mod(ra, 360.0)

def ra_interval_overlap(ra1_min, ra1_max, ra2_min, ra2_max):

    def split_interval(rmin, rmax):

        if rmin <= rmax:

            return [(rmin, rmax)]

        return [(rmin, 360.0), (0.0, rmax)]

    for a1, b1 in split_interval(ra1_min, ra1_max):

        for a2, b2 in split_interval(ra2_min, ra2_max):

            if max(a1, a2) <= min(b1, b2):

                return True

    return False

def exposure_overlaps_tile(exp_wcs, nx, ny, tile_bbox):

    exp_ra, exp_dec = get_image_corners_sky(exp_wcs, nx, ny)

    exp_ra = normalize_ra_deg(exp_ra)

    exp_ra_min, exp_ra_max = exp_ra.min(), exp_ra.max()

    exp_dec_min, exp_dec_max = exp_dec.min(), exp_dec.max()

    tile_ra_min, tile_ra_max, tile_dec_min, tile_dec_max = tile_bbox

    tile_ra_min = normalize_ra_deg(tile_ra_min)

    tile_ra_max = normalize_ra_deg(tile_ra_max)

    return (

        max(exp_dec_min, tile_dec_min) <= min(exp_dec_max, tile_dec_max)

        and

        ra_interval_overlap(exp_ra_min, exp_ra_max, tile_ra_min, tile_ra_max)

    )

# ---------------------------------------------------------------------

# Main drizzle-to-tile pipeline using current drizzle API

# ---------------------------------------------------------------------

def drizzle_coadd_to_tile(

    exposures,

    parquet_path,

    tile_id,

    image_ext=1,

    weight_ext=1,

    flag_ext=1,

    pixfrac=1.0,

    kernel="square",

    fillval=np.nan,

    exptime_key="EXPTIME",

    in_units="counts",

    bad_flag_bits=None,

    do_context=True,

    out_fits=None,

):

    """

    Coadd input exposures onto one predefined tile using drizzle.

    Parameters

    ----------

    exposures : list of dict

        Each dict may contain:

        {

            "image_path": ...,

            "weight_path": ... optional,

            "flag_path": ... optional,

            "exptime": ... optional,

            "iscale": ... optional,

            "exp_id": ... optional

        }

    Notes

    -----

    Current drizzle API expects:

      - Drizzle(out_shape=...)

      - calc_pixmap(input_wcs, output_wcs, shape=...)

      - add_image(data, pixmap=..., inwht=..., exptime=..., ...)

    """

    tile_wcs, shape_out, tile_bbox, tile_center = load_tile_from_parquet(parquet_path, tile_id)

    ny_out, nx_out = shape_out

    driz = Drizzle(

        kernel=kernel,

        fillval=fillval,

        out_shape=shape_out,

        disable_ctx=(not do_context),

    )

    used = []

    skipped = []

    for i, exp in enumerate(exposures):

        image_path = exp["image_path"]

        exp_id = exp.get("exp_id", f"exp_{i}")

        image, in_wcs, nx, ny = read_fits_image_wcs(image_path, ext=image_ext)

        if not exposure_overlaps_tile(in_wcs, nx, ny, tile_bbox):

            skipped.append(exp_id)

            print(f"exp {exp_id} does not overlap with tild {tile_id}")

            continue

        inwht = build_input_weight(

            image=image,

            weight_path=exp.get("weight_path"),

            flag_path=exp.get("flag_path"),

            weight_ext=weight_ext,

            flag_ext=flag_ext,

            bad_flag_bits=bad_flag_bits,

        )

        exptime = exp.get("exptime", None)

        if exptime is None:

            exptime = read_header_value(image_path, exptime_key, ext=image_ext, default=1.0)

        if exptime is None:

            exptime = 1.0

        iscale = exp.get("iscale", 1.0)

        # Current drizzle API uses an explicit pixmap.

        pixmap = calc_pixmap(in_wcs, tile_wcs, shape=(ny, nx))

        driz.add_image(

            image,

            pixmap=pixmap,

            exptime=float(exptime),

            scale=1.0,

            weight_map=inwht,

            wht_scale=1.0,

            pixfrac=float(pixfrac),

            in_units=in_units,

            iscale=float(iscale),

        )

        used.append(exp_id)

    result = {

        "sci": np.array(driz.out_img, copy=True),

        "wht": np.array(driz.out_wht, copy=True),

        "ctx": None if driz.out_ctx is None else np.array(driz.out_ctx, copy=True),

        "tile_wcs": tile_wcs,

        "shape_out": shape_out,

        "tile_bbox": tile_bbox,

        "tile_center": tile_center,

        "used": used,

        "skipped": skipped,

    }

    if out_fits is not None:

        write_drizzle_output(

            out_fits,

            sci=result["sci"],

            wht=result["wht"],

            ctx=result["ctx"],

            wcs=tile_wcs,

            overwrite=True,

        )

    return result

# ---------------------------------------------------------------------

# Output writer

# ---------------------------------------------------------------------

def write_drizzle_output(path, sci, wht, ctx, wcs, overwrite=True):

    hdr = wcs.to_header()

    hdus = [

        fits.PrimaryHDU(),

        fits.ImageHDU(data=np.asarray(sci, dtype=np.float32), header=hdr, name="SCI"),

        fits.ImageHDU(data=np.asarray(wht, dtype=np.float32), header=hdr, name="WHT"),

    ]

    if ctx is not None:

        hdus.append(

            fits.ImageHDU(data=np.asarray(ctx, dtype=np.int32), header=hdr, name="CTX")

        )

    fits.HDUList(hdus).writeto(path, overwrite=overwrite)

    print(f"Wrote {path}")

if __name__ == "__main__":

    exposures = [

        {

            "exp_id": "exp001",

            "image_path": "/path/to/exp001_sci.fits",

            "weight_path": "/path/to/exp001_wht.fits",

            "flag_path": "/path/to/exp001_flg.fits",

            "exptime": 100.0,

        },

        {

            "exp_id": "exp002",

            "image_path": "/path/to/exp002_sci.fits",

            "weight_path": "/path/to/exp002_wht.fits",

            "flag_path": "/path/to/exp002_flg.fits",

        },

    ]

    result = drizzle_coadd_to_tile(

        exposures=exposures,

        parquet_path="tiles.parquet",

        tile_id="tile_000123",

        pixfrac=0.8,

        kernel="square",

        in_units="counts",

        out_fits="tile_000123_coadd.fits",

    )

    coadd = result["sci"]

    weight = result["wht"]

    context = result["ctx"]

 No newline at end of file