jeremy/noaa-goes-visualization
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
noaa-goes-visualization/tests/test_learned_filler.py

656 lines
26 KiB
Python
Raw Blame History

"""The learned filler's wiring: the stack it is handed, and how it is loaded.
The model itself is covered in test_model.py. What matters here is that `bench.py`
hands it the same stack shape `suvi.samples` builds during training, that a band judged
bad is withheld from it, and that damage the judgment missed arrives as *corrupted
pixels* rather than as the pristine original.
"""
import numpy as np
import pytest
import bench
from conftest import solar_disc
from suvi import cases, fillers, paths, samples
WAVELENGTHS = paths.WAVELENGTHS
def overlay_for(satellites=(16, 18), count=40, deleted=()):
"""An Overlay whose slots resolve through in-memory lookups rather than files."""
base = 1715400000 // paths.CADENCE * paths.CADENCE
times = [base + i * paths.CADENCE for i in range(count)]
archive = {(s, w, t): f"truth/{s}/{w}/{t}"
for s in satellites for w in WAVELENGTHS for t in times}
return cases.Overlay(archive=archive, deleted=frozenset(deleted)), times, base
def reader(size=16, missing=()):
"""A six-band reader. `missing` names (satellite, time) pairs with no frame."""
frame = np.stack([solar_disc(size=size, radius=size // 3, peak=1.0 + i)
for i in range(6)]).astype(np.float32)
def read(satellite, when):
if (satellite, when) in missing:
return None
return frame.copy()
return read, frame
def judged_reader(monkeypatch, overlay, bad, size=16):
"""The real `bench._band_reader` over `overlay_for`'s in-memory archive: every
path decodes to the matching band of one six-band frame."""
_, frame = reader(size)
index = {wavelength: i for i, wavelength in enumerate(WAVELENGTHS)}
def read_image(path):
wavelength = int(path.split("/")[2])
return frame[index[wavelength]].copy(), None
monkeypatch.setattr(bench.fitsio, "read_image", read_image)
return bench._band_reader(overlay, WAVELENGTHS, set(bad)), frame
def build(overlay, target, read, bad=frozenset(), satellites=(16, 18)):
"""`bench._build_stack` for GOES-16, with anchors from the judgment `read` uses."""
band_times = bench._band_clean_times(overlay, satellites, WAVELENGTHS, bad)
return bench._build_stack(overlay, 16, target, satellites, WAVELENGTHS, read,
band_times)
# --------------------------------------------------------------------------- stack
def test_stack_matches_the_training_layout():
"""Trained on one stack shape, evaluated on another, would be a silent mismatch.
The window spans 560 slots so even the +/-256 rungs of the exponential ladder
exist; a real window edge simply omits the rungs it cannot reach.
"""
overlay, times, _ = overlay_for(count=560)
read, _ = reader()
stack = build(overlay, times[280], read)
offsets = {(entry["same_satellite"], entry["dt"] / paths.CADENCE) for entry in stack}
for source, offset in samples.stack_layout((16, 18), 16):
assert (source == 16, float(offset)) in offsets, f"missing {(source, offset)}"
def test_stack_carries_all_six_bands():
"""The model is joint across bands; a per-band stack would not fit it."""
overlay, times, _ = overlay_for()
read, _ = reader()
stack = build(overlay, times[20], read)
for entry in stack:
assert entry["image"].shape[0] == len(WAVELENGTHS)
def test_stack_holds_the_targets_undamaged_bands_at_dt_zero(monkeypatch):
"""As in training: the renderer fills only the bands a slot lacks, so the bands
it kept are the model's best evidence -- and the lost ones are withheld."""
overlay, times, _ = overlay_for()
target = times[20]
bad = {(16, 171, target), (16, 304, target)}
read, frame = judged_reader(monkeypatch, overlay, bad)
stack = build(overlay, target, read, bad)
own = [e for e in stack if e["same_satellite"] and e["dt"] == 0.0]
assert len(own) == 1 and own[0]["slot"] == (16, target)
assert own[0]["state"] == "available"
assert own[0]["bands"] == (True, True, False, True, True, False)
assert np.isnan(own[0]["image"][[2, 5]]).all()
np.testing.assert_allclose(own[0]["image"][[0, 1, 3, 4]], frame[[0, 1, 3, 4]])
def test_a_target_with_no_usable_band_is_a_missing_entry(monkeypatch):
overlay, times, _ = overlay_for()
target = times[20]
bad = {(16, w, target) for w in WAVELENGTHS}
read, _ = judged_reader(monkeypatch, overlay, bad)
stack = build(overlay, target, read, bad)
own = next(e for e in stack if e["same_satellite"] and e["dt"] == 0.0)
assert own["state"] == "missing" and own["image"] is None
def test_a_target_absent_from_the_archive_is_no_entry():
overlay, times, _ = overlay_for()
for wavelength in WAVELENGTHS:
del overlay.archive[(16, wavelength, times[20])]
read, _ = reader()
stack = build(overlay, times[20], read)
assert not any(e["same_satellite"] and e["dt"] == 0.0 for e in stack)
def test_stack_includes_the_counterpart_at_the_target_instant():
overlay, times, _ = overlay_for()
read, _ = reader()
stack = build(overlay, times[20], read)
assert any(entry["dt"] == 0.0 and not entry["same_satellite"] for entry in stack)
def test_judged_bad_neighbours_are_withheld(monkeypatch):
"""2024-04-17: shown to v8, N_LONG=0 static still got a small blend weight and
speckled every fill's far field. A judged-bad frame is a hole."""
overlay, times, _ = overlay_for()
target = times[20]
neighbour_time = target - paths.CADENCE
bad = {(16, w, neighbour_time) for w in WAVELENGTHS}
read, _ = judged_reader(monkeypatch, overlay, bad)
stack = build(overlay, target, read, bad)
entry = next(e for e in stack if e["slot"] == (16, neighbour_time))
assert entry["state"] == "missing"
assert entry["image"] is None
assert {e["state"] for e in stack} <= {"available", "missing"}
def test_a_neighbour_bad_in_one_band_keeps_its_other_bands_available(monkeypatch):
"""Per band: a fault in 195 withholds 195, not the frame's five clean bands."""
overlay, times, _ = overlay_for()
target = times[20]
neighbour_time = target - paths.CADENCE
bad = {(16, 195, neighbour_time)}
read, frame = judged_reader(monkeypatch, overlay, bad)
stack = build(overlay, target, read, bad)
entry = next(e for e in stack if e["slot"] == (16, neighbour_time))
assert entry["state"] == "available"
assert entry["bands"] == (True, True, True, False, True, True)
assert np.isnan(entry["image"][3]).all()
np.testing.assert_allclose(entry["image"][[0, 1, 2, 4, 5]], frame[[0, 1, 2, 4, 5]])
def test_unreadable_neighbours_arrive_missing_with_no_pixels():
overlay, times, _ = overlay_for()
target = times[20]
gone = (16, target - paths.CADENCE)
read, _ = reader(missing={gone})
stack = build(overlay, target, read)
entry = next(e for e in stack if e["slot"] == gone)
assert entry["state"] == "missing"
assert entry["image"] is None
def test_undamaged_neighbours_arrive_available():
overlay, times, _ = overlay_for()
read, frame = reader()
stack = build(overlay, times[20], read)
assert {e["state"] for e in stack} == {"available"}
for entry in stack:
np.testing.assert_allclose(entry["image"], frame)
def test_stack_adds_anchors_beyond_the_fixed_offsets(monkeypatch):
"""A long outage must still reach a real frame."""
overlay, times, _ = overlay_for(count=80)
target = times[40]
bad = {(16, w, target + n * paths.CADENCE)
for w in WAVELENGTHS for n in range(-30, 31)}
read, _ = judged_reader(monkeypatch, overlay, bad)
stack = build(overlay, target, read, bad)
reach = [abs(e["dt"]) for e in stack
if e["same_satellite"] and e["state"] == "available"]
assert reach, "no usable same-satellite frame found at all"
assert max(reach) >= 31 * paths.CADENCE
def test_stack_skips_slots_the_window_does_not_contain():
"""Near a window edge the offsets run off the end; those must not appear as data."""
overlay, times, _ = overlay_for(count=40)
read, _ = reader()
stack = build(overlay, times[0], read)
for entry in stack:
satellite, when = entry["slot"]
assert any((satellite, w, when) in overlay.archive for w in WAVELENGTHS)
# ------------------------------------------------------------------- band reader
def test_band_reader_returns_damaged_pixels_the_judgment_missed(tmp_path):
"""A filler handed the pristine frame for a corrupted slot reads the answer key."""
from conftest import write_fits
truth_frame = solar_disc(size=64, radius=20, peak=1.0)
damaged_frame = solar_disc(size=64, radius=20, peak=9.0)
when = 1715400000
archive_map, overrides = {}, {}
for wavelength in WAVELENGTHS:
good = str(tmp_path / f"t{wavelength}.fits")
harmed = str(tmp_path / f"d{wavelength}.fits")
write_fits(good, truth_frame)
write_fits(harmed, damaged_frame)
archive_map[(16, wavelength, when)] = good
overrides[(16, wavelength, when)] = harmed
overlay = cases.Overlay(archive=archive_map, overrides=overrides)
read = bench._band_reader(overlay, WAVELENGTHS, set())
got = read(16, when)
assert got.shape == (6, 64, 64)
# Compared loosely: the FITS tile compression is lossy at the 1e-3 level, and what
# is under test is which *file* was opened, not the codec.
assert got[0].max() == pytest.approx(damaged_frame.max(), rel=1e-2)
assert got[0].max() > truth_frame.max() * 5
def test_band_reader_caches_but_stays_bounded(tmp_path):
"""Six bands of 1280x1280 float32 is 39 MB; an unbounded cache exhausts the VM."""
from conftest import write_fits
frame = solar_disc(size=32, radius=10, peak=1.0)
archive_map = {}
times = [1715400000 + i * paths.CADENCE for i in range(10)]
for when in times:
for wavelength in WAVELENGTHS:
path = str(tmp_path / f"{when}_{wavelength}.fits")
write_fits(path, frame)
archive_map[(16, wavelength, when)] = path
overlay = cases.Overlay(archive=archive_map)
reads = []
real = bench.fitsio.read_image
def counting(path):
reads.append(path)
return real(path)
slot_bytes = 6 * 32 * 32 * 4
read = bench._band_reader(overlay, WAVELENGTHS, set(), budget=3 * slot_bytes)
original = bench.fitsio.read_image
bench.fitsio.read_image = counting
try:
read(16, times[0])
first = len(reads)
read(16, times[0])
assert len(reads) == first, "a repeat read was not served from cache"
for when in times[1:5]:
read(16, when)
before_evicted = len(reads)
read(16, times[0]) # evicted by the 3-slot budget
assert len(reads) > before_evicted, "cache grew past its budget"
# Least recently used goes first: touching a slot keeps it.
for when in times[5:8]:
read(16, when)
read(16, times[5]) # refresh the oldest
read(16, times[8]) # evicts times[6], not times[5]
before = len(reads)
read(16, times[5])
assert len(reads) == before, "a recently used slot was evicted"
read(16, times[6])
assert len(reads) > before, "the least recently used slot was kept"
finally:
bench.fitsio.read_image = original
def test_band_reader_builds_one_float32_array_per_slot(monkeypatch):
overlay, times, _ = overlay_for()
read, frame = judged_reader(monkeypatch, overlay, {(16, 171, times[3])})
image = read(16, times[3])
assert image.dtype == np.float32 and image.shape == (6, 16, 16)
assert np.isnan(image[2]).all()
np.testing.assert_array_equal(image[[0, 1, 3, 4, 5]], frame[[0, 1, 3, 4, 5]])
def _write_bands(tmp_path, wavelengths, when=1715400000, satellite=16):
from conftest import write_fits
frame = solar_disc(size=32, radius=10, peak=1.0)
archive_map = {}
for wavelength in wavelengths:
path = str(tmp_path / f"{satellite}_{wavelength}.fits")
write_fits(path, frame)
archive_map[(satellite, wavelength, when)] = path
return cases.Overlay(archive=archive_map), when
def test_band_reader_keeps_the_bands_a_slot_has(tmp_path):
"""Per band, not all or nothing: a slot without 304 still carries five real
bands, and the missing one is NaN so nothing mistakes it for dark sky."""
overlay, when = _write_bands(tmp_path, WAVELENGTHS[:-1])
image = bench._band_reader(overlay, WAVELENGTHS, set())(16, when)
assert image.shape == (6, 32, 32)
assert np.isnan(image[5]).all()
assert np.isfinite(image[:5]).all()
assert bench.band_presence(image) == (True,) * 5 + (False,)
assert not bench._complete(image)
def test_band_reader_withholds_bands_judged_bad(tmp_path):
"""A judged-bad band reads as absent, like a missing one -- its pixels never
reach the model -- and a slot bad in every band reads as nothing at all."""
overlay, when = _write_bands(tmp_path, WAVELENGTHS)
image = bench._band_reader(overlay, WAVELENGTHS, {(16, 131, when)})(16, when)
assert np.isnan(image[1]).all()
assert bench.band_presence(image) == (True, False, True, True, True, True)
other = bench._band_reader(overlay, WAVELENGTHS, {(18, 131, when)})(16, when)
assert bench._complete(other) # the other satellite's verdict
every = {(16, wavelength, when) for wavelength in WAVELENGTHS}
assert bench._band_reader(overlay, WAVELENGTHS, every)(16, when) is None
def test_band_reader_returns_none_only_when_no_band_reads(tmp_path):
overlay, when = _write_bands(tmp_path, ())
assert bench._band_reader(overlay, WAVELENGTHS, set())(16, when) is None
overlay, when = _write_bands(tmp_path, WAVELENGTHS)
assert bench._complete(bench._band_reader(overlay, WAVELENGTHS, set())(16, when))
assert not bench._complete(None)
# ------------------------------------------------------------ per-band stacks
def partial_reader(size=16, absent=()):
"""A reader whose slots lack the bands named in `absent`:
{(satellite, time): (band index, ...)}."""
read, frame = reader(size)
absent = dict(absent)
def read_partial(satellite, when):
image = read(satellite, when)
for band in absent.get((satellite, when), ()):
image[band] = np.nan
return image
return read_partial, frame
def test_entries_say_which_bands_they_hold():
overlay, times, _ = overlay_for()
read, _ = reader()
stack = build(overlay, times[20], read)
assert all(entry["bands"] == (True,) * 6 for entry in stack)
def test_a_counterpart_missing_a_band_still_joins_the_stack():
"""2024-03-03: GOES-18 had no 304 all day. Its 94 and 131 were exactly what
GOES-16 lacked, and dropping the whole frame left the model nothing."""
overlay, times, _ = overlay_for()
target = times[20]
read, _ = partial_reader(absent={(18, target): (5,)})
stack = build(overlay, target, read)
counterpart = next(e for e in stack if e["slot"] == (18, target))
assert counterpart["state"] == "available"
assert counterpart["bands"] == (True,) * 5 + (False,)
assert np.isfinite(counterpart["image"][:5]).all()
def test_band_clean_times_are_per_satellite_and_band():
overlay, times, _ = overlay_for(count=5)
bad = {(16, 94, times[1]), (18, 304, times[3])}
band_times = bench._band_clean_times(overlay, (16, 18), WAVELENGTHS, bad)
assert band_times[16][94] == [times[i] for i in (0, 2, 3, 4)]
assert band_times[16][304] == times
assert band_times[18][304] == [times[i] for i in (0, 1, 2, 4)]
def test_anchors_reach_each_band_separately(monkeypatch):
"""A band bad for a long stretch gets its own anchor at its nearest clean
slot, though every other band is clean right next to the target."""
overlay, times, _ = overlay_for(count=140)
target = times[70]
step = paths.CADENCE
bad = {(16, 94, target + n * step) for n in range(-45, 46)}
read, _ = judged_reader(monkeypatch, overlay, bad)
stack = build(overlay, target, read, bad)
anchors = {e["slot"][1] - target: e for e in stack
if e["same_satellite"] and e["slot"][1] in (target - 46 * step,
target + 46 * step)}
assert sorted(anchors) == [-46 * step, 46 * step]
assert all(e["bands"][0] and e["state"] == "available" for e in anchors.values())
def test_an_anchor_leaves_out_its_bands_judged_bad(monkeypatch):
overlay, times, _ = overlay_for(count=140)
target = times[70]
step = paths.CADENCE
bad = ({(16, 94, target + n * step) for n in range(-45, 46)}
| {(16, 304, target - 46 * step)})
read, _ = judged_reader(monkeypatch, overlay, bad)
stack = build(overlay, target, read, bad)
anchor = next(e for e in stack if e["slot"] == (16, target - 46 * step))
assert anchor["bands"] == (True,) * 5 + (False,)
assert anchor["state"] == "available"
# -------------------------------------------------------------------------- filler
def test_learned_returns_nothing_without_a_stack():
assert fillers.learned(fillers.FillContext()) is None
def test_learned_returns_nothing_without_a_checkpoint(monkeypatch):
monkeypatch.delenv(fillers.LEARNED_CHECKPOINT_ENV, raising=False)
fillers._LEARNED.clear()
context = fillers.FillContext(stack=[
{"image": np.zeros((6, 16, 16), np.float32), "state": "available",
"dt": -240.0, "same_satellite": True}
])
assert fillers.learned(context) is None
def test_learned_is_registered_alongside_the_others():
assert "learned" in fillers.FILLERS
assert fillers.FILLERS["learned"] is fillers.learned
def test_existing_fillers_ignore_the_stack():
"""Adding `stack` must not perturb any measured baseline."""
before = solar_disc(size=64, radius=20, peak=1.0)
after = solar_disc(size=64, radius=20, peak=1.2)
extra = np.stack([after * 5] * 6)
plain = fillers.FillContext(before=before, after=after, dt_before=240, dt_after=240)
with_stack = fillers.FillContext(
before=before, after=after, dt_before=240, dt_after=240,
stack=[{"image": extra, "state": "available", "dt": -240.0,
"same_satellite": True}],
)
for name in ("hold_last", "linear_blend", "optical_flow", "solar_rotation"):
np.testing.assert_allclose(fillers.FILLERS[name](plain),
fillers.FILLERS[name](with_stack))
def test_learned_runs_end_to_end_against_a_saved_checkpoint(tmp_path, monkeypatch):
"""Checkpoint -> load -> fill, at the archive's native frame size."""
torch = pytest.importorskip("torch")
from suvi import model
net = model.build(base=8, depth=2)
path = tmp_path / "model.pt"
torch.save({"model": net.state_dict(), "args": {"base": 8, "depth": 2, "cond_layout": list(model.COND_LAYOUT)},
"epoch": 0}, path)
monkeypatch.setenv(fillers.LEARNED_CHECKPOINT_ENV, str(path))
fillers._LEARNED.clear()
def six(peak):
return np.stack([solar_disc(size=128, radius=40, peak=peak + b * 0.2)
for b in range(6)]).astype(np.float32)
when = 1715400000
context = fillers.FillContext(
stack=[
{"image": six(1.0), "state": "available", "dt": -240.0,
"same_satellite": True, "slot": (16, when - 240)},
{"image": six(1.1), "state": "available", "dt": 240.0,
"same_satellite": True, "slot": (16, when + 240)},
{"image": six(0.9), "state": "available", "dt": 0.0,
"same_satellite": False, "slot": (18, when)},
{"image": None, "state": "missing", "dt": -960.0,
"same_satellite": True, "slot": (16, when - 960)},
],
calibration=(six(0.9), six(1.05)),
)
filled = fillers.learned(context)
assert filled.shape == (6, 128, 128)
assert np.isfinite(filled).all()
fillers._LEARNED.clear()
def _save_tiny_model(tmp_path, monkeypatch):
torch = pytest.importorskip("torch")
from suvi import model
net = model.build(base=8, depth=2)
path = tmp_path / "model.pt"
torch.save({"model": net.state_dict(),
"args": {"base": 8, "depth": 2, "cond_layout": list(model.COND_LAYOUT)}}, path)
monkeypatch.setenv(fillers.LEARNED_CHECKPOINT_ENV, str(path))
fillers._LEARNED.clear()
def _six(peak, size=64, absent=()):
image = np.stack([solar_disc(size=size, radius=size // 3, peak=peak + b * 0.2)
for b in range(6)]).astype(np.float32)
image[list(absent)] = np.nan
return image
def _entry(image, dt, same, slot):
return {"image": image, "bands": bench.band_presence(image), "state": "available",
"dt": dt, "same_satellite": same, "slot": slot}
def test_assemble_stack_reports_which_bands_each_entry_holds(tmp_path, monkeypatch):
torch = pytest.importorskip("torch")
when = 1715400000
context = fillers.FillContext(stack=[
_entry(_six(1.0, absent=(0, 1)), -240.0, True, (16, when - 240)),
_entry(_six(1.0, absent=(5,)), 0.0, False, (18, when)),
])
aligned, condition, present = fillers.assemble_stack(context, torch,
torch.device("cpu"))
assert present.tolist() == [[[0, 0, 1, 1, 1, 1], [1, 1, 1, 1, 1, 0]]]
assert torch.isfinite(aligned).all()
def test_learned_fills_from_bands_split_across_satellites(tmp_path, monkeypatch):
"""The 2024-03-03 case: the target satellite lacks 94/131, the counterpart
lacks 304. Every band is observed somewhere, so every band is filled."""
_save_tiny_model(tmp_path, monkeypatch)
when = 1715400000
context = fillers.FillContext(stack=[
_entry(_six(1.0, absent=(0, 1)), -240.0, True, (16, when - 240)),
_entry(_six(1.1, absent=(0, 1)), 240.0, True, (16, when + 240)),
_entry(_six(0.9, absent=(5,)), 0.0, False, (18, when)),
])
filled = fillers.learned(context)
assert filled.shape == (6, 64, 64)
assert np.isfinite(filled).all()
fillers._LEARNED.clear()
def test_learned_releases_the_raw_stack_before_the_forward_pass(tmp_path, monkeypatch):
"""The raw six-band images are a second copy of everything the model blends;
held through the forward pass they were ~0.8 GB of a 1280px fill's peak."""
import weakref
_save_tiny_model(tmp_path, monkeypatch)
when = 1715400000
images = [_six(1.0), _six(1.1), _six(0.9)]
refs = [weakref.ref(image) for image in images]
context = fillers.FillContext(
stack=[_entry(images[0], -240.0, True, (16, when - 240)),
_entry(images[1], 240.0, True, (16, when + 240)),
_entry(images[2], 0.0, False, (18, when))],
calibration=(_six(0.9), _six(1.0)),
)
del images
net = fillers.load_learned()[0]
alive_at_forward = []
real_forward = type(net).forward
def forward(self, *args, **kwargs):
alive_at_forward.append(sum(ref() is not None for ref in refs))
return real_forward(self, *args, **kwargs)
monkeypatch.setattr(type(net), "forward", forward)
filled = fillers.learned(context)
assert filled.shape == (6, 64, 64)
assert alive_at_forward == [0]
assert context.stack == [] and context.calibration is None
fillers._LEARNED.clear()
def test_a_band_no_frame_holds_comes_back_nan(tmp_path, monkeypatch):
"""Not a black band: the caller must be able to tell 'unfillable' from 'dark'."""
_save_tiny_model(tmp_path, monkeypatch)
when = 1715400000
context = fillers.FillContext(stack=[
_entry(_six(1.0, absent=(2,)), -240.0, True, (16, when - 240)),
_entry(_six(0.9, absent=(2,)), 0.0, False, (18, when)),
])
filled = fillers.learned(context)
assert np.isnan(filled[2]).all()
assert np.isfinite(np.delete(filled, 2, axis=0)).all()
fillers._LEARNED.clear()
def test_learned_declines_a_stack_with_no_pixels_anywhere(tmp_path, monkeypatch):
"""Silently emitting a black frame here is how a fabricated fill would enter the
archive; the contract is None."""
torch = pytest.importorskip("torch")
from suvi import model
net = model.build(base=8, depth=2)
path = tmp_path / "model.pt"
torch.save({"model": net.state_dict(), "args": {"base": 8, "depth": 2, "cond_layout": list(model.COND_LAYOUT)}}, path)
monkeypatch.setenv(fillers.LEARNED_CHECKPOINT_ENV, str(path))
fillers._LEARNED.clear()
context = fillers.FillContext(stack=[
{"image": None, "state": "missing", "dt": -240.0, "same_satellite": True,
"slot": (16, 1715400000 - 240)},
])
assert fillers.learned(context) is None
fillers._LEARNED.clear()
def test_learned_checkpoint_is_loaded_once(tmp_path, monkeypatch):
"""Thousands of slots per bench run; reloading 14M parameters each time would
dominate the wall clock."""
torch = pytest.importorskip("torch")
from suvi import model
net = model.build(base=8, depth=2)
path = tmp_path / "model.pt"
torch.save({"model": net.state_dict(), "args": {"base": 8, "depth": 2, "cond_layout": list(model.COND_LAYOUT)}}, path)
monkeypatch.setenv(fillers.LEARNED_CHECKPOINT_ENV, str(path))
fillers._LEARNED.clear()
loads = []
real_load = torch.load
monkeypatch.setattr(torch, "load", lambda *a, **k: (loads.append(1), real_load(*a, **k))[1])
for _ in range(3):
fillers.load_learned()
assert len(loads) == 1
fillers._LEARNED.clear()
def test_band_reader_caches_a_repeated_target(tmp_path, monkeypatch):
"""One target pulls one frame per band through the reader, and consecutive targets
share most of their stack. Without the cache each shared frame is re-read, and on
this archive a read is an SMB round trip through two filesystems."""
from conftest import write_fits
frame = solar_disc(size=32, radius=10, peak=1.0)
when = 1715400000
archive_map = {}
for wavelength in WAVELENGTHS:
path = str(tmp_path / f"{wavelength}.fits")
write_fits(path, frame)
archive_map[(16, wavelength, when)] = path
reads = []
real = bench.fitsio.read_image
monkeypatch.setattr(bench.fitsio, "read_image",
lambda p, *a, **k: reads.append(p) or real(p, *a, **k))
overlay = cases.Overlay(archive=archive_map)
read = bench._band_reader(overlay, WAVELENGTHS, set())
read(16, when)
assert len(reads) == len(WAVELENGTHS), "one read per band"
read(16, when)
assert len(reads) == len(WAVELENGTHS), "re-read a slot already in the cache"
Reference in a new issue View git blame Copy permalink