diff --git a/.gitignore b/.gitignore index f0bad73..75c284f 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,6 @@ .venv __pycache__/ *.pyc + +gpu_system_ssh_info +.gpu_paused.json diff --git a/bench.py b/bench.py index ac45592..2834cb7 100644 --- a/bench.py +++ b/bench.py @@ -22,6 +22,7 @@ across runs. """ import argparse +import bisect import datetime import itertools import json @@ -650,6 +651,9 @@ def cmd_fill(args): gap_lengths[(row["satellite"], row["wavelength"], row["t_start"])] = row["gap_length"] wanted = tuple(args.fillers.split(",")) if args.fillers else tuple(fillers.FILLERS) + # The learned filler is joint across all six bands, so it cannot run inside a loop + # that handles one band at a time; it gets its own pass below. + per_band = tuple(name for name in wanted if name != "learned") reliever = vfs.Reliever(label="fill") scored = defaultdict(list) # A slot whose truth frame will not read cannot be scored. Count them: silently @@ -660,7 +664,13 @@ def cmd_fill(args): attempted = 0 payload = [] - for wavelength in wavelengths: + if not per_band: + # Only the learned filler was asked for. Its own pass below reads exactly what + # it needs; walking every band here would pull all 5,304 damaged slots and their + # neighbours off the mount to compute nothing at all. + print(" (skipping per-band pass: no per-band fillers requested)") + + for wavelength in (wavelengths if per_band else ()): truth_cache = {} def truth_image(slot): @@ -670,6 +680,21 @@ def cmd_fill(args): truth_cache[slot] = image return truth_cache[slot] + + # Times where every satellite has a good frame in this band. The two + # spacecraft observe simultaneously, so such a slot isolates the instrument + # difference with none of the Sun's own evolution folded in -- which is the + # only sound way to calibrate one against the other. See fillers.crosssat. + per_satellite = [ + { + s[2] + for s in overlay.series(other, wavelength) + if s not in bad and overlay.path(s) is not None + } + for other in satellites + ] + paired_times = sorted(set.intersection(*per_satellite)) if len(satellites) > 1 else [] + for satellite in satellites: series = overlay.series(satellite, wavelength) good = [s for s in series if s not in bad and overlay.path(s) is not None] @@ -697,16 +722,25 @@ def cmd_fill(args): None, ) + calibration = None + if counterpart_slot is not None and paired_times: + pair_time = _closest(paired_times, slot[2]) + calibration = ( + truth_image((counterpart_slot[0], wavelength, pair_time)), + truth_image((satellite, wavelength, pair_time)), + ) + context = fillers.FillContext( before=truth_image((satellite, wavelength, before)) if before else None, dt_before=slot[2] - before if before else 0.0, after=truth_image((satellite, wavelength, after)) if after else None, dt_after=after - slot[2] if after else 0.0, counterpart=truth_image(counterpart_slot) if counterpart_slot else None, + calibration=calibration, header=_header_for(overlay, slot), ) - for filler_name in wanted: + for filler_name in per_band: filled = fillers.FILLERS[filler_name](context) if filled is None: continue @@ -723,6 +757,15 @@ def cmd_fill(args): truth_cache.clear() print(f" {wavelength}A done") + if "learned" in wanted: + learned_scores, learned_payload, learned_attempted = _fill_learned( + conn, case, overlay, satellites, wavelengths, bad, gap_lengths, source, + reliever, max_targets=args.max_targets, + ) + scored["learned"] = learned_scores + payload.extend(learned_payload) + attempted = max(attempted, learned_attempted) + conn.executemany( """ INSERT INTO bench_fill_result (case_id, frame_id, filler, source, scores_json) @@ -752,6 +795,237 @@ def cmd_fill(args): return 0 +def _stratified_targets(targets, gap_lengths, limit): + """`limit` targets spread evenly over gap length, not the first `limit` found. + + Fill quality is reported against gap length and varies by 10 dB across the range, so + a prefix of the sorted targets would sample whatever gap happened to be placed + earliest and report it as the model's overall quality. + """ + if not limit or limit >= len(targets): + return targets + by_length = defaultdict(list) + for target in targets: + satellite, when = target + lengths = [length for (s, _, t), length in gap_lengths.items() + if s == satellite and t == when] + by_length[max(lengths) if lengths else 0].append(target) + + chosen = [] + groups = [sorted(group) for _, group in sorted(by_length.items())] + position = 0 + while len(chosen) < limit and any(position < len(g) for g in groups): + for group in groups: + if position < len(group) and len(chosen) < limit: + chosen.append(group[position]) + position += 1 + return sorted(chosen) + + +def _fill_learned(conn, case, overlay, satellites, wavelengths, bad, gap_lengths, + source, reliever, max_targets=None): + """Run the learned filler, which reconstructs all six bands in one pass. + + A separate pass because the model is joint across bands while the loop above is not. + Running it inside that loop would mean six forward passes per slot, five of whose + outputs are thrown away, and six times the stack assembly with it. + + Returns (scores, rows, attempted) so the caller can merge it into its own tally. + """ + if fillers.load_learned() is None: + print("\nSkipping 'learned': no checkpoint " + f"(set {fillers.LEARNED_CHECKPOINT_ENV})") + return [], [], 0 + + targets = sorted({(slot[0], slot[2]) for slot in bad}) + if max_targets: + targets = _stratified_targets(targets, gap_lengths, max_targets) + print(f"\nLearned filler: {len(targets)} (satellite, time) targets") + read = _band_reader(overlay, wavelengths, bad, reliever=reliever) + clean_times = _clean_times(overlay, satellites, wavelengths, bad) + # Times where *both* satellites are clean in every band: the calibration pairs + # the cross-satellite photometric transfer is fitted from. A target's own time + # is never in this list -- the target is a bad slot, so it is not clean. + paired_times = (sorted(set(clean_times.get(satellites[0], [])) + & set(clean_times.get(satellites[1], []))) + if len(satellites) > 1 else []) + + scores, rows = [], [] + attempted = 0 + unreadable = 0 + for index, (satellite, when) in enumerate(targets): + stack = _build_stack(overlay, satellite, when, satellites, wavelengths, bad, + read, clean_times) + calibration = None + if paired_times: + other = next(s for s in satellites if s != satellite) + pair_time = _closest(paired_times, when) + pair = (read(other, pair_time), read(satellite, pair_time)) + if pair[0] is not None and pair[1] is not None: + calibration = pair + filled = fillers.learned(fillers.FillContext(stack=stack, + calibration=calibration)) + if filled is None: + unreadable += 1 + continue + for band, wavelength in enumerate(wavelengths): + slot = (satellite, wavelength, when) + if slot not in bad: + continue + attempted += 1 + path = overlay.truth_path(slot) + truth, _ = fitsio.read_image(path) if path else (None, None) + if truth is None: + unreadable += 1 + continue + score = metrics.score_fill(filled[band], truth, wavelength, + gap_lengths.get(slot, 0)) + scores.append(score) + frame_id = db.frame_id_by_slot(conn, *slot) + if frame_id is not None: + rows.append((case["id"], frame_id, "learned", source, + json.dumps(score.as_dict()))) + if index and index % 200 == 0: + print(f" {index}/{len(targets)}", flush=True) + if unreadable: + print(f" {unreadable} targets could not be read or filled") + return scores, rows, attempted + + +def _clean_times(overlay, satellites, wavelengths, bad): + """Per satellite, the times with a usable frame in every band. + + Computed once for the whole case. Deriving it inside :func:`_build_stack` meant + sorting all 31,488 slots afresh for each of ~900 targets, which is most of a minute + spent recomputing the same list. + """ + times = {} + for satellite in satellites: + usable = defaultdict(int) + for slot in overlay.slots(): + if slot[0] != satellite or slot in bad: + continue + if overlay.path(slot) is not None: + usable[slot[2]] += 1 + # *Every* band must be clean. Counting a time as usable when only some + # bands were would hand out anchors marked 'available' that carry damaged + # pixels, and -- worse -- let a damaged band into the calibration pair the + # cross-satellite gain is fitted from. + times[satellite] = sorted(t for t, bands in usable.items() + if bands == len(wavelengths)) + return times + + +def _build_stack(overlay, satellite, target_time, satellites, wavelengths, bad, read, + clean_times=None): + """Assemble the multi-frame stack the learned filler fuses. + + Mirrors `suvi.samples.stack_layout` so a model trained on shards meets the same + stack shape here: the same multi-scale offsets on both satellites, plus the nearest + usable frame in each direction as an anchor so a 300-slot gap is never empty. + + Entries are **all six bands at once**, because the model is joint across them -- + when one band is damaged the others usually are not, and that is much of what makes + the fusion work. `read(satellite, time)` returns a (6, H, W) array or None, and is + responsible for resolving damaged slots to their *damaged* pixels. + + A slot the case damaged is included as **suspect**, carrying those pixels rather + than being dropped. Every hand-written filler discards such a frame outright; that + is the difference this stack exists to measure. + """ + from suvi import samples + + def state_of(source, when): + """'suspect' if any band of this slot was damaged, else 'available'.""" + damaged = any((source, wl, when) in bad for wl in wavelengths) + return "suspect" if damaged else "available" + + entries = [] + for source, offset in samples.stack_layout(satellites, satellite): + when = target_time + offset * paths.CADENCE + if not any((source, wl, when) in overlay.archive for wl in wavelengths): + continue + image = read(source, when) + entries.append({ + "image": image, + "state": "missing" if image is None else state_of(source, when), + "dt": float(when - target_time), + "same_satellite": source == satellite, + "slot": (source, when), + }) + + # Anchors, so the far end of a long outage is still reachable. + seen = {entry["dt"] for entry in entries if entry["same_satellite"]} + if clean_times is None: + clean_times = _clean_times(overlay, (satellite,), wavelengths, bad) + clean = clean_times.get(satellite, []) + for before in (True, False): + anchor = _nearest(clean, target_time, before=before) + if anchor is None or float(anchor - target_time) in seen: + continue + image = read(satellite, anchor) + if image is None: + continue + entries.append({ + "image": image, + "state": "available", + "dt": float(anchor - target_time), + "same_satellite": True, + "slot": (satellite, anchor), + }) + return entries + + +def _band_reader(overlay, wavelengths, bad, limit=48, reliever=None): + """Reader of six-band slots for :func:`_build_stack`, with a bounded cache. + + Damaged slots resolve through the overlay to their *damaged* pixels; a filler handed + the pristine frame for a slot the case corrupted would be reading the answer key. + + Consecutive targets share most of their stack, so a small cache removes almost all + the repeat reads -- and it has to be bounded, because six bands of 1280x1280 float32 + is 39 MB and an unbounded one would run this VM out of memory long before the case + finished. + """ + cache = {} + order = [] + + def read(satellite, when): + key = (satellite, when) + if key in cache: + return cache[key] + bands = [] + for wavelength in wavelengths: + slot = (satellite, wavelength, when) + path = overlay.path(slot) + # Ticked per *read*, not per target. One target pulls ~90 frames through + # here, so ticking once per target would undercount by ninety-fold and the + # reliever would never reach its interval -- which is exactly how the mount + # runs out of file handles. + if reliever is not None: + reliever.tick() + image, _ = fitsio.read_image(path) if path else (None, None) + if image is None: + bands = None + break + bands.append(image) + value = np.stack(bands).astype(np.float32) if bands else None + cache[key] = value + order.append(key) + while len(order) > limit: + cache.pop(order.pop(0), None) + return value + + return read + + +def _closest(times, target): + """Closest time in sorted `times` to `target`, in either direction.""" + index = bisect.bisect_left(times, target) + candidates = times[max(0, index - 1) : index + 1] + return min(candidates, key=lambda t: abs(t - target)) + + def _nearest(times, target, before=True): """Closest time in `times` strictly before or after `target`.""" if before: @@ -769,6 +1043,101 @@ def _header_for(overlay, slot): return values +# -------------------------------------------------------------------------- compare + + +def cmd_compare(args): + """Fill quality by gap length, on one common set of frames. + + Restricting to the frames a chosen filler actually scored is the point. A model + evaluated on 200 sampled targets and a baseline evaluated on 5,304 are not + comparable numbers, and putting them in the same table implies they are. With + --restrict-to the table is one set of frames scored several ways. + """ + conn = db.connect(args.db, readonly=True) + case = conn.execute("SELECT * FROM bench_case WHERE name = ?", (args.case,)).fetchone() + if case is None: + raise SystemExit(f"No case named {args.case!r}") + + gaps = {row["frame_id"]: row["gap_length"] for row in conn.execute( + "SELECT frame_id, gap_length FROM bench_injection WHERE case_id = ?", (case["id"],) + )} + + rows = conn.execute( + "SELECT frame_id, filler, scores_json FROM bench_fill_result " + "WHERE case_id = ? AND source = ?", (case["id"], args.source) + ).fetchall() + if not rows: + raise SystemExit(f"No fill results for case {args.case!r} from {args.source!r}") + + by_filler = defaultdict(dict) + for row in rows: + by_filler[row["filler"]][row["frame_id"]] = json.loads(row["scores_json"]) + + common = None + if args.restrict_to: + if args.restrict_to not in by_filler: + raise SystemExit(f"No results for filler {args.restrict_to!r}") + common = set(by_filler[args.restrict_to]) + for name, scores in by_filler.items(): + common &= set(scores) + print(f"Restricted to the {len(common)} frames scored by " + f"every filler including '{args.restrict_to}'") + + lengths = sorted({gaps.get(frame_id, 0) + for scores in by_filler.values() for frame_id in scores + if common is None or frame_id in common}) + metric = args.metric + + table = {} + for name, scores in by_filler.items(): + cells = defaultdict(list) + for frame_id, values in scores.items(): + if common is not None and frame_id not in common: + continue + cells[gaps.get(frame_id, 0)].append(values[metric]) + table[name] = cells + + print(f"\nCase '{args.case}', source '{args.source}' -- {metric} by gap length\n") + print(f"{'filler':>16} " + " ".join(f"{length:>7}" for length in lengths) + f" {'n':>7}") + order = [n for n in ("hold_last", "linear_blend", "optical_flow", "solar_rotation", + "crosssat", "learned") if n in table] + order += [n for n in sorted(table) if n not in order] + for name in order: + cells = table[name] + printed = [] + for length in lengths: + values = cells.get(length) + printed.append(f"{sum(values) / len(values):>7.2f}" if values else f"{'--':>7}") + total = sum(len(v) for v in cells.values()) + print(f"{name:>16} " + " ".join(printed) + f" {total:>7}") + + # The bar the learned filler has to clear is the best of the others at each gap + # length, not any single one of them: the crossover is real and both sides win + # somewhere. + best = {} + for length in lengths: + options = { + name: sum(cells[length]) / len(cells[length]) + for name, cells in table.items() + if name != "learned" and cells.get(length) + } + if options: + best[length] = max(options.items(), key=lambda kv: kv[1]) + if best: + print("\n best non-learned: " + " ".join( + f"{length}:{name.split('_')[0]}({value:.2f})" + for length, (name, value) in best.items())) + if "learned" in table: + deltas = [] + for length in lengths: + values = table["learned"].get(length) + if values and length in best: + deltas.append(f"{length}:{sum(values) / len(values) - best[length][1]:+.2f}") + print(" learned minus best: " + " ".join(deltas)) + return 0 + + # --------------------------------------------------------------------------- report @@ -1246,6 +1615,9 @@ def build_parser(): fill = subparsers.add_parser("fill", help="reconstruct flagged slots and score them") fill.add_argument("--case", required=True) fill.add_argument("--fillers", default=None, help="comma-separated; default is all") + fill.add_argument("--max-targets", type=int, default=None, + help="score only this many targets, spread evenly over gap " + "lengths -- for a quick read on a model mid-training") group = fill.add_mutually_exclusive_group(required=True) group.add_argument("--oracle", action="store_true", help="use perfect knowledge of which slots are bad") @@ -1253,6 +1625,16 @@ def build_parser(): help="use a detector's output, showing end-to-end error") fill.set_defaults(func=cmd_fill) + compare = subparsers.add_parser( + "compare", help="fill quality by gap length across fillers, on common frames") + compare.add_argument("--case", required=True) + compare.add_argument("--source", default="oracle") + compare.add_argument("--metric", default="psnr", + choices=("psnr", "ssim", "rmse", "log_rmse")) + compare.add_argument("--restrict-to", default=None, + help="score every filler on only the frames this one scored") + compare.set_defaults(func=cmd_compare) + report = subparsers.add_parser("report", help="summarise a case's results") report.add_argument("--case", required=True) report.add_argument("--json", default=None) diff --git a/ceiling.py b/ceiling.py new file mode 100644 index 0000000..37f48f9 --- /dev/null +++ b/ceiling.py @@ -0,0 +1,236 @@ +#!/usr/bin/env python +"""Gate 0: measure the ceiling of the aligned-candidate design before any training. + + ceiling.py --window wide --per-stratum 12 + +The new learned filler can only ever output a per-pixel convex combination of its +*aligned* candidates, times a gain bounded to [0.74, 1.35]. That is a deliberate +structural cage -- it makes the collapse mode of the previous architecture +unreachable -- but it also means the achievable quality is capped by what the +candidates contain. This script measures that cap on the real holdout window, on +the CPU, in under an hour, so a training run that cannot possibly win is never +started. + +For each simulated gap length and counterpart stratum it scores blends of the same +aligned candidates the model would see: + + prior softmax of the fixed blend prior == an untrained model, exactly + uniform equal weight over valid candidates == the floor selection must beat + best1 the single best candidate per band == what perfect *frame* selection earns + +If these sit far below the hand-written baselines (optical_flow 44.79 dB at gap 1 / +34.47 at gap 300; fixed crosssat ~39.9 flat), the candidate set or the alignment is +at fault and no training run can fix it -- which is exactly how the rotation-sign +defect was caught: plain solar_rotation outscored every blend of what should have +been the same warped frames. + +Gaps are simulated by hiding the target satellite's frames inside a centred window +of the given length, exactly as a real outage hides them; `absent` strata hide the +counterpart satellite entirely. Targets never leak: the hidden frames resolve to +`missing`, not to suspiciously pristine `suspect` pixels. +""" + +import argparse +import json +import sys +import time + +import numpy as np + +import bench +from suvi import cases, db, fillers, metrics, paths, samples, vfs + +#: Simulated gap lengths, matching the published baseline table's columns. +GAPS = (1, 3, 30, 100, 300) +#: Margin (slots) a target keeps from the window edges, so anchors always exist. +MARGIN = 310 + + +def band_psnr(shown, shown_truth): + mse = float(np.mean((shown - shown_truth) ** 2)) + return float("inf") if mse == 0 else float(10.0 * np.log10(1.0 / mse)) + + +def blend_psnr(candidates_coded, weights, truth_display): + """Display PSNR of a weighted blend, composed in coded space like the model.""" + composite = np.tensordot(weights, candidates_coded, axes=(0, 0)) + radiance = samples.decode_from_model(composite) + scores = [band_psnr(metrics.to_display(radiance[i], wavelength), truth_display[i]) + for i, wavelength in enumerate(paths.WAVELENGTHS)] + return float(np.mean(scores)) + + +def prior_weights(condition, valid, decay, suspect_penalty=4.0): + distance = condition[:, 5] * 6.0 + logits = -decay * distance - suspect_penalty * condition[:, 2] + logits = np.where(valid > 0.5, logits, -np.inf) + weights = np.exp(logits - logits.max()) + return weights / weights.sum() + + +def score_target(aligned, condition, truth, decays, net=None): + """All the blends for one target. Returns {name: mean-over-band PSNR}.""" + candidates = aligned[0].numpy() # (S, 6, H, W) coded + cond = condition[0].numpy() + valid = np.clip(cond[:, 0] + cond[:, 2], 0, 1) + usable = valid > 0.5 + + truth_display = [metrics.to_display(truth[i], wavelength) + for i, wavelength in enumerate(paths.WAVELENGTHS)] + + out = {} + if net is not None: + import torch + with torch.no_grad(): + predicted = net(aligned, condition)[0].numpy() + radiance = samples.decode_from_model(predicted) + out["model"] = float(np.mean([ + band_psnr(metrics.to_display(radiance[i], wavelength), truth_display[i]) + for i, wavelength in enumerate(paths.WAVELENGTHS) + ])) + uniform = usable / usable.sum() + out["uniform"] = blend_psnr(candidates, uniform, truth_display) + for decay in decays: + out[f"prior[{decay:g}]"] = blend_psnr( + candidates, prior_weights(cond, valid, decay), truth_display + ) + + best_single = [] + for band, wavelength in enumerate(paths.WAVELENGTHS): + shown = np.stack([ + metrics.to_display(samples.decode_from_model(candidates[s, band]), wavelength) + for s in range(candidates.shape[0]) + ]) + best_single.append(max( + band_psnr(shown[s], truth_display[band]) + for s in range(candidates.shape[0]) if usable[s] + )) + out["best1"] = float(np.mean(best_single)) + return out + + +def pick_targets(eligible, count): + """`count` times spread evenly over the eligible list.""" + if len(eligible) <= count: + return list(eligible) + stride = len(eligible) / count + return [eligible[int(i * stride)] for i in range(count)] + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--window", default="wide") + parser.add_argument("--per-stratum", type=int, default=12) + parser.add_argument("--gaps", default=",".join(str(g) for g in GAPS)) + parser.add_argument("--decays", default="1.5", + help="comma-separated prior temporal decays to evaluate") + parser.add_argument("--model", default=None, + help="also score this trained checkpoint, as a 'model' column") + args = parser.parse_args(argv) + + import torch + + gaps = tuple(int(g) for g in args.gaps.split(",")) + decays = tuple(float(d) for d in args.decays.split(",")) + + net = None + if args.model: + loaded = fillers.load_learned(args.model, device="cpu") + if loaded is None: + raise SystemExit(f"could not load checkpoint {args.model}") + net = loaded[0] + + conn = db.connect(readonly=True) + window = bench._load_window(conn, args.window) + satellites = tuple(json.loads(window["satellites"])) + wavelengths = tuple(json.loads(window["wavelengths"])) + print(f"window '{args.window}': {window['t_start']} .. {window['t_end']}, " + f"satellites {satellites}") + + archive = cases.scan_window(paths.data_root(), satellites, wavelengths, + window["t_start"], window["t_end"]) + overlay = cases.Overlay(archive=archive) + reliever = vfs.Reliever(label="ceiling") + reader = bench._band_reader(overlay, wavelengths, set(), limit=64, + reliever=reliever) + clean_times = bench._clean_times(overlay, satellites, wavelengths, set()) + paired = sorted(set(clean_times[satellites[0]]) & set(clean_times[satellites[1]])) + print("clean six-band times: " + + ", ".join(f"G{s}: {len(clean_times[s])}" for s in satellites) + + f", paired: {len(paired)}") + + results = {} + started = time.time() + for gap in gaps: + blocked_offsets = range(-(gap // 2), gap - gap // 2) + for counterpart in (True, False): + label = (gap, "present" if counterpart else "absent") + stratum = [] + for satellite in satellites: + other = next(s for s in satellites if s != satellite) + eligible = [ + t for t in clean_times[satellite][MARGIN:-MARGIN or None] + if not counterpart or t in clean_times[other] + ] + for target in pick_targets(eligible, args.per_stratum // 2 or 1): + blocked = {target + k * paths.CADENCE for k in blocked_offsets} + + def read(source, when): + if source == satellite and when in blocked: + return None # inside the simulated gap + if not counterpart and source == other: + return None # counterpart outage + return reader(source, when) + + bad = {(satellite, wl, t) for t in blocked for wl in wavelengths} + filtered = { + satellite: [t for t in clean_times[satellite] + if t not in blocked], + other: clean_times[other] if counterpart else [], + } + stack = bench._build_stack(overlay, satellite, target, satellites, + wavelengths, bad, read, filtered) + calibration = None + if counterpart: + usable_pairs = [t for t in paired if t not in blocked] + if usable_pairs: + pair_time = bench._closest(usable_pairs, target) + pair = (reader(other, pair_time), + reader(satellite, pair_time)) + if pair[0] is not None and pair[1] is not None: + calibration = pair + + context = fillers.FillContext(stack=stack, calibration=calibration) + assembled = fillers.assemble_stack(context, torch, + torch.device("cpu")) + truth = reader(satellite, target) + if assembled is None or truth is None: + continue + stratum.append(score_target(*assembled, truth, decays, net=net)) + results[label] = stratum + mean = {k: float(np.mean([s[k] for s in stratum])) for k in stratum[0]} \ + if stratum else {} + print(f" gap {gap:>3} {label[1]:>8} n={len(stratum):>2} " + + " ".join(f"{k} {v:6.2f}" for k, v in mean.items()), + flush=True) + + print(f"\n{time.time() - started:.0f} s. Reference: optical_flow " + "44.79/44.36/40.72/36.85/34.47 dB at gaps 1/3/30/100/300; " + "fixed crosssat ~39.5-40.9 flat.") + columns = sorted({k for stratum in results.values() for s in stratum for k in s}) + print(f"\n{'gap':>4} {'counterpart':>11} {'n':>3} " + + " ".join(f"{c:>12}" for c in columns)) + for (gap, counterpart), stratum in results.items(): + if not stratum: + print(f"{gap:>4} {counterpart:>11} 0") + continue + row = " ".join( + f"{float(np.mean([s[c] for s in stratum])):>12.2f}" for c in columns + ) + print(f"{gap:>4} {counterpart:>11} {len(stratum):>3} {row}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/dataset_build.py b/dataset_build.py new file mode 100644 index 0000000..b779338 --- /dev/null +++ b/dataset_build.py @@ -0,0 +1,191 @@ +#!/usr/bin/env python +"""Build training shards from the archive and stream them to the training host. + + dataset_build.py --days 60 --out /tmp/shards --push htpc@192.168.1.66:/var/home/htpc/suvi + +This is a ~260,000-frame traversal of a filesystem that has repeatedly run out of file +handles under exactly this kind of load, on a VM with 184 GB free on the mount and 3.8 GB +on root against a dataset that measures 168 GB. Both facts shape the design: + +* every read goes through :class:`suvi.vfs.Reliever`, which hands handles back as it goes; +* a shard is written to `--out` (point it at tmpfs), pushed, and deleted before the next + one starts, so peak local usage is one day rather than the whole set. + +Days are chosen from the index, not the filesystem, and the split is written first. If +extraction dies half way the manifest still describes the intended experiment, and a +resumed run fills in the shards that are missing rather than choosing different days. +""" + +import argparse +import datetime +import json +import os +import subprocess +import sys +import time +from multiprocessing import Pool + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from suvi import dataset, db, fitsio, paths, vfs + +#: The bench's ground-truth window. Excluded from training so it stays a clean holdout. +HOLDOUT = ("2024-05-08", "2024-05-16") +#: GOES-18 joins the archive in 2022-08; before that there is no second satellite. +DEFAULT_FROM = "2022-09-01" +DEFAULT_TO = "2025-03-20" + + +def _day_bounds(day): + start = datetime.datetime.fromisoformat(day).replace(tzinfo=datetime.timezone.utc) + return int(start.timestamp()), int(start.timestamp()) + 86400 + + +def holdout_days(first=HOLDOUT[0], last=HOLDOUT[1]): + start = datetime.date.fromisoformat(first) + end = datetime.date.fromisoformat(last) + days = [] + while start <= end: + days.append(start.isoformat()) + start += datetime.timedelta(days=1) + return days + + +def _extract(job): + """Worker: one slot -> encoded block. Runs in a subprocess, so it opens its own files.""" + root, time_start, rows = job + fitsio.quiet_astropy() + try: + return time_start, dataset.read_slot(root, rows) + except (OSError, ValueError): + return time_start, None + + +def build_day(conn, root, day, satellite, out_dir, workers, reliever): + """Extract one (day, satellite) shard. Returns (path, bytes, slots_with_frames).""" + start, end = _day_bounds(day) + by_slot = {} + for row in conn.execute( + """ + SELECT wavelength, t_start, path FROM frame + WHERE satellite = ? AND t_start >= ? AND t_start < ? + """, + (satellite, start, end), + ): + by_slot.setdefault(row["t_start"], {})[row["wavelength"]] = row["path"] + + jobs = [(root, t, rows) for t, rows in sorted(by_slot.items())] + records = {} + with Pool(workers) as pool: + for time_start, block in pool.imap_unordered(_extract, jobs, chunksize=1): + records[time_start] = block + reliever.tick(len(paths.WAVELENGTHS)) + + path = os.path.join(out_dir, dataset.shard_name(day, satellite)) + size = dataset.write_shard(path, day, satellite, paths.WAVELENGTHS, records) + return path, size, sum(1 for block in records.values() if block is not None) + + +def push(path, destination, retries=3): + """Copy a shard to the training host and remove the local copy.""" + for attempt in range(retries): + result = subprocess.run( + ["rsync", "-q", "--partial", "--inplace", path, destination], + stderr=subprocess.PIPE, + ) + if result.returncode == 0: + os.remove(path) + return True + print(f" rsync failed (attempt {attempt + 1}): " + f"{result.stderr.decode(errors='replace').strip()[:200]}") + time.sleep(5) + return False + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--days", type=int, default=60) + parser.add_argument("--out", default="/tmp/shards", + help="staging directory; put this on tmpfs") + parser.add_argument("--push", default=None, + help="rsync destination, e.g. host:/path (omit to keep locally)") + parser.add_argument("--from", dest="from_day", default=DEFAULT_FROM) + parser.add_argument("--to", dest="to_day", default=DEFAULT_TO) + parser.add_argument("--satellites", default="16,18") + parser.add_argument("--workers", type=int, default=6) + parser.add_argument("--db", default=None) + parser.add_argument("--manifest", default=None, + help="reuse an existing manifest instead of choosing days afresh") + args = parser.parse_args(argv) + + os.makedirs(args.out, exist_ok=True) + satellites = tuple(int(s) for s in args.satellites.split(",")) + conn = db.connect(args.db, readonly=True) + root = paths.data_root() + + manifest_path = args.manifest or os.path.join(args.out, "manifest.json") + if args.manifest and os.path.exists(args.manifest): + manifest = json.load(open(args.manifest)) + splits = manifest["splits"] + print(f"Reusing manifest {args.manifest} (digest {manifest['digest']})") + else: + chosen = dataset.choose_days( + conn, args.days, + *[_day_bounds(d)[0] for d in (args.from_day, args.to_day)], + exclude=holdout_days(), satellites=satellites, + ) + if len(chosen) < args.days: + print(f"WARNING: only {len(chosen)} complete days available, wanted {args.days}") + splits = dataset.split_days(chosen) + manifest = dataset.manifest(splits, {"satellites": list(satellites), + "holdout": list(HOLDOUT)}) + with open(manifest_path, "w") as handle: + json.dump(manifest, handle, indent=2) + print(f"Wrote {manifest_path} (digest {manifest['digest']})") + + adjacent = dataset.check_split(splits) + for name, group in sorted(splits.items()): + print(f" {name:>5}: {len(group)} days {group[0] if group else '-'}" + f" .. {group[-1] if group else '-'}") + if adjacent: + print(f" NOTE: {len(adjacent)} split boundaries fall on consecutive days; " + "frames four minutes apart sit either side of them.") + + if args.push: + subprocess.run(["rsync", "-q", manifest_path, args.push], check=False) + + days = sorted(sum(splits.values(), [])) + total_bytes = 0 + total_slots = 0 + started = time.time() + reliever = vfs.Reliever(label="dataset") + try: + for index, day in enumerate(days, 1): + for satellite in satellites: + name = dataset.shard_name(day, satellite) + path, size, slots = build_day( + conn, root, day, satellite, args.out, args.workers, reliever + ) + total_bytes += size + total_slots += slots + elapsed = time.time() - started + done = (index - 1) * len(satellites) + satellites.index(satellite) + 1 + remaining = (len(days) * len(satellites) - done) * elapsed / max(done, 1) + print(f" [{done:>3}/{len(days) * len(satellites)}] {name}: " + f"{slots} slots, {size / 1e6:.0f} MB, " + f"{total_bytes / 1e9:.1f} GB total, " + f"eta {remaining / 3600:.1f} h", flush=True) + if args.push and not push(path, args.push): + raise SystemExit(f"could not push {name}; stopping rather than " + "filling the staging disk") + finally: + reliever.finish() + + print(f"\n{total_slots} slots, {total_bytes / 1e9:.1f} GB, " + f"{(time.time() - started) / 3600:.2f} h") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/gpu_train.py b/gpu_train.py new file mode 100644 index 0000000..cd2b69b --- /dev/null +++ b/gpu_train.py @@ -0,0 +1,159 @@ +#!/usr/bin/env python +"""Start a training run on the GPU host and hand the GPU back when it ends. + +Detached by design: a training run is a day long, and tying it to an SSH channel means +it dies with the connection -- or worse, finishes while the client stays blocked on a +pipe that never closes. Both happened while this was being built. The run becomes a +transient systemd unit on the host; this process only starts it, waits, and restores the +LLM services afterwards. + + gpu_train.py --epochs 48 --steps 1500 # start and watch + gpu_train.py --status # what is it doing now + gpu_train.py --stop # stop it and restore the LLMs +""" + +import argparse +import os +import sys +import time + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from suvi import gpubox + +REMOTE = "/var/home/htpc/suvi" +LOG = f"{REMOTE}/train.log" + + +def command(args): + return ( + "bash -c 'cd /work && python -u train.py " + f"--shards /work/shards --manifest /work/shards/manifest.json " + f"--out /work/runs/{args.name} " + f"--epochs {args.epochs} --batch {args.batch} --lr {args.lr} " + f"--base {args.base} --depth {args.depth} --workers {args.workers} " + f"--val-samples {args.val_samples} --eval-every {args.eval_every}" + + (f" --steps-per-epoch {args.steps}" if args.steps else "") + + (f" --max-steps {args.max_steps}" if args.max_steps else "") + + (f" --overfit {args.overfit}" if args.overfit else "") + + (f" --resume /work/runs/{args.name}/last.pt" if args.resume else "") + + "'" + ) + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--name", default="v4") + parser.add_argument("--epochs", type=int, default=20) + parser.add_argument("--steps", type=int, default=None, + help="steps per epoch (default: the whole train split)") + parser.add_argument("--val-samples", type=int, default=64) + parser.add_argument("--eval-every", type=int, default=250) + parser.add_argument("--max-steps", type=int, default=None, + help="stop after this many steps (the 300-step probe gate)") + parser.add_argument("--overfit", type=int, default=None, + help="run the overfit gate on this many frozen samples") + parser.add_argument("--batch", type=int, default=8) + parser.add_argument("--lr", type=float, default=1e-4) + parser.add_argument("--base", type=int, default=32) + parser.add_argument("--depth", type=int, default=3) + parser.add_argument("--workers", type=int, default=6) + parser.add_argument("--resume", action="store_true") + parser.add_argument("--status", action="store_true") + parser.add_argument("--since", type=int, default=None, + help="print history rows with step > SINCE, plus health") + parser.add_argument("--stop", action="store_true") + parser.add_argument("--watch", type=int, default=0, + help="seconds to wait for completion (0 = start and return)") + args = parser.parse_args(argv) + + if args.status: + print(gpubox.tail(LOG, 30)) + print("containers:", gpubox.run("podman ps --format '{{.Names}}'", check=False).strip()) + print(f"GTT {gpubox.gtt_used() / 1024**3:.1f} GiB used") + return 0 + + if args.since is not None: + return _report_since(args) + + if args.stop: + stopped = gpubox.stop_containers() + print(f"stopped {stopped or 'nothing'}") + gpubox.run("systemctl --user start llama-swap.service", check=False) + return 0 + + state = gpubox.stop_llms() + try: + name, log = gpubox.launch(command(args), mounts=[(REMOTE, "/work")], log=LOG) + print(f"launched {name}; log at {log}", flush=True) + if not args.watch: + print("Running detached. gpu_train.py --status to check, --stop to end.") + print("NOTE: the LLM services stay stopped until --stop or the run finishes.") + _paused_by_us(state) + return 0 + deadline = time.time() + args.watch + while time.time() < deadline and gpubox.running(name): + time.sleep(30) + print(gpubox.tail(log, 40)) + finally: + if args.watch: + gpubox.stop_containers() + print("restore failures:", gpubox.start_llms(state) or "none") + return 0 + + +def _report_since(args): + """New history rows past a step, plus the health a loss curve does not show. + + A run can go non-finite, stall with the process still alive, or lose its + container entirely; each of those wastes hours if nobody looks. Returns 1 when + something is wrong, so a supervising loop can react without parsing prose. + """ + import json + + raw = gpubox.run(f"cat {REMOTE}/runs/{args.name}/history.jsonl 2>/dev/null", + check=False) + rows = [json.loads(line) for line in raw.splitlines() if line.strip()] + fresh = [row for row in rows if row.get("step", 0) > args.since] + for row in fresh: + if row.get("untrained"): + print(f"step {row['step']:>6} val {row['val_loss']:.5f} " + f"({row['val_psnr']:.2f} dB) [untrained]") + else: + print(f"step {row['step']:>6} train {row['train_loss']:.5f} " + f"val {row['val_loss']:.5f} ({row['val_psnr']:.2f} dB) " + f"grad {row.get('grad_mean', 0):.3f}/{row.get('grad_max', 0):.2f}") + + problems = [] + for row in fresh: + for key in ("train_loss", "val_loss", "val_psnr"): + if key in row and not (row[key] == row[key] and abs(row[key]) < 1e30): + problems.append(f"step {row['step']}: {key} is not finite") + containers = gpubox.run("podman ps --format '{{.Names}}'", check=False).strip() + if "suvi" not in containers: + problems.append("no suvi container is running") + disk = gpubox.run(f"df --output=pcent {REMOTE} | tail -1", check=False).strip() + if disk.rstrip("%").strip().isdigit() and int(disk.rstrip("%").strip()) > 95: + problems.append(f"host disk at {disk}") + for problem in problems: + print(f"PROBLEM: {problem}") + return 1 if problems else 0 + + +def _paused_by_us(state): + """Record what was stopped, so a later --stop can put it back. + + Written to disk rather than held in memory because the run outlives this process by + a day, and a restore that only works while the launcher is alive is not a restore. + """ + import json + + path = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".gpu_paused.json") + with open(path, "w") as handle: + json.dump(state, handle) + print(f" paused state recorded in {path}") + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/reclaim.py b/reclaim.py index 434c306..53551a3 100644 --- a/reclaim.py +++ b/reclaim.py @@ -1,52 +1,145 @@ -#!/usr/bin/env python +#!/usr/bin/env python3 """Hand the archive's file handles back to the host. -The archive is on a virtiofs share whose daemon holds a host file descriptor per -inode the guest has looked up. Anything that reads a lot of the archive leaves those -inodes cached, and the handles with them, until something forces the guest to evict -them -- at which point the daemon can run out and the mount starts refusing *every* -open, which surfaces across the whole machine as "too many open files in system". +The archive is on a virtiofs share whose daemon holds a host file descriptor per inode +the guest has looked up. Anything that reads a lot of the archive leaves those inodes +cached, and the handles with them, until something forces the guest to evict them -- at +which point the daemon can run out and the mount starts refusing *every* open, which +surfaces across the whole machine as "too many open files in system". -The pipeline's own bulk jobs now reclaim as they go and again when they finish, so -this should rarely be needed. It exists for when something else has loaded the cache, -or to check where things stand: +The pipeline's own bulk jobs reclaim as they go and again when they finish, so this +should rarely be needed. It exists for when something else has loaded the cache: reclaim.py # report, then reclaim if needed reclaim.py --check # report only, change nothing sudo reclaim.py # uses drop_caches: instant, and far gentler -Without root the only lever available is memory pressure, which means briefly -allocating several GiB. Running under sudo avoids that entirely. +**This file deliberately imports nothing from `suvi`, and nothing outside the standard +library.** It has to run in exactly the situation it exists for, and in that situation +the archive mount refuses every open -- which means the venv interpreter (which lives on +the mount), the `suvi` package, and even Python's scan of the working directory all +fail with ENFILE before any of this code runs. Recovery went like this once: + + $ ./.venv/bin/python reclaim.py + bash: ./.venv/bin/python: Too many open files in system + $ /usr/bin/python3 -c ... + OSError: [Errno 23] Too many open files in system: '.../scripts' + +So: no package imports, and run it off the root filesystem with an isolated interpreter, +which keeps the mount out of `sys.path` entirely: + + /usr/bin/python3 -I ~/reclaim.py --force + +`--install` writes that copy for you, to a path that is not on the share. + +Without root the only lever is memory pressure, which means briefly allocating several +GiB. Running under sudo avoids that entirely. """ import argparse import os +import shutil import sys -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - -from suvi import vfs - -#: Roughly the slab size at which the daemon is likely near a 1M descriptor limit. -#: A cached dentry plus inode is on the order of a kilobyte, so a gigabyte of -#: reclaimable slab is on the order of a million inodes. +#: Roughly the slab size at which the daemon is likely near a 1M descriptor limit. A +#: cached dentry plus inode is on the order of a kilobyte, so a gigabyte of reclaimable +#: slab is on the order of a million inodes. CONCERN_GIB = 1.0 +#: Leave at least this much memory free while applying pressure. +RESERVE_GIB = 2.0 +#: Never allocate more than this in total. +MAX_GIB = 48 +#: Below this, the cache is already small enough that pressure achieves nothing. +FLOOR_KB = 1024 * 1024 + + +def meminfo(key): + try: + with open("/proc/meminfo") as handle: + for line in handle: + if line.startswith(key): + return int(line.split()[1]) + except OSError: + pass + return 0 + + +def reclaimable_kb(): + return meminfo("SReclaimable") + + +def available_gib(): + return meminfo("MemAvailable") / (1024 * 1024) + + +def drop_caches(): + """The direct path: ask the kernel to drop dentries and inodes. Needs root.""" + try: + with open("/proc/sys/vm/drop_caches", "w") as handle: + handle.write("2\n") + return True + except OSError: + return False + + +def release_handles(reserve_gib=RESERVE_GIB, max_gib=MAX_GIB, floor_kb=FLOOR_KB): + """Force the guest to evict cached inodes, so the daemon can close their handles.""" + if drop_caches(): + return True + before = reclaimable_kb() + if before <= 0: + return False + if before < floor_kb: + return True + target = before // 2 + blocks = [] + try: + for _ in range(int(max_gib)): + if reclaimable_kb() <= target: + return True + if available_gib() <= reserve_gib: + break + blocks.append(bytearray(1024**3)) + except MemoryError: + pass + finally: + blocks.clear() + return reclaimable_kb() <= target def report(prefix): - slab = vfs.reclaimable_kb() / (1024 * 1024) + slab = reclaimable_kb() / (1024 * 1024) print(f"{prefix:>8}: {slab:.2f} GiB reclaimable slab " - f"(~{slab:.1f}M cached inodes), {vfs.available_gib():.1f} GiB available") + f"(~{slab:.1f}M cached inodes), {available_gib():.1f} GiB available") return slab +def install(destination): + """Copy this script somewhere that is not on the share it rescues.""" + destination = os.path.abspath(os.path.expanduser(destination)) + if os.path.realpath(destination) == os.path.realpath(os.path.abspath(__file__)): + raise SystemExit("refusing to install over the original") + shutil.copyfile(os.path.abspath(__file__), destination) + os.chmod(destination, 0o755) + print(f"Installed to {destination}") + print(f" when the mount is refusing opens, run: /usr/bin/python3 -I {destination}") + return destination + + def main(argv=None): parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("--check", action="store_true", help="report only, reclaim nothing") parser.add_argument("--force", action="store_true", help="reclaim even if it looks fine") + parser.add_argument("--install", nargs="?", const="~/reclaim.py", default=None, + metavar="PATH", + help="copy this script off the share so it survives ENFILE") args = parser.parse_args(argv) + if args.install: + install(args.install) + return 0 + before = report("before") if args.check: print(" (check only; nothing reclaimed)") @@ -58,10 +151,9 @@ def main(argv=None): if os.geteuid() != 0: print(" not root, so using memory pressure; run under sudo for the direct path") - freed = vfs.release_handles() + freed = release_handles() after = report("after") - print(f" {'reclaimed' if freed else 'no change'}: " - f"{before - after:+.2f} GiB") + print(f" {'reclaimed' if freed else 'no change'}: {before - after:+.2f} GiB") return 0 diff --git a/suvi/align.py b/suvi/align.py new file mode 100644 index 0000000..1f0c772 --- /dev/null +++ b/suvi/align.py @@ -0,0 +1,242 @@ +"""Deterministic alignment of a frame stack onto its target instant. + +The learned filler's job is *selection*, not physics. Everything about this archive +that has a closed form is applied here, before any network sees a pixel: + +* **Rotation.** The Sun's differential rotation is known (Snodgrass 1983), so a frame + at dt != 0 is warped forward or back to the target instant -- the same field + :func:`suvi.fillers.solar_rotation` uses, ported to torch so one implementation + serves 640-pixel training shards and 1280-pixel bench frames, on CPU or GPU. +* **Photometry.** GOES-16 and GOES-18 SUVI differ by a band-dependent affine transfer + (gain 0.81-1.49, drifting ~12%/week). A cross-satellite frame is put on the target + instrument's scale using a (gain, offset) fitted from a simultaneous clean pair -- + the corrected-`crosssat` estimator, verified at +16.5 dB over the stale-bracket fit. + +After this stage every stack entry is *an estimate of the target frame*, and the two +views' geometry needs nothing more: the spacecraft sit 0.11 px of parallax apart on +identical WCS grids, so there is no disparity to solve. + +Shards do not store WCS headers, so the geometry (solar B0 angle and apparent disc +radius) comes from the analytic ephemeris below. Both are smooth annual functions of +the date; the accuracy required is loose -- a 0.1 deg B0 error moves a 20-hour warp by +under a tenth of a pixel. +""" + +import math + +import numpy as np + +from . import dataset + +#: Snodgrass (1983) sidereal differential rotation, degrees per day, by latitude. +#: Defined here rather than in :mod:`suvi.fillers` because this module must import +#: inside the ROCm training container, which has torch but not OpenCV; `fillers` +#: re-exports them so its callers see no change. +SNODGRASS_A = 14.713 +SNODGRASS_B = -2.396 +SNODGRASS_C = -1.787 +#: Earth's mean orbital motion, subtracted to get the rotation an Earth-orbiting +#: observer actually sees. +EARTH_ORBIT_DEG_PER_DAY = 0.9856 +SECONDS_PER_DAY = 86400.0 +#: Mean apparent solar radius at 1 AU, arcseconds (IAU 2015 nominal radius). +RADIUS_ARCSEC_1AU = 959.63 +#: SUVI L2 plate scale, arcsec/pixel. Constant across the archive: every frame is +#: reprojected onto the same grid (CDELT1 = CDELT2 = 2.5, CROTA = 0). +PLATE_SCALE = 2.5 +#: Native SUVI L2 frame width, pixels. Geometry is expressed as a *fraction* of the +#: frame so the same numbers serve 640-pixel shards and 1280-pixel archive frames. +NATIVE_SIZE = 1280 +#: Inclination of the solar equator to the ecliptic, degrees (Carrington). +SOLAR_INCLINATION = 7.25 +#: Unix time of the J2000.0 epoch. +J2000_UNIX = 946728000.0 + + +def solar_ephemeris(t_unix): + """(b0 radians, disc radius as a fraction of frame width) for a unix time. + + Low-precision solar position (Meeus, Astronomical Algorithms ch. 25) -- good to + ~0.1 deg in B0 and ~0.1% in distance, far inside what the rotation warp needs. + B0 is the heliographic latitude of the disc centre: the Earth rides 7.25 deg + above and below the solar equator over the year, and ignoring that tilts every + latitude the differential-rotation profile is evaluated at. + """ + n = (t_unix - J2000_UNIX) / SECONDS_PER_DAY + mean_longitude = math.radians((280.460 + 0.9856474 * n) % 360.0) + mean_anomaly = math.radians((357.528 + 0.9856003 * n) % 360.0) + ecliptic_longitude = mean_longitude + math.radians( + 1.915 * math.sin(mean_anomaly) + 0.020 * math.sin(2 * mean_anomaly) + ) + distance_au = 1.00014 - 0.01671 * math.cos(mean_anomaly) \ + - 0.00014 * math.cos(2 * mean_anomaly) + + # Ascending node of the solar equator on the ecliptic, precessing slowly. + node = math.radians(73.6667 + 1.395833 * (n / 36525.0 + 1.5)) + b0 = math.asin( + math.sin(ecliptic_longitude - node) * math.sin(math.radians(SOLAR_INCLINATION)) + ) + + radius_arcsec = RADIUS_ARCSEC_1AU / distance_au + radius_fraction = radius_arcsec / PLATE_SCALE / NATIVE_SIZE + return b0, radius_fraction + + +# ------------------------------------------------------------------ rotation warp + + +def rotation_grid(size, dt_seconds, b0, radius_fraction, synodic=True): + """Sampling grids that undo `dt_seconds` of differential rotation. + + Torch port of :func:`suvi.fillers._rotation_map`, batched: `dt_seconds`, `b0` and + `radius_fraction` are 1-D tensors of N frames, and the result is an (N, H, W, 2) + grid in the normalised align_corners=False convention `grid_sample` expects. + + Where the source point is off-disc or behind the limb the grid holds the pixel's + *own* centre, so sampling returns the unwarped value there -- the corona above the + limb does not co-rotate with the photosphere, matching `solar_rotation`'s + behaviour exactly. + """ + import torch + + height = width = int(size) + dt = dt_seconds.reshape(-1, 1, 1).to(torch.float32) + b0 = b0.reshape(-1, 1, 1).to(torch.float32) + radius = (radius_fraction.reshape(-1, 1, 1) * width).to(torch.float32) + device = dt.device + + centre_x = (width - 1) / 2.0 + centre_y = (height - 1) / 2.0 + grid_y, grid_x = torch.meshgrid( + torch.arange(height, device=device, dtype=torch.float32), + torch.arange(width, device=device, dtype=torch.float32), + indexing="ij", + ) + x = (grid_x - centre_x) / radius + y = (grid_y - centre_y) / radius + + rho2 = x**2 + y**2 + on_disc = rho2 < 1.0 + z = torch.sqrt((1.0 - rho2).clamp(min=0.0)) + + sin_b0, cos_b0 = torch.sin(b0), torch.cos(b0) + sin_lat = (y * cos_b0 + z * sin_b0).clamp(-1.0, 1.0) + latitude = torch.asin(sin_lat) + longitude = torch.atan2(x, z * cos_b0 - y * sin_b0) + + sin2 = sin_lat**2 + rate = SNODGRASS_A + SNODGRASS_B * sin2 + SNODGRASS_C * sin2**2 + if synodic: + rate = rate - EARTH_ORBIT_DEG_PER_DAY + source_longitude = longitude - torch.deg2rad(rate) * (dt / SECONDS_PER_DAY) + + cos_lat = torch.cos(latitude) + source_x = cos_lat * torch.sin(source_longitude) + source_y = sin_lat * cos_b0 - cos_lat * torch.cos(source_longitude) * sin_b0 + source_z = sin_lat * sin_b0 + cos_lat * torch.cos(source_longitude) * cos_b0 + visible = on_disc & (source_z > 0) + + map_x = torch.where(visible, source_x * radius + centre_x, grid_x) + map_y = torch.where(visible, source_y * radius + centre_y, grid_y) + # Pixel-centre normalisation: (2p + 1)/n - 1 is the align_corners=False + # convention; linspace(-1, 1) would shift everything by half a pixel and blur. + grid = torch.stack( + [(2 * map_x + 1) / width - 1, (2 * map_y + 1) / height - 1], dim=-1 + ) + return grid + + +def rotate(frames, dt_seconds, b0, radius_fraction, synodic=True): + """Warp (N, C, H, W) frames by their per-frame dt. Zero dt is the identity.""" + import torch.nn.functional as F + + grid = rotation_grid(frames.shape[-1], dt_seconds, b0, radius_fraction, synodic) + return F.grid_sample(frames, grid.to(frames.dtype), mode="bilinear", + padding_mode="border", align_corners=False) + + +# -------------------------------------------------------------------- photometry + + +def gain_fit(source, reference): + """Least-squares (gain, offset) putting `source` on `reference`'s scale.""" + x = np.nan_to_num(np.asarray(source, dtype=np.float64), nan=0.0, + posinf=0.0, neginf=0.0).ravel() + y = np.nan_to_num(np.asarray(reference, dtype=np.float64), nan=0.0, + posinf=0.0, neginf=0.0).ravel() + variance = float(((x - x.mean()) ** 2).sum()) + if variance <= 0: + return 1.0, 0.0 + gain = float(((x - x.mean()) * (y - y.mean())).sum() / variance) + offset = float(y.mean() - gain * x.mean()) + return gain, offset + + +def apply_photometry(coded, gain, offset): + """Apply a per-band radiance-space affine transfer to asinh-coded frames. + + `coded` is (N, BANDS, H, W) in the model's asinh space; `gain`/`offset` are + (N, BANDS) in radiance units. The transfer is defined on radiance -- a gain + multiplies a physical quantity -- so this computes + ``asinh(sinh(x * R) * g + o / S) / R`` rather than scaling the coded values, + which would model a different (and wrong) transform. + """ + import torch + + scaled = torch.sinh(coded.clamp(-1.0, 1.0) * dataset.ASINH_RANGE) + moved = scaled * gain[..., None, None] + offset[..., None, None] / dataset.ASINH_SCALE + return (torch.asinh(moved) / dataset.ASINH_RANGE).clamp(-1.0, 1.0) + + +def fit_photometry(counterpart, local): + """Per-band (gain, offset) putting `counterpart` on `local`'s radiance scale. + + Least squares on a simultaneous pair of (BANDS, H, W) radiance arrays -- the two + spacecraft observe the same Sun at the same instant, so the fit isolates the + instrument difference with no solar evolution mixed in. numpy, because it runs + in the CPU data path (sampler and bench), once per day rather than per frame. + """ + pairs = [gain_fit(c, l) for c, l in zip(counterpart, local)] + gains = np.array([g for g, _ in pairs], dtype=np.float32) + offsets = np.array([o for _, o in pairs], dtype=np.float32) + return gains, offsets + + +# ------------------------------------------------------------------- whole stacks + + +def align_stack(stack, dts, valid, gains, offsets, b0, radius_fraction, synodic=True): + """Align every frame of a batch of stacks onto its target instant. + + `stack` is (B, S, BANDS, H, W) in asinh space; `dts` (B, S) seconds; `valid` + (B, S) with 1 where a frame carries pixels; `gains`/`offsets` (B, S, BANDS) in + radiance units (identity rows for same-satellite frames); `b0`/`radius_fraction` + (B,) from :func:`solar_ephemeris` at the target instant. + + Photometry first (calibrate the instrument), then rotation (account for time). + Invalid frames are forced to the identity transfer -- an offset applied to a + frame of zeros would manufacture a constant image out of nothing. + """ + batch, stack_depth, bands = stack.shape[:3] + flat = stack.reshape(batch * stack_depth, bands, *stack.shape[-2:]) + keep = valid.reshape(-1, 1).to(flat.dtype) + gain = gains.reshape(-1, bands) * keep + (1.0 - keep) + offset = offsets.reshape(-1, bands) * keep + + needs_transfer = ((gain != 1.0) | (offset != 0.0)).any(dim=1) + if bool(needs_transfer.any()): + moved = apply_photometry(flat[needs_transfer], gain[needs_transfer], + offset[needs_transfer]) + flat = flat.clone() + flat[needs_transfer] = moved.to(flat.dtype) + + expand = lambda values: values.reshape(batch, 1).expand(batch, stack_depth).reshape(-1) + # `dts` is (frame time - target time); the warp must advance each frame by the + # *negation* of that, (target - frame time), to land on the target instant. + # Passing dts unnegated rotates every candidate AWAY from the target, doubling + # the displacement instead of cancelling it -- sub-pixel at short gaps, which is + # how it slipped past the identity and parity tests, and ~7 dB of candidate + # quality at a 300-slot gap, which is how it was caught: the aligned anchors + # scored far below plain solar_rotation on the same frames. + aligned = rotate(flat, -dts.reshape(-1), expand(b0), expand(radius_fraction), synodic) + return aligned.view_as(stack) diff --git a/suvi/corruptions.py b/suvi/corruptions.py index 4f4f132..5e71765 100644 --- a/suvi/corruptions.py +++ b/suvi/corruptions.py @@ -165,12 +165,18 @@ def _salt_pepper(image, rng, severity=DEFAULT_SEVERITY, donor=None): def _translate(image, rng, severity=DEFAULT_SEVERITY, donor=None): - """Mispointing: the solar disc sits off centre.""" - shift = int(200 * severity) + """Mispointing: the solar disc sits off centre. + + The shift is a fraction of the frame, not a fixed pixel count. 200 px was right for + the archive's native 1280, but training samples are half that, and a fixed count + both doubles the apparent severity and -- below 200 px -- indexes past the frame and + raises. Scaling by ``min(h, w) / 1280`` leaves 1280 frames byte-identical. + """ + h, w = image.shape + shift = int(200 * severity * min(h, w) / 1280) dx = int(rng.integers(-shift, shift + 1)) if shift else 0 dy = int(rng.integers(-shift, shift + 1)) if shift else 0 out = np.zeros_like(image) - h, w = image.shape xs, xd = (max(0, -dx), max(0, dx)) ys, yd = (max(0, -dy), max(0, dy)) height, width = h - abs(dy), w - abs(dx) diff --git a/suvi/dataset.py b/suvi/dataset.py new file mode 100644 index 0000000..0c2f58a --- /dev/null +++ b/suvi/dataset.py @@ -0,0 +1,422 @@ +"""Training shards: a compact, portable copy of the archive for the learned filler. + +The archive is 9.2 TB of tile-compressed FITS on a filesystem that cannot survive being +walked repeatedly, and the machine that trains the model is a different machine on the +other end of a gigabit link. So training does not read the archive. It reads *shards*: +one file per (day, satellite) holding that day's frames at half resolution, which is +small enough to ship and to re-read every epoch. + +Two decisions here are load-bearing. + +**Shards store clean frames only.** Corruption is applied on the fly during training +(:mod:`suvi.corruptions`), which costs nothing in storage, gives unlimited variety rather +than one frozen draw, and yields an exact ground-truth class label for the auxiliary +head. Baking damage into the shards would fix the training distribution at extraction +time and make every experiment downstream a re-extraction. + +**Pixels are stored as quantised ``asinh`` radiance.** Solar radiance is heavy-tailed -- +a sample across bands and epochs runs from -3.7 to 1231 -- so linear uint16 would spend +its whole range on the disc and quantise the corona to nothing. ``asinh(x/s)`` is +linear below the knee `s` and logarithmic above it, which is what a signal that is +mostly noise-floor with occasional flares needs. + +``log1p``, the space ``fillers._for_flow`` uses, was the obvious first choice and is +wrong here on two counts. It is effectively linear below 1, and almost every pixel is: +median radiance is 0.004 to 0.19 depending on band. Quantising log1p over [0, 8] gives +a 0.002 radiance pixel **6% error**, not the 0.05% the bright disc gets. And it cannot +represent negative values at all, while 2-18% of pixels are negative after background +subtraction -- clipping those to zero would bias the noise floor the model has to learn +to reproduce. ``asinh`` holds ~0.05% relative across the whole range and is signed. +""" + +import hashlib +import json +import os +import struct + +import numpy as np + +from . import fitsio, paths + +#: Half of the archive's native 1280. Flow, gain and blend fields are smooth, so the +#: model predicts them here and they upsample cleanly to drive full-resolution frames. +SHARD_SIZE = 640 +#: Knee of the asinh transform: linear below, logarithmic above. Relative precision is +#: constant above the knee and decays below it, so the knee sits well under the faintest +#: band's median radiance (0.004 at 131A) -- far enough that the whole corona is in the +#: constant-precision regime, not just the disc. +ASINH_SCALE = 1e-4 +#: Half-range of the stored asinh values, symmetric about zero. +/-18 covers +#: +/-sinh(18)*scale = +/-3284 in radiance, against an observed span of -3.7 to 1231 -- +#: room for a flare well past anything in the archive. Anything beyond that is a +#: corruption, not an observation, and saturates rather than wrapping. +ASINH_RANGE = 18.0 +#: Marker written for a slot whose frame would not read. Distinguishes "we looked and +#: there was nothing usable" from "we never looked", which a plain absence cannot. +UNREADABLE = b"\x00" + +MAGIC = b"SUVISHRD" +VERSION = 1 +SLOTS_PER_DAY = 86400 // paths.CADENCE # 360 + + +def _zstd(): + """zstd compress/decompress, from the stdlib on 3.14+ or the pip package below it. + + The extraction host runs Python 3.14, where zstd is in the standard library; the + training container ships whatever the ROCm image was built against. Both produce + ordinary zstd frames, so shards written by one are read by the other. + """ + try: + from compression import zstd + + return zstd.compress, zstd.decompress + except ImportError: # pragma: no cover - exercised only on the training host + import zstandard + + return ( + lambda data, level=3: zstandard.ZstdCompressor(level=level).compress(data), + lambda data: zstandard.ZstdDecompressor().decompress(data), + ) + + +# ------------------------------------------------------------------- pixel encoding + + +def encode_frames(arrays): + """Six 1280x1280 radiance arrays -> one 6 x 640 x 640 uint16 block. + + Downsampling is a 2x2 mean rather than decimation: SUVI frames carry read noise and + cosmic-ray hits, and taking every other pixel would keep the hits at full amplitude + while throwing away the averaging that suppresses them. + """ + planes = [] + for array in arrays: + image = np.nan_to_num(np.asarray(array, dtype=np.float64), nan=0.0, + posinf=0.0, neginf=0.0) + if image.shape[0] % SHARD_SIZE or image.shape[1] % SHARD_SIZE: + raise ValueError(f"cannot halve {image.shape} to {SHARD_SIZE}") + factor = image.shape[0] // SHARD_SIZE + small = image.reshape(SHARD_SIZE, factor, SHARD_SIZE, factor).mean(axis=(1, 3)) + unit = np.arcsinh(small / ASINH_SCALE) / ASINH_RANGE # -> roughly [-1, 1] + planes.append(np.clip(unit, -1.0, 1.0) * 32767.5 + 32767.5) + return np.rint(np.stack(planes)).astype(np.uint16) + + +def decode_frames(block): + """Inverse of :func:`encode_frames`, back to radiance.""" + return (np.sinh(coded_frames(block).astype(np.float64) * ASINH_RANGE) + * ASINH_SCALE).astype(np.float32) + + +def coded_frames(block): + """The stored block as the roughly-[-1, 1] values the model consumes. + + The uint16 in a shard *is* the model's input space -- ``encode_frames`` already + applied the asinh -- so getting from one to the other is an affine rescale and + nothing more. Going via radiance instead costs a ``sinh`` and an ``arcsinh`` over + 47 million elements per training sample, which measured as a data loader pinning two + cores at 200% while the GPU sat at 0% busy. Only frames that are about to be + corrupted need real radiance, because that is the space corruptions are defined in. + """ + return ((np.asarray(block, dtype=np.float32) - 32767.5) / 32767.5).astype(np.float32) + + +# ------------------------------------------------------------------- shard container +# +# Layout: MAGIC | version | header length | JSON header | records back to back. +# The JSON header carries an (offset, length) per slot, so one record can be read +# without decompressing the rest -- which is what makes random sampling across a +# 60-day dataset cheap. + +_PREFIX = struct.Struct("<8sHI") + + +def write_shard(path, day, satellite, wavelengths, records, level=3): + """Write one (day, satellite) shard. `records` maps t_start -> 6x640x640 uint16.""" + compress, _ = _zstd() + index = {} + blobs = [] + offset = 0 + for t_start in sorted(records): + block = records[t_start] + payload = UNREADABLE if block is None else compress( + np.ascontiguousarray(block, dtype=" archive-relative path. Returns the encoded block, or + None if any band is missing or unreadable -- a partial slot is not a training + sample, and admitting one would teach the model that bands go missing + independently when in practice a satellite drops out whole. + """ + arrays = [] + for wavelength in paths.WAVELENGTHS: + relpath = rows.get(wavelength) + if relpath is None: + return None + image, _ = fitsio.read_image(paths.abspath(relpath, root)) + if image is None or image.shape != (1280, 1280): + return None + arrays.append(image) + return encode_frames(arrays) + + +# ----------------------------------------------------------------------- day choice + + +#: Smallest plausible real frame. A healthy tile-compressed SUVI frame is 1.4-1.8 MB; +#: a failed download leaves a 5,760-byte header-only stub that is indexed like any other +#: frame and reads back as None. There are 125,732 of them in the archive -- 2.3% -- and +#: one chosen day turned out to be 100% stubs, producing an empty shard after a full +#: extraction pass over it. Counting rows is not the same as counting frames. +MIN_FRAME_BYTES = 100_000 + + +def choose_days(conn, count, t_from, t_to, exclude=(), satellites=(16, 18), + wavelengths=paths.WAVELENGTHS, min_bytes=MIN_FRAME_BYTES): + """Days where both satellites have near-complete coverage, spread evenly. + + Spread rather than sampled at random: the point of the training set is to span the + solar cycle, and independent draws over a 2.5-year range clump. Complete rather + than best-effort: a day missing half its slots would silently bias the gap-length + distribution the model trains against. And *real* rather than merely present -- see + :data:`MIN_FRAME_BYTES`. + """ + wanted = len(wavelengths) * len(satellites) * SLOTS_PER_DAY + placeholders = ",".join("?" * len(wavelengths)) + satellite_places = ",".join("?" * len(satellites)) + rows = conn.execute( + f""" + SELECT strftime('%Y-%m-%d', t_start, 'unixepoch') AS day, COUNT(*) AS n + FROM frame + WHERE t_start >= ? AND t_start < ? + AND wavelength IN ({placeholders}) AND satellite IN ({satellite_places}) + AND size_bytes >= ? + GROUP BY day HAVING n >= ? + ORDER BY day + """, + (t_from, t_to, *wavelengths, *satellites, int(min_bytes), int(wanted * 0.98)), + ).fetchall() + + excluded = set(exclude) + days = [row[0] for row in rows if row[0] not in excluded] + if not days: + return [] + if len(days) <= count: + return days + step = len(days) / count + return [days[int(i * step)] for i in range(count)] + + +def block_size(days, when_consecutive=5): + """How many days to keep together when splitting. + + Splitting per *sample* would leak almost perfectly -- consecutive frames are four + minutes apart -- but splitting per *day* only leaks if two days in the set are + themselves adjacent, and then only across the midnight seam. When the chosen days + are spread (60 days over two and a half years puts them ~15 days apart) each day is + an independent view of the Sun and blocks buy nothing while costing resolution in + the split. So: group only when there is something to guard against. + """ + import datetime + + ordered = sorted(days) + adjacent = any( + (datetime.date.fromisoformat(second) - datetime.date.fromisoformat(first)).days <= 1 + for first, second in zip(ordered, ordered[1:]) + ) + if not adjacent: + return 1 + # Never group so coarsely that three splits become impossible. With very few days + # the guard cannot be had at all; `check_split` reports the seams that remain, which + # is more useful than refusing to split. + return max(1, min(when_consecutive, len(ordered) // 3)) + + +def split_days(days, block=None, train=8, val=1, test=1): + """Partition days into disjoint train/val/test sets. + + Assignment is round-robin over blocks so that val and test are *interleaved* through + the date range rather than carved off one end. That matters more than it sounds: + this archive spans the rise of solar cycle 25, so a validation set drawn from one + stretch measures the model on one level of activity and says nothing about the rest. + + `block` defaults to :func:`block_size`, which uses whole-day granularity unless the + chosen days are actually adjacent. Fixing it at 5 with only 12 blocks to hand out + gave val and test one block each -- five consecutive chosen days, four months of + calendar -- which is exactly the failure this is meant to avoid. + """ + days = sorted(days) + block = block_size(days) if block is None else block + blocks = [days[i : i + block] for i in range(0, len(days), block)] + if len(blocks) < 3: + raise ValueError( + f"{len(days)} days in blocks of {block} gives {len(blocks)} blocks; " + "need at least three to make three splits" + ) + + # Take proportions, then place the held-out blocks at evenly spaced positions. + # Walking a repeating ['train'...,'val','test'] cycle instead looks equivalent and + # is not: with 8 blocks and a cycle of 10 it never reaches 'val' or 'test' at all, + # and returns empty held-out sets without complaining. The failure then surfaces + # much later, as a training run that cannot find any validation shards. + total = train + val + test + count = len(blocks) + wanted = {"val": max(1, round(count * val / total)), + "test": max(1, round(count * test / total))} + if wanted["val"] + wanted["test"] >= count: + raise ValueError(f"{count} blocks cannot yield train, val and test") + + held = wanted["val"] + wanted["test"] + stride = count / held + assignment = {} + for position in range(held): + index = min(count - 1, int(position * stride + stride / 2)) + while index in assignment: # collisions when stride is near 1 + index = (index + 1) % count + assignment[index] = "val" if position % 2 == 0 else "test" + + out = {"train": [], "val": [], "test": []} + for number, group in enumerate(blocks): + out[assignment.get(number, "train")].extend(group) + return out + + +def manifest(days_by_split, extra=None): + """A split manifest plus its digest, so a checkpoint can name the data it saw.""" + payload = { + "version": VERSION, + "size": SHARD_SIZE, + "asinh_scale": ASINH_SCALE, + "asinh_range": ASINH_RANGE, + "splits": {name: sorted(days) for name, days in days_by_split.items()}, + } + if extra: + payload.update(extra) + body = json.dumps(payload, sort_keys=True, separators=(",", ":")) + payload["digest"] = hashlib.sha256(body.encode()).hexdigest()[:16] + return payload + + +def check_split(days_by_split): + """Assert the splits are disjoint and no two blocks abut in time. + + Cheap, and it fails loudly. A leak here would not crash anything -- it would just + produce validation numbers that look excellent and mean nothing. + """ + import datetime + + seen = {} + for name, days in days_by_split.items(): + for day in days: + if day in seen: + raise ValueError(f"day {day} is in both {seen[day]!r} and {name!r}") + seen[day] = name + + ordered = sorted(seen) + adjacent = [] + for first, second in zip(ordered, ordered[1:]): + a = datetime.date.fromisoformat(first) + b = datetime.date.fromisoformat(second) + if (b - a).days == 1 and seen[first] != seen[second]: + adjacent.append((first, second)) + return adjacent diff --git a/suvi/fillers.py b/suvi/fillers.py index 59dd042..977d359 100644 --- a/suvi/fillers.py +++ b/suvi/fillers.py @@ -15,16 +15,13 @@ from dataclasses import dataclass, field import cv2 as cv import numpy as np +# Shared with the torch-side alignment; defined there because suvi.align must import +# inside the ROCm container, which has no OpenCV. +from .align import (EARTH_ORBIT_DEG_PER_DAY, SECONDS_PER_DAY, SNODGRASS_A, + SNODGRASS_B, SNODGRASS_C, gain_fit) + #: Nominal solar radius in metres (IAU 2015). R_SUN = 6.957e8 -#: Snodgrass (1983) sidereal differential rotation, degrees per day, by latitude. -SNODGRASS_A = 14.713 -SNODGRASS_B = -2.396 -SNODGRASS_C = -1.787 -#: Earth's mean orbital motion, subtracted to get the rotation an Earth-orbiting -#: observer actually sees. -EARTH_ORBIT_DEG_PER_DAY = 0.9856 -SECONDS_PER_DAY = 86400.0 @dataclass @@ -39,8 +36,24 @@ class FillContext: dt_after: float = 0.0 #: The other satellite's view of this same instant, if it has one. counterpart: np.ndarray | None = None + #: A (counterpart, this satellite) pair from the nearest slot where *both* were + #: good, used to calibrate one instrument against the other. The two spacecraft + #: observe simultaneously, so such a pair isolates the instrument difference with + #: no solar evolution mixed in -- which a bracketing frame from this satellite + #: alone cannot do. See :func:`crosssat`. + calibration: tuple | None = None #: Header of the frame being reconstructed, for the WCS a rotation warp needs. header: dict = field(default_factory=dict) + #: The full stack, for fillers that fuse more than two frames. One entry per input + #: frame: ``{"image", "state", "dt", "same_satellite", "scores", "verdict"}``, where + #: `state` is 'available', 'missing' or 'suspect'. + #: + #: The hand-written fillers ignore this and read only the fields above, which is why + #: adding it changes none of them. A **suspect** entry is the reason it exists: + #: every method in this module discards a flagged frame outright, but 14 of the 20 + #: modes in :mod:`suvi.corruptions` leave one substantially usable, and the learned + #: filler is built to exploit exactly that. + stack: list = field(default_factory=list) @property def alpha(self): @@ -176,53 +189,83 @@ def gain_match(source, reference): GOES-16 and GOES-18 carry different SUVI flight models, so their radiances differ by a roughly affine factor even when both are healthy. """ - x = _finite(source).ravel().astype(np.float64) - y = _finite(reference).ravel().astype(np.float64) - variance = float(((x - x.mean()) ** 2).sum()) - if variance <= 0: - return np.asarray(source, dtype=np.float32) - gain = float(((x - x.mean()) * (y - y.mean())).sum() / variance) - offset = float(y.mean() - gain * x.mean()) + gain, offset = gain_fit(source, reference) return (np.asarray(source, dtype=np.float32) * gain + offset).astype(np.float32) +def _shift(image, dx, dy): + matrix = np.array([[1.0, 0.0, dx], [0.0, 1.0, dy]], dtype=np.float32) + return cv.warpAffine( + np.asarray(image, dtype=np.float32), + matrix, + (image.shape[1], image.shape[0]), + flags=cv.INTER_LINEAR, + borderMode=cv.BORDER_REPLICATE, + ) + + +def _register(source, reference): + """Sub-pixel translation carrying `source` onto `reference`.""" + window = cv.createHanningWindow((source.shape[1], source.shape[0]), cv.CV_64F) + (dx, dy), _ = cv.phaseCorrelate( + source.astype(np.float64), reference.astype(np.float64), window + ) + return dx, dy + + def crosssat(context, align=True): """Substitute the other satellite's view of the same instant. - The two spacecraft see the same Sun from 1 AU, so the substitute is a real - observation of the real Sun at the right time -- not an interpolation. It - should dominate every temporal method whenever it is available, which is the - thing worth quantifying: it is unavailable in the 31% of slots where both - satellites are out simultaneously. + The two spacecraft see the same Sun at the same moment, so the substitute is a + real observation rather than an interpolation, and its quality does not decay + with gap length the way every temporal method does. Measured against truth at + 195A it scores 46.45 dB whether the gap is one slot or three hundred. - Residual differences are instrument calibration (removed by gain matching) and - a few pixels of geostationary parallax (removed by alignment). + What does decay is the *calibration*. The instruments differ by a band-dependent + gain -- 0.81 at 195A, up to 1.49 at 94A, drifting 12% within a single week -- so + the substitute has to be put on this satellite's scale before it is usable, and + that gain has to be estimated from somewhere. + + Estimating it from ``context.before`` is what this used to do, and it is wrong: + at a 300-slot gap that frame is twenty hours old, so the fit absorbs the Sun's + own evolution into what is supposed to be an instrument constant. The result + collapsed from 46.45 dB at gap 1 to 18.75 dB at gap 300 -- which is why the bench + reported this as the worst filler at every length, and why that finding was an + artifact of the estimator rather than a property of the method. + + `context.calibration` instead supplies a *simultaneous* pair from the nearest + slot where both satellites were good. Because the two frames in that pair are of + the same Sun at the same instant, their ratio is the instrument difference and + nothing else, however far away the pair sits in time. Alignment is measured on + the same pair for the same reason. """ if context.counterpart is None: return None counterpart = _finite(context.counterpart) + + if context.calibration is not None: + pair_counterpart, pair_local = context.calibration + if pair_counterpart is not None and pair_local is not None: + pair_counterpart = _finite(pair_counterpart) + pair_local = _finite(pair_local) + if pair_counterpart.shape == pair_local.shape == counterpart.shape: + if align: + dx, dy = _register(pair_counterpart, pair_local) + counterpart = _shift(counterpart, dx, dy) + pair_counterpart = _shift(pair_counterpart, dx, dy) + gain, offset = gain_fit(pair_counterpart, pair_local) + return (counterpart * gain + offset).astype(np.float32) + + # No simultaneous pair anywhere in the series. Fall back to the bracketing + # frame, which is sound at short gaps and degrades as the bracket recedes. reference = context.before if context.before is not None else context.after if reference is None: return counterpart reference = _finite(reference) if counterpart.shape != reference.shape: return counterpart - if align: - window = cv.createHanningWindow( - (counterpart.shape[1], counterpart.shape[0]), cv.CV_64F - ) - (dx, dy), _ = cv.phaseCorrelate( - counterpart.astype(np.float64), reference.astype(np.float64), window - ) - matrix = np.array([[1.0, 0.0, dx], [0.0, 1.0, dy]], dtype=np.float32) - counterpart = cv.warpAffine( - counterpart, - matrix, - (counterpart.shape[1], counterpart.shape[0]), - flags=cv.INTER_LINEAR, - borderMode=cv.BORDER_REPLICATE, - ) + counterpart = _shift(counterpart, *_register(counterpart, reference)) return gain_match(counterpart, reference) @@ -334,16 +377,149 @@ def _warp(image, header, delta_seconds, synodic): return warped, visible +# ------------------------------------------------------------------ learned filler + + +#: Loaded checkpoint, kept between calls. The bench fills thousands of slots one at a +#: time, and reloading 14M parameters per slot would dominate the run. +_LEARNED = {} +#: Where to find the checkpoint, overridable so a bench run can name a specific one. +LEARNED_CHECKPOINT_ENV = "SUVI_MODEL" + + +def load_learned(path=None, device=None): + """Load the trained filler, once. Returns (net, torch, device) or None.""" + import os + + path = path or os.environ.get(LEARNED_CHECKPOINT_ENV) + if not path: + return None + key = (path, device) + if key in _LEARNED: + return _LEARNED[key] + + import torch + + from . import model as model_module + + state = torch.load(path, map_location="cpu", weights_only=False) + settings = state.get("args", {}) + net = model_module.build(base=settings.get("base", 32), depth=settings.get("depth", 3)) + try: + net.load_state_dict(state["model"]) + except RuntimeError as error: + # Say which checkpoint and what changed. Torch's own message names tensor + # shapes and nothing else, which is unhelpful when several runs are on disk and + # only some predate an architecture change. + raise SystemExit( + f"{path} does not match the current model.\n{error}\n" + "This checkpoint was trained against a different architecture; retrain or " + "point SUVI_MODEL at a newer run." + ) from error + resolved = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu")) + net.to(resolved).eval() + _LEARNED[key] = (net, torch, resolved) + return _LEARNED[key] + + +def assemble_stack(context, torch, device): + """Turn a FillContext's stack into the model's aligned inputs. + + Shared by the trained filler and by diagnostics (ceiling.py), so there is + exactly one implementation of the entry-to-tensor path. Returns + ``(aligned, condition)`` -- the (1, S, 6, H, W) aligned candidates and their + (1, S, COND_DIM) conditioning -- or None when nothing in the stack carries + pixels. + """ + from . import align as align_module + from . import model as model_module + from . import samples + + entries = [entry for entry in context.stack if entry.get("image") is not None + or entry.get("state") == "missing"] + shape = next((entry["image"].shape for entry in entries + if entry.get("image") is not None), None) + if shape is None: + return None + + transfer = None + if context.calibration is not None: + pair_counterpart, pair_local = context.calibration + if pair_counterpart is not None and pair_local is not None: + transfer = align_module.fit_photometry(pair_counterpart, pair_local) + + frames, conditions, dts, gains, offsets = [], [], [], [], [] + for entry in entries: + image = entry.get("image") + state = entry.get("state", "available") if image is not None else "missing" + same = bool(entry.get("same_satellite", True)) + frames.append(np.zeros((6, *shape[-2:]), dtype=np.float32) if image is None + else samples.encode_for_model(image)) + conditions.append(model_module.frame_conditioning( + state, same, float(entry.get("dt", 0.0)) + )) + dts.append(float(entry.get("dt", 0.0))) + cross = transfer is not None and not same and image is not None + gains.append(transfer[0] if cross else np.ones(6, dtype=np.float32)) + offsets.append(transfer[1] if cross else np.zeros(6, dtype=np.float32)) + + first = entries[0] + target_time = float(first["slot"][1]) - float(first.get("dt", 0.0)) + b0, radius = align_module.solar_ephemeris(target_time) + + with torch.no_grad(): + stack = torch.from_numpy(np.stack(frames))[None].to(device) + condition = torch.stack(conditions)[None].to(device) + valid = (condition[..., 0] + condition[..., 2]).clamp(0, 1) + if float(valid.sum()) == 0: + return None + aligned = align_module.align_stack( + stack, torch.tensor(dts, device=device)[None], valid, + torch.from_numpy(np.stack(gains))[None].to(device), + torch.from_numpy(np.stack(offsets))[None].to(device), + torch.tensor([b0], device=device), + torch.tensor([radius], device=device), + ) + return aligned, condition + + +def learned(context, path=None, device=None): + """Reconstruct by fusing the whole stack, using the trained model. + + This is the only filler that sees more than three frames, and the only one that + looks at a frame the detector flagged rather than discarding it. The stack is + first aligned deterministically -- rotation-warped to the target instant, and + cross-satellite frames put on this instrument's scale using the simultaneous + pair in ``context.calibration`` (all six bands) -- and the model then chooses, + per pixel and band, which aligned observation to trust. + + Returns None when there is no checkpoint, no stack, or nothing in the stack + carries pixels; the bench reports that as "not applicable" rather than scoring + a fabricated frame. + """ + if not context.stack: + return None + loaded = load_learned(path, device) + if loaded is None: + return None + net, torch, device = loaded + + from . import samples + + assembled = assemble_stack(context, torch, device) + if assembled is None: + return None + aligned, condition = assembled + with torch.no_grad(): + prediction = net(aligned, condition) + return samples.decode_from_model(prediction[0].float().cpu().numpy()) + + FILLERS = { "hold_last": hold_last, "linear_blend": linear_blend, "optical_flow": optical_flow, "crosssat": crosssat, "solar_rotation": solar_rotation, + "learned": learned, } - -# TODO: learned filler. Train a model to predict a frame from its preceding frames, -# following frames, and the other satellite's view, then evaluate it here across -# severities of missing data and prediction horizons (single-frame gaps through -# multi-hour outages, one satellite out versus both). It plugs in as another entry -# in FILLERS and reuses the bench's existing cases and metrics unchanged. diff --git a/suvi/gpubox.py b/suvi/gpubox.py new file mode 100644 index 0000000..2f90595 --- /dev/null +++ b/suvi/gpubox.py @@ -0,0 +1,367 @@ +"""Driving the training host over SSH, including getting the GPU back from the LLMs. + +Training runs on a separate machine -- a Ryzen AI MAX+ 395 with a Radeon 8060S, whose +124 GB of memory is *unified*: the iGPU addresses it through GTT rather than owning +dedicated VRAM. That is why the box can train on frames this large at all, and also why +it can hold nothing else at the time: llama.cpp keeps 115.6 GB of the 124 GB resident, so +a training run does not get a smaller allocation, it gets ``Memory in use`` and dies. + +So something has to give up the GPU for the duration. These services belong to the user, +not to us, which sets the bar for this module: stop as little as will do, put back +whatever it stopped -- including when the training run crashes -- and be able to say +afterwards whether that succeeded. + +Stopping as little as will do is worth the extra code. Measured on this host, almost +all of that 115.6 GB is a 120B model that ``llama-swap`` had loaded on demand; with the +unit stopped, GTT falls to 24.8 GB and roughly 99 GB is free -- ample for training, +without touching anything else. So :func:`stop_llms` works in stages and escalates only +if the first stage leaves too little. + +The two kinds of server are not alike: + +* ``llama-swap.service``, a systemd --user unit. Stopping and starting it is exact, and + it is where the large on-demand models live. Almost always sufficient on its own. +* a **standalone** ``llama-server``, launched by hand and reparented to init. There is + no unit to restart, so the only way back is to record its argv and working directory + before killing it and re-exec them. This is the fragile one, only touched when the + first stage did not free enough, and the reason :func:`llms_paused` verifies the + restore rather than assuming it. +""" + +import json +import os +import shlex +import subprocess +import time +from contextlib import contextmanager + +HOST = os.environ.get("SUVI_GPU_HOST", "htpc@192.168.1.66") +#: Read from a file rather than embedded, so the credential is not in the source tree. +PASSWORD_FILE = os.environ.get( + "SUVI_GPU_PASSWORD_FILE", + os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "gpu_system_ssh_info"), +) +#: systemd --user units that hold the GPU. +UNITS = ("llama-swap.service",) +#: Seconds to wait for GTT to drain / refill before giving up. +SETTLE_TIMEOUT = 120 +#: Unified memory the iGPU can address, from mem_info_gtt_total on this host. +GTT_TOTAL = 115 * 1024**3 +#: Free unified memory a training run needs before it will start. +GTT_FREE_BYTES = 48 * 1024**3 +GTT_USED = "/sys/class/drm/card1/device/mem_info_gtt_used" + + +def _password(): + with open(PASSWORD_FILE) as handle: + lines = [line.strip() for line in handle if line.strip()] + if len(lines) < 2: + raise RuntimeError(f"{PASSWORD_FILE}: expected an ssh line and a password line") + return lines[1] + + +def run(command, check=True, timeout=600, capture=True): + """Run a shell command on the training host.""" + argv = [ + "sshpass", "-p", _password(), "ssh", + "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=10", + HOST, command, + ] + result = subprocess.run( + argv, timeout=timeout, + stdout=subprocess.PIPE if capture else None, + stderr=subprocess.STDOUT if capture else None, + ) + output = result.stdout.decode(errors="replace") if capture else "" + if check and result.returncode != 0: + raise RuntimeError(f"remote command failed ({result.returncode}): {command}\n{output}") + return output + + +def gtt_used(): + """Bytes of unified memory currently mapped to the GPU.""" + try: + return int(run(f"cat {GTT_USED}", check=False).strip() or 0) + except ValueError: + return 0 + + +# --------------------------------------------------------------- pausing the LLMs + + +def _standalone_servers(): + """Every llama-server not owned by one of our units, with enough to restart it. + + A server spawned *by* llama-swap comes back when llama-swap does, so recording it + would restart it twice. Only the ones reparented to init are ours to restore. + """ + script = r""" +for p in $(pgrep -x llama-server); do + ppid=$(ps -o ppid= -p $p 2>/dev/null | tr -d ' ') + [ "$ppid" = "1" ] || continue + cwd=$(readlink /proc/$p/cwd) + argv=$(tr '\0' '\n' < /proc/$p/cmdline | sed 's/"/\\"/g' | awk '{printf "\"%s\",", $0}') + echo "{\"pid\": $p, \"cwd\": \"$cwd\", \"argv\": [${argv%,}]}" +done +""" + found = [] + for line in run(script, check=False).splitlines(): + line = line.strip() + if line.startswith("{"): + try: + found.append(json.loads(line)) + except json.JSONDecodeError: + continue + return found + + +def _wait_for_gtt(below, timeout=SETTLE_TIMEOUT): + deadline = time.time() + timeout + while time.time() < deadline: + if gtt_used() < below: + return True + time.sleep(2) + return False + + +def stop_llms(needed_bytes=None): + """Free the GPU, stopping as little as will do. + + Stage one stops the units, which is where the big on-demand models live and which + restarts exactly. Only if that leaves less than `needed_bytes` free does stage two + kill hand-launched servers, whose restoration is a re-exec rather than a restart. + + Returns the state :func:`start_llms` needs to undo this. + """ + needed = GTT_TOTAL - (needed_bytes or GTT_FREE_BYTES) + state = {"units": [], "standalone": []} + + for unit in UNITS: + active = run(f"systemctl --user is-active {unit}", check=False).strip() + state["units"].append({"unit": unit, "was_active": active == "active"}) + if active == "active": + run(f"systemctl --user stop {unit}", check=False) + if _wait_for_gtt(needed, timeout=30): + return state + + # Stage two: the units alone were not enough. + state["standalone"] = _standalone_servers() + for server in state["standalone"]: + run(f"kill {int(server['pid'])}", check=False) + if not _wait_for_gtt(needed): + for server in state["standalone"]: + run(f"kill -9 {int(server['pid'])}", check=False) + _wait_for_gtt(needed) + return state + + +def start_llms(state): + """Put back exactly what :func:`stop_llms` took away. + + Reports what it could not restore rather than raising: this runs in a `finally`, + and masking the training error with a restore error would lose the more useful of + the two. + """ + failures = [] + for index, server in enumerate(state.get("standalone", [])): + argv = " ".join(shlex.quote(arg) for arg in server["argv"]) + # systemd-run rather than nohup: a backgrounded process still holds the SSH + # channel open, so ssh blocks until it exits -- which for a server is never. + # A transient unit forks away cleanly and returns at once, and --collect means + # it leaves nothing behind when it stops. + started = run( + f"systemd-run --user --collect --unit=suvi-restored-llama-{index} " + f"--property=WorkingDirectory={shlex.quote(server['cwd'])} {argv}", + check=False, timeout=60, + ) + if "Running as unit" not in started and "Failed" in started: + failures.append(f"standalone {os.path.basename(server['argv'][0])}") + + for entry in state.get("units", []): + if not entry["was_active"]: + continue + run(f"systemctl --user start {entry['unit']}", check=False) + if run(f"systemctl --user is-active {entry['unit']}", check=False).strip() != "active": + failures.append(entry["unit"]) + return failures + + +#: States stopped but not yet restored. Module level so the exit hooks can reach them. +_paused = [] +_hooks_installed = False + + +def _restore_all(verbose=True): + """Put back every outstanding pause. Safe to call more than once.""" + while _paused: + state = _paused.pop() + # Our own containers first: restarting the LLMs while a training container still + # holds 60 GB of GTT just moves the memory exhaustion onto the user's services. + try: + stop_containers(verbose=verbose) + except (RuntimeError, OSError): + pass + failures = start_llms(state) + if failures: + print(f" WARNING: could not restart: {', '.join(failures)}") + print(" The GPU host is missing services it had before this run.") + elif verbose: + print(f" restored LLM services; GTT {gtt_used() / 1024**3:.1f} GiB used") + + +def _install_hooks(): + """Restore on the ways out that `finally` does not cover. + + A `finally` block handles a return or an exception. It does not handle SIGTERM, + which is what a timeout, a `kill`, or a parent shell giving up actually sends -- the + default disposition terminates the interpreter without unwinding. That is not + hypothetical: it happened here, and left the user's LLM proxy stopped with nothing + scheduled to start it again. So the same belt-and-braces the archive traversals + use (see :class:`suvi.vfs.Reliever`): atexit for orderly exits, explicit handlers + for the signals. + """ + global _hooks_installed + if _hooks_installed: + return + import atexit + import signal + + atexit.register(_restore_all, verbose=False) + for signum in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP): + previous = signal.getsignal(signum) + + def handler(number, frame, previous=previous): + _restore_all() + if callable(previous) and previous not in (signal.SIG_IGN, signal.SIG_DFL): + return previous(number, frame) + raise SystemExit(128 + number) + + try: + signal.signal(signum, handler) + except (ValueError, OSError): # not the main thread, or no such signal + pass + _hooks_installed = True + + +@contextmanager +def llms_paused(needed_bytes=None, verbose=True): + """Hold the GPU for the duration of the block, then give it back. + + Restoration is registered before the block runs and happens on every exit path -- + return, exception, or signal. What it cannot restore it names, because a + hand-launched server that quietly failed to come back would otherwise be discovered + by the user, later, as a service that is simply gone. + """ + _install_hooks() + state = stop_llms(needed_bytes) + _paused.append(state) + if verbose: + units = [e["unit"] for e in state["units"] if e["was_active"]] + print(f" paused: units {units or 'none'}" + + (f", {len(state['standalone'])} standalone server(s)" + if state["standalone"] else " (no standalone servers touched)")) + print(f" GTT now {gtt_used() / 1024**3:.1f} GiB used") + try: + yield state + finally: + if state in _paused: + _paused.remove(state) + stop_containers(verbose=verbose) + failures = start_llms(state) + if failures: + print(f" WARNING: could not restart: {', '.join(failures)}") + print(" The GPU host is missing services it had before this run.") + elif verbose: + print(f" restored LLM services; GTT {gtt_used() / 1024**3:.1f} GiB used") + + +# ------------------------------------------------------------------ the container + + +#: ROCm+PyTorch plus this project's dependencies. Bazzite's root is read-only and its +#: Python is 3.14, which has no torch wheels, so the toolchain is containerised +#: regardless; the derived tag adds zstandard, opencv, scikit-image and astropy. +IMAGE = os.environ.get("SUVI_GPU_IMAGE", "localhost/suvi-train:latest") +#: Prefix for containers this module starts, so they can be found and stopped again. +CONTAINER_PREFIX = "suvi-run-" + + +def stop_containers(verbose=True): + """Stop every container this module started. + + Necessary because a container does *not* die with the SSH session that launched it. + A run that times out client-side leaves the container running, holding tens of GB of + unified memory; two such orphans were what put GTT at 89.6 GB with the GPU idle. + """ + names = [ + line.strip() + for line in run( + f"podman ps --filter name={CONTAINER_PREFIX} --format '{{{{.Names}}}}'", + check=False, + ).splitlines() + if line.strip().startswith(CONTAINER_PREFIX) + ] + for name in names: + if verbose: + print(f" stopping orphaned container {name}") + run(f"podman stop -t 10 {shlex.quote(name)}", check=False, timeout=60) + return names + + +def launch(command, mounts=(), name=None, log="/var/home/htpc/suvi/run.log"): + """Start a container detached on the host and return (name, log path). + + Long jobs must not be tied to an SSH channel. Streaming a training run through + ``ssh`` means the run dies with the connection, or -- worse, and observed here -- + the container finishes while the local client stays blocked reading a pipe that + will never close. A transient systemd unit is owned by the host, survives the + client entirely, and can be polled with :func:`tail`. + """ + name = name or f"{CONTAINER_PREFIX}{int(time.time())}" + binds = " ".join(f"-v {shlex.quote(src)}:{shlex.quote(dst)}" for src, dst in mounts) + inner = ( + f"podman run --rm --name {shlex.quote(name)} " + f"--device /dev/kfd --device /dev/dri " + f"--security-opt seccomp=unconfined --ipc=host {binds} {IMAGE} {command}" + ) + run(f"rm -f {shlex.quote(log)}", check=False) + run( + f"systemd-run --user --collect --unit={shlex.quote(name)} " + f"--property=StandardOutput=append:{shlex.quote(log)} " + f"--property=StandardError=append:{shlex.quote(log)} " + f"/bin/sh -c {shlex.quote(inner)}", + timeout=60, + ) + return name, log + + +def tail(log, lines=40): + return run(f"tail -n {int(lines)} {shlex.quote(log)} 2>/dev/null", check=False) + + +def running(name): + return run(f"podman ps --filter name={shlex.quote(name)} --format '{{{{.Names}}}}'", + check=False).strip() != "" + + +def podman(command, mounts=(), timeout=None, capture=True, name=None): + """Run `command` inside the training container on the GPU host. + + ``/dev/kfd`` and ``/dev/dri`` are both world-accessible on this host, so no group + mapping is needed; ``seccomp=unconfined`` is what ROCm needs to issue its ioctls. + + The container is named and stopped in a `finally`, so a client-side timeout cannot + leave it running -- ``--rm`` only covers containers that actually exit. + """ + name = name or f"{CONTAINER_PREFIX}{os.getpid()}" + binds = " ".join(f"-v {shlex.quote(src)}:{shlex.quote(dst)}" for src, dst in mounts) + try: + return run( + f"podman run --rm --name {shlex.quote(name)} " + f"--device /dev/kfd --device /dev/dri " + f"--security-opt seccomp=unconfined --ipc=host {binds} {IMAGE} {command}", + timeout=timeout, capture=capture, + ) + finally: + run(f"podman stop -t 10 {shlex.quote(name)}", check=False, timeout=60) diff --git a/suvi/metrics.py b/suvi/metrics.py index 5121b20..0a4a091 100644 --- a/suvi/metrics.py +++ b/suvi/metrics.py @@ -18,7 +18,6 @@ still read as a visible stutter at 60 fps, so flicker is measured explicitly. from dataclasses import dataclass, field import numpy as np -from skimage.metrics import structural_similarity #: Display mapping used by merger_FITS.py to turn radiance into pixels, by band. #: (vmin, vmax, gamma). Fill error is reported through this because it is what the @@ -293,6 +292,12 @@ def score_fill(filled, truth, wavelength, gap_frames=0): log_residual = np.log1p(np.clip(filled, 0, None)) - np.log1p(np.clip(truth, 0, None)) log_rmse = float(np.sqrt(np.mean(log_residual**2))) + # Imported here, not at module scope: DISPLAY_MAPPING and to_display are the only + # parts of this module the training host needs, and scikit-image on that host is + # built against a different numpy ABI than the ROCm image's torch. A constants + # table should not drag an image-processing library in behind it. + from skimage.metrics import structural_similarity + shown_fill, shown_truth = to_display(filled, wavelength), to_display(truth, wavelength) display_mse = float(np.mean((shown_fill - shown_truth) ** 2)) psnr = float("inf") if display_mse == 0 else float(10.0 * np.log10(1.0 / display_mse)) diff --git a/suvi/model.py b/suvi/model.py new file mode 100644 index 0000000..bc8ffdd --- /dev/null +++ b/suvi/model.py @@ -0,0 +1,256 @@ +"""The learned filler: choose, per pixel, which aligned observation to trust. + +The stack this model receives has already been aligned by :mod:`suvi.align`: every +frame is rotation-warped to the target instant and cross-satellite frames are on the +target instrument's radiance scale. Everything with a closed form -- solar rotation, +the photometric transfer -- was applied there, deterministically. What remains is +exactly what has no closed form: + +* **Per-pixel trust.** Which candidate is right varies spatially: on-disc, a warped + temporal neighbour is sharp and co-rotating; above the limb the corona does not + co-rotate and the counterpart's simultaneous view wins; inside a torn or partly + corrupted frame, half the pixels are good and half are not. The weight head emits + per-pixel, per-band blend logits over the stack, and the softmax over candidates + does the selection. +* **A bounded photometric polish.** The daily gain fit is a global affine, and the + measured residual varies 28-61% with radius and drifts between fits. The polish + head emits a per-candidate log-gain field at 20x20, upsampled and bounded to + ``exp(+/-0.3)`` -- enough to fix a level, a radial profile or a flare's broad + brightening, and structurally incapable of synthesising an image. + +Why so little machine learning +------------------------------ +The previous architecture predicted flow, affine, gain, offset, blend, and a free +residual from eight zero-initialised heads, and three training runs failed -- the +last collapsed to an input-independent output because an unbounded additive head is +the cheapest way to satisfy a badly conditioned loss. Here no head can produce an +image on its own: the output is always a convex combination of real aligned +observations, times a gain pinned near one. An untrained model *is* the prior- +weighted blend of its candidates, and every hand-written filler is one softmax +saturation away, so training starts from a sane policy and can only refine it. + +Cross-frame reasoning -- "trust the counterpart *because* the temporal neighbours +disagree" -- needs no attention: each frame's conditioning carries an **agreement** +feature, its RMS deviation from the stack's per-pixel median, computed identically +at training and inference. A frozen, torn or mis-gained frame announces itself +there, whether or not any detector flagged it. +""" + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +#: Resolution the network reasons at. Blend and polish are smooth fields, so they +#: are predicted here and upsampled to the frames' native resolution -- 640 for the +#: training shards, 1280 on the bench -- where they are applied to the real pixels. +#: Sharpness comes from the observations, not the network. +WORK_SIZE = 160 +#: Bands per frame, in suvi.paths.WAVELENGTHS order. +BANDS = 6 +#: Resolution of the polish head's log-gain field. Deliberately coarse: the measured +#: cross-satellite gain residual varies over hundreds of pixels, and a field this +#: smooth cannot carry image content. +POLISH_GRID = 20 +#: Half-width of the polish in log-gain: gains live in [exp(-0.3), exp(0.3)], about +#: [0.74, 1.35]. Covers the measured drift and radial residual with margin, and is +#: the structural bound that keeps the head from becoming a free residual. +POLISH_RANGE = 0.3 + +#: Per-frame conditioning: +#: 3 state one-hot (available / missing / suspect) +#: 1 satellite is the target's +#: 2 dt encoding (sign, log1p|dt| scaled; ~1.0 at a 300-slot gap) +#: 2 agreement (RMS deviation from the stack median: band mean, band max) +#: The last two are written by the model itself, from the aligned stack -- so they +#: exist identically at training and inference, unlike detector scores, which were +#: always zero in training and populated on the bench. +COND_DIM = 8 + +#: Blend prior: fixed logits added to the weight head's output, so the untrained +#: model starts at a sane hand-written policy instead of the uniform stack mean. +#: With candidates aligned and gain-matched there is no cross-satellite penalty any +#: more -- the counterpart arrives on the right scale -- leaving two terms: +#: recency, and distrust of frames known to be damaged. +PRIOR_TEMPORAL_DECAY = 1.5 +PRIOR_SUSPECT_PENALTY = 4.0 + + +def frame_conditioning(state, is_target_satellite, dt_seconds): + """Build one frame's static conditioning vector. See :data:`COND_DIM`.""" + vector = torch.zeros(COND_DIM) + vector[{"available": 0, "missing": 1, "suspect": 2}[state]] = 1.0 + vector[3] = 1.0 if is_target_satellite else 0.0 + dt = float(dt_seconds) + vector[4] = math.copysign(1.0, dt) if dt else 0.0 + vector[5] = math.log1p(abs(dt) / 240.0) / 6.0 + return vector + + +def blend_prior(condition): + """Fixed blend logits from each frame's conditioning. (B, S) -> (B, S).""" + distance = condition[..., 5] * 6.0 + suspect = condition[..., 2] + return -PRIOR_TEMPORAL_DECAY * distance - PRIOR_SUSPECT_PENALTY * suspect + + +# ------------------------------------------------------------------------ backbone + + +class FiLM(nn.Module): + """Per-frame feature modulation from that frame's conditioning vector. + + The stack is processed as a batch of frames, so each frame's `dt`, availability + and agreement have to reach its own features and no other's. Scale-and-shift is + the cheapest thing that does that. + """ + + def __init__(self, channels, cond_dim=COND_DIM): + super().__init__() + self.to_scale_shift = nn.Sequential( + nn.Linear(cond_dim, channels * 2), nn.SiLU(), + nn.Linear(channels * 2, channels * 2), + ) + nn.init.zeros_(self.to_scale_shift[-1].weight) + nn.init.zeros_(self.to_scale_shift[-1].bias) + + def forward(self, features, condition): + scale, shift = self.to_scale_shift(condition).chunk(2, dim=1) + return features * (1 + scale[..., None, None]) + shift[..., None, None] + + +class Block(nn.Module): + def __init__(self, inputs, outputs, cond_dim=COND_DIM): + super().__init__() + self.first = nn.Conv2d(inputs, outputs, 3, padding=1) + self.second = nn.Conv2d(outputs, outputs, 3, padding=1) + self.norm = nn.GroupNorm(8, outputs) + self.film = FiLM(outputs, cond_dim) + self.skip = ( + nn.Identity() if inputs == outputs else nn.Conv2d(inputs, outputs, 1) + ) + + def forward(self, x, condition): + h = F.silu(self.first(x)) + h = self.norm(self.second(h)) + return F.silu(self.film(h, condition) + self.skip(x)) + + +class Encoder(nn.Module): + """Shared per-frame encoder: every frame in the stack goes through one set of + weights, and a frame's role -- which satellite, how far in time, how damaged -- + reaches its features through conditioning rather than through separate branches. + """ + + def __init__(self, in_channels, base=32, depth=3, cond_dim=COND_DIM): + super().__init__() + widths = [base * min(2**i, 8) for i in range(depth + 1)] + self.stem = Block(in_channels, widths[0], cond_dim) + self.down = nn.ModuleList( + [Block(widths[i], widths[i + 1], cond_dim) for i in range(depth)] + ) + self.up = nn.ModuleList( + [Block(widths[i + 1] + widths[i], widths[i], cond_dim) + for i in reversed(range(depth))] + ) + self.width = widths[0] + + def forward(self, x, condition): + h = self.stem(x, condition) + skips = [] + for block in self.down: + skips.append(h) + h = block(F.avg_pool2d(h, 2), condition) + for block, skip in zip(self.up, reversed(skips)): + h = F.interpolate(h, size=skip.shape[-2:], mode="nearest") + h = block(torch.cat([h, skip], dim=1), condition) + return h + + +# --------------------------------------------------------------------------- model + + +class StackFiller(nn.Module): + """Reconstruct one frame as a per-pixel convex combination of aligned candidates. + + Forward takes: + sources (B, S, BANDS, H, W) the **aligned** stack from suvi.align, asinh + space, zeros where a frame is missing. H is 640 in training and + 1280 on the bench; the network downsamples internally. + condition (B, S, COND_DIM) static per-frame conditioning; the agreement + columns are overwritten here from the stack itself. + + Returns the reconstruction at the resolution of `sources`. + """ + + def __init__(self, base=32, depth=3, cond_dim=COND_DIM): + super().__init__() + # +1 input channel: per-frame validity, so the encoder can tell a genuinely + # dark frame (an eclipse) from one that is merely absent. + self.encoder = Encoder(BANDS + 1, base, depth, cond_dim) + self.weight = nn.Conv2d(self.encoder.width, BANDS, 3, padding=1) + self.polish = nn.Conv2d(self.encoder.width, BANDS, 1) + for head in (self.weight, self.polish): + nn.init.zeros_(head.weight) + nn.init.zeros_(head.bias) + + def forward(self, sources, condition): + batch, stack = sources.shape[:2] + size = sources.shape[-2:] + valid = (condition[..., 0] + condition[..., 2]).clamp(0, 1) + + flat = sources.reshape(batch * stack, BANDS, *size) + work = flat + if size[-1] != WORK_SIZE: + work = F.interpolate(flat, size=(WORK_SIZE, WORK_SIZE), + mode="bilinear", align_corners=False) + + condition = self._with_agreement(work, condition, valid, batch, stack) + mask = valid.reshape(batch * stack, 1, 1, 1).expand(-1, 1, *work.shape[-2:]) + features = self.encoder(torch.cat([work, mask], dim=1), + condition.reshape(batch * stack, -1)) + + logits = F.interpolate(self.weight(features), size=size, + mode="bilinear", align_corners=False) + polish = F.interpolate(self.polish(F.adaptive_avg_pool2d(features, POLISH_GRID)), + size=size, mode="bilinear", align_corners=False) + gain = torch.exp(POLISH_RANGE * torch.tanh(polish)) + + logits = logits + blend_prior(condition).reshape(batch * stack, 1, 1, 1) + # A frame with no pixels must not win weight, however confident the head is. + logits = logits.masked_fill( + valid.reshape(batch * stack, 1, 1, 1) < 0.5, float("-inf") + ) + weights = torch.softmax(logits.view(batch, stack, BANDS, *size), dim=1) + weights = torch.nan_to_num(weights, nan=0.0) + + adjusted = (gain * flat).view(batch, stack, BANDS, *size) + return (adjusted * weights).sum(dim=1) + + def _with_agreement(self, work, condition, valid, batch, stack): + """Fill the agreement columns of the conditioning from the stack itself. + + Each frame's RMS deviation from the stack's per-pixel median, over the valid + frames only. A frozen, torn, or mis-gained frame stands out here whether or + not anything flagged it -- this is what replaces both cross-frame attention + and the detector scores of the previous design. float32 throughout: the + deviations are ~0.01-0.1 in asinh units, below bf16's comfort. + """ + with torch.no_grad(): + grouped = work.detach().float().view(batch, stack, BANDS, *work.shape[-2:]) + hidden = torch.where( + valid.reshape(batch, stack, 1, 1, 1) > 0.5, grouped, + torch.full_like(grouped, float("nan")), + ) + median = hidden.nanmedian(dim=1, keepdim=True).values + deviation = torch.sqrt(((grouped - median) ** 2).mean(dim=(-2, -1))) + deviation = torch.nan_to_num(deviation, nan=0.0) * valid[..., None] + condition = condition.clone() + condition[..., 6] = deviation.mean(dim=-1).clamp(0.0, 3.0) + condition[..., 7] = deviation.amax(dim=-1).clamp(0.0, 3.0) + return condition + + +def build(base=32, depth=3): + return StackFiller(base=base, depth=depth) diff --git a/suvi/samples.py b/suvi/samples.py new file mode 100644 index 0000000..a003742 --- /dev/null +++ b/suvi/samples.py @@ -0,0 +1,403 @@ +"""Turning shards into training samples: stacks, masks, and damage applied on the fly. + +One sample is a target slot -- one satellite, one instant, all six bands -- plus the +stack of frames the model may draw on and a per-frame conditioning vector. This module +is the only place that knows how a stack is laid out, so :mod:`suvi.model` and +:mod:`bench` cannot disagree about it. + +Three things here decide what the model can learn. + +**Offsets are exponentially spaced.** Gaps in this archive run from one slot to three +hundred, and a fixed +/-K window either misses the far end of a long outage or wastes +most of its inputs on a short one. Sampling at ``+/-{1, 4, 16, 64, 256}`` slots covers +four minutes to seventeen hours with ten frames per satellite, and an *anchor* -- the +nearest usable frame in each direction, however far -- guarantees the stack is never +empty even beyond that. + +**Damage is applied here, not baked into the shards.** Every epoch draws fresh +corruptions from :mod:`suvi.corruptions`, so the model sees far more variety than a +frozen dataset could hold, and the true mode is known exactly for the auxiliary head. + +**Outages are episodic as well as scattered.** Half of all samples carry a contiguous +simulated outage -- up to 400 slots, on one satellite or both -- because that is how +the real archive fails: measured over the two-satellite era, 11% of slots have both +spacecraft dark at once, in long runs. Independent per-frame drops alone would never +construct that case, and a model cannot learn a regime it has never seen. + +**Context frames get damaged too, not just the target.** This is the requirement that +any supplied frame may itself be invalid. A model trained on clean context and deployed +on an archive where 31% of slots are bad on both satellites would meet, at inference, a +distribution it had never seen. +""" + +import math + +import numpy as np + +from . import align, corruptions, dataset, paths + +#: Slot offsets sampled on each satellite, in both directions. Exponentially spaced, +#: so one fixed-size stack serves every gap length: 4 minutes to 17 hours in five +#: steps, without needing hundreds of input frames to cover a 300-slot outage. The +#: anchors (nearest usable frame either way, at any distance) still guarantee the +#: stack is never empty beyond the ladder's reach. Cost is linear in stack depth in +#: the current model (no attention), so the two extra rungs are affordable where they +#: were not for the attention-based design this replaced. +OFFSETS = (1, 4, 16, 64, 256) +#: How often a sample carries a simulated *episodic* outage -- a contiguous run of +#: missing slots -- on top of the scattered per-frame drops. Real outages are runs, +#: not confetti: measured against the archive, 11% of the two-satellite era has BOTH +#: spacecraft dark simultaneously, in 1,096 contiguous runs of up to 26 days. With +#: independent per-frame drops the probability of reproducing that configuration in +#: a training stack is roughly p^14 -- the model would face at evaluation a regime +#: it had never once seen. +EPISODIC_PROBABILITY = 0.5 +#: Given an episodic outage, how often it takes down *both* satellites at once (a +#: ground-segment or space-weather event) rather than one. +DUAL_OUTAGE_PROBABILITY = 0.4 +#: Longest simulated outage, in slots; matches the anchor search limit, so a stack +#: always retains something real beyond the outage's edge. +MAX_OUTAGE_SLOTS = 400 +#: Modes that leave a frame with no usable signal at all. A frame damaged this way is +#: presented as `missing`; everything else is presented as `suspect` with its pixels. +NO_SIGNAL = frozenset({"all_zero", "drop_image_hdu", "nan_fill", "zblank_fill"}) +#: Auxiliary-head classes: clean plus every catalogued mode, in a fixed order so a +#: checkpoint's class indices stay meaningful. +CLASSES = ("clean",) + tuple(sorted(corruptions.CATALOG)) +CLASS_INDEX = {name: index for index, name in enumerate(CLASSES)} + + +def stack_layout(satellites, target_satellite): + """The (satellite, offset) slots a stack is built from, in a fixed order. + + Deterministic and independent of what happens to be available, so a frame's position + in the stack always means the same thing and the model can rely on it. + + The target satellite's own frame at offset 0 is the frame being reconstructed and is + never an input. The *counterpart's* frame at offset 0 is the single most valuable + entry in the stack -- a real observation of the right Sun at the right instant, worth + ~46 dB on its own once its instrument gain is known. + """ + others = [s for s in satellites if s != target_satellite] + layout = [] + for satellite in [target_satellite] + others: + for offset in OFFSETS: + layout.append((satellite, -offset)) + layout.append((satellite, offset)) + return layout + [(other, 0) for other in others] + + +class Sampler: + """Builds training samples from a set of shards. + + `shards` maps (day, satellite) -> :class:`suvi.dataset.Shard`. + """ + + def __init__(self, shards, satellites=(16, 18), damage_probability=0.30, + drop_probability=0.40, seed=0): + self.shards = shards + self.satellites = tuple(satellites) + #: Ceilings, not rates. Each *sample* draws its own severity uniformly up to + #: these, rather than every frame being damaged at one fixed probability. + #: + #: Fixed rates of 0.25 drop / 0.35 damage left barely half the stack clean in + #: every single sample, so the model never saw the easy case and learned a + #: hedging average instead of "trust the counterpart". It showed up as a model + #: that gained 5 dB on its own harsh validation set and ~1 dB on the bench, + #: where neighbours are mostly clean. Real outages are episodic -- long clean + #: stretches broken by bad runs -- so drawing severity per sample covers both + #: regimes, and roughly a fifth of samples come through almost untouched. + self.damage_probability = damage_probability + self.drop_probability = drop_probability + self.rng = np.random.default_rng(seed) + self._by_satellite = {} + self._day_of = {} + for (day, satellite), shard in shards.items(): + self._day_of[id(shard)] = day + self._by_satellite.setdefault(satellite, {}).update( + {time: shard for time in shard.times()} + ) + #: (day, target satellite, other satellite, pair time) -> per-band (gains, + #: offsets). See :meth:`_calibration`. + self._transfers = {} + + def targets(self): + """Every (satellite, time) that could be reconstructed, in a stable order.""" + found = [] + for satellite in self.satellites: + for time in sorted(self._by_satellite.get(satellite, {})): + found.append((satellite, time)) + return found + + def _read(self, satellite, time): + """One slot in *model* space -- the cheap path, used for every clean frame. + + Shards store the asinh representation already, so this is a rescale rather than + the sinh/arcsinh round trip that going via radiance would cost. + """ + shard = self._by_satellite.get(satellite, {}).get(time) + if shard is None: + return None + return shard.coded(time) if hasattr(shard, "coded") else encode_for_model( + shard.frames(time) + ) + + def _donor(self, satellite, time, rng, reach=60): + """Another real frame, for the modes that splice one in. + + ``frozen`` repeats a neighbour, ``wrong_time`` files one under the wrong stamp + and ``torn_frame`` merges two. All three need a *plausible* second observation, + so the donor is drawn from nearby on the same satellite rather than invented. + """ + available = self._by_satellite.get(satellite, {}) + for _ in range(8): + offset = int(rng.integers(-reach, reach + 1)) + if offset == 0: + continue + candidate = time + offset * paths.CADENCE + if candidate in available: + return self._read(satellite, candidate) + return None + + def _anchor(self, satellite, time, direction, exclude=frozenset(), limit=400): + """Nearest slot with a frame outside `exclude`, up to `limit` slots away. + + `exclude` holds the times inside a simulated outage. The production stack + builder picks anchors from *clean* times, so the sampler must too: an anchor + that lands inside the outage and arrives `missing` would leave the stack with + no real long-range frame at all, which is not what the pipeline would see. + """ + available = self._by_satellite.get(satellite, {}) + for step in range(1, limit + 1): + candidate = time + direction * step * paths.CADENCE + if candidate in available and candidate not in exclude: + return candidate + return None + + def _outages(self, rng, time): + """Per-satellite times inside this sample's simulated episodic outage. + + One contiguous run per affected satellite, log-uniform in length from one + slot to :data:`MAX_OUTAGE_SLOTS`, positioned so the target usually sits + inside it -- reproducing the archive's real failure mode of equipment and + comms dropping out for hours to days, sometimes on both spacecraft at once. + """ + blocked = {satellite: frozenset() for satellite in self.satellites} + # Gated on the drop setting: drop_probability=0 means "no synthetic losses + # of any kind", which diagnostics and tests rely on. + if self.drop_probability <= 0 or rng.random() >= EPISODIC_PROBABILITY: + return blocked + length = int(round(math.exp(rng.uniform(0.0, math.log(MAX_OUTAGE_SLOTS))))) + centre = time + int(rng.integers(-length, length + 1)) * paths.CADENCE + start = centre - (length // 2) * paths.CADENCE + interval = frozenset(start + i * paths.CADENCE for i in range(length)) + if rng.random() < DUAL_OUTAGE_PROBABILITY: + hit = self.satellites + else: + hit = (self.satellites[int(rng.integers(len(self.satellites)))],) + for satellite in hit: + blocked[satellite] = interval + return blocked + + def _calibration(self, day, satellite, other, exclude): + """Per-band (gains, offsets) putting `other`'s radiance on `satellite`'s scale. + + Fitted from one simultaneous pair per day -- the transfer drifts ~12% per + *week*, so a fit at most 24 hours stale is well inside its own noise, and one + pair per day is what a production pipeline could equally afford. + + `exclude` is the time being reconstructed. A pair at that instant would fit + the transfer against the answer itself -- the oracle gain the whole exercise + exists to estimate honestly -- so the fit steps to a neighbouring pair + instead. Returns identity when the day has no usable pair. + """ + local = self.shards.get((day, satellite)) + counterpart = self.shards.get((day, other)) + if local is None or counterpart is None: + return None + candidates = sorted(set(local.times()) & set(counterpart.times()) - {exclude}) + if not candidates: + return None + pair_time = candidates[len(candidates) // 2] + + key = (day, satellite, other, pair_time) + if key not in self._transfers: + self._transfers[key] = align.fit_photometry( + counterpart.frames(pair_time), local.frames(pair_time) + ) + return self._transfers[key] + + def build(self, satellite, time, rng=None): + """One sample: the stack, its per-frame metadata, and the withheld target. + + `frames` is (S, 6, H, W) in *model* space -- roughly [-1, 1] asinh radiance -- + with zeros where a frame is missing; `target` is the withheld frame in the same + space. `gains`/`offsets` are the per-frame radiance transfers that put a + cross-satellite frame on the target instrument's scale (identity rows for + same-satellite and missing frames), and `b0`/`radius` the solar geometry the + rotation warp needs -- both consumed by :func:`suvi.align.align_stack` rather + than here, so the CPU loader stays cheap. + + Returns None if the target itself is absent, or if not one frame in the stack + carries pixels -- there is nothing to reconstruct *from*, and the model's + contract is to decline rather than emit a fabricated frame. + """ + rng = rng if rng is not None else self.rng + target = self._read(satellite, time) + if target is None: + return None + day = self._day_of.get(id(self._by_satellite[satellite][time])) + + # This sample's severity, drawn once and applied to every frame in its stack. + drop_probability = float(rng.uniform(0.0, self.drop_probability)) + damage_probability = float(rng.uniform(0.0, self.damage_probability)) + outages = self._outages(rng, time) + + layout = list(stack_layout(self.satellites, satellite)) + # Anchors: the nearest usable frame either way on the target satellite -- + # outside any simulated outage, as the production stack builder would pick + # them -- so a long gap still has something real to work from. + for direction in (-1, 1): + anchor = self._anchor(satellite, time, direction, exclude=outages[satellite]) + if anchor is not None: + offset = (anchor - time) // paths.CADENCE + if (satellite, offset) not in layout: + layout.append((satellite, offset)) + + frames, states, dts, same, classes, transfers = [], [], [], [], [], [] + for source_satellite, offset in layout: + when = time + offset * paths.CADENCE + image = self._read(source_satellite, when) + state, label = "available", "clean" + + if image is None: + state = "missing" + elif when in outages[source_satellite]: + # Inside this sample's simulated episodic outage. + image, state = None, "missing" + elif rng.random() < drop_probability: + # A scattered single-slot loss, independent of the episodic runs. + image, state = None, "missing" + elif rng.random() < damage_probability: + mode = str(rng.choice(sorted(corruptions.CATALOG))) + corruption = corruptions.CATALOG[mode] + label = mode + if corruption.kind != "array": + # A file-level corruption destroys the container, so by the time a + # frame would reach the model there is nothing left of it. + image, state = None, "missing" + elif mode in NO_SIGNAL: + image, state = None, "missing" + else: + donor = self._donor(source_satellite, when, rng) \ + if corruption.needs_donor else None + if corruption.needs_donor and donor is None: + label, state = "clean", "available" # nothing to draw from + else: + image = _damage(image, mode, int(rng.integers(0, 2**31 - 1)), + float(rng.uniform(0.5, 1.0)), donor) + state = "suspect" + + shape = target.shape + frames.append(np.zeros(shape, dtype=np.float32) if image is None else image) + states.append(state) + dts.append(float(offset * paths.CADENCE)) + same.append(source_satellite == satellite) + classes.append(CLASS_INDEX[label]) + transfer = None + if source_satellite != satellite and image is not None: + transfer = self._calibration(day, satellite, source_satellite, time) + transfers.append(transfer) + + if all(state == "missing" for state in states): + return None + + identity = (np.ones(6, dtype=np.float32), np.zeros(6, dtype=np.float32)) + b0, radius = align.solar_ephemeris(float(time)) + return { + "frames": np.stack(frames), + "states": states, + "dts": np.array(dts, dtype=np.float32), + "same_satellite": np.array(same, dtype=bool), + "classes": np.array(classes, dtype=np.int64), + "gains": np.stack([(t or identity)[0] for t in transfers]), + "offsets": np.stack([(t or identity)[1] for t in transfers]), + "b0": float(b0), + "radius": float(radius), + "target": target, + "slot": (satellite, time), + } + + +def _damage(coded, mode, seed, severity, donor=None): + """Apply one array corruption to all six bands of a frame, in model space. + + Corruptions are *defined* on radiance -- ``gain_shift`` multiplies a physical + quantity, ``saturate`` clips against a physical ceiling -- so the frame is decoded, + damaged, and re-encoded. This is the expensive path, which is why only the frames + actually being corrupted take it. + + Each band gets the same seed, so a fault hits the frame coherently rather than + differently per band, which is how the real ones behave: one instrument produces all + six. ``apply_array`` returns (pixels, header_overrides); the overrides describe a + FITS header this pipeline does not carry, so only the pixels are kept. + """ + radiance = decode_from_model(coded) + donor_radiance = None if donor is None else decode_from_model(donor) + bands = [] + for index, band in enumerate(radiance): + damaged, _ = corruptions.apply_array( + mode, band, seed, severity, + donor=None if donor_radiance is None else donor_radiance[index], + ) + bands.append(damaged) + return encode_for_model(np.stack(bands)) + + +def encode_for_model(radiance): + """Radiance -> the roughly-[-1, 1] asinh space the model works in. + + Corruptions are applied in *radiance*, because that is where they are defined -- a + ``gain_shift`` multiplies a physical quantity and a ``saturate`` clips against a + physical ceiling. Applying them after the transform would model a different fault. + So the encoding happens here, once, on the way into the network. + """ + coded = np.arcsinh( + np.nan_to_num(np.asarray(radiance, dtype=np.float32), nan=0.0, + posinf=0.0, neginf=0.0) / dataset.ASINH_SCALE + ) / dataset.ASINH_RANGE + return np.clip(coded, -1.0, 1.0).astype(np.float32) + + +def decode_from_model(coded): + """Inverse of :func:`encode_for_model`, back to radiance.""" + scaled = np.asarray(coded, dtype=np.float64) * dataset.ASINH_RANGE + return (np.sinh(np.clip(scaled, -dataset.ASINH_RANGE, dataset.ASINH_RANGE)) + * dataset.ASINH_SCALE).astype(np.float32) + + +def to_tensors(sample, torch): + """Pack a sample into tensors for :func:`suvi.align.align_stack` and the model. + + The stack is *unaligned* here: alignment is elementwise math plus a warp, which + the GPU does in milliseconds and the CPU loader should not spend its budget on. + """ + from . import model as model_module + + condition = torch.stack([ + model_module.frame_conditioning( + sample["states"][i], bool(sample["same_satellite"][i]), float(sample["dts"][i]) + ) + for i in range(len(sample["states"])) + ]) + # Already in model space: the sampler works there throughout, so nothing to convert. + return { + "stack": torch.from_numpy(np.ascontiguousarray(sample["frames"])), + "condition": condition, + "dts": torch.from_numpy(np.ascontiguousarray(sample["dts"])), + "gains": torch.from_numpy(np.ascontiguousarray(sample["gains"])), + "offsets": torch.from_numpy(np.ascontiguousarray(sample["offsets"])), + "b0": torch.tensor(sample["b0"], dtype=torch.float32), + "radius": torch.tensor(sample["radius"], dtype=torch.float32), + "target": torch.from_numpy(np.ascontiguousarray(sample["target"])), + } diff --git a/suvi/vfs.py b/suvi/vfs.py index 61efd2c..5a47a83 100644 --- a/suvi/vfs.py +++ b/suvi/vfs.py @@ -21,7 +21,20 @@ import atexit #: Operations between reclaims during a long traversal. Low enough that the mount #: never approaches its ceiling, high enough that the cost is amortised. -RELIEF_INTERVAL = 12_000 +#: +#: Was 12,000, which proved too coarse. The mount refuses opens somewhere around 2 GiB +#: of reclaimable slab -- roughly two million cached inodes -- and a job that touches +#: 25,000 frames with several reads each can cross that between two ticks, which is how +#: a `fill` run left the whole machine unable to exec anything off the share. 6,000 is +#: about a quarter of the way to the ceiling per interval, so a single missed tick is +#: not enough to reach it. +RELIEF_INTERVAL = 6_000 + +#: NOTE: reclaim.py deliberately duplicates the reclaim logic below rather than +#: importing it. That is not an oversight. This module lives on the very share whose +#: exhaustion it addresses, so when the mount starts refusing opens neither this file +#: nor the venv interpreter can be read at all -- the recovery tool has to stand alone +#: on the root filesystem. Keep the two in step by hand; there is not much of either. def drop_caches(): diff --git a/tests/test_align.py b/tests/test_align.py new file mode 100644 index 0000000..4e041f3 --- /dev/null +++ b/tests/test_align.py @@ -0,0 +1,274 @@ +"""The deterministic alignment stage: ephemeris, rotation warp, photometric transfer. + +This is the physics the model no longer has to learn, so its correctness bounds the +whole system: a wrong warp or a wrong gain poisons every candidate the network is +allowed to blend. +""" + +import datetime as dt +import math + +import numpy as np +import pytest + +from conftest import solar_disc +from suvi import align, samples + +torch = pytest.importorskip("torch") + + +def unix(year, month, day): + return dt.datetime(year, month, day, tzinfo=dt.timezone.utc).timestamp() + + +# ---------------------------------------------------------------------- ephemeris + + +def test_ephemeris_b0_reaches_its_annual_extremes(): + """B0 swings +/-7.25 deg, peaking in early March and September.""" + assert math.degrees(align.solar_ephemeris(unix(2024, 3, 7))[0]) == pytest.approx( + -7.25, abs=0.15 + ) + assert math.degrees(align.solar_ephemeris(unix(2024, 9, 8))[0]) == pytest.approx( + 7.25, abs=0.15 + ) + + +def test_ephemeris_b0_crosses_zero_in_june_and_december(): + for when in (unix(2024, 6, 6), unix(2024, 12, 7)): + assert abs(math.degrees(align.solar_ephemeris(when)[0])) < 0.5 + + +def test_ephemeris_radius_tracks_the_orbit(): + """Apparent radius runs ~944 arcsec at aphelion (July) to ~976 at perihelion.""" + def arcsec(when): + return align.solar_ephemeris(when)[1] * align.NATIVE_SIZE * align.PLATE_SCALE + + assert arcsec(unix(2024, 1, 3)) == pytest.approx(975.9, abs=1.0) + assert arcsec(unix(2024, 7, 5)) == pytest.approx(943.9, abs=1.0) + + +def test_ephemeris_matches_a_real_archive_header(): + """Pinned against dr_suvi-l2-ci094_g16_s20240510T000000Z (verified by hand): + DIAM_SUN = 760.2932 px, SOLAR_B0 = -3.205694 deg.""" + when = dt.datetime(2024, 5, 10, 0, 2, tzinfo=dt.timezone.utc).timestamp() + b0, radius_fraction = align.solar_ephemeris(when) + assert 2 * radius_fraction * 1280 == pytest.approx(760.2932, rel=5e-4) + assert math.degrees(b0) == pytest.approx(-3.205694, abs=0.05) + + +# ------------------------------------------------------------------ rotation warp + + +def numpy_reference(size, dt_seconds, b0, radius_fraction): + """The already-validated numpy map from suvi.fillers, on the same geometry.""" + from suvi import fillers + + header = { + "diam_sun": 2 * radius_fraction * size, + "crpix1": (size + 1) / 2, + "crpix2": (size + 1) / 2, + "solar_b0": math.degrees(b0), + } + return fillers._rotation_map((size, size), header, dt_seconds) + + +def test_rotation_grid_matches_the_numpy_map(): + """One implementation trains, the other filled the published baselines; a drift + between them would score the learned filler against different physics.""" + size, lag, b0, radius = 128, 20 * 3600.0, -0.056, 760.29 / 2 / 1280 + map_x, map_y, visible = numpy_reference(size, lag, b0, radius) + + grid = align.rotation_grid(size, torch.tensor([lag]), torch.tensor([b0]), + torch.tensor([radius]))[0].numpy() + got_x = (grid[..., 0] + 1) * size / 2 - 0.5 + got_y = (grid[..., 1] + 1) * size / 2 - 0.5 + assert np.abs(got_x - map_x)[visible].max() < 1e-3 + assert np.abs(got_y - map_y)[visible].max() < 1e-3 + + +def test_rotation_grid_is_identity_where_nothing_co_rotates(): + """Off-disc and behind-the-limb pixels keep their own values: the corona above + the limb does not rotate with the photosphere.""" + size, lag, b0, radius = 64, 20 * 3600.0, 0.02, 0.3 + _, _, visible = numpy_reference(size, lag, b0, radius) + grid = align.rotation_grid(size, torch.tensor([lag]), torch.tensor([b0]), + torch.tensor([radius]))[0].numpy() + grid_y, grid_x = np.mgrid[0:size, 0:size] + got_x = (grid[..., 0] + 1) * size / 2 - 0.5 + got_y = (grid[..., 1] + 1) * size / 2 - 0.5 + assert np.abs(got_x - grid_x)[~visible].max() == 0.0 + assert np.abs(got_y - grid_y)[~visible].max() == 0.0 + + +def test_rotate_at_zero_dt_is_the_identity(): + frames = torch.randn(2, 6, 96, 96) + out = align.rotate(frames, torch.zeros(2), torch.tensor([0.05, -0.05]), + torch.full((2,), 0.3)) + assert float((out - frames).abs().max()) < 1e-3 + + +def test_rotate_moves_on_disc_content(): + disc = torch.from_numpy( + np.stack([solar_disc(size=96, radius=30, peak=2.0)] * 6)[None] + ).float() + # An off-centre bright blob, so rotation has something visible to move. + disc[..., 40:48, 30:38] += 3.0 + out = align.rotate(disc, torch.tensor([12 * 3600.0]), torch.tensor([0.0]), + torch.tensor([30 / 96])) + assert float((out - disc).abs().max()) > 0.5 + + +def test_rotate_round_trips(): + """Forward then back must land where it started, on the visible interior.""" + disc = torch.from_numpy( + np.stack([solar_disc(size=96, radius=30, peak=2.0)] * 6)[None] + ).float() + lag = torch.tensor([6 * 3600.0]) + b0, radius = torch.tensor([0.03]), torch.tensor([30 / 96]) + there = align.rotate(disc, lag, b0, radius) + back = align.rotate(there, -lag, b0, radius) + centre = (slice(None), slice(None), slice(38, 58), slice(38, 58)) + assert float((back[centre] - disc[centre]).abs().max()) < 0.05 + + +# -------------------------------------------------------------------- photometry + + +def test_apply_photometry_matches_radiance_arithmetic(): + """The transfer is defined on radiance; applying it in coded space would model a + different fault entirely.""" + radiance = np.stack([solar_disc(size=32, radius=10, peak=1.0 + b) for b in range(6)]) + coded = torch.from_numpy(samples.encode_for_model(radiance))[None] + gain = torch.tensor([[0.81, 0.885, 1.49, 1.0, 1.2, 0.86]]) + offset = torch.tensor([[0.0, 0.01, -0.02, 0.0, 0.005, 0.0]]) + + moved = align.apply_photometry(coded, gain, offset) + got = samples.decode_from_model(moved[0].numpy()) + expected = radiance * gain[0, :, None, None].numpy() + offset[0, :, None, None].numpy() + np.testing.assert_allclose(got, expected, rtol=3e-3, atol=1e-5) + + +def test_apply_photometry_stays_in_the_coded_range(): + coded = torch.full((1, 6, 8, 8), 1.0) + moved = align.apply_photometry(coded, torch.full((1, 6), 100.0), + torch.full((1, 6), 1e6)) + assert torch.isfinite(moved).all() + assert float(moved.abs().max()) <= 1.0 + + +def test_fit_photometry_recovers_a_known_transfer(): + rng = np.random.default_rng(0) + local = rng.uniform(0.0, 2.0, (6, 32, 32)).astype(np.float32) + gains = np.array([0.81, 0.885, 1.49, 1.0, 1.2, 0.86], dtype=np.float32) + offsets = np.array([0.0, 0.01, -0.02, 0.0, 0.005, 0.03], dtype=np.float32) + counterpart = (local - offsets[:, None, None]) / gains[:, None, None] + + got_gains, got_offsets = align.fit_photometry(counterpart, local) + np.testing.assert_allclose(got_gains, gains, rtol=1e-4) + np.testing.assert_allclose(got_offsets, offsets, atol=1e-4) + + +# ------------------------------------------------------------------- whole stacks + + +def coded_disc(peak=1.0, size=64): + radiance = np.stack([solar_disc(size=size, radius=size // 3, peak=peak + b * 0.2) + for b in range(6)]) + return radiance, torch.from_numpy(samples.encode_for_model(radiance)) + + +def test_align_stack_puts_the_counterpart_on_the_target_scale(): + gains = np.array([0.81, 0.885, 1.49, 1.0, 1.2, 0.86], dtype=np.float32) + local_radiance, local = coded_disc() + counter_radiance = local_radiance / gains[:, None, None] + counter = torch.from_numpy(samples.encode_for_model(counter_radiance)) + + stack = torch.stack([counter])[None] + aligned = align.align_stack( + stack, dts=torch.zeros(1, 1), valid=torch.ones(1, 1), + gains=torch.from_numpy(gains)[None, None], + offsets=torch.zeros(1, 1, 6), + b0=torch.tensor([0.0]), radius_fraction=torch.tensor([0.3]), + ) + assert float((aligned[0, 0] - local).abs().max()) < 5e-3 + + +def test_align_stack_leaves_missing_frames_at_zero(): + """An offset applied to a frame of zeros would manufacture an image from nothing.""" + stack = torch.zeros(1, 1, 6, 32, 32) + aligned = align.align_stack( + stack, dts=torch.zeros(1, 1), valid=torch.zeros(1, 1), + gains=torch.full((1, 1, 6), 3.0), offsets=torch.full((1, 1, 6), 0.5), + b0=torch.tensor([0.0]), radius_fraction=torch.tensor([0.3]), + ) + assert float(aligned.abs().max()) == 0.0 + + +def test_align_stack_leaves_a_simultaneous_local_frame_alone(): + _, local = coded_disc() + stack = torch.stack([local])[None] + aligned = align.align_stack( + stack, dts=torch.zeros(1, 1), valid=torch.ones(1, 1), + gains=torch.ones(1, 1, 6), offsets=torch.zeros(1, 1, 6), + b0=torch.tensor([0.05]), radius_fraction=torch.tensor([0.3]), + ) + assert float((aligned[0, 0] - local).abs().max()) < 1e-3 + + +def test_align_stack_warps_toward_the_target_not_away(): + """The direction test the first training run lacked. + + A 'before' frame (dt < 0) must be rotated *forward* onto the target instant. + The sign error this pins -- warping by dt instead of -dt -- doubled the + misalignment at long gaps while staying sub-pixel at short ones, so only a + test that compares against an independently-warped truth can catch it. + """ + from suvi import fillers + + size, lag, radius = 96, 15 * 3600.0, 30 / 96 + base = np.stack([solar_disc(size=size, radius=30, peak=2.0)] * 6) + base[:, 40:48, 30:38] += 3.0 # feature rotation will move + header = {"diam_sun": 2 * radius * size, "crpix1": (size + 1) / 2, + "crpix2": (size + 1) / 2, "solar_b0": 0.0} + # Independent reference: the validated numpy warp advances `base` by +lag. + truth = np.stack([fillers._warp(band, header, lag, True)[0] for band in base]) + + coded = torch.from_numpy(samples.encode_for_model(base))[None, None] + aligned = align.align_stack( + coded, dts=torch.full((1, 1), -lag), valid=torch.ones(1, 1), + gains=torch.ones(1, 1, 6), offsets=torch.zeros(1, 1, 6), + b0=torch.tensor([0.0]), radius_fraction=torch.tensor([radius]), + ) + got = samples.decode_from_model(aligned[0, 0].numpy()) + interior = (slice(None), slice(30, 66), slice(30, 66)) + aligned_error = float(np.abs(got[interior] - truth[interior]).mean()) + unwarped_error = float(np.abs(base[interior] - truth[interior]).mean()) + assert aligned_error < unwarped_error * 0.35, ( + f"aligned {aligned_error:.4f} vs unwarped {unwarped_error:.4f}: " + "the rotation warp is not moving frames onto the target instant" + ) + + +def test_align_stack_warps_stale_frames_toward_the_target_instant(): + _, local = coded_disc() + marked = local.clone() + marked[..., 20:26, 14:20] += 0.4 # off-centre, on-disc + stack = torch.stack([marked])[None] + aligned = align.align_stack( + stack, dts=torch.full((1, 1), 15 * 3600.0), valid=torch.ones(1, 1), + gains=torch.ones(1, 1, 6), offsets=torch.zeros(1, 1, 6), + b0=torch.tensor([0.0]), radius_fraction=torch.tensor([0.3]), + ) + assert float((aligned[0, 0] - marked).abs().max()) > 0.05 + + +def test_shard_and_ephemeris_share_one_asinh_convention(): + """align.apply_photometry decodes with dataset's constants; if those drift apart + the photometric transfer silently degrades.""" + values = np.array([[-3.7, 0.0, 1e-3, 0.19, 50.0]], dtype=np.float32) + coded = torch.from_numpy(samples.encode_for_model(values)) + identity = align.apply_photometry(coded[None, :, None], torch.ones(1, 1), + torch.zeros(1, 1)) + restored = samples.decode_from_model(identity[0, :, 0].numpy()) + np.testing.assert_allclose(restored, values, rtol=2e-3, atol=1e-7) diff --git a/tests/test_dataset.py b/tests/test_dataset.py new file mode 100644 index 0000000..24b34da --- /dev/null +++ b/tests/test_dataset.py @@ -0,0 +1,396 @@ +import datetime +import os + +import numpy as np +import pytest + +from conftest import solar_disc, write_fits +from suvi import dataset, paths + + +def bands(peak=1.0, size=1280): + """Six bands that differ from each other, so a band mix-up cannot pass.""" + return [solar_disc(size=size, radius=386, peak=peak * (1 + i * 0.3)) for i in range(6)] + + +# ------------------------------------------------------------------ pixel encoding + + +def test_encode_decode_round_trips_radiance(): + """The quantisation has to be tight enough that the model is not learning it.""" + original = bands() + restored = dataset.decode_frames(dataset.encode_frames(original)) + assert restored.shape == (6, dataset.SHARD_SIZE, dataset.SHARD_SIZE) + for index, array in enumerate(original): + small = array.reshape(640, 2, 640, 2).mean(axis=(1, 3)) + relative = np.abs(restored[index] - small) / np.maximum(small, 1e-6) + assert relative.max() < 1e-3 # ~0.055% worst case, one quantisation step + assert relative.mean() < 2e-4 # ~0.014% typical + + +@pytest.mark.parametrize( + "value", [1e-3, 4e-3, 0.05, 1.0, 50.0, 1231.0, -0.01, -0.5, -3.7] +) +def test_relative_precision_holds_across_the_whole_range(value): + """The property log1p failed: faint pixels must be as well resolved as bright ones. + + Median radiance is 0.004 to 0.19 by band, so the faint end is not a corner case -- + it is where most of the frame lives. Values are swept slightly so the test cannot + pass by landing on a quantisation level. + """ + rng = np.random.default_rng(0) + swept = np.full((1280, 1280), value, np.float32) * rng.uniform(1.0, 1.02, (1280, 1280)) + restored = dataset.decode_frames(dataset.encode_frames([swept.astype(np.float32)] * 6)) + small = swept.reshape(640, 2, 640, 2).mean(axis=(1, 3)) + assert (np.abs(restored[0] - small) / np.abs(small)).max() < 1e-3 + + +def test_absolute_precision_holds_below_the_asinh_knee(): + """Below the knee relative precision decays, so pin the absolute error instead. + + A pixel at 1e-6 radiance is five orders of magnitude under the display floor; what + matters there is that it stays negligible, not that its ratio is preserved. + """ + rng = np.random.default_rng(0) + tiny = rng.uniform(-1e-5, 1e-5, (1280, 1280)).astype(np.float32) + restored = dataset.decode_frames(dataset.encode_frames([tiny] * 6)) + small = tiny.reshape(640, 2, 640, 2).mean(axis=(1, 3)) + assert np.abs(restored[0] - small).max() < 1e-7 + + +def test_encoding_preserves_negative_radiance(): + """Background subtraction leaves 2-18% of pixels negative; clipping would bias them.""" + for value in (-3.7, -0.5, -0.01): + flat = np.full((1280, 1280), value, dtype=np.float32) + restored = dataset.decode_frames(dataset.encode_frames([flat] * 6)) + assert restored.max() < 0, f"{value} came back non-negative" + + +def test_encoding_resolves_zero_exactly_enough(): + flat = np.zeros((1280, 1280), dtype=np.float32) + restored = dataset.decode_frames(dataset.encode_frames([flat] * 6)) + assert np.abs(restored).max() < 1e-6 + + +def test_encoding_survives_nan_and_infinity(): + array = solar_disc() + array[10, 10] = np.nan + array[20, 20] = np.inf + array[30, 30] = -np.inf + restored = dataset.decode_frames(dataset.encode_frames([array] * 6)) + assert np.isfinite(restored).all() + + +def test_encoding_clips_rather_than_wrapping_beyond_the_range(): + """A corrupted frame must saturate, never alias into a plausible radiance.""" + for value, expected in ((1e9, 65535), (-1e9, 0)): + block = dataset.encode_frames([np.full((1280, 1280), value, np.float32)] * 6) + assert block.max() == block.min() == expected + assert dataset.decode_frames(np.full((6, 640, 640), 65535, np.uint16)).min() > 1231 + assert dataset.decode_frames(np.zeros((6, 640, 640), np.uint16)).max() < -1231 + + +def test_downsampling_averages_rather_than_decimating(): + """A single hot pixel must be attenuated, not carried through at full amplitude.""" + array = np.zeros((1280, 1280), dtype=np.float32) + array[100, 100] = 50.0 + restored = dataset.decode_frames(dataset.encode_frames([array] * 6)) + assert restored[0, 50, 50] == pytest.approx(12.5, rel=0.01) + + +def test_encode_rejects_a_frame_it_cannot_halve(): + with pytest.raises(ValueError): + dataset.encode_frames([np.zeros((1000, 1000), np.float32)] * 6) + + +# ---------------------------------------------------------------------- containers + + +def test_shard_round_trips_every_slot(tmp_path): + records = {t: dataset.encode_frames(bands(peak=1 + t / 1000.0)) + for t in (1000, 1240, 1480)} + path = str(tmp_path / dataset.shard_name("2024-05-10", 16)) + dataset.write_shard(path, "2024-05-10", 16, paths.WAVELENGTHS, records) + + with dataset.Shard(path) as shard: + assert shard.satellite == 16 + assert shard.day == "2024-05-10" + assert shard.times() == [1000, 1240, 1480] + for t, block in records.items(): + assert t in shard + np.testing.assert_array_equal(shard.raw(t), block) + assert shard.frames(1000).shape == (6, 640, 640) + + +def test_shard_marks_an_unreadable_slot_distinctly(tmp_path): + """Absence and unreadability are different facts and must not be conflated.""" + path = str(tmp_path / "s.shard") + dataset.write_shard(path, "2024-05-10", 16, paths.WAVELENGTHS, + {1000: dataset.encode_frames(bands()), 1240: None}) + with dataset.Shard(path) as shard: + assert 1240 in shard # we looked + assert shard.raw(1240) is None # and there was nothing usable + assert 9999 not in shard # never looked + assert shard.raw(9999) is None + + +def test_shard_rejects_a_foreign_file(tmp_path): + path = tmp_path / "not.shard" + path.write_bytes(b"NOTASHARD" + b"\x00" * 64) + with pytest.raises(ValueError, match="not a shard"): + dataset.Shard(str(path)) + + +def test_shard_rejects_a_future_version(tmp_path): + path = str(tmp_path / "s.shard") + dataset.write_shard(path, "2024-05-10", 16, paths.WAVELENGTHS, {1000: None}) + with open(path, "r+b") as handle: + handle.seek(8) + handle.write((dataset.VERSION + 1).to_bytes(2, "little")) + with pytest.raises(ValueError, match="version"): + dataset.Shard(str(path)) + + +def test_write_shard_leaves_no_partial_file_at_the_real_name(tmp_path): + path = str(tmp_path / "s.shard") + dataset.write_shard(path, "2024-05-10", 16, paths.WAVELENGTHS, {1000: None}) + assert os.path.exists(path) + assert not os.path.exists(path + ".partial") + + +# ---------------------------------------------------------------------- extraction + + +def test_read_slot_matches_a_direct_fits_read(archive): + """The pin on the whole storage path: shard pixels equal archive pixels.""" + arrays = bands() + rows = {} + for index, wavelength in enumerate(paths.WAVELENGTHS): + relpath = f"g16/{wavelength}.fits" + write_fits(os.path.join(str(archive), relpath), arrays[index]) + rows[wavelength] = relpath + + restored = dataset.decode_frames(dataset.read_slot(str(archive), rows)) + for index, array in enumerate(arrays): + expected = array.reshape(640, 2, 640, 2).mean(axis=(1, 3)) + assert np.abs(restored[index] - expected).max() / expected.max() < 0.01 + + +def test_read_slot_refuses_a_partial_slot(archive): + """Five good bands and one missing is not a training sample.""" + rows = {} + for wavelength in paths.WAVELENGTHS[:-1]: + relpath = f"g16/{wavelength}.fits" + write_fits(os.path.join(str(archive), relpath), solar_disc()) + rows[wavelength] = relpath + assert dataset.read_slot(str(archive), rows) is None + + +def test_read_slot_refuses_an_unreadable_band(archive): + rows = {} + for wavelength in paths.WAVELENGTHS: + relpath = f"g16/{wavelength}.fits" + write_fits(os.path.join(str(archive), relpath), solar_disc()) + rows[wavelength] = relpath + with open(os.path.join(str(archive), rows[195]), "r+b") as handle: + handle.truncate(2880) + assert dataset.read_slot(str(archive), rows) is None + + +def test_read_slot_refuses_the_wrong_shape(archive): + rows = {} + for wavelength in paths.WAVELENGTHS: + relpath = f"g16/{wavelength}.fits" + write_fits(os.path.join(str(archive), relpath), solar_disc(size=256, radius=80)) + rows[wavelength] = relpath + assert dataset.read_slot(str(archive), rows) is None + + +def test_read_slot_cannot_be_walked_out_of_the_archive(archive): + """A hostile or corrupt index row must not reach outside the root.""" + with pytest.raises(ValueError): + dataset.read_slot(str(archive), {w: "../../etc/passwd" for w in paths.WAVELENGTHS}) + + +# --------------------------------------------------------------------- day choice + + +def make_index(db_path, days, satellites=(16, 18), per_day=dataset.SLOTS_PER_DAY): + from suvi import db as dbmod + + conn = dbmod.connect(db_path) + rows = [] + for day in days: + base = int(datetime.datetime.fromisoformat(day) + .replace(tzinfo=datetime.timezone.utc).timestamp()) + for satellite in satellites: + for wavelength in paths.WAVELENGTHS: + for slot in range(per_day): + t = base + slot * paths.CADENCE + rows.append((f"{satellite}/{wavelength}/{t}", satellite, wavelength, + t, t + paths.CADENCE, 1_600_000)) + conn.executemany( + "INSERT INTO frame (path, satellite, wavelength, t_start, t_end, size_bytes) " + "VALUES (?,?,?,?,?,?)", + rows, + ) + conn.commit() + return conn + + +def test_choose_days_spreads_across_the_range(db_path): + days = [f"2024-01-{d:02d}" for d in range(1, 31)] + conn = make_index(db_path, days) + chosen = dataset.choose_days(conn, 6, 0, 2**31, satellites=(16, 18)) + assert len(chosen) == 6 + assert chosen == sorted(chosen) + assert chosen[0] < "2024-01-10" and chosen[-1] > "2024-01-20" + + +def test_choose_days_honours_exclusions(db_path): + days = [f"2024-01-{d:02d}" for d in range(1, 31)] + conn = make_index(db_path, days) + banned = {f"2024-01-{d:02d}" for d in range(10, 20)} + chosen = dataset.choose_days(conn, 10, 0, 2**31, exclude=banned) + assert not (set(chosen) & banned) + + +def test_choose_days_skips_incomplete_days(db_path): + """A half-covered day would bias the gap-length distribution silently.""" + conn = make_index(db_path, ["2024-01-01", "2024-01-03"]) + make_index(db_path, ["2024-01-02"], per_day=100) # same db, sparse day + chosen = dataset.choose_days(conn, 10, 0, 2**31) + assert "2024-01-02" not in chosen + assert set(chosen) == {"2024-01-01", "2024-01-03"} + + +def test_choose_days_requires_both_satellites(db_path): + conn = make_index(db_path, ["2024-01-01"], satellites=(16,)) + assert dataset.choose_days(conn, 5, 0, 2**31, satellites=(16, 18)) == [] + + +def test_choose_days_rejects_a_day_of_download_stubs(db_path): + """A failed download leaves a 5,760-byte header with no image, indexed like any + other frame. One chosen day was 100% stubs and produced an empty shard only after + a full extraction pass had read every file on it.""" + conn = make_index(db_path, ["2024-01-01", "2024-01-03"]) + conn.execute("UPDATE frame SET size_bytes = 5760 WHERE t_start < ?", + (int(datetime.datetime(2024, 1, 2, tzinfo=datetime.timezone.utc) + .timestamp()),)) + conn.commit() + chosen = dataset.choose_days(conn, 10, 0, 2**31) + assert "2024-01-01" not in chosen + assert "2024-01-03" in chosen + + +def test_choose_days_accepts_normal_frame_sizes(db_path): + conn = make_index(db_path, ["2024-01-01"]) + conn.execute("UPDATE frame SET size_bytes = 1_600_000") + conn.commit() + assert dataset.choose_days(conn, 5, 0, 2**31) == ["2024-01-01"] + + +# -------------------------------------------------------------------------- splits + + +def test_split_days_partitions_without_overlap(): + days = [f"2024-01-{d:02d}" for d in range(1, 31)] + split = dataset.split_days(days, block=5) + assert sorted(sum(split.values(), [])) == days + assert not (set(split["train"]) & set(split["val"])) + assert not (set(split["val"]) & set(split["test"])) + assert not (set(split["train"]) & set(split["test"])) + + +def test_split_days_keeps_blocks_whole(): + """Days from one block must never straddle two splits.""" + days = [f"2024-01-{d:02d}" for d in range(1, 31)] + split = dataset.split_days(days, block=5) + owner = {day: name for name, group in split.items() for day in group} + for start in range(0, 30, 5): + block = days[start : start + 5] + assert len({owner[day] for day in block}) == 1 + + +def test_split_days_spreads_val_and_test_across_the_range(): + days = [f"2024-{m:02d}-{d:02d}" for m in (1, 4, 7, 10) for d in range(1, 16)] + split = dataset.split_days(days, block=5, train=2, val=1, test=1) + for name in ("val", "test"): + months = {day[:7] for day in split[name]} + assert len(months) > 1, f"{name} landed in a single month" + + +def test_split_never_returns_an_empty_holdout(): + """The silent failure: too few blocks for the cycle and val/test come back empty. + + It surfaces much later as a training run that cannot find validation shards, by + which point the extraction has already run. + """ + for count in range(3, 40): + days = sorted(f"2023-{1 + i // 28:02d}-{1 + i % 28:02d}" for i in range(count)) + split = dataset.split_days(days) + assert split["val"], f"{count} days gave an empty val split" + assert split["test"], f"{count} days gave an empty test split" + assert split["train"], f"{count} days gave an empty train split" + assert sorted(sum(split.values(), [])) == days + + +def test_split_refuses_too_few_days_rather_than_returning_junk(): + with pytest.raises(ValueError, match="at least three"): + dataset.split_days(["2024-01-01", "2024-03-01"]) + + +def test_block_size_is_one_for_days_that_are_far_apart(): + """60 days over two years sit ~15 apart; grouping them guards against nothing.""" + days = [f"2024-{m:02d}-01" for m in range(1, 13)] + assert dataset.block_size(days) == 1 + + +def test_block_size_groups_when_days_are_adjacent(): + days = [f"2024-{1 + i // 28:02d}-{1 + i % 28:02d}" for i in range(30)] + assert dataset.block_size(days, when_consecutive=5) == 5 + + +def test_block_size_never_makes_three_splits_impossible(): + """Three adjacent days cannot have both a guard band and three splits; the split + is the thing that must survive, and check_split still reports the seams.""" + days = ["2024-01-01", "2024-01-02", "2024-01-03"] + assert dataset.block_size(days, when_consecutive=5) == 1 + split = dataset.split_days(days) + assert split["val"] and split["test"] and split["train"] + assert dataset.check_split(split), "adjacency should still be reported" + + +def test_split_of_a_spread_year_covers_the_whole_range(): + """The real failure this caught: a fixed block of 5 over 12 blocks gave val and + test one block each -- four contiguous months out of a two-and-a-half-year archive, + measuring the model at one level of solar activity and nothing else.""" + days = sorted(f"2023-{m:02d}-{d:02d}" for m in range(1, 13) for d in (4, 14, 24)) + assert dataset.block_size(days) == 1, "days are spread; no grouping is needed" + split = dataset.split_days(days) + for name in ("val", "test"): + chosen = split[name] + assert len(chosen) >= 3 + assert chosen[0][:7] < "2023-05", f"{name} starts late: {chosen[0]}" + assert chosen[-1][:7] > "2023-08", f"{name} ends early: {chosen[-1]}" + assert dataset.check_split(split) == [] + + +def test_check_split_catches_a_day_in_two_splits(): + with pytest.raises(ValueError, match="both"): + dataset.check_split({"train": ["2024-01-01"], "val": ["2024-01-01"]}) + + +def test_check_split_reports_abutting_blocks(): + """Adjacent days across a split boundary are four minutes apart at the seam.""" + adjacent = dataset.check_split({"train": ["2024-01-01"], "val": ["2024-01-02"]}) + assert adjacent == [("2024-01-01", "2024-01-02")] + assert dataset.check_split({"train": ["2024-01-01"], "val": ["2024-01-05"]}) == [] + + +def test_manifest_digest_tracks_the_split(): + first = dataset.manifest({"train": ["2024-01-01"], "val": [], "test": []}) + same = dataset.manifest({"train": ["2024-01-01"], "val": [], "test": []}) + other = dataset.manifest({"train": ["2024-01-02"], "val": [], "test": []}) + assert first["digest"] == same["digest"] + assert first["digest"] != other["digest"] diff --git a/tests/test_fillers.py b/tests/test_fillers.py index da02589..b616364 100644 --- a/tests/test_fillers.py +++ b/tests/test_fillers.py @@ -125,6 +125,73 @@ def test_crosssat_aligns_a_parallax_shift(): assert np.abs(filled - truth).mean() < np.abs(unaligned - truth).mean() +def test_crosssat_calibrates_from_the_simultaneous_pair_not_the_bracket(): + """The defect this fixes: a stale bracket poisons the instrument gain. + + The counterpart is a perfect observation of the target instant on the other + instrument's scale. The bracketing frame is twenty hours old, so the Sun in it + is genuinely brighter -- fitting the gain against it folds that evolution into + what should be an instrument constant. A simultaneous pair, however distant in + time, does not have that problem. + """ + truth = solar_disc(size=128, radius=38, peak=2.0) + gain, offset = 0.5, 0.3 # the other instrument's response + counterpart = truth * gain + offset + stale = truth * 1.8 # the Sun, twenty hours earlier + + pair = (stale * gain + offset, stale) # both satellites at that same instant + calibrated = fillers.crosssat( + context(before=stale, dt_before=72000.0, counterpart=counterpart, + calibration=pair), align=False + ) + bracketed = fillers.crosssat( + context(before=stale, dt_before=72000.0, counterpart=counterpart), align=False + ) + np.testing.assert_allclose(calibrated, truth, rtol=1e-3, atol=1e-3) + assert np.abs(bracketed - truth).mean() > 10 * np.abs(calibrated - truth).mean() + + +def test_crosssat_calibration_is_independent_of_how_distant_the_pair_is(): + """A simultaneous pair is equally valid at any separation from the target.""" + truth = solar_disc(size=128, radius=38, peak=2.0) + counterpart = truth * 0.5 + 0.3 + errors = [] + for brightness in (1.0, 1.8, 4.0): # ever more distant, ever more evolved + evolved = truth * brightness + pair = (evolved * 0.5 + 0.3, evolved) + filled = fillers.crosssat( + context(before=evolved, counterpart=counterpart, calibration=pair), + align=False, + ) + errors.append(float(np.abs(filled - truth).mean())) + assert max(errors) < 1e-3 + + +def test_crosssat_falls_back_to_the_bracket_without_a_pair(): + """No slot anywhere had both satellites good; the bracket is all there is.""" + truth = solar_disc(size=128, radius=38, peak=2.0) + counterpart = truth * 0.5 + 0.3 + filled = fillers.crosssat(context(before=truth, counterpart=counterpart), align=False) + np.testing.assert_allclose(filled, truth, atol=1e-3) + + +def test_crosssat_ignores_a_malformed_calibration_pair(): + truth = solar_disc(size=128, radius=38, peak=2.0) + counterpart = truth * 0.5 + 0.3 + for pair in ((None, truth), (truth, None), (np.ones((8, 8), np.float32), truth)): + filled = fillers.crosssat( + context(before=truth, counterpart=counterpart, calibration=pair), align=False + ) + np.testing.assert_allclose(filled, truth, atol=1e-3) + + +def test_gain_fit_returns_identity_for_a_constant_source(): + assert fillers.gain_fit(np.ones((8, 8), np.float32), np.zeros((8, 8), np.float32)) == ( + 1.0, + 0.0, + ) + + def test_gain_match_recovers_an_affine_transform(): source = solar_disc(size=64, radius=20, peak=1.0) reference = source * 3.0 - 0.5 @@ -203,13 +270,31 @@ def test_solar_rotation_gives_up_with_no_brackets(): # ---------------------------------------------------------------------- registry -def test_every_registered_filler_is_callable_and_shape_preserving(): +#: Fillers that reconstruct one band from the two-frame context. `learned` is not one +#: of them: it fuses a whole stack of six-band frames and reports "not applicable" +#: without one, which is why it is excluded here rather than exempted inside the loop. +CONTEXT_FILLERS = ("hold_last", "linear_blend", "optical_flow", "crosssat", + "solar_rotation") + + +def test_every_registered_filler_is_covered_here(): + """A new filler must be classified deliberately, not silently skipped.""" + assert set(fillers.FILLERS) == set(CONTEXT_FILLERS) | {"learned"} + + +def test_every_context_filler_is_callable_and_shape_preserving(): before = solar_disc(size=64, radius=20, peak=2.0) after = solar_disc(size=64, radius=20, peak=2.2) header = dict(HEADER, crpix1=32.5, crpix2=32.5, diam_sun=40.0) ctx = context(before=before, after=after, counterpart=after, header=header) - for name, filler in fillers.FILLERS.items(): - result = filler(ctx) + for name in CONTEXT_FILLERS: + result = fillers.FILLERS[name](ctx) assert result is not None, name assert result.shape == before.shape, name assert np.isfinite(result).all(), name + + +def test_learned_declines_a_two_frame_context(): + """Reporting 'not applicable' beats returning something plausible but untrained.""" + before = solar_disc(size=64, radius=20, peak=2.0) + assert fillers.learned(context(before=before, counterpart=before)) is None diff --git a/tests/test_learned_filler.py b/tests/test_learned_filler.py new file mode 100644 index 0000000..6c794db --- /dev/null +++ b/tests/test_learned_filler.py @@ -0,0 +1,376 @@ +"""The learned filler's wiring: the stack it is handed, and how it is loaded. + +The model itself is covered in test_model.py. What matters here is that `bench.py` +hands it the same stack shape `suvi.samples` builds during training, and that a frame the +case corrupted arrives as *corrupted pixels* rather than as the pristine original. +""" + +import numpy as np +import pytest + +import bench +from conftest import solar_disc +from suvi import cases, fillers, paths, samples + + +WAVELENGTHS = paths.WAVELENGTHS + + +def overlay_for(satellites=(16, 18), count=40, deleted=()): + """An Overlay whose slots resolve through in-memory lookups rather than files.""" + base = 1715400000 // paths.CADENCE * paths.CADENCE + times = [base + i * paths.CADENCE for i in range(count)] + archive = {(s, w, t): f"truth/{s}/{w}/{t}" + for s in satellites for w in WAVELENGTHS for t in times} + return cases.Overlay(archive=archive, deleted=frozenset(deleted)), times, base + + +def reader(size=16, missing=()): + """A six-band reader. `missing` names (satellite, time) pairs with no frame.""" + frame = np.stack([solar_disc(size=size, radius=size // 3, peak=1.0 + i) + for i in range(6)]).astype(np.float32) + + def read(satellite, when): + if (satellite, when) in missing: + return None + return frame.copy() + + return read, frame + + +# --------------------------------------------------------------------------- stack + + +def test_stack_matches_the_training_layout(): + """Trained on one stack shape, evaluated on another, would be a silent mismatch. + + The window spans 560 slots so even the +/-256 rungs of the exponential ladder + exist; a real window edge simply omits the rungs it cannot reach. + """ + overlay, times, _ = overlay_for(count=560) + read, _ = reader() + stack = bench._build_stack(overlay, 16, times[280], (16, 18), WAVELENGTHS, set(), + read) + + offsets = {(entry["same_satellite"], entry["dt"] / paths.CADENCE) for entry in stack} + for source, offset in samples.stack_layout((16, 18), 16): + assert (source == 16, float(offset)) in offsets, f"missing {(source, offset)}" + + +def test_stack_carries_all_six_bands(): + """The model is joint across bands; a per-band stack would not fit it.""" + overlay, times, _ = overlay_for() + read, _ = reader() + stack = bench._build_stack(overlay, 16, times[20], (16, 18), WAVELENGTHS, set(), read) + for entry in stack: + assert entry["image"].shape[0] == len(WAVELENGTHS) + + +def test_stack_excludes_the_target_slot(): + overlay, times, _ = overlay_for() + read, _ = reader() + stack = bench._build_stack(overlay, 16, times[20], (16, 18), WAVELENGTHS, set(), read) + for entry in stack: + assert not (entry["same_satellite"] and entry["dt"] == 0.0) + + +def test_stack_includes_the_counterpart_at_the_target_instant(): + overlay, times, _ = overlay_for() + read, _ = reader() + stack = bench._build_stack(overlay, 16, times[20], (16, 18), WAVELENGTHS, set(), read) + assert any(entry["dt"] == 0.0 and not entry["same_satellite"] for entry in stack) + + +def test_damaged_neighbours_arrive_suspect_rather_than_dropped(): + """The property the whole design rests on: a flagged frame is data, not a hole.""" + overlay, times, _ = overlay_for() + read, _ = reader() + target = times[20] + neighbour_time = target - paths.CADENCE + bad = {(16, w, neighbour_time) for w in WAVELENGTHS} + stack = bench._build_stack(overlay, 16, target, (16, 18), WAVELENGTHS, bad, read) + + entry = next(e for e in stack if e["slot"] == (16, neighbour_time)) + assert entry["state"] == "suspect" + assert entry["image"] is not None + + +def test_a_slot_damaged_in_one_band_marks_the_whole_frame_suspect(): + """One instrument makes all six bands; a fault in one is a reason to distrust all.""" + overlay, times, _ = overlay_for() + read, _ = reader() + target = times[20] + neighbour_time = target - paths.CADENCE + stack = bench._build_stack(overlay, 16, target, (16, 18), WAVELENGTHS, + {(16, 195, neighbour_time)}, read) + entry = next(e for e in stack if e["slot"] == (16, neighbour_time)) + assert entry["state"] == "suspect" + + +def test_unreadable_neighbours_arrive_missing_with_no_pixels(): + overlay, times, _ = overlay_for() + target = times[20] + gone = (16, target - paths.CADENCE) + read, _ = reader(missing={gone}) + stack = bench._build_stack(overlay, 16, target, (16, 18), WAVELENGTHS, set(), read) + entry = next(e for e in stack if e["slot"] == gone) + assert entry["state"] == "missing" + assert entry["image"] is None + + +def test_undamaged_neighbours_arrive_available(): + overlay, times, _ = overlay_for() + read, frame = reader() + stack = bench._build_stack(overlay, 16, times[20], (16, 18), WAVELENGTHS, set(), read) + assert {e["state"] for e in stack} == {"available"} + for entry in stack: + np.testing.assert_allclose(entry["image"], frame) + + +def test_stack_adds_anchors_beyond_the_fixed_offsets(): + """A long outage must still reach a real frame.""" + overlay, times, _ = overlay_for(count=80) + read, _ = reader() + target = times[40] + bad = {(16, w, target + n * paths.CADENCE) + for w in WAVELENGTHS for n in range(-30, 31)} + stack = bench._build_stack(overlay, 16, target, (16, 18), WAVELENGTHS, bad, read) + + reach = [abs(e["dt"]) for e in stack + if e["same_satellite"] and e["state"] == "available"] + assert reach, "no usable same-satellite frame found at all" + assert max(reach) >= 31 * paths.CADENCE + + +def test_stack_skips_slots_the_window_does_not_contain(): + """Near a window edge the offsets run off the end; those must not appear as data.""" + overlay, times, _ = overlay_for(count=40) + read, _ = reader() + stack = bench._build_stack(overlay, 16, times[0], (16, 18), WAVELENGTHS, set(), read) + for entry in stack: + satellite, when = entry["slot"] + assert any((satellite, w, when) in overlay.archive for w in WAVELENGTHS) + + +# ------------------------------------------------------------------- band reader + + +def test_band_reader_returns_damaged_pixels_for_damaged_slots(tmp_path): + """A filler handed the pristine frame for a corrupted slot reads the answer key.""" + from conftest import write_fits + + truth_frame = solar_disc(size=64, radius=20, peak=1.0) + damaged_frame = solar_disc(size=64, radius=20, peak=9.0) + when = 1715400000 + archive_map, overrides = {}, {} + for wavelength in WAVELENGTHS: + good = str(tmp_path / f"t{wavelength}.fits") + harmed = str(tmp_path / f"d{wavelength}.fits") + write_fits(good, truth_frame) + write_fits(harmed, damaged_frame) + archive_map[(16, wavelength, when)] = good + overrides[(16, wavelength, when)] = harmed + + overlay = cases.Overlay(archive=archive_map, overrides=overrides) + read = bench._band_reader(overlay, WAVELENGTHS, set(overrides)) + got = read(16, when) + assert got.shape == (6, 64, 64) + # Compared loosely: the FITS tile compression is lossy at the 1e-3 level, and what + # is under test is which *file* was opened, not the codec. + assert got[0].max() == pytest.approx(damaged_frame.max(), rel=1e-2) + assert got[0].max() > truth_frame.max() * 5 + + +def test_band_reader_caches_but_stays_bounded(tmp_path): + """Six bands of 1280x1280 float32 is 39 MB; an unbounded cache exhausts the VM.""" + from conftest import write_fits + + frame = solar_disc(size=32, radius=10, peak=1.0) + archive_map = {} + times = [1715400000 + i * paths.CADENCE for i in range(10)] + for when in times: + for wavelength in WAVELENGTHS: + path = str(tmp_path / f"{when}_{wavelength}.fits") + write_fits(path, frame) + archive_map[(16, wavelength, when)] = path + + overlay = cases.Overlay(archive=archive_map) + reads = [] + real = bench.fitsio.read_image + + def counting(path): + reads.append(path) + return real(path) + + read = bench._band_reader(overlay, WAVELENGTHS, set(), limit=3) + original = bench.fitsio.read_image + bench.fitsio.read_image = counting + try: + read(16, times[0]) + first = len(reads) + read(16, times[0]) + assert len(reads) == first, "a repeat read was not served from cache" + for when in times[1:5]: + read(16, when) + before_evicted = len(reads) + read(16, times[0]) # evicted by the limit of 3 + assert len(reads) > before_evicted, "cache grew past its limit" + finally: + bench.fitsio.read_image = original + + +def test_band_reader_returns_none_when_a_band_is_missing(tmp_path): + """Five bands is not a frame the joint model can consume.""" + from conftest import write_fits + + frame = solar_disc(size=32, radius=10, peak=1.0) + when = 1715400000 + archive_map = {} + for wavelength in WAVELENGTHS[:-1]: + path = str(tmp_path / f"{wavelength}.fits") + write_fits(path, frame) + archive_map[(16, wavelength, when)] = path + overlay = cases.Overlay(archive=archive_map) + assert bench._band_reader(overlay, WAVELENGTHS, set())(16, when) is None + + +# -------------------------------------------------------------------------- filler + + +def test_learned_returns_nothing_without_a_stack(): + assert fillers.learned(fillers.FillContext()) is None + + +def test_learned_returns_nothing_without_a_checkpoint(monkeypatch): + monkeypatch.delenv(fillers.LEARNED_CHECKPOINT_ENV, raising=False) + fillers._LEARNED.clear() + context = fillers.FillContext(stack=[ + {"image": np.zeros((6, 16, 16), np.float32), "state": "available", + "dt": -240.0, "same_satellite": True} + ]) + assert fillers.learned(context) is None + + +def test_learned_is_registered_alongside_the_others(): + assert "learned" in fillers.FILLERS + assert fillers.FILLERS["learned"] is fillers.learned + + +def test_existing_fillers_ignore_the_stack(): + """Adding `stack` must not perturb any measured baseline.""" + before = solar_disc(size=64, radius=20, peak=1.0) + after = solar_disc(size=64, radius=20, peak=1.2) + extra = np.stack([after * 5] * 6) + plain = fillers.FillContext(before=before, after=after, dt_before=240, dt_after=240) + with_stack = fillers.FillContext( + before=before, after=after, dt_before=240, dt_after=240, + stack=[{"image": extra, "state": "suspect", "dt": -240.0, + "same_satellite": True}], + ) + for name in ("hold_last", "linear_blend", "optical_flow", "solar_rotation"): + np.testing.assert_allclose(fillers.FILLERS[name](plain), + fillers.FILLERS[name](with_stack)) + + +def test_learned_runs_end_to_end_against_a_saved_checkpoint(tmp_path, monkeypatch): + """Checkpoint -> load -> fill, at the archive's native frame size.""" + torch = pytest.importorskip("torch") + from suvi import model + + net = model.build(base=8, depth=2) + path = tmp_path / "model.pt" + torch.save({"model": net.state_dict(), "args": {"base": 8, "depth": 2}, + "epoch": 0}, path) + monkeypatch.setenv(fillers.LEARNED_CHECKPOINT_ENV, str(path)) + fillers._LEARNED.clear() + + def six(peak): + return np.stack([solar_disc(size=128, radius=40, peak=peak + b * 0.2) + for b in range(6)]).astype(np.float32) + + when = 1715400000 + context = fillers.FillContext( + stack=[ + {"image": six(1.0), "state": "available", "dt": -240.0, + "same_satellite": True, "slot": (16, when - 240)}, + {"image": six(1.1), "state": "available", "dt": 240.0, + "same_satellite": True, "slot": (16, when + 240)}, + {"image": six(0.9), "state": "suspect", "dt": 0.0, + "same_satellite": False, "slot": (18, when)}, + {"image": None, "state": "missing", "dt": -960.0, + "same_satellite": True, "slot": (16, when - 960)}, + ], + calibration=(six(0.9), six(1.05)), + ) + filled = fillers.learned(context) + assert filled.shape == (6, 128, 128) + assert np.isfinite(filled).all() + fillers._LEARNED.clear() + + +def test_learned_declines_a_stack_with_no_pixels_anywhere(tmp_path, monkeypatch): + """Silently emitting a black frame here is how a fabricated fill would enter the + archive; the contract is None.""" + torch = pytest.importorskip("torch") + from suvi import model + + net = model.build(base=8, depth=2) + path = tmp_path / "model.pt" + torch.save({"model": net.state_dict(), "args": {"base": 8, "depth": 2}}, path) + monkeypatch.setenv(fillers.LEARNED_CHECKPOINT_ENV, str(path)) + fillers._LEARNED.clear() + context = fillers.FillContext(stack=[ + {"image": None, "state": "missing", "dt": -240.0, "same_satellite": True, + "slot": (16, 1715400000 - 240)}, + ]) + assert fillers.learned(context) is None + fillers._LEARNED.clear() + + +def test_learned_checkpoint_is_loaded_once(tmp_path, monkeypatch): + """Thousands of slots per bench run; reloading 14M parameters each time would + dominate the wall clock.""" + torch = pytest.importorskip("torch") + from suvi import model + + net = model.build(base=8, depth=2) + path = tmp_path / "model.pt" + torch.save({"model": net.state_dict(), "args": {"base": 8, "depth": 2}}, path) + monkeypatch.setenv(fillers.LEARNED_CHECKPOINT_ENV, str(path)) + fillers._LEARNED.clear() + + loads = [] + real_load = torch.load + monkeypatch.setattr(torch, "load", lambda *a, **k: (loads.append(1), real_load(*a, **k))[1]) + for _ in range(3): + fillers.load_learned() + assert len(loads) == 1 + fillers._LEARNED.clear() + + +def test_band_reader_ticks_the_reliever_per_read(tmp_path): + """One target pulls ~90 frames through the reader. Ticking once per target would + undercount ninety-fold, and the mount would exhaust its file handles between two + ticks -- which has taken this machine down more than once.""" + from conftest import write_fits + + class Counter: + def __init__(self): + self.n = 0 + + def tick(self, count=1): + self.n += count + + frame = solar_disc(size=32, radius=10, peak=1.0) + when = 1715400000 + archive_map = {} + for wavelength in WAVELENGTHS: + path = str(tmp_path / f"{wavelength}.fits") + write_fits(path, frame) + archive_map[(16, wavelength, when)] = path + + counter = Counter() + overlay = cases.Overlay(archive=archive_map) + read = bench._band_reader(overlay, WAVELENGTHS, set(), reliever=counter) + read(16, when) + assert counter.n == len(WAVELENGTHS) diff --git a/tests/test_model.py b/tests/test_model.py new file mode 100644 index 0000000..18b4c90 --- /dev/null +++ b/tests/test_model.py @@ -0,0 +1,336 @@ +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +from suvi import model # noqa: E402 + +BANDS = model.BANDS + + +def conditioning(states, satellites=None, dts=None): + """Build a (1, S, COND_DIM) conditioning tensor.""" + count = len(states) + satellites = satellites or [True] * count + dts = dts if dts is not None else [(i - count // 2) * 240.0 for i in range(count)] + rows = [ + model.frame_conditioning(states[i], satellites[i], dts[i]) + for i in range(count) + ] + return torch.stack(rows)[None] + + +def stack(count, size=64, seed=0): + generator = torch.Generator().manual_seed(seed) + return torch.randn(1, count, BANDS, size, size, generator=generator) + + +@pytest.fixture(scope="module") +def net(): + torch.manual_seed(0) + return model.build(base=16, depth=3).eval() + + +# ------------------------------------------------------------------- conditioning + + +def test_conditioning_is_the_declared_width(): + assert model.frame_conditioning("available", True, 0.0).shape == (model.COND_DIM,) + + +def test_conditioning_separates_the_three_states(): + vectors = [model.frame_conditioning(s, True, 240.0)[:3] + for s in ("available", "missing", "suspect")] + for index, vector in enumerate(vectors): + assert vector[index] == 1.0 and vector.sum() == 1.0 + + +def test_conditioning_rejects_an_unknown_state(): + with pytest.raises(KeyError): + model.frame_conditioning("probably-fine", True, 0.0) + + +def test_conditioning_distinguishes_direction_and_distance(): + back = model.frame_conditioning("available", True, -7200.0) + forward = model.frame_conditioning("available", True, 7200.0) + far = model.frame_conditioning("available", True, 72000.0) + assert back[4] == -1.0 and forward[4] == 1.0 + assert back[5] == pytest.approx(forward[5]) # magnitude, not direction + assert far[5] > forward[5] + + +def test_conditioning_leaves_agreement_to_the_model(): + """Columns 6 and 7 are computed from the stack inside forward(), identically at + training and inference -- the fix for detector scores that were always zero in + training and populated on the bench.""" + vector = model.frame_conditioning("available", True, 240.0) + assert vector[6] == 0.0 and vector[7] == 0.0 + + +# ------------------------------------------------------------------- blend prior + + +def test_prior_favours_the_nearest_frame_in_time(): + condition = conditioning(["available"] * 3, satellites=[True] * 3, + dts=[-240.0, -3840.0, -240.0 * 300]) + prior = model.blend_prior(condition)[0] + assert prior[0] > prior[1] > prior[2] + + +def test_prior_no_longer_penalises_the_counterpart(): + """The old prior down-weighted cross-satellite frames because their gain was + wrong at initialisation. Candidates now arrive gain-matched by suvi.align, so + a simultaneous counterpart is as trustworthy as a local frame.""" + same = conditioning(["available"], satellites=[True], dts=[0.0]) + cross = conditioning(["available"], satellites=[False], dts=[0.0]) + assert model.blend_prior(same)[0, 0] == model.blend_prior(cross)[0, 0] + + +def test_prior_downweights_suspect_frames_without_excluding_them(): + clean = conditioning(["available"], satellites=[True], dts=[-240.0]) + suspect = conditioning(["suspect"], satellites=[True], dts=[-240.0]) + assert model.blend_prior(clean)[0, 0] > model.blend_prior(suspect)[0, 0] + assert torch.isfinite(model.blend_prior(suspect)).all() + + +# ------------------------------------------------------------------------- shapes + + +def test_forward_returns_the_source_shape(net): + out = net(stack(5), conditioning(["available"] * 5)) + assert out.shape == (1, BANDS, 64, 64) + assert torch.isfinite(out).all() + + +@pytest.mark.parametrize("size", [160, 320, 640]) +def test_one_model_serves_every_source_resolution(net, size): + """Blend and polish are fields: they upsample, so 1280 needs no retraining.""" + sources = torch.randn(1, 3, BANDS, size, size) + out = net(sources, conditioning(["available"] * 3)) + assert out.shape == (1, BANDS, size, size) + assert torch.isfinite(out).all() + + +@pytest.mark.parametrize("count", [1, 2, 7, 23]) +def test_stack_depth_is_not_fixed(net, count): + """Availability varies per slot, so the stack cannot be a compile-time constant.""" + out = net(stack(count), conditioning(["available"] * count)) + assert out.shape == (1, BANDS, 64, 64) + + +# --------------------------------------------------------------- identity behaviour + + +def expected_blend(sources, condition): + """What an untrained model must output: the prior-weighted blend of its sources. + + Both heads are zero-initialised, so the polish gain is exactly one and the blend + logits are exactly the prior.""" + valid = (condition[..., 0] + condition[..., 2]).clamp(0, 1) + logits = model.blend_prior(condition).masked_fill(valid < 0.5, float("-inf")) + weights = torch.nan_to_num(torch.softmax(logits, dim=1), nan=0.0) + return (sources * weights[..., None, None, None]).sum(dim=1) + + +def test_untrained_model_is_the_prior_weighted_blend(net): + """The wiring gate: any scaling or orientation bug in the composition shows up + here, and is invisible once training starts moving the loss.""" + frames = stack(5) + condition = conditioning(["available"] * 5) + out = net(frames, condition) + assert torch.allclose(out, expected_blend(frames, condition), atol=1e-4) + + +def test_identity_holds_when_fields_are_upsampled(net): + """A 160-resolution field applied to larger sources must still be the identity.""" + sources = torch.randn(1, 3, BANDS, 320, 320) + condition = conditioning(["available"] * 3) + out = net(sources, condition) + assert torch.allclose(out, expected_blend(sources, condition), atol=1e-3) + + +def test_missing_frames_take_no_weight(net): + """A frame with no pixels must not contribute, however confident the head is.""" + frames = stack(4) + frames[:, 0] = 999.0 # would dominate any blend that included it + condition = conditioning(["missing", "available", "available", "available"]) + out = net(frames, condition) + assert torch.allclose(out, expected_blend(frames, condition), atol=1e-4) + assert out.abs().max() < 100, "the masked frame leaked into the blend" + + +def test_suspect_frames_do_take_weight(net): + """The design decision under test: a flagged frame is data, not a hole.""" + frames = stack(3) + excluded = conditioning(["missing", "available", "available"]) + included = conditioning(["suspect", "available", "available"]) + a = net(frames, excluded) + b = net(frames, included) + assert not torch.allclose(a, b, atol=1e-3) + assert torch.allclose(b, expected_blend(frames, included), atol=1e-4) + + +def test_a_stack_with_nothing_usable_returns_zeros_not_nan(net): + """The sampler declines such targets; the network itself must still not NaN.""" + out = net(stack(4), conditioning(["missing"] * 4)) + assert torch.isfinite(out).all() + assert out.abs().max() == 0.0 + + +def test_heads_start_at_zero(): + """Training must start exactly at the hand-written policy; every step afterwards + is a departure the data paid for.""" + torch.manual_seed(0) + fresh = model.build(base=8, depth=2) + for head in (fresh.weight, fresh.polish): + with torch.no_grad(): + assert float(head.weight.abs().sum()) == 0.0 + assert float(head.bias.abs().sum()) == 0.0 + + +# --------------------------------------------------------- structural non-collapse + + +def test_output_is_a_bounded_combination_whatever_the_weights(): + """The invariant that makes the old failure mode unreachable: with *arbitrary* + parameters, the output of a single-frame stack stays within the polish bound of + that frame. No parameter setting can emit an image that is not built from the + inputs -- which is exactly what the collapsed run did.""" + torch.manual_seed(3) + wild = model.build(base=8, depth=2) + with torch.no_grad(): + for parameter in wild.parameters(): + parameter.normal_(0.0, 2.0) + frames = torch.full((1, 1, BANDS, 48, 48), 0.5) + with torch.no_grad(): + out = wild(frames, conditioning(["available"], dts=[240.0])) + bound = 0.5 * float(np.exp(model.POLISH_RANGE)) + lower = 0.5 * float(np.exp(-model.POLISH_RANGE)) + assert float(out.max()) <= bound + 1e-4 + assert float(out.min()) >= lower - 1e-4 + + +def test_polish_gain_saturates_at_its_bound(): + torch.manual_seed(0) + fresh = model.build(base=8, depth=2).eval() + with torch.no_grad(): + fresh.polish.bias.fill_(1000.0) + frames = torch.ones(1, 1, BANDS, 48, 48) + out = fresh(frames, conditioning(["available"], dts=[240.0])) + assert float(out.mean()) == pytest.approx(np.exp(model.POLISH_RANGE), rel=1e-3) + + +# --------------------------------------------------------------------- agreement + + +def test_agreement_flags_the_frame_that_disagrees(): + """A frozen or torn frame announces itself as deviation from the stack median -- + the deterministic replacement for both cross-attention and detector scores.""" + torch.manual_seed(0) + fresh = model.build(base=8, depth=2) + frames = torch.zeros(1, 4, BANDS, 32, 32) + frames[:, 3] = 5.0 # the odd one out + condition = conditioning(["available"] * 4) + valid = (condition[..., 0] + condition[..., 2]).clamp(0, 1) + got = fresh._with_agreement(frames.reshape(4, BANDS, 32, 32), condition, valid, + batch=1, stack=4) + assert float(got[0, 3, 6]) > float(got[0, 0, 6]) + 0.5 + assert float(got[0, 0, 6]) < 0.1 + + +def test_agreement_is_zero_for_a_unanimous_stack(): + torch.manual_seed(0) + fresh = model.build(base=8, depth=2) + frames = torch.full((1, 3, BANDS, 32, 32), 0.25) + condition = conditioning(["available"] * 3) + valid = (condition[..., 0] + condition[..., 2]).clamp(0, 1) + got = fresh._with_agreement(frames.reshape(3, BANDS, 32, 32), condition, valid, + batch=1, stack=3) + assert float(got[..., 6:].abs().max()) < 1e-6 + + +def test_agreement_ignores_missing_frames(): + """A zeroed missing frame must not drag the median toward black.""" + torch.manual_seed(0) + fresh = model.build(base=8, depth=2) + frames = torch.full((1, 3, BANDS, 32, 32), 0.5) + frames[:, 0] = 0.0 # missing, zero-filled + condition = conditioning(["missing", "available", "available"]) + valid = (condition[..., 0] + condition[..., 2]).clamp(0, 1) + got = fresh._with_agreement(frames.reshape(3, BANDS, 32, 32), condition, valid, + batch=1, stack=3) + assert float(got[0, 1, 6]) < 1e-6 + assert float(got[0, 2, 6]) < 1e-6 + + +# ------------------------------------------------------------------------ training + + +def test_gradients_reach_every_head(): + """A head with no gradient path never learns, and zero-init hides it. Checked + after three optimiser steps, because the zero initialisation is nested two deep: + at step 1 only the output heads move, at step 2 the FiLM output layers behind + them, and only then does gradient reach everything. Asserting earlier would + fail on a model that is perfectly healthy.""" + torch.manual_seed(0) + fresh = model.build(base=8, depth=2) + frames = stack(3, size=48) + condition = conditioning(["available", "suspect", "available"]) + optimiser = torch.optim.Adam(fresh.parameters(), lr=1e-3) + + for _ in range(3): + optimiser.zero_grad() + out = fresh(frames, condition) + out.square().mean().backward() + optimiser.step() + + for name, parameter in fresh.named_parameters(): + assert parameter.grad is not None, f"{name} has no gradient" + assert parameter.grad.abs().sum() > 0, f"{name} gradient is identically zero" + + +def smooth_stack(count, size=64, seed=1): + """Spatially correlated frames; white noise tests nothing the real task needs.""" + generator = torch.Generator().manual_seed(seed) + coarse = torch.randn(1, count * BANDS, size // 8, size // 8, generator=generator) + smooth = torch.nn.functional.interpolate(coarse, size=(size, size), mode="bicubic", + align_corners=False) + return smooth.view(1, count, BANDS, size, size) + + +def train_a_little(net, frames, condition, target, steps=120, lr=3e-3): + optimiser = torch.optim.Adam(net.parameters(), lr=lr) + first = None + for step in range(steps): + optimiser.zero_grad() + out = net(frames, condition) + loss = (out - target).abs().mean() + loss.backward() + optimiser.step() + if step == 0: + first = loss.item() + return first, loss.item() + + +def test_model_learns_to_select_one_frame_from_the_stack(): + """The weight head's core job: when one candidate is the answer, put the weight + there. Fails in seconds if the softmax, masking or composition is miswired.""" + torch.manual_seed(0) + fresh = model.build(base=16, depth=2) + frames = smooth_stack(3) + first, last = train_a_little( + fresh, frames, conditioning(["available"] * 3), frames[:, 1].clone() + ) + assert last < first * 0.2, f"loss barely moved: {first:.4f} -> {last:.4f}" + + +def test_model_learns_a_per_band_gain_within_the_polish_bound(): + """The polish head's core job: the residual drift the daily fit leaves behind.""" + torch.manual_seed(0) + fresh = model.build(base=16, depth=2) + frames = smooth_stack(1) + per_band = torch.tensor([1.2, 0.85, 1.1, 0.9, 1.25, 0.8]).view(1, BANDS, 1, 1) + target = frames[:, 0] * per_band + first, last = train_a_little(fresh, frames, conditioning(["available"], + dts=[240.0]), target) + assert last < first * 0.3, f"loss barely moved: {first:.4f} -> {last:.4f}" diff --git a/tests/test_samples.py b/tests/test_samples.py new file mode 100644 index 0000000..356f42f --- /dev/null +++ b/tests/test_samples.py @@ -0,0 +1,415 @@ +import numpy as np +import pytest + +from conftest import solar_disc +from suvi import corruptions, dataset, paths, samples + +SIZE = dataset.SHARD_SIZE + + +class FakeShard: + """A shard-shaped object holding frames in memory, so tests need no files.""" + + def __init__(self, satellite, times, size=32): + self.satellite = satellite + self.size = size + self._frames = { + t: np.stack([ + solar_disc(size=size, radius=size // 3, peak=1.0 + band * 0.2 + t / 1e6) + for band in range(6) + ]).astype(np.float32) + for t in times + } + + def times(self): + return sorted(self._frames) + + def frames(self, time): + got = self._frames.get(time) + return None if got is None else got.copy() + + +def make_sampler(count=120, satellites=(16, 18), size=256, **kwargs): + base = 1715400000 // paths.CADENCE * paths.CADENCE + times = [base + i * paths.CADENCE for i in range(count)] + shards = {("d", s): FakeShard(s, times, size) for s in satellites} + return samples.Sampler(shards, satellites=satellites, **kwargs), times + + +# --------------------------------------------------------------------------- layout + + +def test_layout_excludes_the_frame_being_reconstructed(): + layout = samples.stack_layout((16, 18), 16) + assert (16, 0) not in layout + + +def test_layout_includes_the_counterpart_at_the_target_instant(): + """The single most valuable entry: a real observation of the right Sun, right time.""" + assert (18, 0) in samples.stack_layout((16, 18), 16) + assert (16, 0) in samples.stack_layout((16, 18), 18) + + +def test_layout_is_multi_scale_in_both_directions(): + layout = samples.stack_layout((16, 18), 16) + for offset in samples.OFFSETS: + assert (16, offset) in layout and (16, -offset) in layout + assert (18, offset) in layout and (18, -offset) in layout + + +def test_layout_is_stable(): + """A frame's position in the stack must always mean the same thing.""" + assert samples.stack_layout((16, 18), 16) == samples.stack_layout((16, 18), 16) + assert len(set(samples.stack_layout((16, 18), 16))) == len( + samples.stack_layout((16, 18), 16) + ) + + +def test_classes_cover_every_catalogued_mode_plus_clean(): + assert samples.CLASSES[0] == "clean" + assert set(samples.CLASSES[1:]) == set(corruptions.CATALOG) + assert len(samples.CLASS_INDEX) == len(corruptions.CATALOG) + 1 + + +# -------------------------------------------------------------------------- samples + + +def test_build_returns_a_full_stack(): + sampler, times = make_sampler() + sample = sampler.build(16, times[60]) + expected = len(samples.stack_layout((16, 18), 16)) + assert sample["frames"].shape[0] >= expected + assert sample["frames"].shape[1:] == (6, 256, 256) + assert sample["target"].shape == (6, 256, 256) + assert len(sample["states"]) == sample["frames"].shape[0] + assert len(sample["dts"]) == sample["frames"].shape[0] + assert len(sample["classes"]) == sample["frames"].shape[0] + + +def test_build_returns_nothing_without_a_target(): + sampler, times = make_sampler() + assert sampler.build(16, times[-1] + 99 * paths.CADENCE) is None + + +def test_target_is_never_in_its_own_stack(): + """Leakage of the most direct kind: the answer among the inputs.""" + sampler, times = make_sampler(damage_probability=0.0, drop_probability=0.0) + sample = sampler.build(16, times[60]) + for index, (dt, same) in enumerate(zip(sample["dts"], sample["same_satellite"])): + assert not (dt == 0.0 and same), f"stack entry {index} is the target itself" + + +def test_undamaged_sampler_marks_everything_available(): + # 600 slots, so even the +/-256 rungs of the exponential ladder land on frames. + sampler, times = make_sampler(count=600, size=32, + damage_probability=0.0, drop_probability=0.0) + sample = sampler.build(16, times[300]) + assert set(sample["states"]) == {"available"} + assert set(sample["classes"].tolist()) == {samples.CLASS_INDEX["clean"]} + + +def test_missing_frames_are_zeroed_and_labelled(): + """Off the end of the shard there is genuinely nothing.""" + sampler, times = make_sampler(damage_probability=0.0, drop_probability=0.0) + sample = sampler.build(16, times[1]) # offsets -2, -4, -16 fall off the start + missing = [i for i, s in enumerate(sample["states"]) if s == "missing"] + assert missing, "expected some offsets to fall outside the shard" + for index in missing: + assert np.all(sample["frames"][index] == 0.0) + + +def test_anchors_reach_past_the_sampled_offsets(): + """A long gap must still find something *real* to work from -- reach is measured + over frames that carry pixels, not over empty ladder rungs.""" + base = 1715400000 // paths.CADENCE * paths.CADENCE + times = [base] + [base + i * paths.CADENCE for i in range(100, 130)] + shards = {("d", 16): FakeShard(16, times, 256)} + sampler = samples.Sampler(shards, satellites=(16,), damage_probability=0.0, + drop_probability=0.0) + sample = sampler.build(16, base) + reach = max(abs(dt) for dt, state in zip(sample["dts"], sample["states"]) + if state == "available") + assert reach >= 100 * paths.CADENCE, "no anchor was added beyond the fixed offsets" + + +def test_damage_produces_suspect_frames_that_keep_their_pixels(): + """The design claim under test: a flagged frame is data, not a hole.""" + sampler, times = make_sampler(damage_probability=1.0, drop_probability=0.0, seed=3) + seen = set() + for time in times[20:80]: + sample = sampler.build(16, time) + seen.update(sample["states"]) + for index, state in enumerate(sample["states"]): + if state == "suspect": + assert np.any(sample["frames"][index] != 0.0) + assert sample["classes"][index] != samples.CLASS_INDEX["clean"] + assert "suspect" in seen + + +def test_no_signal_modes_are_presented_as_missing(): + """all_zero and friends leave nothing; calling them 'suspect' would be a lie.""" + sampler, times = make_sampler(damage_probability=1.0, drop_probability=0.0, seed=5) + for time in times[20:80]: + sample = sampler.build(16, time) + for index, state in enumerate(sample["states"]): + label = samples.CLASSES[sample["classes"][index]] + if label in samples.NO_SIGNAL: + assert state == "missing" + assert np.all(sample["frames"][index] == 0.0) + + +def test_every_catalogued_mode_can_be_applied(): + """Including the three that need a donor frame, which raise if given none.""" + sampler, times = make_sampler(damage_probability=1.0, drop_probability=0.0, seed=11) + applied = set() + for time in times[10:110]: + sample = sampler.build(16, time) + applied.update(samples.CLASSES[c] for c in sample["classes"].tolist()) + for mode, corruption in corruptions.CATALOG.items(): + if corruption.needs_donor: + assert mode in applied, f"{mode} never applied; donors may be unavailable" + + +def test_damage_is_coherent_across_bands(): + """One instrument makes all six bands, so a fault hits them together.""" + radiance = np.stack([solar_disc(size=32, radius=10, peak=1.0 + i) for i in range(6)]) + flipped = samples.decode_from_model( + samples._damage(samples.encode_for_model(radiance), "yaw_flip", seed=1, + severity=1.0) + ) + for band in range(6): + expected, _ = corruptions.apply_array("yaw_flip", radiance[band], 1, 1.0) + np.testing.assert_allclose(flipped[band], expected, rtol=2e-3, atol=1e-6) + + +def test_damage_operates_on_radiance_not_model_space(): + """A gain_shift multiplies a physical quantity; applying it to asinh values would + model an entirely different fault.""" + radiance = np.stack([np.full((16, 16), 2.0, np.float32) for _ in range(6)]) + damaged = samples.decode_from_model( + samples._damage(samples.encode_for_model(radiance), "gain_shift", seed=4, + severity=1.0) + ) + expected, _ = corruptions.apply_array("gain_shift", radiance[0], 4, 1.0) + ratio = float(damaged[0].mean() / radiance[0].mean()) + np.testing.assert_allclose(damaged[0], expected, rtol=2e-3) + assert abs(ratio - float(expected.mean() / 2.0)) < 1e-3 + + +def test_sampling_is_reproducible_from_its_seed(): + first, times = make_sampler(seed=7) + second, _ = make_sampler(seed=7) + a = first.build(16, times[60]) + b = second.build(16, times[60]) + assert a["states"] == b["states"] + np.testing.assert_array_equal(a["classes"], b["classes"]) + np.testing.assert_allclose(a["frames"], b["frames"]) + + +def test_targets_lists_both_satellites(): + sampler, times = make_sampler() + found = sampler.targets() + assert {s for s, _ in found} == {16, 18} + assert len(found) == 2 * len(times) + + +# ------------------------------------------------------------------------ encoding + + +def test_encode_decode_round_trips(): + values = np.array([[-3.7, -0.01, 0.0, 1e-3, 0.19, 50.0, 1231.0]], dtype=np.float32) + restored = samples.decode_from_model(samples.encode_for_model(values)) + np.testing.assert_allclose(restored, values, rtol=2e-3, atol=1e-7) + + +def test_encoding_matches_the_shard_transform(): + """Shards and live frames must land in the same space, or inference sees a shift.""" + frames = np.stack([solar_disc(size=1280, radius=386, peak=1.0 + i) for i in range(6)]) + from_shard = dataset.decode_frames(dataset.encode_frames(frames)) + direct = samples.decode_from_model(samples.encode_for_model(from_shard)) + np.testing.assert_allclose(direct, from_shard, rtol=2e-3, atol=1e-6) + + +def test_encoding_stays_in_range_for_absurd_input(): + coded = samples.encode_for_model(np.array([1e12, -1e12, np.nan, np.inf])) + assert np.isfinite(coded).all() + assert coded.max() <= 1.0 and coded.min() >= -1.0 + + +# -------------------------------------------------------------------------- tensors + + +def test_to_tensors_produces_what_the_model_expects(): + torch = pytest.importorskip("torch") + from suvi import align, model + + sampler, times = make_sampler(size=64) + sample = sampler.build(16, times[60]) + packed = samples.to_tensors(sample, torch) + + stack = sample["frames"].shape[0] + assert packed["stack"].shape == (stack, 6, 64, 64) + assert packed["condition"].shape == (stack, model.COND_DIM) + assert packed["gains"].shape == (stack, 6) + assert packed["offsets"].shape == (stack, 6) + assert packed["target"].shape == (6, 64, 64) + + condition = packed["condition"][None] + valid = (condition[..., 0] + condition[..., 2]).clamp(0, 1) + aligned = align.align_stack( + packed["stack"][None], packed["dts"][None], valid, packed["gains"][None], + packed["offsets"][None], packed["b0"][None], packed["radius"][None], + ) + net = model.build(base=8, depth=2) + out = net(aligned, condition) + assert out.shape == (1, 6, 64, 64) + assert torch.isfinite(out).all() + + +# ---------------------------------------------------------------------- transfers + + +def test_same_satellite_frames_carry_the_identity_transfer(): + sampler, times = make_sampler(damage_probability=0.0, drop_probability=0.0) + sample = sampler.build(16, times[60]) + for index, same in enumerate(sample["same_satellite"]): + if same: + np.testing.assert_array_equal(sample["gains"][index], np.ones(6)) + np.testing.assert_array_equal(sample["offsets"][index], np.zeros(6)) + + +def test_cross_satellite_frames_carry_a_fitted_transfer(): + """The counterpart must arrive with the day's calibration attached, so alignment + can put it on the target instrument's scale before the model sees it.""" + + class Scaled(FakeShard): + def frames(self, time): + got = super().frames(time) + return None if got is None else got * 0.8 + 0.01 + + base = 1715400000 // paths.CADENCE * paths.CADENCE + times = [base + i * paths.CADENCE for i in range(40)] + shards = {("d", 16): FakeShard(16, times, 32), ("d", 18): Scaled(18, times, 32)} + sampler = samples.Sampler(shards, damage_probability=0.0, drop_probability=0.0) + sample = sampler.build(16, times[20]) + + cross = [i for i, same in enumerate(sample["same_satellite"]) + if not same and sample["states"][i] == "available"] + assert cross, "no cross-satellite frame in the stack" + for index in cross: + np.testing.assert_allclose(sample["gains"][index], np.full(6, 1 / 0.8), + rtol=1e-3) + np.testing.assert_allclose(sample["offsets"][index], np.full(6, -0.01 / 0.8), + atol=1e-4) + + +def test_calibration_never_fits_against_the_target_itself(): + """A pair at the target instant would fit the transfer against the answer -- + the oracle gain the whole exercise exists to estimate honestly.""" + sampler, times = make_sampler() + for target in (times[0], times[60], times[-1]): + sampler._transfers.clear() + assert sampler._calibration("d", 16, 18, exclude=target) is not None + for (_, _, _, pair_time) in sampler._transfers: + assert pair_time != target + + +def test_outages_are_contiguous_runs(): + sampler, times = make_sampler() + rng = np.random.default_rng(2) + seen = 0 + for _ in range(50): + blocked = sampler._outages(rng, times[60]) + for satellite, interval in blocked.items(): + if not interval: + continue + seen += 1 + ordered = sorted(interval) + gaps = {b - a for a, b in zip(ordered, ordered[1:])} + assert gaps <= {paths.CADENCE}, "outage is not a contiguous run" + assert len(ordered) <= samples.MAX_OUTAGE_SLOTS + assert seen, "no episodic outage was ever drawn" + + +def test_dual_outages_take_down_both_satellites_sometimes(): + sampler, times = make_sampler() + rng = np.random.default_rng(3) + dual = single = 0 + for _ in range(200): + blocked = sampler._outages(rng, times[60]) + affected = [s for s, interval in blocked.items() if interval] + if len(affected) == 2: + assert blocked[16] == blocked[18], "dual outage must share one interval" + dual += 1 + elif len(affected) == 1: + single += 1 + assert dual > 10, f"dual outages too rare to train on ({dual}/200)" + assert single > 10, f"single-satellite outages too rare ({single}/200)" + + +def test_the_long_dual_outage_case_actually_occurs_in_training(): + """The regime the archive says is 11% of reality: both satellites dark around + the target, nearest real frame far away. With independent per-frame drops this + configuration had probability ~p^14 and was never trained.""" + sampler, times = make_sampler(count=600, size=32, seed=9) + starved = 0 + for time in times[280:380]: + sample = sampler.build(16, time) + if sample is None: + continue + near_same = [i for i, (dt, same) in enumerate(zip(sample["dts"], + sample["same_satellite"])) + if same and abs(dt) <= 16 * paths.CADENCE] + cross = [i for i, (dt, same) in enumerate(zip(sample["dts"], + sample["same_satellite"])) + if not same and abs(dt) <= 16 * paths.CADENCE] + if all(sample["states"][i] == "missing" for i in near_same) and \ + all(sample["states"][i] == "missing" for i in cross): + starved += 1 + assert starved >= 2, f"long dual-outage stacks essentially absent ({starved}/100)" + + +def test_anchors_land_outside_a_simulated_outage(): + sampler, times = make_sampler(count=600, size=32, + damage_probability=0.0, drop_probability=0.0) + exclude = frozenset(times[300 + k] for k in range(-150, 150)) + anchor = sampler._anchor(16, times[300], direction=1, exclude=exclude) + assert anchor is not None and anchor not in exclude + assert anchor >= times[300] + 150 * paths.CADENCE + + +def test_build_declines_when_nothing_in_the_stack_has_pixels(): + """Both satellites out for the whole window: the contract is None, not a + fabricated frame.""" + base = 1715400000 // paths.CADENCE * paths.CADENCE + shards = {("d", 16): FakeShard(16, [base], 32)} + sampler = samples.Sampler(shards, satellites=(16, 18), damage_probability=0.0, + drop_probability=0.0) + assert sampler.build(16, base) is None + + +def test_build_reports_solar_geometry(): + sampler, times = make_sampler() + sample = sampler.build(16, times[60]) + assert -0.13 < sample["b0"] < 0.13 # +/-7.25 deg in radians + assert 0.28 < sample["radius"] < 0.32 # disc fraction of the frame + + +def test_severity_is_drawn_per_sample_not_fixed_per_frame(): + """Fixed rates left barely half of every stack clean, so the model never saw the + easy case and learned to hedge. Real outages are episodic: some stacks should come + through almost untouched and some should be wrecked.""" + # A wide window, so every ladder rung exists and the clean fraction measures + # injected damage rather than the shard's edges. + sampler, times = make_sampler(count=600, size=32, damage_probability=0.30, + drop_probability=0.40, seed=17) + clean_fractions = [] + for time in times[280:380]: + sample = sampler.build(16, time) + states = sample["states"] + clean_fractions.append(sum(s == "available" for s in states) / len(states)) + + assert max(clean_fractions) > 0.95, "no sample came through nearly clean" + assert min(clean_fractions) < 0.75, "no sample was substantially degraded" + spread = max(clean_fractions) - min(clean_fractions) + assert spread > 0.3, f"severity barely varied between samples (spread {spread:.2f})" diff --git a/tests/test_train.py b/tests/test_train.py new file mode 100644 index 0000000..989e41c --- /dev/null +++ b/tests/test_train.py @@ -0,0 +1,289 @@ +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +import train # noqa: E402 +from suvi import dataset, metrics, model, paths, samples # noqa: E402 + + +# ------------------------------------------------------------------ display mapping + + +def test_hard_display_map_matches_the_metric_the_bench_scores(): + """With eps=0 the torch map must *be* metrics.to_display, or the reported PSNR + stops being comparable to the published filler table.""" + rng = np.random.default_rng(0) + radiance = rng.uniform(-1.0, 60.0, (1, 6, 16, 16)).astype(np.float32) + coded = torch.from_numpy(samples.encode_for_model(radiance)) + + got = train.display_map(coded, eps=0.0).numpy() + for index, wavelength in enumerate(paths.WAVELENGTHS): + expected = metrics.to_display(radiance[0, index], wavelength) + np.testing.assert_allclose(got[0, index], expected, atol=2e-3) + + +def test_display_map_still_sends_zero_to_zero_when_softened(): + """The softening must not lift the black level: (0+eps)^g - eps^g == 0.""" + coded = samples.encode_for_model(np.zeros((1, 6, 4, 4), np.float32)) + out = train.display_map(torch.from_numpy(coded)) + assert float(out.abs().max()) < 1e-6 # fp32 rounding only; 1/255 is 4e-3 + + +def test_display_map_saturates_without_nan(): + for value in (-1.0, 1.0): + coded = torch.full((1, 6, 4, 4), value, requires_grad=True) + out = train.display_map(coded) + out.sum().backward() + assert torch.isfinite(out).all() and torch.isfinite(coded.grad).all() + + +# --------------------------------------------------------------- loss conditioning + + +def test_the_loss_gradient_is_bounded_everywhere(): + """The property whose absence sank three training runs. + + The old display-space loss had |dL/dpred| of exactly 0 across most of the frame + (below the vmin clamp) and a 25x spike with a hard discontinuity at the + threshold; clip_grad_norm then renormalised every update onto that annulus. + Here the whole coded range, including the sub-vmin corona and the threshold + itself, must see a per-pixel gradient that is neither zero over wide regions nor + orders of magnitude apart. + """ + # Sweep the entire coded range, so the threshold annulus of every band is in, + # with a *constant* residual: the per-pixel gradient then measures the loss + # landscape's Jacobian, not the accident of which pixels happened to be right. + target = torch.linspace(-1.0, 1.0, 64 * 64).reshape(1, 1, 64, 64).repeat(1, 6, 1, 1) + prediction = (target + 0.05).requires_grad_(True) + + train.criterion(prediction, target).backward() + per_pixel = prediction.grad.abs() * prediction.numel() + assert torch.isfinite(prediction.grad).all() + assert float(per_pixel.max()) < 100.0, "a gradient spike survived the redesign" + # The asinh term supervises everything: no dead zones anywhere. + assert float(per_pixel.min()) > 0.1, "part of the frame is unsupervised" + + +def test_charbonnier_is_zero_for_a_perfect_match(): + x = torch.rand(4, 6, 8, 8) + assert train.charbonnier(x, x) == pytest.approx(train.CHARBONNIER_EPS, abs=1e-6) + + +def test_charbonnier_is_less_outlier_driven_than_squared_error(): + """A saturated pixel must not be able to dominate the whole frame's loss.""" + truth = torch.zeros(1, 1, 8, 8) + spike = truth.clone() + spike[0, 0, 0, 0] = 50.0 + spread = truth + 0.1 + + charbonnier_ratio = float(train.charbonnier(spike, truth) + / train.charbonnier(spread, truth)) + squared_ratio = float(((spike - truth) ** 2).mean() / ((spread - truth) ** 2).mean()) + assert charbonnier_ratio < 10 + assert squared_ratio > 1000 + + +def test_gradient_loss_punishes_blur(): + """The term that makes a sharp answer beat a hedged smooth one.""" + truth = torch.zeros(1, 1, 32, 32) + truth[:, :, 8:24, 8:24] = 1.0 + blurred = torch.nn.functional.avg_pool2d(truth, 5, stride=1, padding=2) + assert train.gradient_loss(blurred, truth) > train.gradient_loss(truth, truth) + + +def test_display_psnr_matches_the_bench_definition(): + """Per band, then averaged -- the pooled-MSE PSNR reads ~4 dB lower (Jensen) and + would not be comparable to the published baselines.""" + rng = np.random.default_rng(1) + radiance = rng.uniform(0.0, 30.0, (6, 16, 16)).astype(np.float32) + noisy = radiance + rng.normal(0, 0.5, radiance.shape).astype(np.float32) + + expected = [] + for index, wavelength in enumerate(paths.WAVELENGTHS): + a = metrics.to_display(radiance[index], wavelength) + b = metrics.to_display(noisy[index], wavelength) + expected.append(10.0 * np.log10(1.0 / max(((a - b) ** 2).mean(), 1e-12))) + + got = train.display_psnr( + torch.from_numpy(samples.encode_for_model(noisy))[None], + torch.from_numpy(samples.encode_for_model(radiance))[None], + ) + assert got == pytest.approx(float(np.mean(expected)), abs=0.1) + + +def test_display_psnr_of_identical_inputs_is_finite(): + x = torch.rand(1, 6, 8, 8) * 0.5 + assert np.isfinite(train.display_psnr(x, x)) + + +# ------------------------------------------------------------------------- collate + + +def item(depth, bands=6, size=16): + return { + "stack": torch.randn(depth, bands, size, size), + "condition": torch.stack( + [model.frame_conditioning("available", True, 240.0) for _ in range(depth)] + ), + "dts": torch.full((depth,), 240.0), + "gains": torch.ones(depth, bands), + "offsets": torch.zeros(depth, bands), + "b0": torch.tensor(0.05), + "radius": torch.tensor(0.3), + "target": torch.randn(bands, size, size), + } + + +def test_collate_pads_ragged_stacks(): + """Stack depth varies with how many anchors were found near a shard's edge.""" + batch = train.collate([item(5), item(7), item(6)]) + assert batch["stack"].shape[:2] == (3, 7) + assert batch["condition"].shape == (3, 7, model.COND_DIM) + assert batch["gains"].shape == (3, 7, 6) + assert batch["target"].shape == (3, 6, 16, 16) + assert batch["b0"].shape == (3,) + + +def test_collate_marks_padding_as_missing_with_identity_transfer(): + """Padding must take zero weight and pass alignment untouched, or the batch + shape would change the answer.""" + batch = train.collate([item(3), item(6)]) + assert torch.all(batch["condition"][0, 3:, 1] == 1.0) # 'missing' one-hot + assert torch.all(batch["stack"][0, 3:] == 0.0) + assert torch.all(batch["gains"][0, 3:] == 1.0) + assert torch.all(batch["offsets"][0, 3:] == 0.0) + + +def test_padding_does_not_change_the_prediction(): + """The property that makes ragged batching safe at all.""" + torch.manual_seed(0) + net = model.build(base=8, depth=2).eval() + short = item(3, size=32) + alone = train.collate([short]) + padded = train.collate([short, item(6, size=32)]) + with torch.no_grad(): + a = net(alone["stack"], alone["condition"]) + b = net(padded["stack"][:1], padded["condition"][:1]) + assert torch.allclose(a, b, atol=1e-5) + + +# ---------------------------------------------------------------------- data feed + + +class MemoryShard: + def __init__(self, satellite, times, size=64): + from conftest import solar_disc + + self.satellite = satellite + self._frames = { + t: np.stack([solar_disc(size=size, radius=size // 3, peak=1.0 + b * 0.2) + for b in range(6)]).astype(np.float32) + for t in times + } + + def times(self): + return sorted(self._frames) + + def frames(self, time): + got = self._frames.get(time) + return None if got is None else got.copy() + + +def shard_set(monkeypatch, tmp_path, days, deterministic=False, seed=0): + base = 1715400000 // paths.CADENCE * paths.CADENCE + times = [base + i * paths.CADENCE for i in range(40)] + monkeypatch.setattr(dataset, "Shard", lambda path: MemoryShard(16, times)) + real_exists = train.os.path.exists + # Pretend only the shard files exist; a blanket True breaks cv2's lazy loader. + monkeypatch.setattr( + train.os.path, "exists", + lambda p: True if str(p).startswith(str(tmp_path)) else real_exists(p), + ) + return train.ShardSet(str(tmp_path), days, satellites=(16, 18), + deterministic=deterministic, seed=seed) + + +def test_shard_set_yields_model_ready_tensors(monkeypatch, tmp_path): + data = shard_set(monkeypatch, tmp_path, ["2024-01-01"]) + sample = data[0] + assert sample["stack"].shape[1:] == (6, 64, 64) + assert sample["condition"].shape[1] == model.COND_DIM + assert sample["gains"].shape == (sample["stack"].shape[0], 6) + assert torch.isfinite(sample["stack"]).all() + assert torch.isfinite(sample["target"]).all() + + +def test_validation_damage_is_reproducible(monkeypatch, tmp_path): + """A val number that moves because the damage moved measures nothing.""" + first = shard_set(monkeypatch, tmp_path, ["2024-01-01"], deterministic=True, seed=1) + second = shard_set(monkeypatch, tmp_path, ["2024-01-01"], deterministic=True, seed=1) + for index in (0, 5, 11): + assert torch.allclose(first[index]["stack"], second[index]["stack"]) + + +def test_training_damage_varies_between_epochs(monkeypatch, tmp_path): + """The reason shards store clean frames: unlimited corruption variety.""" + data = shard_set(monkeypatch, tmp_path, ["2024-01-01"], deterministic=False) + assert not torch.allclose(data[3]["stack"], data[3]["stack"]) + + +def test_shard_set_refuses_to_run_on_nothing(tmp_path): + with pytest.raises(SystemExit, match="no shards"): + train.ShardSet(str(tmp_path), ["2024-01-01"]) + + +# ------------------------------------------------------------------ training steps + + +def loader_for(monkeypatch, tmp_path): + data = shard_set(monkeypatch, tmp_path, ["2024-01-01"]) + return torch.utils.data.DataLoader(data, batch_size=2, num_workers=0, + collate_fn=train.collate) + + +def test_a_training_step_is_finite_end_to_end(monkeypatch, tmp_path): + """Data feed, alignment, forward, loss, backward: all wired together.""" + loader = loader_for(monkeypatch, tmp_path) + net = model.build(base=8, depth=2) + optimiser = torch.optim.AdamW(net.parameters(), lr=1e-4) + device = torch.device("cpu") + + for index, batch in enumerate(loader): + if index >= 2: + break + aligned, condition, target = train.aligned_batch(batch, device) + prediction = net(aligned, condition) + loss = train.criterion(prediction, target) + assert torch.isfinite(loss) + optimiser.zero_grad() + loss.backward() + grad = float(torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)) + assert np.isfinite(grad) + optimiser.step() + + +def test_evaluate_leaves_weights_alone(monkeypatch, tmp_path): + """Validation must not train, or the held-out split stops being held out.""" + loader = loader_for(monkeypatch, tmp_path) + net = model.build(base=8, depth=2) + before = [p.detach().clone() for p in net.parameters()] + stats = train.evaluate(net, loader, torch.device("cpu")) + assert np.isfinite(stats["loss"]) and np.isfinite(stats["psnr"]) + for old, new in zip(before, net.parameters()): + assert torch.equal(old, new) + + +def test_validation_subset_spans_the_whole_set(monkeypatch, tmp_path): + """A shortened val pass must not become 'the first N targets': `targets()` is + sorted by (satellite, time), so a prefix is the earliest days of one satellite.""" + data = shard_set(monkeypatch, tmp_path, ["2024-01-01"]) + wanted = 8 + stride = len(data) / wanted + subset = torch.utils.data.Subset(data, [int(i * stride) for i in range(wanted)]) + indices = subset.indices + assert len(set(indices)) == wanted + assert max(indices) > len(data) * 0.8, "subset does not reach the end of the set" + assert min(indices) < len(data) * 0.2, "subset does not start near the beginning" + satellites = {data.items[i][0] for i in indices} + assert len(satellites) > 1, "subset covers only one satellite" diff --git a/train.py b/train.py new file mode 100644 index 0000000..9523e23 --- /dev/null +++ b/train.py @@ -0,0 +1,422 @@ +#!/usr/bin/env python +"""Train the learned filler. Runs on the training host, inside the ROCm container. + + train.py --shards /data/shards --manifest /data/shards/manifest.json --out /data/runs/v4 + +The loss lives in **asinh space**, where the shards already are: it is bounded, +smooth, and supervises every pixel. The previous trainer optimised the display +mapping instead, which clamps at a per-band `vmin` sitting *above* the median pixel +-- most of the frame had exactly zero gradient, and the annulus straddling the +threshold carried a 25x spike with a hard discontinuity (measured: |dL/dpred| of 0 / +49 / 2-7 below, at, and above the threshold). Three runs diverged on that loss. +A display-space term is kept, softened so its worst slope is ~10, because the bench +scores display PSNR and the loss should feel the same per-band weighting; reported +PSNR uses the *exact* hard bench mapping so the numbers are comparable to the +baseline table (optical_flow 44.79 dB at gap 1, 34.47 at gap 300). + +Validation runs on a fixed, evenly spaced subset of held-out days **every 250 +steps**, not per epoch: the last collapse happened entirely inside one 41-minute +epoch and was invisible until it ended. Training aborts automatically if validation +PSNR falls well below its running best twice in a row. +""" + +import argparse +import json +import os +import random +import sys +import time + +import numpy as np +import torch + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from suvi import align, dataset, metrics, model, samples + +#: Weight of the gradient-matching term (asinh space): penalises blur, not just error. +GRADIENT_WEIGHT = 0.5 +#: Weight of the softened display-space term. +DISPLAY_WEIGHT = 0.25 +#: Charbonnier knee; small enough not to blunt real errors. +CHARBONNIER_EPS = 1e-3 +#: Softening of the display gamma for the *loss*. At 5e-3 the worst-case slope of +#: the mapping is ~10 against the 37,500 of the 1e-8 the old trainer used; the +#: mapping still sends 0 to 0 exactly. The reported PSNR never uses this -- metrics +#: use the hard bench mapping. +DISPLAY_EPS = 5e-3 +#: Steps between validation passes. +EVAL_EVERY = 250 +#: Abort when validation PSNR sits this far (dB) below its best, twice in a row. +ABORT_DROP_DB = 1.5 + + +def display_map(coded, eps=DISPLAY_EPS): + """Model space -> display space for all six bands, differentiably when eps > 0. + + The torch equivalent of :func:`suvi.metrics.to_display`; with ``eps=0.0`` it *is* + that mapping, for metrics under no_grad. + """ + radiance = torch.sinh(coded.clamp(-1, 1) * dataset.ASINH_RANGE) * dataset.ASINH_SCALE + out = [] + for index, wavelength in enumerate(samples.paths.WAVELENGTHS): + vmin, vmax, gamma = metrics.DISPLAY_MAPPING[wavelength] + scaled = ((radiance[:, index] - vmin) / vmax).clamp(0.0, 1.0) + out.append((scaled + eps) ** gamma - eps**gamma) + return torch.stack(out, dim=1) + + +def charbonnier(prediction, truth, eps=CHARBONNIER_EPS): + return torch.sqrt((prediction - truth) ** 2 + eps**2).mean() + + +def gradient_loss(prediction, truth): + """Match image gradients, so a blurred answer is penalised as well as a wrong one.""" + loss = 0.0 + for axis in (-1, -2): + loss = loss + charbonnier(prediction.diff(dim=axis), truth.diff(dim=axis)) + return loss / 2 + + +def criterion(prediction, target): + """The full training loss, in spaces whose gradients are bounded everywhere.""" + loss = charbonnier(prediction, target) + loss = loss + GRADIENT_WEIGHT * gradient_loss(prediction, target) + loss = loss + DISPLAY_WEIGHT * charbonnier( + display_map(prediction), display_map(target) + ) + return loss + + +def display_psnr(prediction, target): + """Mean per-band display PSNR, the bench's definition, in dB. + + Per band then averaged -- PSNR of the pooled MSE reads ~4 dB lower (Jensen) and + would not be comparable to the published filler table. + """ + with torch.no_grad(): + shown = display_map(prediction.float(), eps=0.0) + truth = display_map(target.float(), eps=0.0) + mse = ((shown - truth) ** 2).mean(dim=(-2, -1)).clamp(min=1e-12) + return float((10.0 * torch.log10(1.0 / mse)).mean()) + + +# ----------------------------------------------------------------------- data feed + + +class ShardSet(torch.utils.data.Dataset): + """Samples drawn from the shards of one split. + + Shards hold clean frames; damage is drawn fresh here on every access, so two + epochs never see the same corruption of the same slot. Validation pins its seed + to the slot index instead, because a validation number that moves because the + *damage* moved says nothing about the model. + """ + + def __init__(self, directory, days, satellites=(16, 18), deterministic=False, + seed=0): + self.deterministic = deterministic + self.seed = seed + shards = {} + for day in days: + for satellite in satellites: + path = os.path.join(directory, dataset.shard_name(day, satellite)) + if os.path.exists(path): + shards[(day, satellite)] = dataset.Shard(path) + if not shards: + raise SystemExit(f"no shards for {len(days)} days under {directory}") + self.sampler = samples.Sampler(shards, satellites=satellites, seed=seed) + self.items = self.sampler.targets() + + def __len__(self): + return len(self.items) + + def __getitem__(self, index): + satellite, when = self.items[index] + rng = (np.random.default_rng(self.seed * 1_000_003 + index) + if self.deterministic else self.sampler.rng) + sample = self.sampler.build(satellite, when, rng) + if sample is None: # target absent; fall back to a neighbour + return self[(index + 1) % len(self.items)] + return samples.to_tensors(sample, torch) + + +def collate(batch): + """Pad stacks to a common depth. + + Depth varies because the anchors are the nearest usable frames, and near the + edge of a shard there may be one or none. Padding with `missing` frames is + exactly right -- the model gives them zero weight -- and keeps the batch + rectangular. Padded rows carry the identity transfer so alignment ignores them. + """ + depth = max(item["stack"].shape[0] for item in batch) + out = {} + for item in batch: + short = depth - item["stack"].shape[0] + if short: + item = dict(item) + item["stack"] = torch.cat( + [item["stack"], torch.zeros(short, *item["stack"].shape[1:])] + ) + filler = model.frame_conditioning("missing", True, 0.0) + item["condition"] = torch.cat([item["condition"], filler.repeat(short, 1)]) + item["dts"] = torch.cat([item["dts"], torch.zeros(short)]) + item["gains"] = torch.cat([item["gains"], torch.ones(short, model.BANDS)]) + item["offsets"] = torch.cat([item["offsets"], + torch.zeros(short, model.BANDS)]) + for key, value in item.items(): + out.setdefault(key, []).append(value) + return {key: torch.stack(value) for key, value in out.items()} + + +def aligned_batch(batch, device): + """Move a batch to the device and run the deterministic alignment, in fp32. + + Alignment is physics, not learning: it runs outside autocast because a bf16 + ``sinh`` at the top of the asinh range keeps only two decimal digits of a + radiance the gain is about to multiply. + """ + stack = batch["stack"].to(device, non_blocking=True).float() + condition = batch["condition"].to(device, non_blocking=True) + target = batch["target"].to(device, non_blocking=True).float() + valid = (condition[..., 0] + condition[..., 2]).clamp(0, 1) + aligned = align.align_stack( + stack, batch["dts"].to(device, non_blocking=True), valid, + batch["gains"].to(device, non_blocking=True), + batch["offsets"].to(device, non_blocking=True), + batch["b0"].to(device, non_blocking=True), + batch["radius"].to(device, non_blocking=True), + ) + return aligned, condition, target + + +def evaluate(net, loader, device): + """Loss and bench-comparable PSNR over a whole loader.""" + net.eval() + totals = {"loss": 0.0, "psnr": 0.0, "n": 0} + with torch.no_grad(): + for batch in loader: + aligned, condition, target = aligned_batch(batch, device) + with torch.autocast(device.type, dtype=torch.bfloat16, + enabled=device.type == "cuda"): + prediction = net(aligned, condition) + prediction = prediction.float() + count = aligned.shape[0] + totals["loss"] += float(criterion(prediction, target)) * count + totals["psnr"] += display_psnr(prediction, target) * count + totals["n"] += count + net.train() + n = max(totals["n"], 1) + return {"loss": totals["loss"] / n, "psnr": totals["psnr"] / n} + + +# ------------------------------------------------------------------------ training + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--shards", required=True) + parser.add_argument("--manifest", required=True) + parser.add_argument("--out", required=True) + parser.add_argument("--epochs", type=int, default=20) + parser.add_argument("--batch", type=int, default=8) + parser.add_argument("--lr", type=float, default=1e-4) + parser.add_argument("--base", type=int, default=32) + parser.add_argument("--depth", type=int, default=3) + parser.add_argument("--workers", type=int, default=8) + parser.add_argument("--steps-per-epoch", type=int, default=None) + parser.add_argument("--val-samples", type=int, default=64) + parser.add_argument("--eval-every", type=int, default=EVAL_EVERY) + parser.add_argument("--max-steps", type=int, default=None, + help="stop after this many optimiser steps (gates and probes)") + parser.add_argument("--overfit", type=int, default=None, + help="train and validate on this many frozen held-in samples " + "-- the overfit gate; loss must approach zero") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--resume", default=None) + args = parser.parse_args(argv) + + torch.manual_seed(args.seed) + random.seed(args.seed) + np.random.seed(args.seed) + torch.backends.cudnn.benchmark = True # worth ~2x on this GPU + + manifest = json.load(open(args.manifest)) + splits = manifest["splits"] + leaks = dataset.check_split(splits) # raises if a day is in two splits + satellites = tuple(manifest.get("satellites", (16, 18))) + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(f"device: {device}" + + (f" ({torch.cuda.get_device_properties(0).gcnArchName})" + if device.type == "cuda" else "")) + print(f"manifest {manifest['digest']}: " + + ", ".join(f"{k} {len(v)}d" for k, v in sorted(splits.items()))) + if leaks: + print(f" {len(leaks)} split boundaries on consecutive days: {leaks[:3]}") + + if args.overfit: + # The overfit gate: a handful of *frozen* samples (deterministic damage), and + # validation on those same samples. A model that cannot drive this loss + # toward zero has a wiring or conditioning problem that a long run will not + # cure -- two of the three failed runs would have failed this in minutes. + held_in = ShardSet(args.shards, splits["train"], satellites, + deterministic=True, seed=args.seed) + stride = max(len(held_in) // args.overfit, 1) + train_set = torch.utils.data.Subset( + held_in, [i * stride for i in range(args.overfit)] + ) + val_set = train_set + else: + train_set = ShardSet(args.shards, splits["train"], satellites, seed=args.seed) + val_set = ShardSet(args.shards, splits["val"], satellites, deterministic=True, + seed=1) + # An evenly spaced subset, so a small validation pass still spans every val + # day and both satellites, and is the same subset every time. + if args.val_samples and args.val_samples < len(val_set): + stride = len(val_set) / args.val_samples + val_set = torch.utils.data.Subset( + val_set, [int(i * stride) for i in range(args.val_samples)] + ) + print(f" train {len(train_set)} targets, val {len(val_set)}") + + loader_args = dict(batch_size=args.batch, num_workers=args.workers, + collate_fn=collate, pin_memory=device.type == "cuda", + persistent_workers=args.workers > 0) + train_loader = torch.utils.data.DataLoader(train_set, shuffle=True, drop_last=True, + **loader_args) + val_loader = torch.utils.data.DataLoader(val_set, shuffle=False, **loader_args) + + net = model.build(base=args.base, depth=args.depth).to(device) + net = net.to(memory_format=torch.channels_last) + print(f" {sum(p.numel() for p in net.parameters()) / 1e6:.1f}M parameters") + + optimiser = torch.optim.AdamW(net.parameters(), lr=args.lr, weight_decay=1e-4) + steps_per_epoch = args.steps_per_epoch or len(train_loader) + total_steps = args.epochs * steps_per_epoch + schedule = torch.optim.lr_scheduler.OneCycleLR( + optimiser, max_lr=args.lr, total_steps=total_steps, pct_start=0.2 + ) + + step = 0 + start_epoch = 0 + if args.resume and os.path.exists(args.resume): + state = torch.load(args.resume, map_location=device) + net.load_state_dict(state["model"]) + optimiser.load_state_dict(state["optimiser"]) + start_epoch = state["epoch"] + 1 + step = state.get("step", start_epoch * steps_per_epoch) + for _ in range(min(step, total_steps - 1)): + schedule.step() + print(f" resumed from {args.resume} at epoch {start_epoch}, step {step}") + + os.makedirs(args.out, exist_ok=True) + history_path = os.path.join(args.out, "history.jsonl") + + def checkpoint(name, epoch): + state = {"model": net.state_dict(), "optimiser": optimiser.state_dict(), + "epoch": epoch, "step": step, "args": vars(args), + "manifest": manifest["digest"]} + torch.save(state, os.path.join(args.out, name)) + + def record(row): + with open(history_path, "a") as handle: + handle.write(json.dumps(row) + "\n") + + # The untrained model is the prior-weighted blend of aligned candidates -- a + # meaningful policy, and the floor every later eval is judged against. + baseline = evaluate(net, val_loader, device) + print(f"step {step:>6} val {baseline['loss']:.5f} ({baseline['psnr']:.2f} dB) " + f"[untrained]", flush=True) + record({"step": step, "val_loss": baseline["loss"], "val_psnr": baseline["psnr"], + "untrained": True}) + best = baseline["psnr"] + if not os.path.exists(os.path.join(args.out, "best.pt")): + checkpoint("best.pt", epoch=-1) # never clobber a resumed run's best + drops = 0 + + window = {"loss": 0.0, "grad": 0.0, "grad_max": 0.0, "n": 0} + started = time.time() + for epoch in range(start_epoch, args.epochs): + epoch_steps = 0 + for batch in train_loader: + # --steps caps the *loop*, not just the LR schedule. The first run + # honoured it only in the schedule, so an "epoch" silently became the + # whole 4,312-batch split and last.pt was never written before the + # machine stalled at step ~1,700 -- leaving nothing to resume from. + if epoch_steps >= steps_per_epoch: + break + if step >= total_steps or (args.max_steps and step >= args.max_steps): + break + epoch_steps += 1 + aligned, condition, target = aligned_batch(batch, device) + with torch.autocast(device.type, dtype=torch.bfloat16, + enabled=device.type == "cuda"): + prediction = net(aligned, condition) + loss = criterion(prediction.float(), target) + + optimiser.zero_grad(set_to_none=True) + loss.backward() + grad = float(torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)) + optimiser.step() + schedule.step() + step += 1 + + window["loss"] += float(loss.detach()) + window["grad"] += grad + window["grad_max"] = max(window["grad_max"], grad) + window["n"] += 1 + + if step % args.eval_every == 0: + val = evaluate(net, val_loader, device) + n = max(window["n"], 1) + row = { + "step": step, "epoch": epoch, + "train_loss": window["loss"] / n, + "val_loss": val["loss"], "val_psnr": val["psnr"], + "grad_mean": window["grad"] / n, "grad_max": window["grad_max"], + "lr": optimiser.param_groups[0]["lr"], + "seconds": time.time() - started, + } + record(row) + print(f"step {step:>6} train {row['train_loss']:.5f} " + f"val {val['loss']:.5f} ({val['psnr']:.2f} dB) " + f"grad {row['grad_mean']:.3f}/{row['grad_max']:.2f} " + f"lr {row['lr']:.2e} {row['seconds'] / 60:.1f} min", flush=True) + window = {"loss": 0.0, "grad": 0.0, "grad_max": 0.0, "n": 0} + started = time.time() + if device.type == "cuda": + # The composition holds multi-GB full-resolution tensors, and on + # unified GTT memory the caching allocator's fragmentation growth + # is host RAM disappearing: the first full run crept to 75 GB + # over ~1,700 steps and stalled the machine into swap. Releasing + # the cache every eval costs milliseconds and caps the creep. + torch.cuda.empty_cache() + + # A resumable point every eval, not just every epoch: the stall cost + # 1,700 steps because the only last.pt cadence was the epoch end. + checkpoint("last.pt", epoch) + if val["psnr"] > best: + best = val["psnr"] + checkpoint("best.pt", epoch) + drops = 0 + elif val["psnr"] < best - ABORT_DROP_DB: + drops += 1 + if drops >= 2: + print(f"ABORT: validation {val['psnr']:.2f} dB has sat " + f">{ABORT_DROP_DB} dB below the best ({best:.2f} dB) " + "on two consecutive evals.", flush=True) + checkpoint("last.pt", epoch) + return 1 + else: + drops = 0 + checkpoint("last.pt", epoch) + if step >= total_steps or (args.max_steps and step >= args.max_steps): + break + print(f"best validation display PSNR {best:.2f} dB") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/verify_shards.py b/verify_shards.py new file mode 100644 index 0000000..0d1dfbf --- /dev/null +++ b/verify_shards.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python3 +"""Check a shard set before spending a day training on it. + +Cheap insurance. The extraction is a four-hour pass over a filesystem that has failed +mid-run before, and every failure mode it has -- a truncated push, a day of download +stubs, a split that leaks -- is silent at training time and expensive to discover from a +loss curve. +""" + +import argparse +import json +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from suvi import dataset + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--shards", required=True) + parser.add_argument("--manifest", default=None) + parser.add_argument("--satellites", default="16,18") + args = parser.parse_args(argv) + + satellites = tuple(int(s) for s in args.satellites.split(",")) + manifest_path = args.manifest or os.path.join(args.shards, "manifest.json") + manifest = json.load(open(manifest_path)) + splits = manifest["splits"] + + print(f"manifest {manifest['digest']} size={manifest['size']} " + f"asinh={manifest.get('asinh_scale')}/{manifest.get('asinh_range')}") + + leaks = dataset.check_split(splits) + print("split: " + ", ".join(f"{k} {len(v)}d" for k, v in sorted(splits.items()))) + print(f" adjacent boundaries: {len(leaks)}{' ' + str(leaks[:3]) if leaks else ''}") + + problems = [] + totals = {} + for name, days in sorted(splits.items()): + slots = empty = missing = 0 + for day in days: + for satellite in satellites: + path = os.path.join(args.shards, dataset.shard_name(day, satellite)) + if not os.path.exists(path): + missing += 1 + problems.append(f"missing shard {os.path.basename(path)}") + continue + try: + with dataset.Shard(path) as shard: + times = shard.times() + usable = sum(1 for t in times if shard.raw(t) is not None) + slots += usable + if usable == 0: + empty += 1 + problems.append(f"empty shard {os.path.basename(path)}") + # Spot-check one decode per shard. + for when in times[:1]: + frames = shard.coded(when) + if frames is None: + continue + if frames.shape != (6, dataset.SHARD_SIZE, dataset.SHARD_SIZE): + problems.append(f"{os.path.basename(path)}: {frames.shape}") + if abs(frames).max() > 1.0001: + problems.append( + f"{os.path.basename(path)}: out of range " + f"{abs(frames).max():.3f}") + except (ValueError, OSError) as error: + problems.append(f"{os.path.basename(path)}: {error}") + totals[name] = slots + print(f" {name:>5}: {slots:>6} usable slots" + + (f", {empty} empty shards" if empty else "") + + (f", {missing} MISSING" if missing else "")) + + print(f"\ntotal usable slots: {sum(totals.values())}") + if problems: + print(f"\n{len(problems)} problems:") + for line in problems[:20]: + print(f" {line}") + return 1 + print("no problems found") + return 0 + + +if __name__ == "__main__": + sys.exit(main())