290 lines
12 KiB
Python
290 lines
12 KiB
Python
|
import numpy as np
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
torch = pytest.importorskip("torch")
|
||
|
|
|
||
|
|
import train # noqa: E402
|
||
|
|
from suvi import dataset, metrics, model, paths, samples # noqa: E402
|
||
|
|
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------ display mapping
|
||
|
|
|
||
|
|
|
||
|
|
def test_hard_display_map_matches_the_metric_the_bench_scores():
|
||
|
|
"""With eps=0 the torch map must *be* metrics.to_display, or the reported PSNR
|
||
|
|
stops being comparable to the published filler table."""
|
||
|
|
rng = np.random.default_rng(0)
|
||
|
|
radiance = rng.uniform(-1.0, 60.0, (1, 6, 16, 16)).astype(np.float32)
|
||
|
|
coded = torch.from_numpy(samples.encode_for_model(radiance))
|
||
|
|
|
||
|
|
got = train.display_map(coded, eps=0.0).numpy()
|
||
|
|
for index, wavelength in enumerate(paths.WAVELENGTHS):
|
||
|
|
expected = metrics.to_display(radiance[0, index], wavelength)
|
||
|
|
np.testing.assert_allclose(got[0, index], expected, atol=2e-3)
|
||
|
|
|
||
|
|
|
||
|
|
def test_display_map_still_sends_zero_to_zero_when_softened():
|
||
|
|
"""The softening must not lift the black level: (0+eps)^g - eps^g == 0."""
|
||
|
|
coded = samples.encode_for_model(np.zeros((1, 6, 4, 4), np.float32))
|
||
|
|
out = train.display_map(torch.from_numpy(coded))
|
||
|
|
assert float(out.abs().max()) < 1e-6 # fp32 rounding only; 1/255 is 4e-3
|
||
|
|
|
||
|
|
|
||
|
|
def test_display_map_saturates_without_nan():
|
||
|
|
for value in (-1.0, 1.0):
|
||
|
|
coded = torch.full((1, 6, 4, 4), value, requires_grad=True)
|
||
|
|
out = train.display_map(coded)
|
||
|
|
out.sum().backward()
|
||
|
|
assert torch.isfinite(out).all() and torch.isfinite(coded.grad).all()
|
||
|
|
|
||
|
|
|
||
|
|
# --------------------------------------------------------------- loss conditioning
|
||
|
|
|
||
|
|
|
||
|
|
def test_the_loss_gradient_is_bounded_everywhere():
|
||
|
|
"""The property whose absence sank three training runs.
|
||
|
|
|
||
|
|
The old display-space loss had |dL/dpred| of exactly 0 across most of the frame
|
||
|
|
(below the vmin clamp) and a 25x spike with a hard discontinuity at the
|
||
|
|
threshold; clip_grad_norm then renormalised every update onto that annulus.
|
||
|
|
Here the whole coded range, including the sub-vmin corona and the threshold
|
||
|
|
itself, must see a per-pixel gradient that is neither zero over wide regions nor
|
||
|
|
orders of magnitude apart.
|
||
|
|
"""
|
||
|
|
# Sweep the entire coded range, so the threshold annulus of every band is in,
|
||
|
|
# with a *constant* residual: the per-pixel gradient then measures the loss
|
||
|
|
# landscape's Jacobian, not the accident of which pixels happened to be right.
|
||
|
|
target = torch.linspace(-1.0, 1.0, 64 * 64).reshape(1, 1, 64, 64).repeat(1, 6, 1, 1)
|
||
|
|
prediction = (target + 0.05).requires_grad_(True)
|
||
|
|
|
||
|
|
train.criterion(prediction, target).backward()
|
||
|
|
per_pixel = prediction.grad.abs() * prediction.numel()
|
||
|
|
assert torch.isfinite(prediction.grad).all()
|
||
|
|
assert float(per_pixel.max()) < 100.0, "a gradient spike survived the redesign"
|
||
|
|
# The asinh term supervises everything: no dead zones anywhere.
|
||
|
|
assert float(per_pixel.min()) > 0.1, "part of the frame is unsupervised"
|
||
|
|
|
||
|
|
|
||
|
|
def test_charbonnier_is_zero_for_a_perfect_match():
|
||
|
|
x = torch.rand(4, 6, 8, 8)
|
||
|
|
assert train.charbonnier(x, x) == pytest.approx(train.CHARBONNIER_EPS, abs=1e-6)
|
||
|
|
|
||
|
|
|
||
|
|
def test_charbonnier_is_less_outlier_driven_than_squared_error():
|
||
|
|
"""A saturated pixel must not be able to dominate the whole frame's loss."""
|
||
|
|
truth = torch.zeros(1, 1, 8, 8)
|
||
|
|
spike = truth.clone()
|
||
|
|
spike[0, 0, 0, 0] = 50.0
|
||
|
|
spread = truth + 0.1
|
||
|
|
|
||
|
|
charbonnier_ratio = float(train.charbonnier(spike, truth)
|
||
|
|
/ train.charbonnier(spread, truth))
|
||
|
|
squared_ratio = float(((spike - truth) ** 2).mean() / ((spread - truth) ** 2).mean())
|
||
|
|
assert charbonnier_ratio < 10
|
||
|
|
assert squared_ratio > 1000
|
||
|
|
|
||
|
|
|
||
|
|
def test_gradient_loss_punishes_blur():
|
||
|
|
"""The term that makes a sharp answer beat a hedged smooth one."""
|
||
|
|
truth = torch.zeros(1, 1, 32, 32)
|
||
|
|
truth[:, :, 8:24, 8:24] = 1.0
|
||
|
|
blurred = torch.nn.functional.avg_pool2d(truth, 5, stride=1, padding=2)
|
||
|
|
assert train.gradient_loss(blurred, truth) > train.gradient_loss(truth, truth)
|
||
|
|
|
||
|
|
|
||
|
|
def test_display_psnr_matches_the_bench_definition():
|
||
|
|
"""Per band, then averaged -- the pooled-MSE PSNR reads ~4 dB lower (Jensen) and
|
||
|
|
would not be comparable to the published baselines."""
|
||
|
|
rng = np.random.default_rng(1)
|
||
|
|
radiance = rng.uniform(0.0, 30.0, (6, 16, 16)).astype(np.float32)
|
||
|
|
noisy = radiance + rng.normal(0, 0.5, radiance.shape).astype(np.float32)
|
||
|
|
|
||
|
|
expected = []
|
||
|
|
for index, wavelength in enumerate(paths.WAVELENGTHS):
|
||
|
|
a = metrics.to_display(radiance[index], wavelength)
|
||
|
|
b = metrics.to_display(noisy[index], wavelength)
|
||
|
|
expected.append(10.0 * np.log10(1.0 / max(((a - b) ** 2).mean(), 1e-12)))
|
||
|
|
|
||
|
|
got = train.display_psnr(
|
||
|
|
torch.from_numpy(samples.encode_for_model(noisy))[None],
|
||
|
|
torch.from_numpy(samples.encode_for_model(radiance))[None],
|
||
|
|
)
|
||
|
|
assert got == pytest.approx(float(np.mean(expected)), abs=0.1)
|
||
|
|
|
||
|
|
|
||
|
|
def test_display_psnr_of_identical_inputs_is_finite():
|
||
|
|
x = torch.rand(1, 6, 8, 8) * 0.5
|
||
|
|
assert np.isfinite(train.display_psnr(x, x))
|
||
|
|
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------------- collate
|
||
|
|
|
||
|
|
|
||
|
|
def item(depth, bands=6, size=16):
|
||
|
|
return {
|
||
|
|
"stack": torch.randn(depth, bands, size, size),
|
||
|
|
"condition": torch.stack(
|
||
|
|
[model.frame_conditioning("available", True, 240.0) for _ in range(depth)]
|
||
|
|
),
|
||
|
|
"dts": torch.full((depth,), 240.0),
|
||
|
|
"gains": torch.ones(depth, bands),
|
||
|
|
"offsets": torch.zeros(depth, bands),
|
||
|
|
"b0": torch.tensor(0.05),
|
||
|
|
"radius": torch.tensor(0.3),
|
||
|
|
"target": torch.randn(bands, size, size),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def test_collate_pads_ragged_stacks():
|
||
|
|
"""Stack depth varies with how many anchors were found near a shard's edge."""
|
||
|
|
batch = train.collate([item(5), item(7), item(6)])
|
||
|
|
assert batch["stack"].shape[:2] == (3, 7)
|
||
|
|
assert batch["condition"].shape == (3, 7, model.COND_DIM)
|
||
|
|
assert batch["gains"].shape == (3, 7, 6)
|
||
|
|
assert batch["target"].shape == (3, 6, 16, 16)
|
||
|
|
assert batch["b0"].shape == (3,)
|
||
|
|
|
||
|
|
|
||
|
|
def test_collate_marks_padding_as_missing_with_identity_transfer():
|
||
|
|
"""Padding must take zero weight and pass alignment untouched, or the batch
|
||
|
|
shape would change the answer."""
|
||
|
|
batch = train.collate([item(3), item(6)])
|
||
|
|
assert torch.all(batch["condition"][0, 3:, 1] == 1.0) # 'missing' one-hot
|
||
|
|
assert torch.all(batch["stack"][0, 3:] == 0.0)
|
||
|
|
assert torch.all(batch["gains"][0, 3:] == 1.0)
|
||
|
|
assert torch.all(batch["offsets"][0, 3:] == 0.0)
|
||
|
|
|
||
|
|
|
||
|
|
def test_padding_does_not_change_the_prediction():
|
||
|
|
"""The property that makes ragged batching safe at all."""
|
||
|
|
torch.manual_seed(0)
|
||
|
|
net = model.build(base=8, depth=2).eval()
|
||
|
|
short = item(3, size=32)
|
||
|
|
alone = train.collate([short])
|
||
|
|
padded = train.collate([short, item(6, size=32)])
|
||
|
|
with torch.no_grad():
|
||
|
|
a = net(alone["stack"], alone["condition"])
|
||
|
|
b = net(padded["stack"][:1], padded["condition"][:1])
|
||
|
|
assert torch.allclose(a, b, atol=1e-5)
|
||
|
|
|
||
|
|
|
||
|
|
# ---------------------------------------------------------------------- data feed
|
||
|
|
|
||
|
|
|
||
|
|
class MemoryShard:
|
||
|
|
def __init__(self, satellite, times, size=64):
|
||
|
|
from conftest import solar_disc
|
||
|
|
|
||
|
|
self.satellite = satellite
|
||
|
|
self._frames = {
|
||
|
|
t: np.stack([solar_disc(size=size, radius=size // 3, peak=1.0 + b * 0.2)
|
||
|
|
for b 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()
|
||
|
|
|
||
|
|
|
||
|
|
def shard_set(monkeypatch, tmp_path, days, deterministic=False, seed=0):
|
||
|
|
base = 1715400000 // paths.CADENCE * paths.CADENCE
|
||
|
|
times = [base + i * paths.CADENCE for i in range(40)]
|
||
|
|
monkeypatch.setattr(dataset, "Shard", lambda path: MemoryShard(16, times))
|
||
|
|
real_exists = train.os.path.exists
|
||
|
|
# Pretend only the shard files exist; a blanket True breaks cv2's lazy loader.
|
||
|
|
monkeypatch.setattr(
|
||
|
|
train.os.path, "exists",
|
||
|
|
lambda p: True if str(p).startswith(str(tmp_path)) else real_exists(p),
|
||
|
|
)
|
||
|
|
return train.ShardSet(str(tmp_path), days, satellites=(16, 18),
|
||
|
|
deterministic=deterministic, seed=seed)
|
||
|
|
|
||
|
|
|
||
|
|
def test_shard_set_yields_model_ready_tensors(monkeypatch, tmp_path):
|
||
|
|
data = shard_set(monkeypatch, tmp_path, ["2024-01-01"])
|
||
|
|
sample = data[0]
|
||
|
|
assert sample["stack"].shape[1:] == (6, 64, 64)
|
||
|
|
assert sample["condition"].shape[1] == model.COND_DIM
|
||
|
|
assert sample["gains"].shape == (sample["stack"].shape[0], 6)
|
||
|
|
assert torch.isfinite(sample["stack"]).all()
|
||
|
|
assert torch.isfinite(sample["target"]).all()
|
||
|
|
|
||
|
|
|
||
|
|
def test_validation_damage_is_reproducible(monkeypatch, tmp_path):
|
||
|
|
"""A val number that moves because the damage moved measures nothing."""
|
||
|
|
first = shard_set(monkeypatch, tmp_path, ["2024-01-01"], deterministic=True, seed=1)
|
||
|
|
second = shard_set(monkeypatch, tmp_path, ["2024-01-01"], deterministic=True, seed=1)
|
||
|
|
for index in (0, 5, 11):
|
||
|
|
assert torch.allclose(first[index]["stack"], second[index]["stack"])
|
||
|
|
|
||
|
|
|
||
|
|
def test_training_damage_varies_between_epochs(monkeypatch, tmp_path):
|
||
|
|
"""The reason shards store clean frames: unlimited corruption variety."""
|
||
|
|
data = shard_set(monkeypatch, tmp_path, ["2024-01-01"], deterministic=False)
|
||
|
|
assert not torch.allclose(data[3]["stack"], data[3]["stack"])
|
||
|
|
|
||
|
|
|
||
|
|
def test_shard_set_refuses_to_run_on_nothing(tmp_path):
|
||
|
|
with pytest.raises(SystemExit, match="no shards"):
|
||
|
|
train.ShardSet(str(tmp_path), ["2024-01-01"])
|
||
|
|
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------ training steps
|
||
|
|
|
||
|
|
|
||
|
|
def loader_for(monkeypatch, tmp_path):
|
||
|
|
data = shard_set(monkeypatch, tmp_path, ["2024-01-01"])
|
||
|
|
return torch.utils.data.DataLoader(data, batch_size=2, num_workers=0,
|
||
|
|
collate_fn=train.collate)
|
||
|
|
|
||
|
|
|
||
|
|
def test_a_training_step_is_finite_end_to_end(monkeypatch, tmp_path):
|
||
|
|
"""Data feed, alignment, forward, loss, backward: all wired together."""
|
||
|
|
loader = loader_for(monkeypatch, tmp_path)
|
||
|
|
net = model.build(base=8, depth=2)
|
||
|
|
optimiser = torch.optim.AdamW(net.parameters(), lr=1e-4)
|
||
|
|
device = torch.device("cpu")
|
||
|
|
|
||
|
|
for index, batch in enumerate(loader):
|
||
|
|
if index >= 2:
|
||
|
|
break
|
||
|
|
aligned, condition, target = train.aligned_batch(batch, device)
|
||
|
|
prediction = net(aligned, condition)
|
||
|
|
loss = train.criterion(prediction, target)
|
||
|
|
assert torch.isfinite(loss)
|
||
|
|
optimiser.zero_grad()
|
||
|
|
loss.backward()
|
||
|
|
grad = float(torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0))
|
||
|
|
assert np.isfinite(grad)
|
||
|
|
optimiser.step()
|
||
|
|
|
||
|
|
|
||
|
|
def test_evaluate_leaves_weights_alone(monkeypatch, tmp_path):
|
||
|
|
"""Validation must not train, or the held-out split stops being held out."""
|
||
|
|
loader = loader_for(monkeypatch, tmp_path)
|
||
|
|
net = model.build(base=8, depth=2)
|
||
|
|
before = [p.detach().clone() for p in net.parameters()]
|
||
|
|
stats = train.evaluate(net, loader, torch.device("cpu"))
|
||
|
|
assert np.isfinite(stats["loss"]) and np.isfinite(stats["psnr"])
|
||
|
|
for old, new in zip(before, net.parameters()):
|
||
|
|
assert torch.equal(old, new)
|
||
|
|
|
||
|
|
|
||
|
|
def test_validation_subset_spans_the_whole_set(monkeypatch, tmp_path):
|
||
|
|
"""A shortened val pass must not become 'the first N targets': `targets()` is
|
||
|
|
sorted by (satellite, time), so a prefix is the earliest days of one satellite."""
|
||
|
|
data = shard_set(monkeypatch, tmp_path, ["2024-01-01"])
|
||
|
|
wanted = 8
|
||
|
|
stride = len(data) / wanted
|
||
|
|
subset = torch.utils.data.Subset(data, [int(i * stride) for i in range(wanted)])
|
||
|
|
indices = subset.indices
|
||
|
|
assert len(set(indices)) == wanted
|
||
|
|
assert max(indices) > len(data) * 0.8, "subset does not reach the end of the set"
|
||
|
|
assert min(indices) < len(data) * 0.2, "subset does not start near the beginning"
|
||
|
|
satellites = {data.items[i][0] for i in indices}
|
||
|
|
assert len(satellites) > 1, "subset covers only one satellite"
|