noaa-goes-visualization/tests/test_samples.py

"""Training samples: the stack, which bands it holds, and the bands to reconstruct.

What must hold: damage is withheld, never shown -- every band in a stack is the
pristine shard frame or absent; the target is a real fill case, its own undamaged
bands in the stack at dt=0 and at least one band lost; damage patterns follow the
real archive band by band; and the layout, anchors and transfers match what the
bench and renderer build.
"""

import numpy as np
import pytest

from conftest import solar_disc
from suvi import availability, dataset, paths, samples

BASE = 1715400000 // paths.CADENCE * paths.CADENCE
BANDS = len(paths.WAVELENGTHS)


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


class HoleyShard(FakeShard):
    """A vetted shard: the slot is listed, but its payload is the unreadable
    sentinel -- times() promises nothing about pixels."""

    def __init__(self, satellite, times, size=32, unreadable=()):
        super().__init__(satellite, times, size)
        self.unreadable = set(unreadable)

    def frames(self, time):
        return None if time in self.unreadable else super().frames(time)


def make_sampler(count=120, satellites=(16, 18), size=256, **kwargs):
    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


def make_availability(count=600, satellites=(16, 18)):
    """An all-GOOD per-band trace; tests spoil (satellite row, band, slots) in it."""
    states = np.full((len(satellites), BANDS, count), availability.GOOD, dtype=np.uint8)
    return availability.Availability(states=states, t0=BASE, cadence=paths.CADENCE,
                                     satellites=satellites)


def entry(sample, satellite_is_target, dt_slots):
    """Index of the stack entry with this satellite flag and offset."""
    for index, (dt, same) in enumerate(zip(sample["dts"], sample["same_satellite"])):
        if bool(same) == satellite_is_target and dt == dt_slots * paths.CADENCE:
            return index
    raise AssertionError(f"no entry {satellite_is_target, dt_slots}")


def pristine(satellite, time, size):
    return samples.encode_for_model(FakeShard(satellite, [time], size).frames(time))


# --------------------------------------------------------------------------- layout


def test_layout_leaves_the_target_slot_to_the_sampler():
    """stack_layout is the fixed ladder; the target's own entry is added by build."""
    assert (16, 0) not in samples.stack_layout((16, 18), 16)


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


# -------------------------------------------------------------------------- samples


def test_build_returns_a_full_stack():
    sampler, times = make_sampler()
    sample = sampler.build(16, times[60])
    depth = sample["frames"].shape[0]
    assert depth >= len(samples.stack_layout((16, 18), 16)) + 1
    assert sample["frames"].shape[1:] == (6, 256, 256)
    assert sample["present"].shape == (depth, 6) and sample["present"].dtype == bool
    assert sample["target"].shape == (6, 256, 256)
    assert sample["lost"].shape == (6,) and sample["lost"].any()
    assert len(sample["states"]) == len(sample["dts"]) == depth


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_the_target_slot_holds_exactly_its_undamaged_bands():
    """The render fills the bands a slot lacks and has the rest: so does training."""
    sampler, times = make_sampler(size=32, seed=4)
    for time in times[40:60]:
        sample = sampler.build(16, time)
        own = entry(sample, True, 0)
        np.testing.assert_array_equal(sample["present"][own], ~sample["lost"])
        for band in range(BANDS):
            expected = sample["target"][band] if not sample["lost"][band] else 0.0
            np.testing.assert_array_equal(sample["frames"][own][band], expected)


def test_damage_is_withheld_never_shown():
    """Every band in the stack is the pristine shard frame or zeros marked absent."""
    sampler, times = make_sampler(size=32, loss_probability=1.0, seed=9)
    absent_seen = 0
    for time in times[30:90:6]:
        sample = sampler.build(16, time)
        for index, dt in enumerate(sample["dts"]):
            source = 16 if sample["same_satellite"][index] else 18
            when = time + int(dt)
            if when not in sampler._by_satellite[source]:
                assert not sample["present"][index].any()
                continue
            truth = pristine(source, when, 32)
            for band in range(BANDS):
                if sample["present"][index][band]:
                    np.testing.assert_array_equal(sample["frames"][index][band],
                                                  truth[band])
                else:
                    absent_seen += 1
                    assert np.all(sample["frames"][index][band] == 0.0)
    assert absent_seen > 50


def test_no_synthetic_loss_leaves_the_context_whole():
    sampler, times = make_sampler(count=600, size=32, loss_probability=0.0)
    sample = sampler.build(16, times[300])
    own = entry(sample, True, 0)
    for index, held in enumerate(sample["present"]):
        if index != own:
            assert held.all()
    assert set(sample["states"]) == {"available"}


def test_frames_off_the_shard_are_missing():
    sampler, times = make_sampler(loss_probability=0.0)
    sample = sampler.build(16, times[1])          # -4, -16, ... fall off the start
    missing = [i for i, s in enumerate(sample["states"]) if s == "missing"]
    assert missing
    for index in missing:
        assert not sample["present"][index].any()
        assert np.all(sample["frames"][index] == 0.0)


def test_build_declines_when_no_band_survives(monkeypatch):
    shards = {("d", 16): FakeShard(16, [BASE], 32)}
    sampler = samples.Sampler(shards, satellites=(16, 18), loss_probability=0.0)
    monkeypatch.setattr(sampler, "_pattern", lambda rng: np.ones(BANDS, dtype=bool))
    assert sampler.build(16, BASE) is None


def test_anchors_reach_past_the_sampled_offsets():
    """A long gap must still find something *real* to work from."""
    times = [BASE] + [BASE + i * paths.CADENCE for i in range(100, 130)]
    shards = {("d", 16): FakeShard(16, times, 64)}
    sampler = samples.Sampler(shards, satellites=(16,), loss_probability=0.0)
    sample = sampler.build(16, BASE)
    reach = max(abs(dt) for dt, held in zip(sample["dts"], sample["present"]) if held.any())
    assert reach >= 100 * paths.CADENCE


def test_an_anchor_skips_unreadable_slots():
    times = [BASE + i * paths.CADENCE for i in range(10)]
    shards = {("d", 16): HoleyShard(16, times, 32, unreadable=[times[1]])}
    sampler = samples.Sampler(shards, satellites=(16,), loss_probability=0.0)
    offset = sampler._anchor(sampler._read, 16, times[0], 1, 0, None)
    assert offset == 2


def test_anchors_are_found_band_by_band_outside_the_traced_outage():
    """Production anchors each band at its nearest clean slot; so does the sampler."""
    trace = make_availability()
    trace.states[0, 0, 250:350] = availability.MISSING        # GOES-16 94A out
    sampler, times = make_sampler(count=600, size=32, availability=trace)
    window = availability.TraceWindow(trace, 300)
    read = sampler._read
    assert sampler._anchor(read, 16, times[300], 1, 0, window) == 50
    assert sampler._anchor(read, 16, times[300], -1, 0, window) == -51
    assert sampler._anchor(read, 16, times[300], 1, 3, window) == 1   # 195A is fine


def test_sampling_is_reproducible_from_its_seed():
    first, times = make_sampler(seed=7)
    second, _ = make_sampler(seed=7)
    a, b = first.build(16, times[60]), second.build(16, times[60])
    assert a["states"] == b["states"]
    np.testing.assert_array_equal(a["present"], b["present"])
    np.testing.assert_array_equal(a["lost"], b["lost"])
    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)


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_loss_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.  Some stacks should come through almost
    untouched and some should be wrecked."""
    sampler, times = make_sampler(count=600, size=32, loss_probability=0.5, seed=17)
    whole = []
    for time in times[280:380]:
        sample = sampler.build(16, time)
        own = entry(sample, True, 0)
        held = [h.all() for i, h in enumerate(sample["present"]) if i != own]
        whole.append(sum(held) / len(held))
    assert max(whole) > 0.95 and min(whole) < 0.75
    assert max(whole) - min(whole) > 0.3


def test_synthetic_losses_follow_the_archives_band_patterns():
    """Without a trace window, a lost frame still loses bands the way real ones do."""
    trace = make_availability()
    trace.states[0, 1, 0:300] = availability.BAD               # 131A alone, often
    trace.states[1, :, 0:100] = availability.MISSING           # whole slots, sometimes
    masks, odds = trace.damage_patterns()
    assert {tuple(m) for m in masks} == {(False, True, False, False, False, False),
                                         (True,) * 6}
    sampler, times = make_sampler(count=600, size=32, availability=trace, seed=3)
    drawn = {tuple(sampler._pattern(np.random.default_rng(k))) for k in range(50)}
    assert drawn <= {tuple(m) for m in masks}


def test_sampler_targets_respect_segments_and_span():
    """With a segment corpus, a target must sit a full span inside a pristine run."""
    times = [BASE + i * paths.CADENCE for i in range(200)]
    shards = {("d", s): FakeShard(s, times, 24) for s in (16, 18)}
    segment = (times[20], times[180])
    sampler = samples.Sampler(shards, offsets=(1, 4, 16, 64), segments=[segment],
                              loss_probability=0.0)
    times_used = sorted({t for _, t in sampler.targets()})
    assert times_used[0] == times[20 + 64]
    assert times_used[-1] == times[180 - 64 - 1]
    layout = samples.stack_layout((16, 18), 16, (1, 4, 16, 64))
    assert (16, 256) not in layout and (16, 64) in layout


# ------------------------------------------------------------------- real traces


def test_a_trace_window_is_a_real_fill_case(monkeypatch):
    """Centred on the one slot GOES-16 lost 131A in; GOES-18 lost 131A there too,
    94A+171A four slots later and everything sixteen slots earlier."""
    monkeypatch.setattr(samples, "TRACE_PROBABILITY", 1.0)
    trace = make_availability()
    trace.states[0, 1, 300] = availability.BAD
    trace.states[1, 1, 300] = availability.MISSING
    trace.states[1, [0, 2], 304] = availability.MISSING
    trace.states[1, :, 284] = availability.MISSING
    sampler, times = make_sampler(count=600, size=32, availability=trace, seed=1)
    for time in times[250:260]:
        sample = sampler.build(16, time)
        np.testing.assert_array_equal(sample["lost"], [False, True] + [False] * 4)
        np.testing.assert_array_equal(sample["present"][entry(sample, False, 0)],
                                      [True, False, True, True, True, True])
        np.testing.assert_array_equal(sample["present"][entry(sample, False, 4)],
                                      [False, True, False, True, True, True])
        assert not sample["present"][entry(sample, False, -16)].any()
        assert sample["present"][entry(sample, True, -4)].all()


def test_overlong_dual_traces_are_never_overlaid(monkeypatch):
    """A window whose centre sits in a dual outage longer than MAX_DUAL_GAP is
    redrawn away; with nothing else to draw, synthetic losses take over -- visible
    because the all-missing trace would otherwise empty the stack."""
    monkeypatch.setattr(samples, "TRACE_PROBABILITY", 1.0)
    trace = make_availability()
    trace.states[:] = availability.MISSING
    sampler, times = make_sampler(count=600, size=32, availability=trace, seed=2)
    for time in times[295:305]:
        sample = sampler.build(16, time)
        assert sample is not None and "available" in sample["states"]


# -------------------------------------------------------------------- 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["present"].shape == (stack, 6)
    assert packed["lost"].shape == (6,) and packed["lost"].sum() >= 1
    assert packed["target"].shape == (6, 64, 64)

    condition = packed["condition"][None]
    valid = condition[..., model.AVAILABLE]
    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, packed["present"][None])
    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(loss_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."""

    class Scaled(FakeShard):
        def frames(self, time):
            got = super().frames(time)
            return None if got is None else got * 0.8 + 0.01

    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, loss_probability=0.0)
    sample = sampler.build(16, times[20])
    cross = [i for i, same in enumerate(sample["same_satellite"])
             if not same and sample["present"][i].any()]
    assert cross
    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_steps_past_unreadable_pair_slots():
    times = [BASE + i * paths.CADENCE for i in range(41)]
    hole = times[20]
    shards = {("d", 16): FakeShard(16, times, 32),
              ("d", 18): HoleyShard(18, times, 32, unreadable=[hole])}
    sampler = samples.Sampler(shards, loss_probability=0.0)
    assert sampler._calibration("d", 16, 18, exclude=None) is not None
    assert all(pair_time != hole for (_, _, _, pair_time) in sampler._transfers)


class DarkAtShard(FakeShard):
    """A shard whose frames at given times are near-black -- a real dropout."""

    def __init__(self, satellite, times, size=32, dark=()):
        super().__init__(satellite, times, size)
        self.dark = set(dark)

    def frames(self, time):
        got = super().frames(time)
        return got * 1e-4 if time in self.dark else got


def test_calibration_walks_past_a_photometrically_degenerate_pair():
    times = [BASE + i * paths.CADENCE for i in range(41)]
    shards = {("d", 16): FakeShard(16, times, 32),
              ("d", 18): DarkAtShard(18, times, 32, dark=[times[20]])}
    sampler = samples.Sampler(shards, loss_probability=0.0)
    fit = sampler._calibration("d", 16, 18, exclude=None)
    assert fit is not None and float(np.abs(fit[0]).max()) < 5.0
    assert all(pair_time != times[20] for (_, _, _, pair_time) in sampler._transfers)


def test_calibration_declines_when_no_pair_reads():
    times = [BASE + i * paths.CADENCE for i in range(10)]
    shards = {("d", 16): FakeShard(16, times, 32),
              ("d", 18): HoleyShard(18, times, 32, unreadable=times)}
    sampler = samples.Sampler(shards, loss_probability=0.0)
    assert sampler._calibration("d", 16, 18, exclude=None) is None


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