396 lines
16 KiB
Python
396 lines
16 KiB
Python
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"]
|