404 lines
20 KiB
Python
404 lines
20 KiB
Python
|
"""Turning shards into training samples: stacks, masks, and damage applied on the fly.
|
||
|
|
|
||
|
|
One sample is a target slot -- one satellite, one instant, all six bands -- plus the
|
||
|
|
stack of frames the model may draw on and a per-frame conditioning vector. This module
|
||
|
|
is the only place that knows how a stack is laid out, so :mod:`suvi.model` and
|
||
|
|
:mod:`bench` cannot disagree about it.
|
||
|
|
|
||
|
|
Three things here decide what the model can learn.
|
||
|
|
|
||
|
|
**Offsets are exponentially spaced.** Gaps in this archive run from one slot to three
|
||
|
|
hundred, and a fixed +/-K window either misses the far end of a long outage or wastes
|
||
|
|
most of its inputs on a short one. Sampling at ``+/-{1, 4, 16, 64, 256}`` slots covers
|
||
|
|
four minutes to seventeen hours with ten frames per satellite, and an *anchor* -- the
|
||
|
|
nearest usable frame in each direction, however far -- guarantees the stack is never
|
||
|
|
empty even beyond that.
|
||
|
|
|
||
|
|
**Damage is applied here, not baked into the shards.** Every epoch draws fresh
|
||
|
|
corruptions from :mod:`suvi.corruptions`, so the model sees far more variety than a
|
||
|
|
frozen dataset could hold, and the true mode is known exactly for the auxiliary head.
|
||
|
|
|
||
|
|
**Outages are episodic as well as scattered.** Half of all samples carry a contiguous
|
||
|
|
simulated outage -- up to 400 slots, on one satellite or both -- because that is how
|
||
|
|
the real archive fails: measured over the two-satellite era, 11% of slots have both
|
||
|
|
spacecraft dark at once, in long runs. Independent per-frame drops alone would never
|
||
|
|
construct that case, and a model cannot learn a regime it has never seen.
|
||
|
|
|
||
|
|
**Context frames get damaged too, not just the target.** This is the requirement that
|
||
|
|
any supplied frame may itself be invalid. A model trained on clean context and deployed
|
||
|
|
on an archive where 31% of slots are bad on both satellites would meet, at inference, a
|
||
|
|
distribution it had never seen.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import math
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from . import align, corruptions, dataset, paths
|
||
|
|
|
||
|
|
#: Slot offsets sampled on each satellite, in both directions. Exponentially spaced,
|
||
|
|
#: so one fixed-size stack serves every gap length: 4 minutes to 17 hours in five
|
||
|
|
#: steps, without needing hundreds of input frames to cover a 300-slot outage. The
|
||
|
|
#: anchors (nearest usable frame either way, at any distance) still guarantee the
|
||
|
|
#: stack is never empty beyond the ladder's reach. Cost is linear in stack depth in
|
||
|
|
#: the current model (no attention), so the two extra rungs are affordable where they
|
||
|
|
#: were not for the attention-based design this replaced.
|
||
|
|
OFFSETS = (1, 4, 16, 64, 256)
|
||
|
|
#: How often a sample carries a simulated *episodic* outage -- a contiguous run of
|
||
|
|
#: missing slots -- on top of the scattered per-frame drops. Real outages are runs,
|
||
|
|
#: not confetti: measured against the archive, 11% of the two-satellite era has BOTH
|
||
|
|
#: spacecraft dark simultaneously, in 1,096 contiguous runs of up to 26 days. With
|
||
|
|
#: independent per-frame drops the probability of reproducing that configuration in
|
||
|
|
#: a training stack is roughly p^14 -- the model would face at evaluation a regime
|
||
|
|
#: it had never once seen.
|
||
|
|
EPISODIC_PROBABILITY = 0.5
|
||
|
|
#: Given an episodic outage, how often it takes down *both* satellites at once (a
|
||
|
|
#: ground-segment or space-weather event) rather than one.
|
||
|
|
DUAL_OUTAGE_PROBABILITY = 0.4
|
||
|
|
#: Longest simulated outage, in slots; matches the anchor search limit, so a stack
|
||
|
|
#: always retains something real beyond the outage's edge.
|
||
|
|
MAX_OUTAGE_SLOTS = 400
|
||
|
|
#: Modes that leave a frame with no usable signal at all. A frame damaged this way is
|
||
|
|
#: presented as `missing`; everything else is presented as `suspect` with its pixels.
|
||
|
|
NO_SIGNAL = frozenset({"all_zero", "drop_image_hdu", "nan_fill", "zblank_fill"})
|
||
|
|
#: Auxiliary-head classes: clean plus every catalogued mode, in a fixed order so a
|
||
|
|
#: checkpoint's class indices stay meaningful.
|
||
|
|
CLASSES = ("clean",) + tuple(sorted(corruptions.CATALOG))
|
||
|
|
CLASS_INDEX = {name: index for index, name in enumerate(CLASSES)}
|
||
|
|
|
||
|
|
|
||
|
|
def stack_layout(satellites, target_satellite):
|
||
|
|
"""The (satellite, offset) slots a stack is built from, in a fixed order.
|
||
|
|
|
||
|
|
Deterministic and independent of what happens to be available, so a frame's position
|
||
|
|
in the stack always means the same thing and the model can rely on it.
|
||
|
|
|
||
|
|
The target satellite's own frame at offset 0 is the frame being reconstructed and is
|
||
|
|
never an input. The *counterpart's* frame at offset 0 is the single most valuable
|
||
|
|
entry in the stack -- a real observation of the right Sun at the right instant, worth
|
||
|
|
~46 dB on its own once its instrument gain is known.
|
||
|
|
"""
|
||
|
|
others = [s for s in satellites if s != target_satellite]
|
||
|
|
layout = []
|
||
|
|
for satellite in [target_satellite] + others:
|
||
|
|
for offset in OFFSETS:
|
||
|
|
layout.append((satellite, -offset))
|
||
|
|
layout.append((satellite, offset))
|
||
|
|
return layout + [(other, 0) for other in others]
|
||
|
|
|
||
|
|
|
||
|
|
class Sampler:
|
||
|
|
"""Builds training samples from a set of shards.
|
||
|
|
|
||
|
|
`shards` maps (day, satellite) -> :class:`suvi.dataset.Shard`.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, shards, satellites=(16, 18), damage_probability=0.30,
|
||
|
|
drop_probability=0.40, seed=0):
|
||
|
|
self.shards = shards
|
||
|
|
self.satellites = tuple(satellites)
|
||
|
|
#: Ceilings, not rates. Each *sample* draws its own severity uniformly up to
|
||
|
|
#: these, rather than every frame being damaged at one fixed probability.
|
||
|
|
#:
|
||
|
|
#: Fixed rates of 0.25 drop / 0.35 damage left barely half the stack clean in
|
||
|
|
#: every single sample, so the model never saw the easy case and learned a
|
||
|
|
#: hedging average instead of "trust the counterpart". It showed up as a model
|
||
|
|
#: that gained 5 dB on its own harsh validation set and ~1 dB on the bench,
|
||
|
|
#: where neighbours are mostly clean. Real outages are episodic -- long clean
|
||
|
|
#: stretches broken by bad runs -- so drawing severity per sample covers both
|
||
|
|
#: regimes, and roughly a fifth of samples come through almost untouched.
|
||
|
|
self.damage_probability = damage_probability
|
||
|
|
self.drop_probability = drop_probability
|
||
|
|
self.rng = np.random.default_rng(seed)
|
||
|
|
self._by_satellite = {}
|
||
|
|
self._day_of = {}
|
||
|
|
for (day, satellite), shard in shards.items():
|
||
|
|
self._day_of[id(shard)] = day
|
||
|
|
self._by_satellite.setdefault(satellite, {}).update(
|
||
|
|
{time: shard for time in shard.times()}
|
||
|
|
)
|
||
|
|
#: (day, target satellite, other satellite, pair time) -> per-band (gains,
|
||
|
|
#: offsets). See :meth:`_calibration`.
|
||
|
|
self._transfers = {}
|
||
|
|
|
||
|
|
def targets(self):
|
||
|
|
"""Every (satellite, time) that could be reconstructed, in a stable order."""
|
||
|
|
found = []
|
||
|
|
for satellite in self.satellites:
|
||
|
|
for time in sorted(self._by_satellite.get(satellite, {})):
|
||
|
|
found.append((satellite, time))
|
||
|
|
return found
|
||
|
|
|
||
|
|
def _read(self, satellite, time):
|
||
|
|
"""One slot in *model* space -- the cheap path, used for every clean frame.
|
||
|
|
|
||
|
|
Shards store the asinh representation already, so this is a rescale rather than
|
||
|
|
the sinh/arcsinh round trip that going via radiance would cost.
|
||
|
|
"""
|
||
|
|
shard = self._by_satellite.get(satellite, {}).get(time)
|
||
|
|
if shard is None:
|
||
|
|
return None
|
||
|
|
return shard.coded(time) if hasattr(shard, "coded") else encode_for_model(
|
||
|
|
shard.frames(time)
|
||
|
|
)
|
||
|
|
|
||
|
|
def _donor(self, satellite, time, rng, reach=60):
|
||
|
|
"""Another real frame, for the modes that splice one in.
|
||
|
|
|
||
|
|
``frozen`` repeats a neighbour, ``wrong_time`` files one under the wrong stamp
|
||
|
|
and ``torn_frame`` merges two. All three need a *plausible* second observation,
|
||
|
|
so the donor is drawn from nearby on the same satellite rather than invented.
|
||
|
|
"""
|
||
|
|
available = self._by_satellite.get(satellite, {})
|
||
|
|
for _ in range(8):
|
||
|
|
offset = int(rng.integers(-reach, reach + 1))
|
||
|
|
if offset == 0:
|
||
|
|
continue
|
||
|
|
candidate = time + offset * paths.CADENCE
|
||
|
|
if candidate in available:
|
||
|
|
return self._read(satellite, candidate)
|
||
|
|
return None
|
||
|
|
|
||
|
|
def _anchor(self, satellite, time, direction, exclude=frozenset(), limit=400):
|
||
|
|
"""Nearest slot with a frame outside `exclude`, up to `limit` slots away.
|
||
|
|
|
||
|
|
`exclude` holds the times inside a simulated outage. The production stack
|
||
|
|
builder picks anchors from *clean* times, so the sampler must too: an anchor
|
||
|
|
that lands inside the outage and arrives `missing` would leave the stack with
|
||
|
|
no real long-range frame at all, which is not what the pipeline would see.
|
||
|
|
"""
|
||
|
|
available = self._by_satellite.get(satellite, {})
|
||
|
|
for step in range(1, limit + 1):
|
||
|
|
candidate = time + direction * step * paths.CADENCE
|
||
|
|
if candidate in available and candidate not in exclude:
|
||
|
|
return candidate
|
||
|
|
return None
|
||
|
|
|
||
|
|
def _outages(self, rng, time):
|
||
|
|
"""Per-satellite times inside this sample's simulated episodic outage.
|
||
|
|
|
||
|
|
One contiguous run per affected satellite, log-uniform in length from one
|
||
|
|
slot to :data:`MAX_OUTAGE_SLOTS`, positioned so the target usually sits
|
||
|
|
inside it -- reproducing the archive's real failure mode of equipment and
|
||
|
|
comms dropping out for hours to days, sometimes on both spacecraft at once.
|
||
|
|
"""
|
||
|
|
blocked = {satellite: frozenset() for satellite in self.satellites}
|
||
|
|
# Gated on the drop setting: drop_probability=0 means "no synthetic losses
|
||
|
|
# of any kind", which diagnostics and tests rely on.
|
||
|
|
if self.drop_probability <= 0 or rng.random() >= EPISODIC_PROBABILITY:
|
||
|
|
return blocked
|
||
|
|
length = int(round(math.exp(rng.uniform(0.0, math.log(MAX_OUTAGE_SLOTS)))))
|
||
|
|
centre = time + int(rng.integers(-length, length + 1)) * paths.CADENCE
|
||
|
|
start = centre - (length // 2) * paths.CADENCE
|
||
|
|
interval = frozenset(start + i * paths.CADENCE for i in range(length))
|
||
|
|
if rng.random() < DUAL_OUTAGE_PROBABILITY:
|
||
|
|
hit = self.satellites
|
||
|
|
else:
|
||
|
|
hit = (self.satellites[int(rng.integers(len(self.satellites)))],)
|
||
|
|
for satellite in hit:
|
||
|
|
blocked[satellite] = interval
|
||
|
|
return blocked
|
||
|
|
|
||
|
|
def _calibration(self, day, satellite, other, exclude):
|
||
|
|
"""Per-band (gains, offsets) putting `other`'s radiance on `satellite`'s scale.
|
||
|
|
|
||
|
|
Fitted from one simultaneous pair per day -- the transfer drifts ~12% per
|
||
|
|
*week*, so a fit at most 24 hours stale is well inside its own noise, and one
|
||
|
|
pair per day is what a production pipeline could equally afford.
|
||
|
|
|
||
|
|
`exclude` is the time being reconstructed. A pair at that instant would fit
|
||
|
|
the transfer against the answer itself -- the oracle gain the whole exercise
|
||
|
|
exists to estimate honestly -- so the fit steps to a neighbouring pair
|
||
|
|
instead. Returns identity when the day has no usable pair.
|
||
|
|
"""
|
||
|
|
local = self.shards.get((day, satellite))
|
||
|
|
counterpart = self.shards.get((day, other))
|
||
|
|
if local is None or counterpart is None:
|
||
|
|
return None
|
||
|
|
candidates = sorted(set(local.times()) & set(counterpart.times()) - {exclude})
|
||
|
|
if not candidates:
|
||
|
|
return None
|
||
|
|
pair_time = candidates[len(candidates) // 2]
|
||
|
|
|
||
|
|
key = (day, satellite, other, pair_time)
|
||
|
|
if key not in self._transfers:
|
||
|
|
self._transfers[key] = align.fit_photometry(
|
||
|
|
counterpart.frames(pair_time), local.frames(pair_time)
|
||
|
|
)
|
||
|
|
return self._transfers[key]
|
||
|
|
|
||
|
|
def build(self, satellite, time, rng=None):
|
||
|
|
"""One sample: the stack, its per-frame metadata, and the withheld target.
|
||
|
|
|
||
|
|
`frames` is (S, 6, H, W) in *model* space -- roughly [-1, 1] asinh radiance --
|
||
|
|
with zeros where a frame is missing; `target` is the withheld frame in the same
|
||
|
|
space. `gains`/`offsets` are the per-frame radiance transfers that put a
|
||
|
|
cross-satellite frame on the target instrument's scale (identity rows for
|
||
|
|
same-satellite and missing frames), and `b0`/`radius` the solar geometry the
|
||
|
|
rotation warp needs -- both consumed by :func:`suvi.align.align_stack` rather
|
||
|
|
than here, so the CPU loader stays cheap.
|
||
|
|
|
||
|
|
Returns None if the target itself is absent, or if not one frame in the stack
|
||
|
|
carries pixels -- there is nothing to reconstruct *from*, and the model's
|
||
|
|
contract is to decline rather than emit a fabricated frame.
|
||
|
|
"""
|
||
|
|
rng = rng if rng is not None else self.rng
|
||
|
|
target = self._read(satellite, time)
|
||
|
|
if target is None:
|
||
|
|
return None
|
||
|
|
day = self._day_of.get(id(self._by_satellite[satellite][time]))
|
||
|
|
|
||
|
|
# This sample's severity, drawn once and applied to every frame in its stack.
|
||
|
|
drop_probability = float(rng.uniform(0.0, self.drop_probability))
|
||
|
|
damage_probability = float(rng.uniform(0.0, self.damage_probability))
|
||
|
|
outages = self._outages(rng, time)
|
||
|
|
|
||
|
|
layout = list(stack_layout(self.satellites, satellite))
|
||
|
|
# Anchors: the nearest usable frame either way on the target satellite --
|
||
|
|
# outside any simulated outage, as the production stack builder would pick
|
||
|
|
# them -- so a long gap still has something real to work from.
|
||
|
|
for direction in (-1, 1):
|
||
|
|
anchor = self._anchor(satellite, time, direction, exclude=outages[satellite])
|
||
|
|
if anchor is not None:
|
||
|
|
offset = (anchor - time) // paths.CADENCE
|
||
|
|
if (satellite, offset) not in layout:
|
||
|
|
layout.append((satellite, offset))
|
||
|
|
|
||
|
|
frames, states, dts, same, classes, transfers = [], [], [], [], [], []
|
||
|
|
for source_satellite, offset in layout:
|
||
|
|
when = time + offset * paths.CADENCE
|
||
|
|
image = self._read(source_satellite, when)
|
||
|
|
state, label = "available", "clean"
|
||
|
|
|
||
|
|
if image is None:
|
||
|
|
state = "missing"
|
||
|
|
elif when in outages[source_satellite]:
|
||
|
|
# Inside this sample's simulated episodic outage.
|
||
|
|
image, state = None, "missing"
|
||
|
|
elif rng.random() < drop_probability:
|
||
|
|
# A scattered single-slot loss, independent of the episodic runs.
|
||
|
|
image, state = None, "missing"
|
||
|
|
elif rng.random() < damage_probability:
|
||
|
|
mode = str(rng.choice(sorted(corruptions.CATALOG)))
|
||
|
|
corruption = corruptions.CATALOG[mode]
|
||
|
|
label = mode
|
||
|
|
if corruption.kind != "array":
|
||
|
|
# A file-level corruption destroys the container, so by the time a
|
||
|
|
# frame would reach the model there is nothing left of it.
|
||
|
|
image, state = None, "missing"
|
||
|
|
elif mode in NO_SIGNAL:
|
||
|
|
image, state = None, "missing"
|
||
|
|
else:
|
||
|
|
donor = self._donor(source_satellite, when, rng) \
|
||
|
|
if corruption.needs_donor else None
|
||
|
|
if corruption.needs_donor and donor is None:
|
||
|
|
label, state = "clean", "available" # nothing to draw from
|
||
|
|
else:
|
||
|
|
image = _damage(image, mode, int(rng.integers(0, 2**31 - 1)),
|
||
|
|
float(rng.uniform(0.5, 1.0)), donor)
|
||
|
|
state = "suspect"
|
||
|
|
|
||
|
|
shape = target.shape
|
||
|
|
frames.append(np.zeros(shape, dtype=np.float32) if image is None else image)
|
||
|
|
states.append(state)
|
||
|
|
dts.append(float(offset * paths.CADENCE))
|
||
|
|
same.append(source_satellite == satellite)
|
||
|
|
classes.append(CLASS_INDEX[label])
|
||
|
|
transfer = None
|
||
|
|
if source_satellite != satellite and image is not None:
|
||
|
|
transfer = self._calibration(day, satellite, source_satellite, time)
|
||
|
|
transfers.append(transfer)
|
||
|
|
|
||
|
|
if all(state == "missing" for state in states):
|
||
|
|
return None
|
||
|
|
|
||
|
|
identity = (np.ones(6, dtype=np.float32), np.zeros(6, dtype=np.float32))
|
||
|
|
b0, radius = align.solar_ephemeris(float(time))
|
||
|
|
return {
|
||
|
|
"frames": np.stack(frames),
|
||
|
|
"states": states,
|
||
|
|
"dts": np.array(dts, dtype=np.float32),
|
||
|
|
"same_satellite": np.array(same, dtype=bool),
|
||
|
|
"classes": np.array(classes, dtype=np.int64),
|
||
|
|
"gains": np.stack([(t or identity)[0] for t in transfers]),
|
||
|
|
"offsets": np.stack([(t or identity)[1] for t in transfers]),
|
||
|
|
"b0": float(b0),
|
||
|
|
"radius": float(radius),
|
||
|
|
"target": target,
|
||
|
|
"slot": (satellite, time),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def _damage(coded, mode, seed, severity, donor=None):
|
||
|
|
"""Apply one array corruption to all six bands of a frame, in model space.
|
||
|
|
|
||
|
|
Corruptions are *defined* on radiance -- ``gain_shift`` multiplies a physical
|
||
|
|
quantity, ``saturate`` clips against a physical ceiling -- so the frame is decoded,
|
||
|
|
damaged, and re-encoded. This is the expensive path, which is why only the frames
|
||
|
|
actually being corrupted take it.
|
||
|
|
|
||
|
|
Each band gets the same seed, so a fault hits the frame coherently rather than
|
||
|
|
differently per band, which is how the real ones behave: one instrument produces all
|
||
|
|
six. ``apply_array`` returns (pixels, header_overrides); the overrides describe a
|
||
|
|
FITS header this pipeline does not carry, so only the pixels are kept.
|
||
|
|
"""
|
||
|
|
radiance = decode_from_model(coded)
|
||
|
|
donor_radiance = None if donor is None else decode_from_model(donor)
|
||
|
|
bands = []
|
||
|
|
for index, band in enumerate(radiance):
|
||
|
|
damaged, _ = corruptions.apply_array(
|
||
|
|
mode, band, seed, severity,
|
||
|
|
donor=None if donor_radiance is None else donor_radiance[index],
|
||
|
|
)
|
||
|
|
bands.append(damaged)
|
||
|
|
return encode_for_model(np.stack(bands))
|
||
|
|
|
||
|
|
|
||
|
|
def encode_for_model(radiance):
|
||
|
|
"""Radiance -> the roughly-[-1, 1] asinh space the model works in.
|
||
|
|
|
||
|
|
Corruptions are applied in *radiance*, because that is where they are defined -- a
|
||
|
|
``gain_shift`` multiplies a physical quantity and a ``saturate`` clips against a
|
||
|
|
physical ceiling. Applying them after the transform would model a different fault.
|
||
|
|
So the encoding happens here, once, on the way into the network.
|
||
|
|
"""
|
||
|
|
coded = np.arcsinh(
|
||
|
|
np.nan_to_num(np.asarray(radiance, dtype=np.float32), nan=0.0,
|
||
|
|
posinf=0.0, neginf=0.0) / dataset.ASINH_SCALE
|
||
|
|
) / dataset.ASINH_RANGE
|
||
|
|
return np.clip(coded, -1.0, 1.0).astype(np.float32)
|
||
|
|
|
||
|
|
|
||
|
|
def decode_from_model(coded):
|
||
|
|
"""Inverse of :func:`encode_for_model`, back to radiance."""
|
||
|
|
scaled = np.asarray(coded, dtype=np.float64) * dataset.ASINH_RANGE
|
||
|
|
return (np.sinh(np.clip(scaled, -dataset.ASINH_RANGE, dataset.ASINH_RANGE))
|
||
|
|
* dataset.ASINH_SCALE).astype(np.float32)
|
||
|
|
|
||
|
|
|
||
|
|
def to_tensors(sample, torch):
|
||
|
|
"""Pack a sample into tensors for :func:`suvi.align.align_stack` and the model.
|
||
|
|
|
||
|
|
The stack is *unaligned* here: alignment is elementwise math plus a warp, which
|
||
|
|
the GPU does in milliseconds and the CPU loader should not spend its budget on.
|
||
|
|
"""
|
||
|
|
from . import model as model_module
|
||
|
|
|
||
|
|
condition = torch.stack([
|
||
|
|
model_module.frame_conditioning(
|
||
|
|
sample["states"][i], bool(sample["same_satellite"][i]), float(sample["dts"][i])
|
||
|
|
)
|
||
|
|
for i in range(len(sample["states"]))
|
||
|
|
])
|
||
|
|
# Already in model space: the sampler works there throughout, so nothing to convert.
|
||
|
|
return {
|
||
|
|
"stack": torch.from_numpy(np.ascontiguousarray(sample["frames"])),
|
||
|
|
"condition": condition,
|
||
|
|
"dts": torch.from_numpy(np.ascontiguousarray(sample["dts"])),
|
||
|
|
"gains": torch.from_numpy(np.ascontiguousarray(sample["gains"])),
|
||
|
|
"offsets": torch.from_numpy(np.ascontiguousarray(sample["offsets"])),
|
||
|
|
"b0": torch.tensor(sample["b0"], dtype=torch.float32),
|
||
|
|
"radius": torch.tensor(sample["radius"], dtype=torch.float32),
|
||
|
|
"target": torch.from_numpy(np.ascontiguousarray(sample["target"])),
|
||
|
|
}
|