noaa-goes-visualization/tests/test_availability.py

"""The per-band availability model: the archive's own damage, as data.

What must hold: states are per (satellite, band, slot); reading outside the era is
MISSING; windows can be centred on a real fill case; the damage-pattern
distribution counts slots, not bands; the dual-gap policy is the renderer's,
applied per band; and an old slot-level model is refused, not misread.
"""

import numpy as np
import pytest

from suvi import availability, paths

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


def model(count=600, satellites=(16, 18)):
    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 test_states_must_be_per_band():
    with pytest.raises(ValueError, match="bands"):
        availability.Availability(states=np.zeros((2, 10), np.uint8), t0=0, cadence=240,
                                  satellites=(16, 18))


def test_reading_outside_the_era_or_an_unknown_satellite_is_missing():
    m = model(count=10)
    assert (m.state(16, -1) == availability.MISSING).all()
    assert (m.state(16, 10) == availability.MISSING).all()
    assert (m.state(17, 3) == availability.MISSING).all()
    assert m.usable(16, 3).all()


def test_pristine_needs_every_band_of_every_satellite():
    m = model(count=5)
    m.states[1, 4, 2] = availability.BAD
    m.states[0, 0, 3] = availability.MISSING
    assert m.pristine().tolist() == [True, True, False, False, True]


def test_windows_for_a_satellite_centre_on_its_damaged_slots():
    m = model()
    m.states[0, 1, [10, 400]] = availability.BAD
    m.states[1, :, 200] = availability.MISSING        # the other satellite's damage
    rng = np.random.default_rng(0)
    centres = {m.sample_window(rng, 16).centre for _ in range(50)}
    assert centres == {10, 400}
    assert {m.sample_window(rng, 18).centre for _ in range(10)} == {200}


def test_a_satellite_with_no_damage_cannot_centre_a_window():
    with pytest.raises(ValueError):
        model().sample_window(np.random.default_rng(0), 16)


def test_damage_patterns_count_slots_and_their_band_sets():
    m = model(count=10)
    m.states[0, 1, 0:6] = availability.BAD                 # 131 alone, six slots
    m.states[1, :, 0:2] = availability.MISSING             # whole slots, two
    m.states[1, [0, 4], 5] = availability.BAD              # 94+284, one
    masks, odds = m.damage_patterns()
    found = {tuple(np.flatnonzero(mask)): p for mask, p in zip(masks, odds)}
    assert found == pytest.approx({(1,): 6 / 9, (0, 1, 2, 3, 4, 5): 2 / 9, (0, 4): 1 / 9})


def test_fractions_report_bands_and_whole_slots():
    m = model(count=4)
    m.states[0, 0, 0] = availability.BAD
    rates = m.fractions()[16]
    assert rates["bad"] == pytest.approx(1 / (4 * BANDS))
    assert rates["pristine"] == pytest.approx(3 / 4)


def test_save_and_load_round_trip(tmp_path):
    m = model(count=7)
    m.states[1, 2, 3] = availability.BAD
    m.save(str(tmp_path / "a.npz"))
    back = availability.Availability.load(str(tmp_path / "a.npz"))
    np.testing.assert_array_equal(back.states, m.states)
    assert (back.t0, back.cadence, back.satellites) == (m.t0, m.cadence, m.satellites)


def test_a_slot_level_model_is_refused(tmp_path):
    np.savez_compressed(str(tmp_path / "old.npz"), states=np.zeros((2, 10), np.uint8),
                        t0=np.int64(0), cadence=np.int64(240),
                        satellites=np.asarray([16, 18]))
    with pytest.raises(ValueError, match="profile_archive"):
        availability.Availability.load(str(tmp_path / "old.npz"))


# ------------------------------------------------------------------ dual gaps


def test_dual_gap_measures_the_run_containing_the_centre_per_band():
    m = model()
    m.states[:, 0, 200:300] = availability.MISSING         # 94A on neither satellite
    window = availability.TraceWindow(m, 250)
    assert window.dual_gap() == 100
    assert availability.TraceWindow(m, 299).dual_gap() == 100
    assert availability.TraceWindow(m, 100).dual_gap() == 0


def test_dual_gap_is_zero_while_one_satellite_has_every_band():
    m = model()
    m.states[0, :, :] = availability.MISSING
    assert availability.TraceWindow(m, 300).dual_gap() == 0


def test_dual_gap_takes_the_longest_band():
    m = model()
    m.states[:, 0, 290:310] = availability.MISSING
    m.states[:, 3, 250:350] = availability.BAD
    assert availability.TraceWindow(m, 300).dual_gap() == 100


def test_dual_gap_reads_beyond_the_era_as_dark_and_caps_the_walk():
    m = model()
    m.states[:, 2, :] = availability.BAD
    assert availability.TraceWindow(m, 300).dual_gap(limit=50) == 51
    assert availability.TraceWindow(m, 0).dual_gap(limit=50) == 51


def test_a_band_split_between_satellites_is_no_dual_gap():
    """2024-03-03: GOES-16 had no 94/131 and GOES-18 no 304, all day."""
    m = model()
    m.states[0, [0, 1], :] = availability.MISSING
    m.states[1, 5, :] = availability.MISSING
    assert availability.TraceWindow(m, 300).dual_gap() == 0


# ------------------------------------------------------------------ helpers


def test_run_lengths_finds_contiguous_runs():
    states = np.array([2, 2, 0, 0, 0, 2, 1, 1, 2, 2, 2])
    np.testing.assert_array_equal(availability.run_lengths(states, 2), [2, 1, 3])
    np.testing.assert_array_equal(availability.run_lengths(states, 1), [2])
    np.testing.assert_array_equal(availability.run_lengths(states, 0), [3])


def test_unfillable_times_flags_only_overlong_dual_gaps():
    times = list(range(0, 100 * 240, 240))
    usable16 = set(times) - set(times[10:40])
    usable18 = set(times) - set(times[20:30])
    assert availability.unfillable_times(times, {16: usable16, 18: usable18},
                                         limit=5) == set(times[20:30])
    assert availability.unfillable_times(times, {16: usable16, 18: usable18},
                                         limit=10) == set()
    assert availability.unfillable_times(times, {16: set(times), 18: set()},
                                         limit=5) == set()