noaa-goes-visualization/tests/test_learned_filler.py

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