noaa-goes-visualization/tests/test_dataset.py

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_reads_survive_deletion_of_the_file(tmp_path):
    """The pool rotator deletes rotated-out shards while an epoch may still be
    reading them; the handle must be open from construction, or the first lazy
    read after the delete dies (it did, mid-run)."""
    records = {1000: dataset.encode_frames(bands())}
    path = str(tmp_path / dataset.shard_name("2024-05-10", 16))
    dataset.write_shard(path, "2024-05-10", 16, paths.WAVELENGTHS, records)

    shard = dataset.Shard(path)
    os.unlink(path)
    np.testing.assert_array_equal(shard.raw(1000), records[1000])
    shard.close()


def test_shard_reads_ignore_the_shared_file_offset(tmp_path):
    """Forked DataLoader workers share the eagerly-opened file description; a read
    must neither depend on nor move the shared offset, or two workers' reads
    interleave (zstd screamed 'data corruption' mid-run)."""
    records = {t: dataset.encode_frames(bands(peak=1 + t / 1000.0))
               for t in (1000, 1240)}
    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:
        shard._handle.seek(3)               # another process moved the offset
        np.testing.assert_array_equal(shard.raw(1240), records[1240])
        assert shard._handle.tell() == 3    # ...and this read must not move it


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


# ------------------------------------------------- streaming the container out
#
# Shards are produced straight into a network stream, so two properties that used
# to be free now have to be checked: that nothing larger than one record is ever
# held, and that a stream cut off part way is refused rather than misread.


def test_streaming_a_shard_never_runs_ahead_of_the_writer(tmp_path):
    """The memory property, stated as a test.

    A shard is ~1 GB; the point of streaming is that it never exists at once.
    Feeding a generator that records how far it has been consumed each time the
    stream is written to proves the writer pulls one record at a time rather than
    draining the iterable first.
    """
    produced = []
    depth = []

    class Watching:
        def __init__(self):
            self.buffer = bytearray()

        def write(self, data):
            depth.append(len(produced))
            self.buffer.extend(data)

    def records():
        for t in (1000, 1240, 1480, 1720):
            produced.append(t)
            yield t, dataset.encode_frames(bands(peak=1 + t / 1000.0))

    handle = Watching()
    dataset.write_shard_stream(handle, "2024-05-10", 16, paths.WAVELENGTHS,
                               records())
    assert produced == [1000, 1240, 1480, 1720]
    # Every write happens with at most one record produced beyond those already
    # written; a buffering implementation would show all four before the first.
    for index, seen in enumerate(depth):
        assert seen <= index + 1, f"writer ran {seen} records ahead at write {index}"


def test_streamed_and_file_written_shards_are_byte_identical(tmp_path):
    """One implementation of the format, whichever end it is written to."""
    records = {t: dataset.encode_frames(bands(peak=1 + t / 1000.0))
               for t in (1000, 1240, 1480)}
    path = str(tmp_path / "file.shard")
    dataset.write_shard(path, "2024-05-10", 16, paths.WAVELENGTHS, records)

    import io
    stream = io.BytesIO()
    written = dataset.write_shard_stream(
        stream, "2024-05-10", 16, paths.WAVELENGTHS,
        ((t, records[t]) for t in sorted(records)))
    assert stream.getvalue() == open(path, "rb").read()
    assert written == os.path.getsize(path)


def test_the_same_records_always_produce_the_same_bytes(tmp_path):
    """A regenerated corpus must be reproducible, not merely equivalent."""
    records = {t: dataset.encode_frames(bands(peak=1 + t / 1000.0))
               for t in (1000, 1240)}
    first = str(tmp_path / "a.shard")
    second = str(tmp_path / "b.shard")
    for path in (first, second):
        dataset.write_shard(path, "2024-05-10", 16, paths.WAVELENGTHS, records)
    assert open(first, "rb").read() == open(second, "rb").read()


def test_streaming_rejects_records_out_of_order():
    """Offsets are cumulative, so an out-of-order record would build an index that
    does not describe the file -- silently, and only detectable much later."""
    import io

    block = dataset.encode_frames(bands())
    with pytest.raises(ValueError, match="ascend"):
        dataset.write_shard_stream(io.BytesIO(), "2024-05-10", 16,
                                   paths.WAVELENGTHS,
                                   [(1240, block), (1000, block)])


@pytest.mark.parametrize("cut", [1, dataset._TRAILER.size, None])
def test_a_shard_cut_off_mid_stream_is_refused_at_open(tmp_path, cut):
    """The failure the streaming writer makes possible.

    A file-written shard was renamed into place only once complete, so truncation
    was near-impossible; a streamed one can die mid-flight.  `cut=None` is the
    whole file's tail removed down to half, which loses the trailer.
    """
    records = {t: dataset.encode_frames(bands(peak=1 + t / 1000.0))
               for t in (1000, 1240, 1480)}
    path = tmp_path / "s.shard"
    dataset.write_shard(str(path), "2024-05-10", 16, paths.WAVELENGTHS, records)
    whole = path.read_bytes()
    path.write_bytes(whole[:len(whole) // 2] if cut is None else whole[:-cut])
    with pytest.raises(ValueError, match="trailer|too short"):
        dataset.Shard(str(path))


def test_a_shard_whose_records_are_short_is_refused_at_open(tmp_path):
    """The case a valid trailer cannot catch.

    If bytes are lost from the *records* while the header and trailer survive,
    every offset in the index still parses and every read still returns data --
    from the wrong place.  Only cross-checking that the last record ends exactly
    where the header begins catches it, which is why that check is load-bearing
    rather than belt-and-braces.
    """
    records = {t: dataset.encode_frames(bands(peak=1 + t / 1000.0))
               for t in (1000, 1240, 1480)}
    path = tmp_path / "s.shard"
    dataset.write_shard(str(path), "2024-05-10", 16, paths.WAVELENGTHS, records)
    whole = path.read_bytes()
    # Splice 500 bytes out of the record region, leaving header and trailer intact.
    path.write_bytes(whole[:dataset._PREFIX.size]
                     + whole[dataset._PREFIX.size + 500:])
    with pytest.raises(ValueError, match="records end"):
        dataset.Shard(str(path))


def test_a_shard_of_the_previous_format_is_refused(tmp_path):
    """v1 put the index in front.  Read as v2 its header length would be parsed as
    record bytes, so it has to fail on the version rather than on the contents."""
    path = tmp_path / "old.shard"
    path.write_bytes(dataset.MAGIC + (1).to_bytes(2, "little") + b"\x00" * 64)
    with pytest.raises(ValueError, match="version"):
        dataset.Shard(str(path))


def test_index_offsets_are_contiguous_from_the_first_record(tmp_path):
    """Record offsets are relative to a fixed base now, not to the header's end."""
    records = {t: dataset.encode_frames(bands(peak=1 + t / 1000.0))
               for t in (1000, 1240, 1480)}
    path = str(tmp_path / "s.shard")
    dataset.write_shard(path, "2024-05-10", 16, paths.WAVELENGTHS, records)
    with dataset.Shard(path) as shard:
        assert shard._base == dataset._PREFIX.size
        expected = 0
        for t in shard.times():
            offset, length = shard._index[t]
            assert offset == expected
            expected += length


# ---------------------------------------------------------------------- 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_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_shard_usable_reads_the_index_not_the_payload(tmp_path):
    """The segment verifier asks this for every corpus slot; the one-byte
    UNREADABLE sentinel must answer from the index alone."""
    records = {1000: dataset.encode_frames(bands()), 1240: None}
    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.usable(1000)
        assert not shard.usable(1240)          # vetting-dropped sentinel
        assert not shard.usable(1480)          # absent entirely