274 lines
12 KiB
Python
274 lines
12 KiB
Python
"""The deterministic alignment stage: ephemeris, rotation warp, photometric transfer.
|
|
|
|
This is the physics the model no longer has to learn, so its correctness bounds the
|
|
whole system: a wrong warp or a wrong gain poisons every candidate the network is
|
|
allowed to blend.
|
|
"""
|
|
|
|
import datetime as dt
|
|
import math
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from conftest import solar_disc
|
|
from suvi import align, samples
|
|
|
|
torch = pytest.importorskip("torch")
|
|
|
|
|
|
def unix(year, month, day):
|
|
return dt.datetime(year, month, day, tzinfo=dt.timezone.utc).timestamp()
|
|
|
|
|
|
# ---------------------------------------------------------------------- ephemeris
|
|
|
|
|
|
def test_ephemeris_b0_reaches_its_annual_extremes():
|
|
"""B0 swings +/-7.25 deg, peaking in early March and September."""
|
|
assert math.degrees(align.solar_ephemeris(unix(2024, 3, 7))[0]) == pytest.approx(
|
|
-7.25, abs=0.15
|
|
)
|
|
assert math.degrees(align.solar_ephemeris(unix(2024, 9, 8))[0]) == pytest.approx(
|
|
7.25, abs=0.15
|
|
)
|
|
|
|
|
|
def test_ephemeris_b0_crosses_zero_in_june_and_december():
|
|
for when in (unix(2024, 6, 6), unix(2024, 12, 7)):
|
|
assert abs(math.degrees(align.solar_ephemeris(when)[0])) < 0.5
|
|
|
|
|
|
def test_ephemeris_radius_tracks_the_orbit():
|
|
"""Apparent radius runs ~944 arcsec at aphelion (July) to ~976 at perihelion."""
|
|
def arcsec(when):
|
|
return align.solar_ephemeris(when)[1] * align.NATIVE_SIZE * align.PLATE_SCALE
|
|
|
|
assert arcsec(unix(2024, 1, 3)) == pytest.approx(975.9, abs=1.0)
|
|
assert arcsec(unix(2024, 7, 5)) == pytest.approx(943.9, abs=1.0)
|
|
|
|
|
|
def test_ephemeris_matches_a_real_archive_header():
|
|
"""Pinned against dr_suvi-l2-ci094_g16_s20240510T000000Z (verified by hand):
|
|
DIAM_SUN = 760.2932 px, SOLAR_B0 = -3.205694 deg."""
|
|
when = dt.datetime(2024, 5, 10, 0, 2, tzinfo=dt.timezone.utc).timestamp()
|
|
b0, radius_fraction = align.solar_ephemeris(when)
|
|
assert 2 * radius_fraction * 1280 == pytest.approx(760.2932, rel=5e-4)
|
|
assert math.degrees(b0) == pytest.approx(-3.205694, abs=0.05)
|
|
|
|
|
|
# ------------------------------------------------------------------ rotation warp
|
|
|
|
|
|
def numpy_reference(size, dt_seconds, b0, radius_fraction):
|
|
"""The already-validated numpy map from suvi.fillers, on the same geometry."""
|
|
from suvi import fillers
|
|
|
|
header = {
|
|
"diam_sun": 2 * radius_fraction * size,
|
|
"crpix1": (size + 1) / 2,
|
|
"crpix2": (size + 1) / 2,
|
|
"solar_b0": math.degrees(b0),
|
|
}
|
|
return fillers._rotation_map((size, size), header, dt_seconds)
|
|
|
|
|
|
def test_rotation_grid_matches_the_numpy_map():
|
|
"""One implementation trains, the other filled the published baselines; a drift
|
|
between them would score the learned filler against different physics."""
|
|
size, lag, b0, radius = 128, 20 * 3600.0, -0.056, 760.29 / 2 / 1280
|
|
map_x, map_y, visible = numpy_reference(size, lag, b0, radius)
|
|
|
|
grid = align.rotation_grid(size, torch.tensor([lag]), torch.tensor([b0]),
|
|
torch.tensor([radius]))[0].numpy()
|
|
got_x = (grid[..., 0] + 1) * size / 2 - 0.5
|
|
got_y = (grid[..., 1] + 1) * size / 2 - 0.5
|
|
assert np.abs(got_x - map_x)[visible].max() < 1e-3
|
|
assert np.abs(got_y - map_y)[visible].max() < 1e-3
|
|
|
|
|
|
def test_rotation_grid_is_identity_where_nothing_co_rotates():
|
|
"""Off-disc and behind-the-limb pixels keep their own values: the corona above
|
|
the limb does not rotate with the photosphere."""
|
|
size, lag, b0, radius = 64, 20 * 3600.0, 0.02, 0.3
|
|
_, _, visible = numpy_reference(size, lag, b0, radius)
|
|
grid = align.rotation_grid(size, torch.tensor([lag]), torch.tensor([b0]),
|
|
torch.tensor([radius]))[0].numpy()
|
|
grid_y, grid_x = np.mgrid[0:size, 0:size]
|
|
got_x = (grid[..., 0] + 1) * size / 2 - 0.5
|
|
got_y = (grid[..., 1] + 1) * size / 2 - 0.5
|
|
assert np.abs(got_x - grid_x)[~visible].max() == 0.0
|
|
assert np.abs(got_y - grid_y)[~visible].max() == 0.0
|
|
|
|
|
|
def test_rotate_at_zero_dt_is_the_identity():
|
|
frames = torch.randn(2, 6, 96, 96)
|
|
out = align.rotate(frames, torch.zeros(2), torch.tensor([0.05, -0.05]),
|
|
torch.full((2,), 0.3))
|
|
assert float((out - frames).abs().max()) < 1e-3
|
|
|
|
|
|
def test_rotate_moves_on_disc_content():
|
|
disc = torch.from_numpy(
|
|
np.stack([solar_disc(size=96, radius=30, peak=2.0)] * 6)[None]
|
|
).float()
|
|
# An off-centre bright blob, so rotation has something visible to move.
|
|
disc[..., 40:48, 30:38] += 3.0
|
|
out = align.rotate(disc, torch.tensor([12 * 3600.0]), torch.tensor([0.0]),
|
|
torch.tensor([30 / 96]))
|
|
assert float((out - disc).abs().max()) > 0.5
|
|
|
|
|
|
def test_rotate_round_trips():
|
|
"""Forward then back must land where it started, on the visible interior."""
|
|
disc = torch.from_numpy(
|
|
np.stack([solar_disc(size=96, radius=30, peak=2.0)] * 6)[None]
|
|
).float()
|
|
lag = torch.tensor([6 * 3600.0])
|
|
b0, radius = torch.tensor([0.03]), torch.tensor([30 / 96])
|
|
there = align.rotate(disc, lag, b0, radius)
|
|
back = align.rotate(there, -lag, b0, radius)
|
|
centre = (slice(None), slice(None), slice(38, 58), slice(38, 58))
|
|
assert float((back[centre] - disc[centre]).abs().max()) < 0.05
|
|
|
|
|
|
# -------------------------------------------------------------------- photometry
|
|
|
|
|
|
def test_apply_photometry_matches_radiance_arithmetic():
|
|
"""The transfer is defined on radiance; applying it in coded space would model a
|
|
different fault entirely."""
|
|
radiance = np.stack([solar_disc(size=32, radius=10, peak=1.0 + b) for b in range(6)])
|
|
coded = torch.from_numpy(samples.encode_for_model(radiance))[None]
|
|
gain = torch.tensor([[0.81, 0.885, 1.49, 1.0, 1.2, 0.86]])
|
|
offset = torch.tensor([[0.0, 0.01, -0.02, 0.0, 0.005, 0.0]])
|
|
|
|
moved = align.apply_photometry(coded, gain, offset)
|
|
got = samples.decode_from_model(moved[0].numpy())
|
|
expected = radiance * gain[0, :, None, None].numpy() + offset[0, :, None, None].numpy()
|
|
np.testing.assert_allclose(got, expected, rtol=3e-3, atol=1e-5)
|
|
|
|
|
|
def test_apply_photometry_stays_in_the_coded_range():
|
|
coded = torch.full((1, 6, 8, 8), 1.0)
|
|
moved = align.apply_photometry(coded, torch.full((1, 6), 100.0),
|
|
torch.full((1, 6), 1e6))
|
|
assert torch.isfinite(moved).all()
|
|
assert float(moved.abs().max()) <= 1.0
|
|
|
|
|
|
def test_fit_photometry_recovers_a_known_transfer():
|
|
rng = np.random.default_rng(0)
|
|
local = rng.uniform(0.0, 2.0, (6, 32, 32)).astype(np.float32)
|
|
gains = np.array([0.81, 0.885, 1.49, 1.0, 1.2, 0.86], dtype=np.float32)
|
|
offsets = np.array([0.0, 0.01, -0.02, 0.0, 0.005, 0.03], dtype=np.float32)
|
|
counterpart = (local - offsets[:, None, None]) / gains[:, None, None]
|
|
|
|
got_gains, got_offsets = align.fit_photometry(counterpart, local)
|
|
np.testing.assert_allclose(got_gains, gains, rtol=1e-4)
|
|
np.testing.assert_allclose(got_offsets, offsets, atol=1e-4)
|
|
|
|
|
|
# ------------------------------------------------------------------- whole stacks
|
|
|
|
|
|
def coded_disc(peak=1.0, size=64):
|
|
radiance = np.stack([solar_disc(size=size, radius=size // 3, peak=peak + b * 0.2)
|
|
for b in range(6)])
|
|
return radiance, torch.from_numpy(samples.encode_for_model(radiance))
|
|
|
|
|
|
def test_align_stack_puts_the_counterpart_on_the_target_scale():
|
|
gains = np.array([0.81, 0.885, 1.49, 1.0, 1.2, 0.86], dtype=np.float32)
|
|
local_radiance, local = coded_disc()
|
|
counter_radiance = local_radiance / gains[:, None, None]
|
|
counter = torch.from_numpy(samples.encode_for_model(counter_radiance))
|
|
|
|
stack = torch.stack([counter])[None]
|
|
aligned = align.align_stack(
|
|
stack, dts=torch.zeros(1, 1), valid=torch.ones(1, 1),
|
|
gains=torch.from_numpy(gains)[None, None],
|
|
offsets=torch.zeros(1, 1, 6),
|
|
b0=torch.tensor([0.0]), radius_fraction=torch.tensor([0.3]),
|
|
)
|
|
assert float((aligned[0, 0] - local).abs().max()) < 5e-3
|
|
|
|
|
|
def test_align_stack_leaves_missing_frames_at_zero():
|
|
"""An offset applied to a frame of zeros would manufacture an image from nothing."""
|
|
stack = torch.zeros(1, 1, 6, 32, 32)
|
|
aligned = align.align_stack(
|
|
stack, dts=torch.zeros(1, 1), valid=torch.zeros(1, 1),
|
|
gains=torch.full((1, 1, 6), 3.0), offsets=torch.full((1, 1, 6), 0.5),
|
|
b0=torch.tensor([0.0]), radius_fraction=torch.tensor([0.3]),
|
|
)
|
|
assert float(aligned.abs().max()) == 0.0
|
|
|
|
|
|
def test_align_stack_leaves_a_simultaneous_local_frame_alone():
|
|
_, local = coded_disc()
|
|
stack = torch.stack([local])[None]
|
|
aligned = align.align_stack(
|
|
stack, dts=torch.zeros(1, 1), valid=torch.ones(1, 1),
|
|
gains=torch.ones(1, 1, 6), offsets=torch.zeros(1, 1, 6),
|
|
b0=torch.tensor([0.05]), radius_fraction=torch.tensor([0.3]),
|
|
)
|
|
assert float((aligned[0, 0] - local).abs().max()) < 1e-3
|
|
|
|
|
|
def test_align_stack_warps_toward_the_target_not_away():
|
|
"""The direction test the first training run lacked.
|
|
|
|
A 'before' frame (dt < 0) must be rotated *forward* onto the target instant.
|
|
The sign error this pins -- warping by dt instead of -dt -- doubled the
|
|
misalignment at long gaps while staying sub-pixel at short ones, so only a
|
|
test that compares against an independently-warped truth can catch it.
|
|
"""
|
|
from suvi import fillers
|
|
|
|
size, lag, radius = 96, 15 * 3600.0, 30 / 96
|
|
base = np.stack([solar_disc(size=size, radius=30, peak=2.0)] * 6)
|
|
base[:, 40:48, 30:38] += 3.0 # feature rotation will move
|
|
header = {"diam_sun": 2 * radius * size, "crpix1": (size + 1) / 2,
|
|
"crpix2": (size + 1) / 2, "solar_b0": 0.0}
|
|
# Independent reference: the validated numpy warp advances `base` by +lag.
|
|
truth = np.stack([fillers._warp(band, header, lag, True)[0] for band in base])
|
|
|
|
coded = torch.from_numpy(samples.encode_for_model(base))[None, None]
|
|
aligned = align.align_stack(
|
|
coded, dts=torch.full((1, 1), -lag), valid=torch.ones(1, 1),
|
|
gains=torch.ones(1, 1, 6), offsets=torch.zeros(1, 1, 6),
|
|
b0=torch.tensor([0.0]), radius_fraction=torch.tensor([radius]),
|
|
)
|
|
got = samples.decode_from_model(aligned[0, 0].numpy())
|
|
interior = (slice(None), slice(30, 66), slice(30, 66))
|
|
aligned_error = float(np.abs(got[interior] - truth[interior]).mean())
|
|
unwarped_error = float(np.abs(base[interior] - truth[interior]).mean())
|
|
assert aligned_error < unwarped_error * 0.35, (
|
|
f"aligned {aligned_error:.4f} vs unwarped {unwarped_error:.4f}: "
|
|
"the rotation warp is not moving frames onto the target instant"
|
|
)
|
|
|
|
|
|
def test_align_stack_warps_stale_frames_toward_the_target_instant():
|
|
_, local = coded_disc()
|
|
marked = local.clone()
|
|
marked[..., 20:26, 14:20] += 0.4 # off-centre, on-disc
|
|
stack = torch.stack([marked])[None]
|
|
aligned = align.align_stack(
|
|
stack, dts=torch.full((1, 1), 15 * 3600.0), valid=torch.ones(1, 1),
|
|
gains=torch.ones(1, 1, 6), offsets=torch.zeros(1, 1, 6),
|
|
b0=torch.tensor([0.0]), radius_fraction=torch.tensor([0.3]),
|
|
)
|
|
assert float((aligned[0, 0] - marked).abs().max()) > 0.05
|
|
|
|
|
|
def test_shard_and_ephemeris_share_one_asinh_convention():
|
|
"""align.apply_photometry decodes with dataset's constants; if those drift apart
|
|
the photometric transfer silently degrades."""
|
|
values = np.array([[-3.7, 0.0, 1e-3, 0.19, 50.0]], dtype=np.float32)
|
|
coded = torch.from_numpy(samples.encode_for_model(values))
|
|
identity = align.apply_photometry(coded[None, :, None], torch.ones(1, 1),
|
|
torch.zeros(1, 1))
|
|
restored = samples.decode_from_model(identity[0, :, 0].numpy())
|
|
np.testing.assert_allclose(restored, values, rtol=2e-3, atol=1e-7)
|