776 lines
34 KiB
Python
776 lines
34 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
torch = pytest.importorskip("torch")
|
|
|
|
from suvi import model # noqa: E402
|
|
|
|
BANDS = model.BANDS
|
|
|
|
|
|
def conditioning(states, satellites=None, dts=None):
|
|
"""Build a (1, S, COND_DIM) conditioning tensor."""
|
|
count = len(states)
|
|
satellites = satellites or [True] * count
|
|
dts = dts if dts is not None else [(i - count // 2) * 240.0 for i in range(count)]
|
|
rows = [
|
|
model.frame_conditioning(states[i], satellites[i], dts[i])
|
|
for i in range(count)
|
|
]
|
|
return torch.stack(rows)[None]
|
|
|
|
|
|
def expand(net, frames, condition, present=None):
|
|
"""The stack with the model's trend candidates appended, as the blend sees it."""
|
|
trends, condition, present = net._trends(frames, condition, present)
|
|
return torch.cat([frames, trends], dim=1), condition, present
|
|
|
|
|
|
def stack(count, size=64, seed=0):
|
|
generator = torch.Generator().manual_seed(seed)
|
|
return torch.randn(1, count, BANDS, size, size, generator=generator)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def net():
|
|
torch.manual_seed(0)
|
|
return model.build(base=16, depth=3).eval()
|
|
|
|
|
|
# ------------------------------------------------------------------- conditioning
|
|
|
|
|
|
def test_conditioning_is_the_declared_width():
|
|
assert model.frame_conditioning("available", True, 0.0).shape == (model.COND_DIM,)
|
|
|
|
|
|
def test_the_layout_names_every_column_once():
|
|
assert model.COND_DIM == len(model.COND_LAYOUT) == len(set(model.COND_LAYOUT))
|
|
assert [model.AVAILABLE, model.MISSING, model.TARGET_SATELLITE, model.DT_SIGN,
|
|
model.DT_LOG, model.AGREE_MEAN, model.AGREE_MAX, model.DERIVED] == list(
|
|
range(model.COND_DIM))
|
|
|
|
|
|
def test_conditioning_separates_the_two_states():
|
|
available = model.frame_conditioning("available", True, 240.0)
|
|
missing = model.frame_conditioning("missing", True, 240.0)
|
|
assert available[model.AVAILABLE] == 1.0 and available[model.MISSING] == 0.0
|
|
assert missing[model.AVAILABLE] == 0.0 and missing[model.MISSING] == 1.0
|
|
|
|
|
|
@pytest.mark.parametrize("state", ["suspect", "probably-fine"])
|
|
def test_conditioning_rejects_an_unknown_state(state):
|
|
"""v9 has no `suspect`: a judged-bad band is withheld, never shown."""
|
|
with pytest.raises(KeyError):
|
|
model.frame_conditioning(state, True, 0.0)
|
|
|
|
|
|
def test_conditioning_distinguishes_direction_and_distance():
|
|
back = model.frame_conditioning("available", True, -7200.0)
|
|
forward = model.frame_conditioning("available", True, 7200.0)
|
|
far = model.frame_conditioning("available", True, 72000.0)
|
|
assert back[model.DT_SIGN] == -1.0 and forward[model.DT_SIGN] == 1.0
|
|
assert back[model.DT_LOG] == pytest.approx(forward[model.DT_LOG]) # magnitude only
|
|
assert far[model.DT_LOG] > forward[model.DT_LOG]
|
|
|
|
|
|
def test_conditioning_leaves_agreement_to_the_model():
|
|
"""The agreement columns are computed from the stack inside forward(),
|
|
identically at training and inference -- the fix for detector scores that were
|
|
always zero in training and populated on the bench."""
|
|
vector = model.frame_conditioning("available", True, 240.0)
|
|
assert vector[model.AGREE_MEAN] == 0.0 and vector[model.AGREE_MAX] == 0.0
|
|
|
|
|
|
# ------------------------------------------------------------------- blend prior
|
|
|
|
|
|
def test_prior_favours_the_nearest_frame_in_time():
|
|
condition = conditioning(["available"] * 3, satellites=[True] * 3,
|
|
dts=[-240.0, -3840.0, -240.0 * 300])
|
|
prior = model.blend_prior(condition)[0]
|
|
assert prior[0] > prior[1] > prior[2]
|
|
|
|
|
|
def test_prior_no_longer_penalises_the_counterpart():
|
|
"""The old prior down-weighted cross-satellite frames because their gain was
|
|
wrong at initialisation. Candidates now arrive gain-matched by suvi.align, so
|
|
a simultaneous counterpart is as trustworthy as a local frame."""
|
|
same = conditioning(["available"], satellites=[True], dts=[0.0])
|
|
cross = conditioning(["available"], satellites=[False], dts=[0.0])
|
|
assert model.blend_prior(same)[0, 0] == model.blend_prior(cross)[0, 0]
|
|
|
|
|
|
def test_prior_disfavours_derived_candidates():
|
|
real = conditioning(["available"], dts=[-240.0])
|
|
derived = real.clone()
|
|
derived[..., model.DERIVED] = 1.0
|
|
gap = model.blend_prior(real)[0, 0] - model.blend_prior(derived)[0, 0]
|
|
assert float(gap) == pytest.approx(model.PRIOR_DERIVED_PENALTY)
|
|
|
|
|
|
# ------------------------------------------------------------------------- shapes
|
|
|
|
|
|
def test_forward_returns_the_source_shape(net):
|
|
out = net(stack(5), conditioning(["available"] * 5))
|
|
assert out.shape == (1, BANDS, 64, 64)
|
|
assert torch.isfinite(out).all()
|
|
|
|
|
|
@pytest.mark.parametrize("size", [160, 320, 640])
|
|
def test_one_model_serves_every_source_resolution(net, size):
|
|
"""Blend and polish are fields: they upsample, so 1280 needs no retraining."""
|
|
sources = torch.randn(1, 3, BANDS, size, size)
|
|
out = net(sources, conditioning(["available"] * 3))
|
|
assert out.shape == (1, BANDS, size, size)
|
|
assert torch.isfinite(out).all()
|
|
|
|
|
|
@pytest.mark.parametrize("count", [1, 2, 7, 23])
|
|
def test_stack_depth_is_not_fixed(net, count):
|
|
"""Availability varies per slot, so the stack cannot be a compile-time constant."""
|
|
out = net(stack(count), conditioning(["available"] * count))
|
|
assert out.shape == (1, BANDS, 64, 64)
|
|
|
|
|
|
# --------------------------------------------------------------- identity behaviour
|
|
|
|
|
|
def expected_blend(sources, condition):
|
|
"""What an untrained model must output: the prior-weighted blend of its sources.
|
|
|
|
Both heads are zero-initialised, so the polish gain is exactly one and the blend
|
|
logits are exactly the prior."""
|
|
valid = condition[..., model.AVAILABLE]
|
|
logits = model.blend_prior(condition).masked_fill(valid < 0.5, float("-inf"))
|
|
weights = torch.nan_to_num(torch.softmax(logits, dim=1), nan=0.0)
|
|
return (sources * weights[..., None, None, None]).sum(dim=1)
|
|
|
|
|
|
def test_untrained_model_is_the_prior_weighted_blend(net):
|
|
"""The wiring gate: any scaling or orientation bug in the composition shows up
|
|
here, and is invisible once training starts moving the loss. The stack is
|
|
expanded through the trend builder first: derived candidates are inputs to the
|
|
blend like any other, and their prior penalty is part of the wiring."""
|
|
frames = stack(5)
|
|
condition = conditioning(["available"] * 5)
|
|
out = net(frames, condition)
|
|
expanded_frames, expanded_condition, _ = expand(net, frames, condition)
|
|
assert torch.allclose(out, expected_blend(expanded_frames, expanded_condition),
|
|
atol=1e-4)
|
|
|
|
|
|
def test_identity_holds_when_fields_are_upsampled(net):
|
|
"""A 160-resolution field applied to larger sources must still be the identity."""
|
|
sources = torch.randn(1, 3, BANDS, 320, 320)
|
|
condition = conditioning(["available"] * 3)
|
|
out = net(sources, condition)
|
|
assert torch.allclose(out, expected_blend(sources, condition), atol=1e-3)
|
|
|
|
|
|
def test_missing_frames_take_no_weight(net):
|
|
"""A frame with no pixels must not contribute, however confident the head is."""
|
|
frames = stack(4)
|
|
frames[:, 0] = 999.0 # would dominate any blend that included it
|
|
condition = conditioning(["missing", "available", "available", "available"])
|
|
out = net(frames, condition)
|
|
assert torch.allclose(out, expected_blend(frames, condition), atol=1e-4)
|
|
assert out.abs().max() < 100, "the masked frame leaked into the blend"
|
|
|
|
|
|
def test_the_targets_own_bands_win_only_where_it_holds_them(net):
|
|
"""The target slot is an entry at dt=0 holding its undamaged bands: the prior
|
|
trusts it most where it has pixels, and it takes no weight in a lost band."""
|
|
frames = torch.stack([torch.full((BANDS, 32, 32), v) for v in (0.9, 0.2, 0.3)])[None]
|
|
condition = conditioning(["available"] * 3, satellites=[True, True, True],
|
|
dts=[0.0, -240.0, 240.0])
|
|
present = torch.ones(1, 3, BANDS)
|
|
present[0, 0, 2] = 0.0 # the target lost band 2
|
|
out = net(frames, condition, present)
|
|
assert torch.allclose(out, expected_band_blend(frames, condition, present),
|
|
atol=1e-4)
|
|
assert float(out[0, 2].detach().max()) < 0.31 # nothing of 0.9 in band 2
|
|
assert float(out[0, 0].detach().min()) > 0.5 # mostly the target's own
|
|
|
|
|
|
def test_a_stack_with_nothing_usable_returns_zeros_not_nan(net):
|
|
"""The sampler declines such targets; the network itself must still not NaN."""
|
|
out = net(stack(4), conditioning(["missing"] * 4))
|
|
assert torch.isfinite(out).all()
|
|
assert out.abs().max() == 0.0
|
|
|
|
|
|
def test_heads_start_at_zero():
|
|
"""Training must start exactly at the hand-written policy; every step afterwards
|
|
is a departure the data paid for."""
|
|
torch.manual_seed(0)
|
|
fresh = model.build(base=8, depth=2)
|
|
for head in (fresh.weight, fresh.polish):
|
|
with torch.no_grad():
|
|
assert float(head.weight.abs().sum()) == 0.0
|
|
assert float(head.bias.abs().sum()) == 0.0
|
|
|
|
|
|
# --------------------------------------------------------- structural non-collapse
|
|
|
|
|
|
def test_output_is_a_bounded_combination_whatever_the_weights():
|
|
"""The invariant that makes the old failure mode unreachable: with *arbitrary*
|
|
parameters, the output of a single-frame stack stays within the polish bound of
|
|
that frame. No parameter setting can emit an image that is not built from the
|
|
inputs -- which is exactly what the collapsed run did."""
|
|
torch.manual_seed(3)
|
|
wild = model.build(base=8, depth=2)
|
|
with torch.no_grad():
|
|
for parameter in wild.parameters():
|
|
parameter.normal_(0.0, 2.0)
|
|
frames = torch.full((1, 1, BANDS, 48, 48), 0.5)
|
|
with torch.no_grad():
|
|
out = wild(frames, conditioning(["available"], dts=[240.0]))
|
|
bound = 0.5 * float(np.exp(model.POLISH_RANGE))
|
|
lower = 0.5 * float(np.exp(-model.POLISH_RANGE))
|
|
assert float(out.max()) <= bound + 1e-4
|
|
assert float(out.min()) >= lower - 1e-4
|
|
|
|
|
|
def test_polish_gain_saturates_at_its_bound():
|
|
torch.manual_seed(0)
|
|
fresh = model.build(base=8, depth=2).eval()
|
|
with torch.no_grad():
|
|
fresh.polish.bias.fill_(1000.0)
|
|
frames = torch.ones(1, 1, BANDS, 48, 48)
|
|
out = fresh(frames, conditioning(["available"], dts=[240.0]))
|
|
assert float(out.mean()) == pytest.approx(np.exp(model.POLISH_RANGE), rel=1e-3)
|
|
|
|
|
|
# --------------------------------------------------------------------- agreement
|
|
|
|
|
|
def test_agreement_flags_the_frame_that_disagrees():
|
|
"""A frozen or torn frame announces itself as deviation from the stack median --
|
|
the deterministic replacement for both cross-attention and detector scores."""
|
|
torch.manual_seed(0)
|
|
fresh = model.build(base=8, depth=2)
|
|
frames = torch.zeros(1, 4, BANDS, 32, 32)
|
|
frames[:, 3] = 5.0 # the odd one out
|
|
condition = conditioning(["available"] * 4)
|
|
valid = condition[..., model.AVAILABLE]
|
|
got = fresh._with_agreement(frames.reshape(4, BANDS, 32, 32), condition, valid,
|
|
batch=1, stack=4)
|
|
assert float(got[0, 3, model.AGREE_MEAN]) > float(got[0, 0, model.AGREE_MEAN]) + 0.5
|
|
assert float(got[0, 0, model.AGREE_MEAN]) < 0.1
|
|
|
|
|
|
# ------------------------------------------------------------------------- flow
|
|
|
|
|
|
def test_flow_is_zero_for_a_simultaneous_frame():
|
|
"""dt=0 means a real observation of the target instant: the flow allowance is
|
|
zero there, however loud the head, so the most trusted candidate stays put."""
|
|
torch.manual_seed(0)
|
|
fresh = model.build(base=8, depth=2).eval()
|
|
with torch.no_grad():
|
|
fresh.flow.bias.fill_(1000.0)
|
|
frames = stack(1, size=48)
|
|
out = fresh(frames, conditioning(["available"], dts=[0.0]))
|
|
assert torch.allclose(out, frames[:, 0], atol=1e-4)
|
|
|
|
|
|
def test_flow_displaces_by_the_dt_scaled_bound_at_saturation():
|
|
"""A saturated flow head moves a ramp by exactly the allowance for its dt --
|
|
and no further. Displacement is a fraction of frame width, so the same model
|
|
moves features the same physical distance at 640 and 1280."""
|
|
torch.manual_seed(0)
|
|
fresh = model.build(base=8, depth=2).eval()
|
|
with torch.no_grad():
|
|
fresh.flow.bias.fill_(1000.0) # tanh saturates: full allowance, +x and +y
|
|
size = 64
|
|
ramp = torch.linspace(-1.0, 1.0, size).view(1, 1, 1, size)
|
|
frames = ramp.expand(1, 1, BANDS, size, size).clone()
|
|
condition = conditioning(["available"], dts=[240.0])
|
|
out = fresh(frames, condition)
|
|
allowance = model.MAX_FLOW_FRACTION * float(condition[0, 0, model.DT_LOG])
|
|
interior = out[..., 8:-8, 8:-8] - frames[:, 0, :, 8:-8, 8:-8]
|
|
assert float(interior.mean()) == pytest.approx(2.0 * allowance, rel=0.05)
|
|
|
|
|
|
def test_untrained_flow_is_the_identity():
|
|
frames = stack(3, size=32)
|
|
warped = model._displace(frames[0], torch.zeros(3, 2, 32, 32))
|
|
assert torch.allclose(warped, frames[0], atol=1e-5)
|
|
|
|
|
|
# ----------------------------------------------------------------------- trends
|
|
|
|
|
|
def test_dt_encoding_round_trips_exactly():
|
|
"""The trend builder recovers real seconds from the conditioning encoding."""
|
|
dts = [-72000.0, -960.0, -240.0, 0.0, 240.0, 61440.0]
|
|
condition = conditioning(["available"] * len(dts), dts=dts)
|
|
recovered = model.dt_from_conditioning(condition)[0]
|
|
assert torch.allclose(recovered, torch.tensor(dts), rtol=1e-5, atol=1e-3)
|
|
|
|
|
|
def trend_setup(dts, values, states=None):
|
|
"""A stack of uniform frames at given dts/values, expanded through the builder."""
|
|
count = len(dts)
|
|
frames = torch.zeros(1, count, BANDS, 16, 16)
|
|
for i, value in enumerate(values):
|
|
frames[:, i] = value
|
|
condition = conditioning(states or ["available"] * count, dts=dts)
|
|
net = model.build(base=8, depth=2)
|
|
return expand(net, frames, condition)[:2]
|
|
|
|
|
|
def test_trend_candidates_recover_a_linear_ramp():
|
|
"""Brightness rising linearly into the gap: the extrapolation must land exactly
|
|
on the line's value at the target instant -- the reach a convex blend of the
|
|
observations alone cannot have."""
|
|
expanded, condition = trend_setup(
|
|
dts=[-480.0, -240.0, 240.0, 480.0], values=[0.2, 0.3, 0.5, 0.6])
|
|
past, future = expanded[0, -2], expanded[0, -1]
|
|
assert torch.allclose(past, torch.full_like(past, 0.4), atol=1e-6)
|
|
assert torch.allclose(future, torch.full_like(future, 0.4), atol=1e-6)
|
|
assert bool(condition[0, -2, model.DERIVED]) and bool(condition[0, -1, model.DERIVED])
|
|
assert float(condition[0, -2, model.AVAILABLE]) == 1.0 # available, not missing
|
|
|
|
|
|
def test_trend_reach_is_capped():
|
|
"""A short-baseline slope far from the target must not extrapolate freely: at
|
|
6x its baseline the line would say 0.95, but reach caps at 4x -> 0.75."""
|
|
expanded, _ = trend_setup(dts=[-1680.0, -1440.0], values=[0.25, 0.35])
|
|
past = expanded[0, -2]
|
|
assert torch.allclose(past, torch.full_like(past, 0.75), atol=1e-6)
|
|
|
|
|
|
def test_trend_output_is_clamped_to_model_range():
|
|
expanded, _ = trend_setup(dts=[-480.0, -240.0], values=[0.2, 0.9])
|
|
assert float(expanded[0, -2].max()) <= 1.0
|
|
|
|
|
|
def test_trends_missing_when_a_side_lacks_two_frames():
|
|
expanded, condition = trend_setup(dts=[-240.0, 240.0, 480.0], values=[0.1, 0.2, 0.3])
|
|
assert float(condition[0, -2, model.MISSING]) == 1.0 # only one past frame
|
|
assert float(condition[0, -1, model.AVAILABLE]) == 1.0 # future trend exists
|
|
assert float(expanded[0, -2].abs().max()) == 0.0
|
|
|
|
|
|
def test_trends_ignore_invalid_and_cross_satellite_frames():
|
|
"""Extrapolation slopes must come from clean target-satellite data only: a
|
|
missing frame has no pixels, and a counterpart's residual gain error would be
|
|
amplified by the lever arm."""
|
|
frames = torch.full((1, 4, BANDS, 16, 16), 0.5)
|
|
condition = conditioning(["missing", "available", "available", "available"],
|
|
satellites=[True, False, True, True],
|
|
dts=[-480.0, -240.0, 240.0, 480.0])
|
|
net = model.build(base=8, depth=2)
|
|
_, expanded, _ = net._trends(frames, condition)
|
|
assert float(expanded[0, -2, model.MISSING]) == 1.0 # one missing, one cross-sat
|
|
assert float(expanded[0, -1, model.AVAILABLE]) == 1.0 # two clean same-sat frames
|
|
|
|
|
|
def test_agreement_is_zero_for_a_unanimous_stack():
|
|
torch.manual_seed(0)
|
|
fresh = model.build(base=8, depth=2)
|
|
frames = torch.full((1, 3, BANDS, 32, 32), 0.25)
|
|
condition = conditioning(["available"] * 3)
|
|
valid = condition[..., model.AVAILABLE]
|
|
got = fresh._with_agreement(frames.reshape(3, BANDS, 32, 32), condition, valid,
|
|
batch=1, stack=3)
|
|
assert float(got[..., model.AGREE_MEAN:model.AGREE_MAX + 1].abs().max()) < 1e-6
|
|
|
|
|
|
def test_agreement_ignores_missing_frames():
|
|
"""A zeroed missing frame must not drag the median toward black."""
|
|
torch.manual_seed(0)
|
|
fresh = model.build(base=8, depth=2)
|
|
frames = torch.full((1, 3, BANDS, 32, 32), 0.5)
|
|
frames[:, 0] = 0.0 # missing, zero-filled
|
|
condition = conditioning(["missing", "available", "available"])
|
|
valid = condition[..., model.AVAILABLE]
|
|
got = fresh._with_agreement(frames.reshape(3, BANDS, 32, 32), condition, valid,
|
|
batch=1, stack=3)
|
|
assert float(got[0, 1, model.AGREE_MEAN]) < 1e-6
|
|
assert float(got[0, 2, model.AGREE_MEAN]) < 1e-6
|
|
|
|
|
|
# ------------------------------------------------------------- per-band presence
|
|
|
|
|
|
def expected_band_blend(sources, condition, present):
|
|
"""The untrained model's output when some frames lack some bands: per band,
|
|
the prior blend over the valid frames that hold that band."""
|
|
valid = condition[..., model.AVAILABLE]
|
|
logits = model.blend_prior(condition)[..., None].expand(-1, -1, BANDS)
|
|
logits = logits.masked_fill((valid[..., None] < 0.5) | (present < 0.5),
|
|
float("-inf"))
|
|
weights = torch.nan_to_num(torch.softmax(logits, dim=1), nan=0.0)
|
|
return (sources * weights[..., None, None]).sum(dim=1)
|
|
|
|
|
|
def partial_setup():
|
|
"""Three frames, none a trend source pair, the first missing bands 0 and 4."""
|
|
frames = torch.stack([torch.full((BANDS, 32, 32), v) for v in (0.2, 0.4, 0.6)])[None]
|
|
condition = conditioning(["available"] * 3, satellites=[False, True, False],
|
|
dts=[0.0, -240.0, 480.0])
|
|
present = torch.ones(1, 3, BANDS)
|
|
present[0, 0, [0, 4]] = 0.0
|
|
return frames, condition, present
|
|
|
|
|
|
def test_all_bands_present_is_the_default(net):
|
|
"""`present` of all ones is exactly what None means, bit for bit."""
|
|
frames = stack(5)
|
|
condition = conditioning(["available", "available", "available", "missing",
|
|
"available"])
|
|
assert torch.equal(net(frames, condition),
|
|
net(frames, condition, torch.ones(1, 5, BANDS)))
|
|
|
|
|
|
def test_a_frame_missing_a_band_still_contributes_its_other_bands(net):
|
|
frames, condition, present = partial_setup()
|
|
out = net(frames, condition, present)
|
|
assert torch.allclose(out, expected_band_blend(frames, condition, present),
|
|
atol=1e-4)
|
|
without_first = expected_band_blend(frames, condition,
|
|
torch.tensor([[[0.0] * BANDS, [1.0] * BANDS,
|
|
[1.0] * BANDS]]))
|
|
# Where the first frame holds the band, it pulls the blend toward its 0.2...
|
|
assert float(out[0, 1].mean()) < float(without_first[0, 1].mean()) - 0.01
|
|
# ...and where it does not, the band is the blend of the other two alone.
|
|
assert torch.allclose(out[0, 0], without_first[0, 0], atol=1e-4)
|
|
assert torch.allclose(out[0, 4], without_first[0, 4], atol=1e-4)
|
|
|
|
|
|
def test_a_missing_bands_pixels_never_reach_the_output(net):
|
|
frames, condition, present = partial_setup()
|
|
reference = net(frames, condition, present)
|
|
frames[0, 0, [0, 4]] = 999.0 # whatever the stack carries there
|
|
out = net(frames, condition, present)
|
|
assert torch.allclose(out, reference, atol=1e-4)
|
|
|
|
|
|
def test_a_band_no_frame_holds_comes_back_zero_not_nan(net):
|
|
frames, condition, present = partial_setup()
|
|
present[0, :, 2] = 0.0
|
|
out = net(frames, condition, present)
|
|
assert torch.isfinite(out).all()
|
|
assert float(out[0, 2].abs().max()) == 0.0
|
|
assert float(out[0, 3].abs().max()) > 0.1
|
|
|
|
|
|
def test_trends_come_only_from_frames_holding_every_band():
|
|
"""A frame missing a band has no line to continue in it."""
|
|
frames = torch.full((1, 4, BANDS, 16, 16), 0.5)
|
|
condition = conditioning(["available"] * 4, dts=[-480.0, -240.0, 240.0, 480.0])
|
|
present = torch.ones(1, 4, BANDS)
|
|
present[0, 1, 3] = 0.0 # one of the past pair is partial
|
|
net = model.build(base=8, depth=2)
|
|
_, expanded, held = net._trends(frames, condition, present)
|
|
assert float(expanded[0, -2, model.MISSING]) == 1.0 # past trend: missing
|
|
assert float(expanded[0, -1, model.AVAILABLE]) == 1.0 # future trend: built
|
|
assert held.shape == (1, 6, BANDS) and bool((held[0, 4:] == 1.0).all())
|
|
|
|
|
|
def test_agreement_ignores_the_bands_a_frame_lacks():
|
|
"""A partial frame is judged on the bands it holds, not on whatever stands in
|
|
for the rest."""
|
|
torch.manual_seed(0)
|
|
fresh = model.build(base=8, depth=2)
|
|
frames = torch.full((1, 3, BANDS, 32, 32), 0.5)
|
|
frames[0, 0, 0] = 0.9 # absent band, stand-in
|
|
condition = conditioning(["available"] * 3)
|
|
valid = condition[..., model.AVAILABLE]
|
|
present = torch.ones(1, 3, BANDS)
|
|
present[0, 0, 0] = 0.0
|
|
got = fresh._with_agreement(frames.reshape(3, BANDS, 32, 32), condition, valid,
|
|
batch=1, stack=3, present=present)
|
|
assert float(got[0, :, model.AGREE_MEAN:model.AGREE_MAX + 1].abs().max()) < 1e-6
|
|
|
|
|
|
def test_the_median_skips_frames_without_the_band():
|
|
torch.manual_seed(0)
|
|
fresh = model.build(base=8, depth=2)
|
|
frames = torch.full((1, 3, BANDS, 8, 8), 0.5)
|
|
frames[0, 0, 2] = -1.0
|
|
valid = torch.ones(1, 3)
|
|
present = torch.ones(1, 3, BANDS)
|
|
present[0, 0, 2] = 0.0
|
|
median = fresh._stack_median(frames.reshape(3, BANDS, 8, 8), valid, 1, 3, present)
|
|
assert float(median[0, 0, 2].min()) == pytest.approx(0.5)
|
|
|
|
|
|
# ------------------------------------------------------------------------ training
|
|
|
|
|
|
def test_gradients_reach_every_head():
|
|
"""A head with no gradient path never learns, and zero-init hides it. Checked
|
|
after three optimiser steps, because the zero initialisation is nested two deep:
|
|
at step 1 only the output heads move, at step 2 the FiLM output layers behind
|
|
them, and only then does gradient reach everything. Asserting earlier would
|
|
fail on a model that is perfectly healthy."""
|
|
torch.manual_seed(0)
|
|
fresh = model.build(base=8, depth=2)
|
|
frames = stack(3, size=48)
|
|
condition = conditioning(["available", "missing", "available"])
|
|
optimiser = torch.optim.Adam(fresh.parameters(), lr=1e-3)
|
|
|
|
for _ in range(3):
|
|
optimiser.zero_grad()
|
|
out = fresh(frames, condition)
|
|
out.square().mean().backward()
|
|
optimiser.step()
|
|
|
|
for name, parameter in fresh.named_parameters():
|
|
assert parameter.grad is not None, f"{name} has no gradient"
|
|
assert parameter.grad.abs().sum() > 0, f"{name} gradient is identically zero"
|
|
|
|
|
|
def smooth_stack(count, size=64, seed=1):
|
|
"""Spatially correlated frames; white noise tests nothing the real task needs."""
|
|
generator = torch.Generator().manual_seed(seed)
|
|
coarse = torch.randn(1, count * BANDS, size // 8, size // 8, generator=generator)
|
|
smooth = torch.nn.functional.interpolate(coarse, size=(size, size), mode="bicubic",
|
|
align_corners=False)
|
|
return smooth.view(1, count, BANDS, size, size)
|
|
|
|
|
|
def train_a_little(net, frames, condition, target, steps=120, lr=3e-3):
|
|
optimiser = torch.optim.Adam(net.parameters(), lr=lr)
|
|
first = None
|
|
for step in range(steps):
|
|
optimiser.zero_grad()
|
|
out = net(frames, condition)
|
|
loss = (out - target).abs().mean()
|
|
loss.backward()
|
|
optimiser.step()
|
|
if step == 0:
|
|
first = loss.item()
|
|
return first, loss.item()
|
|
|
|
|
|
def test_model_learns_to_select_one_frame_from_the_stack():
|
|
"""The weight head's core job: when one candidate is the answer, put the weight
|
|
there. Fails in seconds if the softmax, masking or composition is miswired."""
|
|
torch.manual_seed(0)
|
|
fresh = model.build(base=16, depth=2)
|
|
frames = smooth_stack(3)
|
|
first, last = train_a_little(
|
|
fresh, frames, conditioning(["available"] * 3), frames[:, 1].clone()
|
|
)
|
|
assert last < first * 0.2, f"loss barely moved: {first:.4f} -> {last:.4f}"
|
|
|
|
|
|
def test_model_learns_a_per_band_gain_within_the_polish_bound():
|
|
"""The polish head's core job: the residual drift the daily fit leaves behind."""
|
|
torch.manual_seed(0)
|
|
fresh = model.build(base=16, depth=2)
|
|
frames = smooth_stack(1)
|
|
per_band = torch.tensor([1.2, 0.85, 1.1, 0.9, 1.25, 0.8]).view(1, BANDS, 1, 1)
|
|
target = frames[:, 0] * per_band
|
|
first, last = train_a_little(fresh, frames, conditioning(["available"],
|
|
dts=[240.0]), target)
|
|
assert last < first * 0.3, f"loss barely moved: {first:.4f} -> {last:.4f}"
|
|
|
|
|
|
# ---------------------------------------------------------------- variant arms
|
|
|
|
|
|
@pytest.mark.parametrize("flags", [
|
|
{"relative": True}, {"corr_flow": True}, {"attention": True},
|
|
{"relative": True, "corr_flow": True, "attention": True},
|
|
])
|
|
def test_every_variant_is_the_prior_blend_untrained(flags):
|
|
"""Each screened cross-frame pathway must start silent: zero-init heads and
|
|
gates mean an untrained variant is exactly the same sane policy as the
|
|
baseline, and the v1-v3 collapse mode stays structurally unreachable."""
|
|
torch.manual_seed(0)
|
|
net = model.build(base=8, depth=2, **flags).eval()
|
|
frames = stack(5)
|
|
condition = conditioning(["available"] * 5)
|
|
out = net(frames, condition)
|
|
expanded_frames, expanded_condition, _ = expand(net, frames, condition)
|
|
assert torch.allclose(out, expected_blend(expanded_frames, expanded_condition),
|
|
atol=1e-4)
|
|
|
|
|
|
@pytest.mark.parametrize("flags", [
|
|
{"relative": True}, {"corr_flow": True}, {"attention": True},
|
|
])
|
|
def test_variant_forward_and_backward_are_finite(flags):
|
|
torch.manual_seed(1)
|
|
net = model.build(base=8, depth=2, **flags)
|
|
frames = stack(6, size=32)
|
|
condition = conditioning(["available", "available", "available", "missing",
|
|
"available", "available"])
|
|
out = net(frames, condition)
|
|
assert torch.isfinite(out).all()
|
|
out.square().mean().backward()
|
|
for name, parameter in net.named_parameters():
|
|
if parameter.grad is not None:
|
|
assert torch.isfinite(parameter.grad).all(), name
|
|
|
|
|
|
def test_variant_checkpoints_round_trip_through_load_learned(tmp_path, monkeypatch):
|
|
"""A variant's flags ride in the checkpoint; loading must rebuild the same
|
|
shapes or state_dict loading fails loudly."""
|
|
import torch as torch_module
|
|
|
|
from suvi import fillers
|
|
|
|
net = model.build(base=8, depth=2, relative=True, corr_flow=True)
|
|
path = str(tmp_path / "variant.pt")
|
|
torch_module.save({"model": net.state_dict(),
|
|
"args": {"base": 8, "depth": 2,
|
|
"cond_layout": list(model.COND_LAYOUT),
|
|
"relative": True, "corr_flow": True,
|
|
"offsets": [1, 4, 16, 64]}}, path)
|
|
monkeypatch.setenv(fillers.LEARNED_CHECKPOINT_ENV, path)
|
|
fillers._LEARNED.clear()
|
|
loaded, _, _ = fillers.load_learned()
|
|
assert loaded.relative and loaded.corr_flow and loaded.attend is None
|
|
assert loaded.offsets == (1, 4, 16, 64)
|
|
fillers._LEARNED.clear()
|
|
|
|
|
|
@pytest.mark.parametrize("args", [
|
|
{"base": 8, "depth": 2, "cond_dim": 9}, # v8: a suspect column
|
|
{"base": 8, "depth": 2, "cond_layout": ["available", "missing", "suspect"]},
|
|
])
|
|
def test_a_checkpoint_under_another_conditioning_layout_is_refused(
|
|
tmp_path, monkeypatch, args):
|
|
"""Its weights would read every column shifted: refuse by name, not by a
|
|
shape error -- or worse, a v5 checkpoint whose shapes happen to fit."""
|
|
import torch as torch_module
|
|
|
|
from suvi import fillers
|
|
|
|
path = str(tmp_path / "old.pt")
|
|
torch_module.save({"model": model.build(base=8, depth=2).state_dict(),
|
|
"args": args}, path)
|
|
monkeypatch.setenv(fillers.LEARNED_CHECKPOINT_ENV, path)
|
|
fillers._LEARNED.clear()
|
|
with pytest.raises(SystemExit, match="conditioning layout"):
|
|
fillers.load_learned()
|
|
fillers._LEARNED.clear()
|
|
|
|
|
|
def test_correlation_peaks_at_the_true_shift():
|
|
"""The cost volume must actually measure correspondence: a frame shifted by
|
|
+2 px correlates best with the median at the matching displacement."""
|
|
torch.manual_seed(0)
|
|
base_img = torch.randn(1, 6, 32, 32)
|
|
shifted = torch.roll(base_img, shifts=2, dims=-1)
|
|
volume = model._correlation(shifted, base_img, radius=3)
|
|
centre = volume.shape[1] // 2
|
|
interior = volume[..., 8:-8, 8:-8].mean(dim=(-2, -1))[0]
|
|
assert int(interior.argmax()) != centre
|
|
best_dx = int(interior.argmax()) % 7 - 3
|
|
assert best_dx == -2 or best_dx == 2
|
|
|
|
|
|
# ------------------------------------------------- inference path memory refactor
|
|
|
|
|
|
def _trained_looking(attention):
|
|
"""A model whose heads are non-zero, so the blend, gain and flow all act."""
|
|
torch.manual_seed(1)
|
|
built = model.build(base=8, depth=2, attention=attention)
|
|
with torch.no_grad():
|
|
for head in (built.weight, built.polish, built.flow):
|
|
head.weight.normal_(0.0, 0.2)
|
|
head.bias.normal_(0.0, 0.2)
|
|
if built.attn_gate is not None:
|
|
built.attn_gate.fill_(0.5)
|
|
return built.eval()
|
|
|
|
|
|
@pytest.mark.parametrize("attention", [True, False])
|
|
def test_band_at_a_time_inference_matches_the_batched_path_exactly(
|
|
attention, monkeypatch):
|
|
"""Inference blends one band at a time to bound memory (a 24-candidate stack
|
|
at 1280px peaked at 9.4 GB all at once); it must not change one bit. Both
|
|
sides run without gradients: attention itself takes a different fused kernel
|
|
when gradients are off, so a grad-vs-no-grad comparison would test torch."""
|
|
net_ = _trained_looking(attention)
|
|
sources = stack(6, size=96, seed=3).clamp(-1, 1)
|
|
cond = conditioning(["available", "available", "available", "missing",
|
|
"available", "available"],
|
|
satellites=[True, True, True, True, False, True],
|
|
dts=[-240.0, 240.0, -960.0, 960.0, 0.0, -3840.0])
|
|
with torch.no_grad():
|
|
monkeypatch.setattr(model, "BAND_AT_A_TIME", False)
|
|
batched = net_(sources, cond)
|
|
monkeypatch.setattr(model, "BAND_AT_A_TIME", True)
|
|
banded = net_(sources, cond)
|
|
assert banded.shape == batched.shape == (1, BANDS, 96, 96)
|
|
assert torch.equal(banded, batched)
|
|
|
|
|
|
def test_inference_matches_the_gradient_path(monkeypatch):
|
|
"""Without gradients the encoder and the blend run a few frames at a time and the
|
|
softmax is streamed; the same computation, so equal up to float rounding."""
|
|
net_ = _trained_looking(False)
|
|
sources = stack(6, size=96, seed=4).clamp(-1, 1)
|
|
cond = conditioning(["available", "available", "available", "missing",
|
|
"available", "available"],
|
|
satellites=[True, True, True, True, False, True],
|
|
dts=[-240.0, 240.0, -960.0, 960.0, 0.0, -3840.0])
|
|
present = torch.ones(1, 6, BANDS)
|
|
present[0, 4, 1] = 0.0
|
|
monkeypatch.setattr(model, "BAND_AT_A_TIME", False)
|
|
with torch.enable_grad():
|
|
reference = net_(sources, cond, present).detach()
|
|
with torch.no_grad():
|
|
assert torch.allclose(net_(sources, cond, present), reference,
|
|
rtol=0, atol=5e-5)
|
|
|
|
|
|
@pytest.mark.parametrize("band_at_a_time", [True, False])
|
|
def test_chunked_inference_handles_batches_ragged_chunks_and_unheld_bands(
|
|
monkeypatch, band_at_a_time):
|
|
"""Two stacks of seven frames -- not a multiple of the chunk -- with one band no
|
|
frame of the second stack holds: that band blends to zero, not NaN."""
|
|
monkeypatch.setattr(model, "BAND_AT_A_TIME", band_at_a_time)
|
|
net_ = _trained_looking(False)
|
|
sources = torch.cat([stack(7, size=48, seed=6), stack(7, size=48, seed=7)]).clamp(-1, 1)
|
|
states = ["available", "available", "missing", "available", "available",
|
|
"available", "available"]
|
|
dts = [-240.0, 240.0, -960.0, 960.0, 0.0, -3840.0, 3840.0]
|
|
cond = torch.cat([conditioning(states, dts=dts),
|
|
conditioning(states, satellites=[True] * 6 + [False], dts=dts)])
|
|
present = torch.ones(2, 7, BANDS)
|
|
present[1, :, 3] = 0.0
|
|
with torch.enable_grad():
|
|
reference = net_(sources, cond, present).detach()
|
|
with torch.no_grad():
|
|
out = net_(sources, cond, present)
|
|
assert torch.isfinite(out).all()
|
|
assert float(out[1, 3].abs().max()) == 0.0
|
|
assert torch.allclose(out, reference, rtol=0, atol=5e-5)
|
|
|
|
|
|
def test_trends_are_returned_apart_from_the_stack():
|
|
frames = torch.full((1, 4, BANDS, 16, 16), 0.5)
|
|
condition = conditioning(["available"] * 4, dts=[-480.0, -240.0, 240.0, 480.0])
|
|
trends, expanded, present = model.build(base=8, depth=2)._trends(frames, condition)
|
|
assert trends.shape == (1, 2, BANDS, 16, 16)
|
|
assert expanded.shape[1] == present.shape[1] == 6
|
|
|
|
|
|
def test_a_one_frame_stack_has_no_trends_and_still_blends(net):
|
|
frames = stack(1, size=48)
|
|
trends, condition, _ = net._trends(frames, conditioning(["available"]))
|
|
assert trends.shape[1] == 0 and condition.shape[1] == 1
|
|
out = net(frames, conditioning(["available"], dts=[240.0]))
|
|
assert torch.isfinite(out).all()
|
|
|
|
|
|
def test_training_keeps_the_batched_path_and_its_gradients():
|
|
net_ = _trained_looking(False).train()
|
|
sources = stack(4, size=48, seed=5).clamp(-1, 1).requires_grad_(True)
|
|
cond = conditioning(["available"] * 4)
|
|
net_(sources, cond).square().mean().backward()
|
|
assert sources.grad is not None and float(sources.grad.abs().sum()) > 0
|
|
assert net_.weight.weight.grad is not None
|