astrophotography/session-scripts/mo_sensitivity.py

"""Moving-object search, pass 3: measure what the search could have found.

A null result is only worth reading if the sensitivity behind it is known, so
this injects synthetic moving point sources into the real pixel data and runs
the identical detection -> registration -> static-rejection -> linking chain
over them.

Sources are laid out on a grid of magnitude against apparent rate. Each is
given a random position and direction, is placed at p0 + v*t in the reference
frame (mapped back through each sub's own affine into native pixels before
injection), and is trailed across the 300 s exposure. Its flux is set from the
per-frame zero point, which is tied to the master's by the stars the two have
in common, so a quoted magnitude means the same thing as it does for the
master's photometry.

Recovery is scored by matching linked tracklets back to the injection truth.
The same run also reports the false-alarm rate, measured by permuting the
frame times of the real residual detections: that keeps the real spatial
density of cosmic rays and noise peaks but destroys any genuine temporal
ordering, so anything the linker still produces is by construction spurious.

Output: _mo_sensitivity.npz
"""
import os
import sys

import numpy as np
from astropy.io import fits

import mo_common as C

import layout

MAGS = np.array([17.0, 17.5, 18.0, 18.5, 18.75, 19.0, 19.25])
RATES_PX_HR = np.array([3.0, 5.0, 6.5, 8.0, 10.0, 15.0, 25.0, 40.0,
                        70.0, 120.0, 180.0, 250.0])
NREP = 4                 # injections per (magnitude, rate) cell
MARGIN = 350             # px kept clear of the frame edge
MINSEP = 120             # px between injected tracks at t=0
EXPHR = 300.0 / 3600.0   # exposure length in hours


def plan(times, shape, rng):
    """Random start positions and directions for the injection grid."""
    ny, nx = shape
    tspan = times[-1]
    rows = []
    placed = []
    for mag in MAGS:
        for rate in RATES_PX_HR:
            for _ in range(NREP):
                for _try in range(200):
                    ang = rng.uniform(0, 2 * np.pi)
                    vx, vy = rate * np.cos(ang), rate * np.sin(ang)
                    x0 = rng.uniform(MARGIN, nx - MARGIN)
                    y0 = rng.uniform(MARGIN, ny - MARGIN)
                    x1, y1 = x0 + vx * tspan, y0 + vy * tspan
                    if not (MARGIN < x1 < nx - MARGIN and
                            MARGIN < y1 < ny - MARGIN):
                        continue
                    if placed and min(np.hypot(x0 - p[0], y0 - p[1])
                                      for p in placed) < MINSEP:
                        continue
                    placed.append((x0, y0))
                    rows.append(dict(mag=float(mag), rate=float(rate),
                                     x0=x0, y0=y0, vx=vx, vy=vy))
                    break
    return rows


def run_pass(inject, times, tforms, mtree, zps, rng, thresh=3.5):
    """Detect on every sub (optionally with synthetics added) and link."""
    frames = C.lum_frames()
    pts, nats, mags = [], [], []
    for i, (key, path) in enumerate(frames):
        with fits.open(path, memmap=False) as hd:
            img = hd[0].data.astype(np.float32)
        tf = tforms[key]
        inv = tf.inverse
        lin = np.linalg.inv(tf.params[:2, :2])       # ref -> native, linear
        for s in inject:
            xr = s["x0"] + s["vx"] * times[i]
            yr = s["y0"] + s["vy"] * times[i]
            xn, yn = inv(np.array([[xr, yr]]))[0]
            d = lin @ np.array([s["vx"] * EXPHR, s["vy"] * EXPHR])
            ap_flux = 10 ** ((zps[key] - s["mag"]) / 2.5)
            C.add_source(img, xn, yn, ap_flux / C.APFRAC, d[0], d[1])
        objs, ap, _rms = C.detect(img, thresh=thresh)
        del img
        keep = (objs["npix"] > 3) & (objs["npix"] < 3000) & (ap > 0)
        objs, ap = objs[keep], ap[keep]
        nat = np.column_stack([objs["x"], objs["y"]]).astype(float)
        nats.append(nat)
        pts.append(tf(nat))
        mags.append(-2.5 * np.log10(ap) + zps[key])
    return pts, nats, mags


def score(cands, inject, times, tol=8.0):
    """Match linked tracklets back to the injection truth.

    Every candidate is assigned to its nearest compatible injection, so a
    source recovered as two overlapping tracklets counts once as a detection
    and does not also count as a false positive. Candidates that match no
    injection at all are the genuinely spurious ones.
    """
    tmid = times.mean()
    tx = np.array([s["x0"] + s["vx"] * tmid for s in inject])
    ty = np.array([s["y0"] + s["vy"] * tmid for s in inject])
    trate = np.array([s["rate"] for s in inject])
    found = np.zeros(len(inject), dtype=bool)
    spurious = []
    for ci, c in enumerate(cands):
        cxm = c["cx"][0] + c["cx"][1] * tmid
        cym = c["cy"][0] + c["cy"][1] * tmid
        d = np.hypot(tx - cxm, ty - cym)
        ok = (d < tol) & (np.abs(trate - c["rate"]) < 0.25 * trate + 3)
        if ok.any():
            found[np.argmin(np.where(ok, d, np.inf))] = True
        else:
            spurious.append(ci)
    return found, spurious


def main():
    times = C.lum_times()
    tforms = C.frame_transforms()
    mobjs, map_, _, _, mimg = C.master_sources()
    shape = mimg.shape
    del mimg
    from scipy.spatial import cKDTree
    mtree = cKDTree(np.column_stack([mobjs["x"], mobjs["y"]]))
    print(f"master: {len(mobjs)} static sources, field {shape}")

    # Per-frame zero points from the stars the sub and the master share.
    _meta, dets = C.load_dets()
    zps = {}
    for key, path in C.lum_frames():
        with fits.open(path, memmap=False) as hd:
            img = hd[0].data.astype(np.float32)
        objs, ap, _ = C.detect(img)
        del img
        z, n = C.frame_zeropoint(objs["x"], objs["y"], ap, mobjs, map_,
                                 tforms[key])
        zps[key] = z
        print(f"  {key}: zero point {z:.3f} from {n} matched stars")

    seed = int(sys.argv[1]) if len(sys.argv) > 1 else 20260721
    rng = np.random.default_rng(seed)
    inject = plan(times, shape, rng)
    print(f"\ninjecting {len(inject)} synthetic movers "
          f"({len(MAGS)} mags x {len(RATES_PX_HR)} rates x {NREP})")

    pts, nats, mags = run_pass(inject, times, tforms, mtree, zps, rng)
    print("detections per frame:", " ".join(str(len(p)) for p in pts))
    keep = C.residual_mask(pts, nats, mtree)
    rpts = [p[m] for p, m in zip(pts, keep)]
    rmags = [g[m] for g, m in zip(mags, keep)]
    rnat = [n[m] for n, m in zip(nats, keep)]
    print("residuals per frame:", " ".join(str(len(p)) for p in rpts))

    cands = C.link(times, rpts, rmags, nat=rnat)
    print(f"{len(cands)} tracklets linked")
    found, spurious = score(cands, inject, times)
    nspur = len(spurious)
    print(f"recovered {found.sum()}/{len(inject)}, "
          f"{nspur} unmatched (spurious) tracklets")

    grid = np.zeros((len(MAGS), len(RATES_PX_HR)))
    for k, s in enumerate(inject):
        i = int(np.where(MAGS == s["mag"])[0][0])
        j = int(np.where(RATES_PX_HR == s["rate"])[0][0])
        grid[i, j] += found[k]
    grid /= NREP
    print("\nrecovery fraction (rows = G mag, cols = rate px/hr):")
    print("        " + " ".join(f"{r:6.0f}" for r in RATES_PX_HR))
    for i, m in enumerate(MAGS):
        print(f"  {m:5.1f} " + " ".join(f"{v:6.2f}" for v in grid[i]))

    # ---- false-alarm control on the real (uninjected) data ----------
    real = np.load(layout.path("_mo_residuals.npz"),
                   allow_pickle=True)
    rp = [real[f"p{i}"] for i in range(len(times))]
    rf = [real[f"m{i}"] for i in range(len(times))]
    rn = [real[f"n{i}"] for i in range(len(times))]
    rm = rf
    fa = []
    for trial in range(20):
        perm = np.random.default_rng(1000 + trial).permutation(len(times))
        cs = C.link(times, [rp[k] for k in perm], [rm[k] for k in perm],
                    nat=[rn[k] for k in perm])
        fa.append(len(cs))
    print(f"\nfalse-alarm control: time-permuted real residuals, "
          f"20 trials -> {np.sum(fa)} tracklets total "
          f"({np.mean(fa):.2f} per trial)")

    np.savez(layout.path(f"_mo_sensitivity_{seed}.npz"),
             mags=MAGS, rates=RATES_PX_HR, grid=grid, nrep=NREP,
             inject=np.array([(s["mag"], s["rate"], s["x0"], s["y0"],
                               s["vx"], s["vy"]) for s in inject]),
             found=found, nspur=nspur, falsealarm=np.array(fa),
             times=times, scale=C.SCALE)
    print(f"saved _mo_sensitivity_{seed}.npz")


OUT_NPZ_DIR = C.OUT

if __name__ == "__main__":
    main()