453 lines
19 KiB
Python
453 lines
19 KiB
Python
"""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
|