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 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_conditioning_separates_the_three_states(): vectors = [model.frame_conditioning(s, True, 240.0)[:3] for s in ("available", "missing", "suspect")] for index, vector in enumerate(vectors): assert vector[index] == 1.0 and vector.sum() == 1.0 def test_conditioning_rejects_an_unknown_state(): with pytest.raises(KeyError): model.frame_conditioning("probably-fine", 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[4] == -1.0 and forward[4] == 1.0 assert back[5] == pytest.approx(forward[5]) # magnitude, not direction assert far[5] > forward[5] def test_conditioning_leaves_agreement_to_the_model(): """Columns 6 and 7 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[6] == 0.0 and vector[7] == 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_downweights_suspect_frames_without_excluding_them(): clean = conditioning(["available"], satellites=[True], dts=[-240.0]) suspect = conditioning(["suspect"], satellites=[True], dts=[-240.0]) assert model.blend_prior(clean)[0, 0] > model.blend_prior(suspect)[0, 0] assert torch.isfinite(model.blend_prior(suspect)).all() # ------------------------------------------------------------------------- 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[..., 0] + condition[..., 2]).clamp(0, 1) 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.""" frames = stack(5) condition = conditioning(["available"] * 5) out = net(frames, condition) assert torch.allclose(out, expected_blend(frames, 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_suspect_frames_do_take_weight(net): """The design decision under test: a flagged frame is data, not a hole.""" frames = stack(3) excluded = conditioning(["missing", "available", "available"]) included = conditioning(["suspect", "available", "available"]) a = net(frames, excluded) b = net(frames, included) assert not torch.allclose(a, b, atol=1e-3) assert torch.allclose(b, expected_blend(frames, included), atol=1e-4) 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[..., 0] + condition[..., 2]).clamp(0, 1) got = fresh._with_agreement(frames.reshape(4, BANDS, 32, 32), condition, valid, batch=1, stack=4) assert float(got[0, 3, 6]) > float(got[0, 0, 6]) + 0.5 assert float(got[0, 0, 6]) < 0.1 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[..., 0] + condition[..., 2]).clamp(0, 1) got = fresh._with_agreement(frames.reshape(3, BANDS, 32, 32), condition, valid, batch=1, stack=3) assert float(got[..., 6:].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[..., 0] + condition[..., 2]).clamp(0, 1) got = fresh._with_agreement(frames.reshape(3, BANDS, 32, 32), condition, valid, batch=1, stack=3) assert float(got[0, 1, 6]) < 1e-6 assert float(got[0, 2, 6]) < 1e-6 # ------------------------------------------------------------------------ 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", "suspect", "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}"