noaa-goes-visualization/tests/test_ceiling.py

76 lines
3.2 KiB
Python

"""ceiling.py's reference blends, which are per band: a frame may lack some bands."""
import numpy as np
import pytest
torch = pytest.importorskip("torch")
import ceiling # noqa: E402
from suvi import model, samples # noqa: E402
def conditioning(dts, states=None):
states = states or ["available"] * len(dts)
return torch.stack([model.frame_conditioning(state, True, dt)
for state, dt in zip(states, dts)])[None]
def test_prior_weights_normalise_per_band_over_the_frames_holding_it():
cond = conditioning([-240.0, -2400.0, 240.0])[0].numpy()
usable = np.ones((3, 6), dtype=bool)
usable[0, 2] = False
weights = ceiling.prior_weights(cond, usable, decay=1.5)
assert weights.shape == (3, 6)
np.testing.assert_allclose(weights.sum(axis=0), 1.0)
assert weights[0, 2] == 0.0
assert weights[0, 0] > weights[1, 0] # recency still rules
np.testing.assert_allclose(weights[:, 0], weights[:, 1])
def test_prior_weights_ignore_satellite_and_direction():
"""Recency alone: the counterpart arrives gain-matched, and past and future are
equally good evidence."""
cond = torch.stack([model.frame_conditioning("available", same, dt)
for same, dt in ((True, -240.0), (False, 240.0))]).numpy()
weights = ceiling.prior_weights(cond, np.ones((2, 6), dtype=bool), decay=1.5)
np.testing.assert_allclose(weights[0], weights[1])
def test_blend_psnr_blends_each_band_with_its_own_weights():
"""Truth is frame 1 in every band but 3, where it is frame 0 (slightly off,
so every PSNR is finite): the per-band mix must beat taking frame 1 whole."""
from suvi import metrics, paths
radiance = (1.0, 2.0)
candidates = np.stack([samples.encode_for_model(np.full((6, 8, 8), r, np.float32))
for r in radiance])
truth = np.full((6, 8, 8), radiance[1] * 1.02, np.float32)
truth[3] = radiance[0] * 1.02
truth_display = [metrics.to_display(t, w) for t, w in zip(truth, paths.WAVELENGTHS)]
whole = np.zeros((2, 6))
whole[1] = 1.0
mixed = whole.copy()
mixed[:, 3] = (1.0, 0.0)
assert (ceiling.blend_psnr(candidates, mixed, truth_display)
> ceiling.blend_psnr(candidates, whole, truth_display) + 5.0)
def test_score_target_declines_when_a_band_has_no_usable_frame():
aligned = torch.zeros(1, 2, 6, 8, 8)
condition = conditioning([-240.0, 240.0])
present = torch.ones(1, 2, 6)
present[0, :, 4] = 0.0
truth = np.ones((6, 8, 8), np.float32)
assert ceiling.score_target(aligned, condition, present, truth, (1.5,)) is None
def test_score_target_scores_a_partial_stack_on_the_bands_held():
aligned = torch.full((1, 2, 6, 8, 8), 0.2)
aligned[0, 0, 4] = 0.9 # absent band, stand-in pixels
condition = conditioning([-240.0, 240.0])
present = torch.ones(1, 2, 6)
present[0, 0, 4] = 0.0
truth = samples.decode_from_model(np.full((6, 8, 8), 0.2, np.float32))
scores = ceiling.score_target(aligned, condition, present, truth, (1.5,))
assert scores["best1"] == float("inf")
assert scores["uniform"] == float("inf") # the 0.9 never blended
assert scores["prior[1.5]"] == float("inf")