412 lines
18 KiB
Python
412 lines
18 KiB
Python
"""Training shards: a compact, portable copy of the archive for the learned filler.
|
|
|
|
The archive is 9.2 TB of tile-compressed FITS reached over three hops -- a virtiofs
|
|
re-export of an SMB share -- and the machine that trains the model is a different machine
|
|
on the other end of a gigabit link. So training does not read the archive. It reads
|
|
*shards*:
|
|
one file per (day, satellite) holding that day's frames at half resolution, which is
|
|
small enough to ship and to re-read every epoch.
|
|
|
|
Two decisions here are load-bearing.
|
|
|
|
**Shards store clean frames only.** Corruption is applied on the fly during training
|
|
(:mod:`suvi.corruptions`), which costs nothing in storage, gives unlimited variety rather
|
|
than one frozen draw, and yields an exact ground-truth class label for the auxiliary
|
|
head. Baking damage into the shards would fix the training distribution at extraction
|
|
time and make every experiment downstream a re-extraction.
|
|
|
|
**Pixels are stored as quantised ``asinh`` radiance.** Solar radiance is heavy-tailed --
|
|
a sample across bands and epochs runs from -3.7 to 1231 -- so linear uint16 would spend
|
|
its whole range on the disc and quantise the corona to nothing. ``asinh(x/s)`` is
|
|
linear below the knee `s` and logarithmic above it, which is what a signal that is
|
|
mostly noise-floor with occasional flares needs.
|
|
|
|
``log1p``, the space ``fillers._for_flow`` uses, was the obvious first choice and is
|
|
wrong here on two counts. It is effectively linear below 1, and almost every pixel is:
|
|
median radiance is 0.004 to 0.19 depending on band. Quantising log1p over [0, 8] gives
|
|
a 0.002 radiance pixel **6% error**, not the 0.05% the bright disc gets. And it cannot
|
|
represent negative values at all, while 2-18% of pixels are negative after background
|
|
subtraction -- clipping those to zero would bias the noise floor the model has to learn
|
|
to reproduce. ``asinh`` holds ~0.05% relative across the whole range and is signed.
|
|
"""
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import struct
|
|
|
|
import numpy as np
|
|
|
|
from . import fitsio, paths
|
|
|
|
#: Half of the archive's native 1280. Flow, gain and blend fields are smooth, so the
|
|
#: model predicts them here and they upsample cleanly to drive full-resolution frames.
|
|
SHARD_SIZE = 640
|
|
#: Knee of the asinh transform: linear below, logarithmic above. Relative precision is
|
|
#: constant above the knee and decays below it, so the knee sits well under the faintest
|
|
#: band's median radiance (0.004 at 131A) -- far enough that the whole corona is in the
|
|
#: constant-precision regime, not just the disc.
|
|
ASINH_SCALE = 1e-4
|
|
#: Half-range of the stored asinh values, symmetric about zero. +/-18 covers
|
|
#: +/-sinh(18)*scale = +/-3284 in radiance, against an observed span of -3.7 to 1231 --
|
|
#: room for a flare well past anything in the archive. Anything beyond that is a
|
|
#: corruption, not an observation, and saturates rather than wrapping.
|
|
ASINH_RANGE = 18.0
|
|
#: Marker written for a slot whose frame would not read. Distinguishes "we looked and
|
|
#: there was nothing usable" from "we never looked", which a plain absence cannot.
|
|
UNREADABLE = b"\x00"
|
|
|
|
MAGIC = b"SUVISHRD"
|
|
VERSION = 2
|
|
SLOTS_PER_DAY = 86400 // paths.CADENCE # 360
|
|
|
|
|
|
def _zstd():
|
|
"""zstd compress/decompress, from the stdlib on 3.14+ or the pip package below it.
|
|
|
|
The extraction host runs Python 3.14, where zstd is in the standard library; the
|
|
training container ships whatever the ROCm image was built against. Both produce
|
|
ordinary zstd frames, so shards written by one are read by the other.
|
|
"""
|
|
try:
|
|
from compression import zstd
|
|
|
|
return zstd.compress, zstd.decompress
|
|
except ImportError: # pragma: no cover - exercised only on the training host
|
|
import zstandard
|
|
|
|
return (
|
|
lambda data, level=3: zstandard.ZstdCompressor(level=level).compress(data),
|
|
lambda data: zstandard.ZstdDecompressor().decompress(data),
|
|
)
|
|
|
|
|
|
# ------------------------------------------------------------------- pixel encoding
|
|
|
|
|
|
def encode_frames(arrays):
|
|
"""Six 1280x1280 radiance arrays -> one 6 x 640 x 640 uint16 block.
|
|
|
|
Downsampling is a 2x2 mean rather than decimation: SUVI frames carry read noise and
|
|
cosmic-ray hits, and taking every other pixel would keep the hits at full amplitude
|
|
while throwing away the averaging that suppresses them.
|
|
"""
|
|
planes = []
|
|
for array in arrays:
|
|
image = np.nan_to_num(np.asarray(array, dtype=np.float64), nan=0.0,
|
|
posinf=0.0, neginf=0.0)
|
|
if image.shape[0] % SHARD_SIZE or image.shape[1] % SHARD_SIZE:
|
|
raise ValueError(f"cannot halve {image.shape} to {SHARD_SIZE}")
|
|
factor = image.shape[0] // SHARD_SIZE
|
|
small = image.reshape(SHARD_SIZE, factor, SHARD_SIZE, factor).mean(axis=(1, 3))
|
|
unit = np.arcsinh(small / ASINH_SCALE) / ASINH_RANGE # -> roughly [-1, 1]
|
|
planes.append(np.clip(unit, -1.0, 1.0) * 32767.5 + 32767.5)
|
|
return np.rint(np.stack(planes)).astype(np.uint16)
|
|
|
|
|
|
def decode_frames(block):
|
|
"""Inverse of :func:`encode_frames`, back to radiance."""
|
|
return (np.sinh(coded_frames(block).astype(np.float64) * ASINH_RANGE)
|
|
* ASINH_SCALE).astype(np.float32)
|
|
|
|
|
|
def coded_frames(block):
|
|
"""The stored block as the roughly-[-1, 1] values the model consumes.
|
|
|
|
The uint16 in a shard *is* the model's input space -- ``encode_frames`` already
|
|
applied the asinh -- so getting from one to the other is an affine rescale and
|
|
nothing more. Going via radiance instead costs a ``sinh`` and an ``arcsinh`` over
|
|
47 million elements per training sample, which measured as a data loader pinning two
|
|
cores at 200% while the GPU sat at 0% busy. Only frames that are about to be
|
|
corrupted need real radiance, because that is the space corruptions are defined in.
|
|
"""
|
|
return ((np.asarray(block, dtype=np.float32) - 32767.5) / 32767.5).astype(np.float32)
|
|
|
|
|
|
# ------------------------------------------------------------------- shard container
|
|
#
|
|
# Layout:
|
|
# offset 0 _PREFIX MAGIC, VERSION (10 bytes)
|
|
# offset 10 records back to back, sorted by t_start
|
|
# JSON header, carrying (offset, length) per slot
|
|
# EOF - 12 _TRAILER header length, MAGIC (12 bytes)
|
|
#
|
|
# The index gives one record per slot, so a single record can be read without
|
|
# decompressing the rest -- which is what makes random sampling across the corpus
|
|
# cheap. It lives at the *end* because its contents are not known until every
|
|
# payload has been compressed: with the index in front, writing a shard means
|
|
# buffering the whole ~1 GB of it, and shards are produced straight into a network
|
|
# stream that should never hold more than one record at a time.
|
|
#
|
|
# Record offsets are relative to the first record, which is therefore at a fixed
|
|
# position rather than one that depends on the header's length.
|
|
|
|
_PREFIX = struct.Struct("<8sH")
|
|
_TRAILER = struct.Struct("<I8s")
|
|
|
|
|
|
def _header_bytes(day, satellite, wavelengths, index):
|
|
return json.dumps({
|
|
"day": day,
|
|
"satellite": int(satellite),
|
|
"wavelengths": list(wavelengths),
|
|
"size": SHARD_SIZE,
|
|
"asinh_scale": ASINH_SCALE,
|
|
"asinh_range": ASINH_RANGE,
|
|
"index": index,
|
|
}, separators=(",", ":")).encode()
|
|
|
|
|
|
def write_shard_stream(handle, day, satellite, wavelengths, records, level=3):
|
|
"""Write one shard to a binary stream. Returns the bytes written.
|
|
|
|
`records` is an *iterable* of (t_start, block) in ascending t_start order --
|
|
an iterable, not a mapping, because the point of this function is that a
|
|
shard can be produced without ever existing in memory or on disk. Each
|
|
payload is compressed and written as it arrives, so the caller may hand over
|
|
a generator that reads and encodes one slot at a time and nothing larger than
|
|
a single record is ever live.
|
|
|
|
Ascending order is the caller's responsibility and is checked: offsets are
|
|
cumulative, so an out-of-order record would silently produce an index that
|
|
does not describe the file.
|
|
"""
|
|
compress, _ = _zstd()
|
|
index = {}
|
|
offset = 0
|
|
previous = None
|
|
handle.write(_PREFIX.pack(MAGIC, VERSION))
|
|
for t_start, block in records:
|
|
if previous is not None and t_start <= previous:
|
|
raise ValueError(
|
|
f"records must ascend by t_start: {t_start} after {previous}")
|
|
previous = t_start
|
|
payload = UNREADABLE if block is None else compress(
|
|
np.ascontiguousarray(block, dtype="<u2").tobytes(), level
|
|
)
|
|
handle.write(payload)
|
|
index[str(t_start)] = [offset, len(payload)]
|
|
offset += len(payload)
|
|
|
|
header = _header_bytes(day, satellite, wavelengths, index)
|
|
handle.write(header)
|
|
handle.write(_TRAILER.pack(len(header), MAGIC))
|
|
return _PREFIX.size + offset + len(header) + _TRAILER.size
|
|
|
|
|
|
def write_shard(path, day, satellite, wavelengths, records, level=3):
|
|
"""Write one (day, satellite) shard to a file.
|
|
|
|
`records` maps t_start -> 6x640x640 uint16 block (or None for a slot that was
|
|
looked at and found unusable). A thin wrapper over
|
|
:func:`write_shard_stream` so there is only one implementation of the format.
|
|
"""
|
|
temporary = f"{path}.partial"
|
|
with open(temporary, "wb") as handle:
|
|
write_shard_stream(handle, day, satellite, wavelengths,
|
|
((t, records[t]) for t in sorted(records)), level)
|
|
os.replace(temporary, path) # never leave a half-written shard at the real name
|
|
return os.path.getsize(path)
|
|
|
|
|
|
class Shard:
|
|
"""Random access to one shard. Holds a single file handle, opened eagerly.
|
|
|
|
Eagerly, because the pool rotator's consistency protocol depends on it: a
|
|
rotated-out shard is *deleted* on the training host, and POSIX only keeps the
|
|
bytes readable for handles that already exist. A lazily-opened shard whose
|
|
first read came after the delete once killed a run mid-epoch."""
|
|
|
|
def __init__(self, path):
|
|
self.path = path
|
|
self._handle = open(path, "rb")
|
|
size = os.fstat(self._handle.fileno()).st_size
|
|
if size < _PREFIX.size + _TRAILER.size:
|
|
raise ValueError(f"{path}: too short to be a shard ({size} bytes)")
|
|
|
|
magic, version = _PREFIX.unpack(self._handle.read(_PREFIX.size))
|
|
if magic != MAGIC:
|
|
raise ValueError(f"{path}: not a shard")
|
|
if version != VERSION:
|
|
raise ValueError(f"{path}: shard version {version}, expected {VERSION}")
|
|
|
|
length, tail_magic = _TRAILER.unpack(
|
|
os.pread(self._handle.fileno(), _TRAILER.size, size - _TRAILER.size))
|
|
# A shard is now streamed to its destination rather than written whole and
|
|
# renamed, so a truncated one is a real possibility rather than a
|
|
# theoretical one. Every check below exists to make that fail here, at
|
|
# open, instead of returning plausible-looking pixels from the wrong offset.
|
|
if tail_magic != MAGIC:
|
|
raise ValueError(f"{path}: truncated or not a shard (bad trailer)")
|
|
start = size - _TRAILER.size - length
|
|
if not _PREFIX.size <= start <= size - _TRAILER.size:
|
|
raise ValueError(f"{path}: header length {length} does not fit the file")
|
|
self.header = json.loads(os.pread(self._handle.fileno(), length, start))
|
|
|
|
self._base = _PREFIX.size
|
|
self.satellite = self.header["satellite"]
|
|
self.day = self.header["day"]
|
|
self.wavelengths = tuple(self.header["wavelengths"])
|
|
self._index = {int(k): tuple(v) for k, v in self.header["index"].items()}
|
|
end = max((o + n for o, n in self._index.values()), default=0)
|
|
if self._base + end != start:
|
|
raise ValueError(
|
|
f"{path}: records end at {self._base + end}, header starts at {start}")
|
|
|
|
def times(self):
|
|
return sorted(self._index)
|
|
|
|
def __contains__(self, t_start):
|
|
return t_start in self._index
|
|
|
|
def usable(self, t_start):
|
|
"""Whether the slot holds real frames, from the index alone.
|
|
|
|
A vetting-dropped slot is stored as the one-byte UNREADABLE sentinel, so
|
|
the index's length column answers without decompressing anything -- the
|
|
segment verifier asks this for every slot of the corpus.
|
|
"""
|
|
entry = self._index.get(t_start)
|
|
return entry is not None and entry[1] > len(UNREADABLE)
|
|
|
|
def raw(self, t_start):
|
|
"""The stored uint16 block for a slot, or None if it holds no usable frame."""
|
|
entry = self._index.get(t_start)
|
|
if entry is None:
|
|
return None
|
|
offset, length = entry
|
|
if self._handle is None: # reopened after an explicit close()
|
|
self._handle = open(self.path, "rb")
|
|
# pread, not seek+read: DataLoader workers fork *after* the shard opens
|
|
# (they must -- the handle has to predate any pool-rotation delete), so
|
|
# they share one file description. A seek from one worker lands in the
|
|
# middle of another's read; positional reads share nothing.
|
|
payload = os.pread(self._handle.fileno(), length, self._base + offset)
|
|
if payload == UNREADABLE:
|
|
return None
|
|
_, decompress = _zstd()
|
|
flat = np.frombuffer(decompress(payload), dtype="<u2")
|
|
return flat.reshape(len(self.wavelengths), SHARD_SIZE, SHARD_SIZE)
|
|
|
|
def frames(self, t_start):
|
|
"""Radiance for a slot, 6 x 640 x 640 float32, or None."""
|
|
block = self.raw(t_start)
|
|
return None if block is None else decode_frames(block)
|
|
|
|
def coded(self, t_start):
|
|
"""Model-space values for a slot, 6 x 640 x 640 float32, or None.
|
|
|
|
The cheap path, and the one training uses. See :func:`coded_frames`.
|
|
"""
|
|
block = self.raw(t_start)
|
|
return None if block is None else coded_frames(block)
|
|
|
|
def close(self):
|
|
if self._handle is not None:
|
|
self._handle.close()
|
|
self._handle = None
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *exc):
|
|
self.close()
|
|
|
|
|
|
def shard_name(day, satellite):
|
|
return f"{day}_g{satellite:d}.shard"
|
|
|
|
|
|
# ------------------------------------------------------------------------ extraction
|
|
|
|
|
|
def read_slot(root, rows):
|
|
"""Decode one slot's six bands from the archive.
|
|
|
|
`rows` maps wavelength -> archive-relative path. Returns the encoded block, or
|
|
None if any band is missing or unreadable -- a partial slot is not a training
|
|
sample, and admitting one would teach the model that bands go missing
|
|
independently when in practice a satellite drops out whole.
|
|
"""
|
|
arrays = []
|
|
for wavelength in paths.WAVELENGTHS:
|
|
relpath = rows.get(wavelength)
|
|
if relpath is None:
|
|
return None
|
|
image, _ = fitsio.read_image(paths.abspath(relpath, root))
|
|
if image is None or image.shape != (1280, 1280):
|
|
return None
|
|
arrays.append(image)
|
|
return encode_frames(arrays)
|
|
|
|
|
|
# ----------------------------------------------------------------------- day choice
|
|
|
|
|
|
#: Smallest plausible real frame. A healthy tile-compressed SUVI frame is 1.4-1.8 MB;
|
|
#: a failed download leaves a 5,760-byte header-only stub that is indexed like any other
|
|
#: frame and reads back as None. There are 125,732 of them in the archive -- 2.3% -- and
|
|
#: one chosen day turned out to be 100% stubs, producing an empty shard after a full
|
|
#: extraction pass over it. Counting rows is not the same as counting frames.
|
|
MIN_FRAME_BYTES = 100_000
|
|
|
|
|
|
def choose_days(conn, count, t_from, t_to, exclude=(), satellites=(16, 18),
|
|
wavelengths=paths.WAVELENGTHS, min_bytes=MIN_FRAME_BYTES):
|
|
"""Days where both satellites have near-complete coverage, spread evenly.
|
|
|
|
Spread rather than sampled at random: the point of the training set is to span the
|
|
solar cycle, and independent draws over a 2.5-year range clump. Complete rather
|
|
than best-effort: a day missing half its slots would silently bias the gap-length
|
|
distribution the model trains against. And *real* rather than merely present -- see
|
|
:data:`MIN_FRAME_BYTES`.
|
|
"""
|
|
wanted = len(wavelengths) * len(satellites) * SLOTS_PER_DAY
|
|
placeholders = ",".join("?" * len(wavelengths))
|
|
satellite_places = ",".join("?" * len(satellites))
|
|
rows = conn.execute(
|
|
f"""
|
|
SELECT strftime('%Y-%m-%d', t_start, 'unixepoch') AS day, COUNT(*) AS n
|
|
FROM frame
|
|
WHERE t_start >= ? AND t_start < ?
|
|
AND wavelength IN ({placeholders}) AND satellite IN ({satellite_places})
|
|
AND size_bytes >= ?
|
|
GROUP BY day HAVING n >= ?
|
|
ORDER BY day
|
|
""",
|
|
(t_from, t_to, *wavelengths, *satellites, int(min_bytes), int(wanted * 0.98)),
|
|
).fetchall()
|
|
|
|
excluded = set(exclude)
|
|
days = [row[0] for row in rows if row[0] not in excluded]
|
|
if not days:
|
|
return []
|
|
if len(days) <= count:
|
|
return days
|
|
step = len(days) / count
|
|
return [days[int(i * step)] for i in range(count)]
|
|
|
|
|
|
def check_split(days_by_split):
|
|
"""Assert the splits are disjoint and no two blocks abut in time.
|
|
|
|
Cheap, and it fails loudly. A leak here would not crash anything -- it would just
|
|
produce validation numbers that look excellent and mean nothing.
|
|
"""
|
|
import datetime
|
|
|
|
seen = {}
|
|
for name, days in days_by_split.items():
|
|
for day in days:
|
|
if day in seen:
|
|
raise ValueError(f"day {day} is in both {seen[day]!r} and {name!r}")
|
|
seen[day] = name
|
|
|
|
ordered = sorted(seen)
|
|
adjacent = []
|
|
for first, second in zip(ordered, ordered[1:]):
|
|
a = datetime.date.fromisoformat(first)
|
|
b = datetime.date.fromisoformat(second)
|
|
if (b - a).days == 1 and seen[first] != seen[second]:
|
|
adjacent.append((first, second))
|
|
return adjacent
|