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