diff --git a/pipeline/measure.py b/pipeline/measure.py new file mode 100644 index 0000000..2593ee1 --- /dev/null +++ b/pipeline/measure.py @@ -0,0 +1,127 @@ +"""Stage 2: measure every frame and cache its star list. + +Two jobs. It produces the quality numbers used to choose a registration +reference and to weight the stack - sky level, noise, star sharpness - and it +caches each frame's detected stars, because both registration and the plate +solve need them and detection costs far more than reading a small array back. + +Frames are opened one at a time. A 4096x4096 float32 frame is 67 MB and a +session can hold ninety of them; there is no reason to have more than one in +memory at once. +""" +import os + +import numpy as np +import sep + +from astropy.io import fits + +import layout + +MAX_STARS = 500 # cached per frame, brightest first +DETECT_SIGMA = 5.0 + + +def measure_frame(path, max_stars=MAX_STARS): + """Sky, noise, seeing and a star list for one frame.""" + with fits.open(path, memmap=False) as hd: + data = hd[0].data.astype(np.float32) + + bkg = sep.Background(data, bw=64, bh=64, fw=3, fh=3) + sky = float(np.median(bkg.back())) + rms = float(bkg.globalrms) + sub = data - bkg.back() + + objs = sep.extract(sub, DETECT_SIGMA, err=rms, minarea=9, + deblend_cont=0.005) + # Saturated cores and cosmic-ray hits both make poor registration anchors + # and poor seeing estimates, so they are excluded before anything is + # measured from them. + objs = objs[(objs["flag"] == 0) & (objs["npix"] > 12) & + (objs["npix"] < 3000)] + if len(objs) == 0: + del data, sub, bkg + return dict(sky=sky, rms=rms, fwhm=np.nan, nstars=0, + xy=np.zeros((0, 2)), flux=np.zeros(0)) + + order = np.argsort(objs["flux"])[::-1][:max_stars] + objs = objs[order] + radius, _ = sep.flux_radius(sub, objs["x"], objs["y"], 6.0 * objs["a"], + 0.5, normflux=objs["flux"]) + fwhm = float(np.median(radius) * 2.0) + + result = dict( + sky=sky, rms=rms, fwhm=fwhm, nstars=int(len(objs)), + xy=np.column_stack([objs["x"], objs["y"]]).astype(np.float64), + flux=objs["flux"].astype(np.float64), + ) + del data, sub, bkg, objs + return result + + +def run(session, force=False, verbose=True): + """Measure every frame in a session, caching to intermediates/.""" + cache_path = os.path.join(session.root, layout.INTERMEDIATES, + "_measure.npz") + os.makedirs(os.path.dirname(cache_path), exist_ok=True) + + cached = {} + if os.path.exists(cache_path) and not force: + z = np.load(cache_path, allow_pickle=True) + cached = {k: z[k] for k in z.files} + + store, rows = {}, [] + for frame in session.frames: + key = frame.name + if f"{key}|xy" in cached and not force: + xy = cached[f"{key}|xy"] + stats = cached[f"{key}|stats"].item() + else: + stats = measure_frame(frame.path) + xy = stats.pop("xy") + flux = stats.pop("flux") + store[f"{key}|flux"] = flux + store[f"{key}|xy"] = xy + store[f"{key}|stats"] = np.array(stats, dtype=object) + rows.append((frame, stats)) + if verbose: + print(f" {frame.filter:10s} {frame.name[-28:]:28s} " + f"sky={stats['sky']:8.1f} rms={stats['rms']:7.2f} " + f"fwhm={stats['fwhm']:5.2f}px stars={stats['nstars']:4d}") + + # Carry forward anything already cached that was not re-measured. + for k, v in cached.items(): + store.setdefault(k, v) + np.savez_compressed(cache_path, **store) + return rows + + +def load(session): + """Star lists and stats for a session, as measured earlier.""" + path = os.path.join(session.root, layout.INTERMEDIATES, "_measure.npz") + if not os.path.exists(path): + raise FileNotFoundError( + f"{path} not found - run the measure stage first") + z = np.load(path, allow_pickle=True) + xy = {k.split("|")[0]: z[k] for k in z.files if k.endswith("|xy")} + stats = {k.split("|")[0]: z[k].item() for k in z.files + if k.endswith("|stats")} + return xy, stats + + +def choose_reference(session, stats): + """Pick the frame everything else is aligned to. + + The reference sets the output grid, so it should be the sharpest frame of + the filter with the most signal - which is also usually the filter with the + finest sampling. Choosing a poor reference costs resolution in every other + frame, permanently, because registration can only resample onto it. + """ + best_filter = session.filters[0] + candidates = [f for f in session.by_filter(best_filter) + if stats.get(f.name, {}).get("nstars", 0) >= 8] + if not candidates: + candidates = session.by_filter(best_filter) + ref = min(candidates, + key=lambda f: stats.get(f.name, {}).get("fwhm", np.inf)) + return ref diff --git a/pipeline/register.py b/pipeline/register.py new file mode 100644 index 0000000..c4a522d --- /dev/null +++ b/pipeline/register.py @@ -0,0 +1,159 @@ +"""Stage 3: register every frame onto one grid and combine per filter. + +Registration matches asterisms between the cached star lists rather than +cross-correlating pixels: matching a few hundred coordinates is far cheaper +than comparing 15 megapixel frames, and it copes with the field rotating +between the east and west sides of the meridian. + +**One reference for all filters, not one per filter.** That is what makes the +masters pixel-aligned, so the colour composite needs no further registration. + +The awkward case this has to survive: NGC 6744 was shot with luminance at bin1 +(4096x4096) and colour at bin2 (2048x2048). Frames do not share a shape or a +pixel scale. Because astroalign returns a similarity transform - rotation, +translation AND scale - the maths already handles it; what does not handle it +is assuming the output is the same shape as the input. Every warp is therefore +given the reference's shape explicitly, and the bin2 frames are resampled up +onto the bin1 grid. +""" +import os + +import astroalign as aa +import numpy as np +from astropy.io import fits +from astropy.stats import sigma_clip +from skimage.transform import warp + +import layout +import measure as measure_mod + +aa.MIN_MATCHES_FRACTION = 0.6 +aa.NUM_NEAREST_NEIGHBORS = 8 + + +def combine(cube, weights, sigma=3.0): + """Weighted sigma-clipped mean, in row blocks to bound peak memory. + + sigma_clip allocates a mask and float64 intermediates; on a full cube that + can triple peak usage, which matters on a machine with a few GB free. + """ + n, ny, nx = cube.shape + out = np.zeros((ny, nx), np.float32) + w = np.asarray(weights, np.float32)[:, None, None] + step = 256 + for y0 in range(0, ny, step): + y1 = min(y0 + step, ny) + block = cube[:, y0:y1, :] + if n >= 3: + clipped = sigma_clip(block, sigma=sigma, maxiters=1, axis=0, + masked=True, copy=True) + good = ~clipped.mask + else: + # With one or two frames there is nothing to reject against, and + # clipping would just throw away signal. + good = np.ones(block.shape, bool) + finite = np.isfinite(block) + good &= finite + vals = np.where(finite, block, 0.0) + wb = np.broadcast_to(w, block.shape) * good + denom = wb.sum(axis=0) + denom[denom == 0] = np.nan + out[y0:y1, :] = np.nansum(vals * wb, axis=0) / denom + del block, finite, good, vals, wb, denom + return np.nan_to_num(out, nan=0.0) + + +def run(session, verbose=True): + """Register and stack every filter; write masters to stacks/masters/.""" + xy, stats = measure_mod.load(session) + ref = measure_mod.choose_reference(session, stats) + ref_xy = xy[ref.name] + ref_shape = ref.shape + if verbose: + print(f" reference {ref.name[-32:]} ({ref.filter}, " + f"fwhm {stats[ref.name]['fwhm']:.2f} px, " + f"{ref_shape[1]}x{ref_shape[0]})") + + out_dir = os.path.join(session.root, layout.MASTERS) + os.makedirs(out_dir, exist_ok=True) + written = {} + + for filt in session.filters: + frames = session.by_filter(filt) + planes, weights, used = [], [], [] + header0 = None + for frame in frames: + with fits.open(frame.path, memmap=False) as hd: + data = hd[0].data.astype(np.float32) + if header0 is None: + header0 = hd[0].header.copy() + if frame.name == ref.name: + reg = data + else: + src = xy.get(frame.name) + if src is None or len(src) < 8: + print(f" {frame.name[-28:]}: too few stars - " + f"dropped") + del data + continue + try: + tform, _ = aa.find_transform(src, ref_xy) + except Exception as exc: # noqa: BLE001 + print(f" {frame.name[-28:]}: registration failed " + f"({type(exc).__name__}) - dropped") + del data + continue + # output_shape is the whole point: without it a bin2 frame + # would be written onto a bin2-sized grid and silently fail to + # line up with a bin1 reference. + reg = warp(data, inverse_map=tform.inverse, + output_shape=ref_shape, order=3, mode="constant", + cval=np.nan, preserve_range=True).astype(np.float32) + del data + sky = float(np.nanmedian(reg)) + reg -= sky + planes.append(reg) + rms = stats.get(frame.name, {}).get("rms", 1.0) or 1.0 + weights.append(1.0 / (rms * rms)) + used.append(frame) + + if not planes: + print(f" {filt}: no frames survived registration - skipped") + continue + + cube = np.stack(planes, axis=0) + del planes + # Normalise for transparency: scale each frame so its bright signal + # matches the group, so a frame through thin cloud cannot drag the + # mean down. + levels = np.array([np.nanpercentile(p, 99.5) for p in cube]) + ref_level = np.nanmedian(levels) + for i, lv in enumerate(levels): + if lv > 0: + cube[i] *= float(ref_level / lv) + + master = combine(cube, weights) + nframes = cube.shape[0] + total_exp = sum(f.exptime for f in used) + del cube + + hdr = header0 + for key in ("CBLACK", "CWHITE", "PEDESTAL", "HISTORY"): + hdr.remove(key, ignore_missing=True, remove_all=True) + hdr["FILTER"] = filt + hdr["NCOMBINE"] = (nframes, "frames in this master") + hdr["EXPTOTAL"] = (total_exp, "[s] total integration") + hdr["STACKREF"] = (ref.name[:60], "registration reference frame") + hdr["STACKALG"] = ("sigma-clipped weighted mean" if nframes >= 3 + else "weighted mean (too few frames to clip)", + "combine method") + hdr["IMAGETYP"] = "Master Light" + path = os.path.join(out_dir, f"master-{filt}.fit") + fits.PrimaryHDU(master.astype(np.float32), hdr).writeto( + path, overwrite=True) + written[filt] = path + if verbose: + print(f" {filt:10s} {nframes:3d} frames " + f"{total_exp/60:6.1f} min -> master-{filt}.fit") + del master + return written, ref