94 lines
4 KiB
Python
94 lines
4 KiB
Python
"""Choosing which slices of the pristine corpus to ship to the training host.
|
|
|
|
Training targets come from ``segment_length - 2*span`` positions, so a seven-day
|
|
segment yields ~2,300 targets and a minimum-length one yields a handful.
|
|
Admitting segments longest-first would therefore fill the host with a few long
|
|
stretches of Sun; admitting them in random order samples the whole era.
|
|
|
|
A segment's contribution is capped at :func:`suvi.segments.decorrelation_slots`
|
|
(~4,720 frames, half a synodic rotation). That is the length at which a segment
|
|
stops being one evolving view of one hemisphere: past it the far side has rotated
|
|
into view and the frames share no surface features, so the later part is
|
|
effectively an independent segment and capping it would be arbitrary. Below that
|
|
scale every extra frame buys real diversity -- the Sun turns ~13 degrees a day --
|
|
so truncating short of it discards new Sun rather than redundancy.
|
|
"""
|
|
|
|
import numpy as np
|
|
|
|
from . import paths, segments as segments_lib
|
|
|
|
#: Bytes per slot per satellite, measured on the archive (six bands, 640^2,
|
|
#: uint16, zstd-3). Compression is already at the entropy limit here -- level 19
|
|
#: buys 4.8% for twenty times the CPU -- so this is a fixed cost, and the only
|
|
#: lever on corpus size is how many slots are shipped.
|
|
SLOT_BYTES = 4.0e6
|
|
|
|
|
|
def choose_slices(segments, cap, budget_bytes, satellites, seed,
|
|
val_count, test_count):
|
|
"""Pick the slices to ship: val and test first, then train until full.
|
|
|
|
val/test segment choice is seeded off the *whole* segment list rather than
|
|
the admitted subset, so the evaluation sets stay identical when the corpus
|
|
is rebuilt at a different budget or cap.
|
|
"""
|
|
rng = np.random.default_rng(seed)
|
|
order = rng.permutation(len(segments))
|
|
roles = {}
|
|
for rank, index in enumerate(order):
|
|
if rank < val_count:
|
|
roles[int(index)] = "val"
|
|
elif rank < val_count + test_count:
|
|
roles[int(index)] = "test"
|
|
|
|
per_slot = SLOT_BYTES * len(satellites)
|
|
chosen = {"train": [], "val": [], "test": []}
|
|
owner = {}
|
|
used = 0.0
|
|
# Evaluation slices are admitted first: a corpus that cannot be measured is
|
|
# worth nothing, however much training data it holds.
|
|
for wanted in ("val", "test", "train"):
|
|
for index in order:
|
|
index = int(index)
|
|
role = roles.get(index, "train")
|
|
if role != wanted:
|
|
continue
|
|
t0, t1 = segments[index]
|
|
length = (t1 - t0) // paths.CADENCE
|
|
take = min(length, cap)
|
|
if used + take * per_slot > budget_bytes:
|
|
continue
|
|
start = t0 + int(rng.integers(0, length - take + 1)) * paths.CADENCE
|
|
slice_ = (start, start + take * paths.CADENCE)
|
|
# A shard is day-keyed, so a day shared between two splits would put
|
|
# validation frames inside the training sampler's reach as *context*
|
|
# -- never as targets, but close enough to memorise. One role per
|
|
# day keeps the splits genuinely disjoint.
|
|
days = segments_lib.days_touched([slice_])
|
|
if any(owner.get(day, role) != role for day in days):
|
|
continue
|
|
for day in days:
|
|
owner[day] = role
|
|
chosen[role].append(slice_)
|
|
used += take * per_slot
|
|
return chosen, used
|
|
|
|
|
|
def wanted_slots(chosen, satellites, cadence=paths.CADENCE):
|
|
"""The slots each (day, satellite) shard must hold, from the chosen slices.
|
|
|
|
Returns {(day, satellite): [t_start, ...]} with each list ascending, which is
|
|
the order the shard writer requires.
|
|
"""
|
|
import datetime
|
|
|
|
wanted = {}
|
|
for slices in chosen.values():
|
|
for t0, t1 in slices:
|
|
for t in range(t0, t1, cadence):
|
|
day = datetime.datetime.fromtimestamp(
|
|
t, datetime.timezone.utc).strftime("%Y-%m-%d")
|
|
for satellite in satellites:
|
|
wanted.setdefault((day, satellite), []).append(t)
|
|
return {key: sorted(times) for key, times in wanted.items()}
|