Initial ML model training progress. Commit before moving to next stage with larger training corpus and longer train time.

.gitignore

@@ -2,3 +2,6 @@
 .venv
 __pycache__/
 *.pyc
+
+gpu_system_ssh_info
+.gpu_paused.json

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)

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())

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())

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())

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
+The archive is on a virtiofs share whose daemon holds a host file descriptor per inode
-inode the guest has looked up.  Anything that reads a lot of the archive leaves those
+the guest has looked up.  Anything that reads a lot of the archive leaves those inodes
-inodes cached, and the handles with them, until something forces the guest to evict
+cached, and the handles with them, until something forces the guest to evict them -- at
-them -- at which point the daemon can run out and the mount starts refusing *every*
+which point the daemon can run out and the mount starts refusing *every* open, which
-open, which surfaces across the whole machine as "too many open files in system".
+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
+The pipeline's own bulk jobs reclaim as they go and again when they finish, so this
-this should rarely be needed.  It exists for when something else has loaded the cache,
+should rarely be needed.  It exists for when something else has loaded the cache:
-or to check where things stand:
 
     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
+**This file deliberately imports nothing from `suvi`, and nothing outside the standard
-allocating several GiB.  Running under sudo avoids that entirely.
+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__)))
+#: 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
-from suvi import vfs
+#: 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'}: "
+    print(f"  {'reclaimed' if freed else 'no change'}: {before - after:+.2f} GiB")
-          f"{before - after:+.2f} GiB")
     return 0
 
 

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)

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."""
+    """Mispointing: the solar disc sits off centre.
-    shift = int(200 * severity)
+
+    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)

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="<u2").tobytes(), level
+        )
+        index[str(t_start)] = [offset, len(payload)]
+        blobs.append(payload)
+        offset += len(payload)
+
+    header = json.dumps({
+        "day": day,
+        "satellite": int(satellite),
+        "wavelengths": list(wavelengths),
+        "size": SHARD_SIZE,
+        "asinh_scale": ASINH_SCALE,
+        "asinh_range": ASINH_RANGE,
+        "index": index,
+    }, separators=(",", ":")).encode()
+
+    temporary = f"{path}.partial"
+    with open(temporary, "wb") as handle:
+        handle.write(_PREFIX.pack(MAGIC, VERSION, len(header)))
+        handle.write(header)
+        for blob in blobs:
+            handle.write(blob)
+    os.replace(temporary, path)   # never leave a half-written shard at the real name
+    return os.path.getsize(path)
+
+
+class Shard:
+    """Random access to one shard.  Holds a single file handle, opened lazily."""
+
+    def __init__(self, path):
+        self.path = path
+        self._handle = None
+        with open(path, "rb") as handle:
+            magic, version, length = _PREFIX.unpack(handle.read(_PREFIX.size))
+            if magic != MAGIC:
+                raise ValueError(f"{path}: not a shard")
+            if version != VERSION:
+                raise ValueError(f"{path}: shard version {version}, expected {VERSION}")
+            self.header = json.loads(handle.read(length))
+            self._base = _PREFIX.size + length
+        self.satellite = self.header["satellite"]
+        self.day = self.header["day"]
+        self.wavelengths = tuple(self.header["wavelengths"])
+        self._index = {int(k): tuple(v) for k, v in self.header["index"].items()}
+
+    def times(self):
+        return sorted(self._index)
+
+    def __contains__(self, t_start):
+        return t_start in self._index
+
+    def raw(self, t_start):
+        """The stored uint16 block for a slot, or None if it holds no usable frame."""
+        entry = self._index.get(t_start)
+        if entry is None:
+            return None
+        offset, length = entry
+        if self._handle is None:
+            self._handle = open(self.path, "rb")
+        self._handle.seek(self._base + offset)
+        payload = self._handle.read(length)
+        if payload == UNREADABLE:
+            return None
+        _, decompress = _zstd()
+        flat = np.frombuffer(decompress(payload), dtype="<u2")
+        return flat.reshape(len(self.wavelengths), SHARD_SIZE, SHARD_SIZE)
+
+    def frames(self, t_start):
+        """Radiance for a slot, 6 x 640 x 640 float32, or None."""
+        block = self.raw(t_start)
+        return None if block is None else decode_frames(block)
+
+    def coded(self, t_start):
+        """Model-space values for a slot, 6 x 640 x 640 float32, or None.
+
+        The cheap path, and the one training uses.  See :func:`coded_frames`.
+        """
+        block = self.raw(t_start)
+        return None if block is None else coded_frames(block)
+
+    def close(self):
+        if self._handle is not None:
+            self._handle.close()
+            self._handle = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        self.close()
+
+
+def shard_name(day, satellite):
+    return f"{day}_g{satellite:d}.shard"
+
+
+# ------------------------------------------------------------------------ extraction
+
+
+def read_slot(root, rows):
+    """Decode one slot's six bands from the archive.
+
+    `rows` maps wavelength -> 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

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)
+    gain, offset = gain_fit(source, reference)
-    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())
     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
+    The two spacecraft see the same Sun at the same moment, so the substitute is a
-    observation of the real Sun at the right time -- not an interpolation.  It
+    real observation rather than an interpolation, and its quality does not decay
-    should dominate every temporal method whenever it is available, which is the
+    with gap length the way every temporal method does.  Measured against truth at
-    thing worth quantifying: it is unavailable in the 31% of slots where both
+    195A it scores 46.45 dB whether the gap is one slot or three hundred.
-    satellites are out simultaneously.
 
-    Residual differences are instrument calibration (removed by gain matching) and
+    What does decay is the *calibration*.  The instruments differ by a band-dependent
-    a few pixels of geostationary parallax (removed by alignment).
+    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 = _shift(counterpart, *_register(counterpart, reference))
-            (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,
-        )
     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.

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)

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))

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)

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"])),
+    }

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():

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)

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"]

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():
+    for name in CONTEXT_FILLERS:
-        result = filler(ctx)
+        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

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)

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}"

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})"

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"

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())

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())