111 lines
4.3 KiB
Python
111 lines
4.3 KiB
Python
|
import os
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from suvi import availability, paths, segments
|
||
|
|
|
||
|
|
CAD = paths.CADENCE
|
||
|
|
BASE = 1715400000 // CAD * CAD
|
||
|
|
|
||
|
|
|
||
|
|
def make_model(pattern16, pattern18):
|
||
|
|
# Slot patterns, applied to every band: segments need all six GOOD.
|
||
|
|
slots = np.stack([np.asarray(pattern16), np.asarray(pattern18)])
|
||
|
|
states = np.repeat(slots[:, None, :], len(availability.BANDS), axis=1)
|
||
|
|
return availability.Availability(states=states, t0=BASE, cadence=CAD,
|
||
|
|
satellites=(16, 18))
|
||
|
|
|
||
|
|
|
||
|
|
def test_segments_are_jointly_pristine_runs_of_at_least_the_minimum():
|
||
|
|
g16 = np.full(100, availability.GOOD)
|
||
|
|
g18 = np.full(100, availability.GOOD)
|
||
|
|
g16[30] = availability.BAD # splits the run
|
||
|
|
g18[70:80] = availability.MISSING # second split, other satellite
|
||
|
|
found = segments.from_availability(make_model(g16, g18), minimum_slots=10)
|
||
|
|
assert found == [
|
||
|
|
(BASE, BASE + 30 * CAD), # 30 slots
|
||
|
|
(BASE + 31 * CAD, BASE + 70 * CAD), # 39 slots
|
||
|
|
(BASE + 80 * CAD, BASE + 100 * CAD), # 20 slots
|
||
|
|
]
|
||
|
|
# A tighter minimum drops the tail fragment.
|
||
|
|
found = segments.from_availability(make_model(g16, g18), minimum_slots=35)
|
||
|
|
assert found == [(BASE + 31 * CAD, BASE + 70 * CAD)]
|
||
|
|
|
||
|
|
|
||
|
|
def test_days_touched_spans_midnights():
|
||
|
|
start = BASE # 2024-05-11 04:00 UTC
|
||
|
|
segment = (start, start + 350 * CAD) # ~23.3 h -> crosses one midnight
|
||
|
|
days = segments.days_touched([segment])
|
||
|
|
assert days == ["2024-05-11", "2024-05-12"]
|
||
|
|
|
||
|
|
|
||
|
|
def test_verify_splits_where_extraction_disagrees():
|
||
|
|
"""disc_v1 can veto a slot the header profile called pristine; the segment
|
||
|
|
must split there, and fragments below the model input span must go."""
|
||
|
|
segment = (BASE, BASE + 100 * CAD)
|
||
|
|
bad = {BASE + 50 * CAD}
|
||
|
|
out = segments.verify([segment], lambda t: t not in bad, minimum_slots=30)
|
||
|
|
assert out == [(BASE, BASE + 50 * CAD), (BASE + 51 * CAD, BASE + 100 * CAD)]
|
||
|
|
out = segments.verify([segment], lambda t: t not in bad, minimum_slots=60)
|
||
|
|
assert out == [] # both fragments too short
|
||
|
|
|
||
|
|
|
||
|
|
def test_targets_keep_a_full_span_inside_the_segment():
|
||
|
|
segment = (BASE, BASE + 100 * CAD)
|
||
|
|
eligible = segments.targets([segment], span_slots=30)
|
||
|
|
assert eligible[0] == BASE + 30 * CAD
|
||
|
|
assert eligible[-1] == BASE + 69 * CAD
|
||
|
|
assert len(eligible) == 40
|
||
|
|
|
||
|
|
|
||
|
|
def test_split_assigns_whole_segments_and_is_deterministic():
|
||
|
|
runs = [(BASE + i * 1000 * CAD, BASE + (i * 1000 + 200) * CAD)
|
||
|
|
for i in range(10)]
|
||
|
|
first = segments.split(runs, seed=3, val_targets=100, test_targets=100,
|
||
|
|
span_slots=50)
|
||
|
|
again = segments.split(runs, seed=3, val_targets=100, test_targets=100,
|
||
|
|
span_slots=50)
|
||
|
|
assert first == again
|
||
|
|
everything = first["train"] + first["val"] + first["test"]
|
||
|
|
assert sorted(everything) == sorted(map(tuple, runs))
|
||
|
|
for name in ("val", "test"):
|
||
|
|
mass = sum((b - a) // CAD - 100 for a, b in first[name])
|
||
|
|
assert mass >= 100
|
||
|
|
assert first["train"], "everything landed in val/test"
|
||
|
|
|
||
|
|
|
||
|
|
def test_save_load_round_trip(tmp_path):
|
||
|
|
runs = [(BASE, BASE + 200 * CAD), (BASE + 500 * CAD, BASE + 900 * CAD)]
|
||
|
|
path = str(tmp_path / "segments.json")
|
||
|
|
segments.save(path, runs)
|
||
|
|
assert segments.load(path) == runs
|
||
|
|
|
||
|
|
|
||
|
|
def test_containing():
|
||
|
|
runs = [(BASE, BASE + 100 * CAD)]
|
||
|
|
assert segments.containing(runs, BASE + 50 * CAD) == runs[0]
|
||
|
|
assert segments.containing(runs, BASE - CAD) is None
|
||
|
|
|
||
|
|
|
||
|
|
def test_ladder_for_matches_spans():
|
||
|
|
from suvi import samples
|
||
|
|
|
||
|
|
assert samples.ladder_for(64) == (1, 4, 16, 64)
|
||
|
|
assert samples.ladder_for(128) == (1, 4, 16, 64, 128)
|
||
|
|
assert samples.ladder_for(512) == (1, 4, 16, 64, 128, 256, 512)
|
||
|
|
with pytest.raises(ValueError):
|
||
|
|
samples.ladder_for(100)
|
||
|
|
|
||
|
|
|
||
|
|
def test_saved_segments_are_plain_json_ints(tmp_path):
|
||
|
|
"""numpy int64 is not JSON serializable, and the failure lands at the very
|
||
|
|
end of a multi-hour corpus build."""
|
||
|
|
model = make_model(np.full(600, availability.GOOD), np.full(600, availability.GOOD))
|
||
|
|
found = segments.from_availability(model, minimum_slots=10)
|
||
|
|
assert all(type(v) is int for pair in found for v in pair)
|
||
|
|
path = str(tmp_path / "segments.json")
|
||
|
|
segments.save(path, found)
|
||
|
|
assert segments.load(path) == found
|
||
|
|
assert not os.path.exists(path + ".tmp")
|