noaa-goes-visualization/suvi/samples.py

"""Turning shards into training samples: stacks, masks, and damage drawn on the fly.

One sample is a target slot -- one satellite, one instant, some of its six bands
lost -- 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.

Four things here decide what the model can learn.

**Damage is withheld, never shown.**  Production gives the model nothing for a band
that is missing or judged bad -- it reads as absent -- so training does exactly
that: simulated damage removes bands, and the pristine shard frames are the ground
truth.  The model never sees a damaged pixel in either.

**The target is a real fill case.**  The renderer fills the bands a slot lacks and
keeps the rest, and four in five damaged slots keep one to five good bands.  So the
target's own undamaged bands are an input (the stack entry at dt=0 on its own
satellite), the lost ones are reconstructed, and the loss is taken on those alone.

**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 *anchors* -- per
band, the nearest usable frame in each direction, however far -- guarantee the stack
is never empty even beyond that.

**Damage is drawn from the real archive, band by band.**  Most samples
(:data:`TRACE_PROBABILITY`) take their joint two-satellite, six-band availability
pattern verbatim from a window of the profiled state sequence
(:mod:`suvi.availability`) centred on a slot where the target satellite lost bands:
the target's loss, the missing runs around it, the bands that fail together and the
correlation between the two spacecraft all arrive exactly as the archive contains
them.  Windows whose centre falls in a run longer than
:data:`suvi.availability.MAX_DUAL_GAP` where no satellite has some band are redrawn,
because the fill policy declines such gaps and the model should never train on work
it will not be asked to do.  The remainder lose frames at a per-sample rate, each
loss taking a band pattern drawn from the archive's own distribution, so the model
also meets stacks more and less damaged than any one window.
"""

import numpy as np

from . import align, availability as availability_module, dataset, paths, quality

#: 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's damage is a real archive window rather than a synthetic
#: draw.  Not 1.0: one window holds one outage's severity, and the synthetic draws
#: spread the stack's damage from untouched to heavy, which the model needs to
#: learn when to trust the counterpart and when its own neighbours.
TRACE_PROBABILITY = 0.8
#: How far the anchor search reaches, in slots.
ANCHOR_REACH = 400


#: The full rung ladder; a model's `span` selects the prefix it trains with.
LADDER = (1, 4, 16, 64, 128, 256, 512)


def ladder_for(span):
    """The exponential offsets a model of input span +/-`span` uses."""
    offsets = tuple(offset for offset in LADDER if offset <= span)
    if not offsets or offsets[-1] != span:
        raise ValueError(f"span {span} is not on the ladder {LADDER}")
    return offsets


def stack_layout(satellites, target_satellite, offsets=OFFSETS):
    """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), loss_probability=0.5,
                 availability=None, seed=0, offsets=OFFSETS, segments=None):
        self.shards = shards
        self.satellites = tuple(satellites)
        #: The rung ladder this model trains with; the maximum rung is the input
        #: span, and with `segments` set it is also the target margin.
        self.offsets = tuple(offsets)
        #: Pristine contiguous (t0, t1) ranges.  When present, only slots at
        #: least one full span inside a segment are targets, so every rung of
        #: the ladder lands on jointly-pristine ground truth and every hole the
        #: model sees within reach is simulated, never real.
        self.segments = None if segments is None else sorted(map(tuple, segments))
        #: The archive's profiled state sequences (:class:`suvi.availability
        #: .Availability`), or None to use synthetic draws only.  When present,
        #: :data:`TRACE_PROBABILITY` of samples overlay a real availability window.
        self.availability = availability
        #: A ceiling, not a rate.  Each synthetic sample draws its own loss rate
        #: uniformly up to it, rather than every frame failing 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 a fair share of samples come through almost untouched.
        self.loss_probability = loss_probability
        #: Which bands a lost frame loses together: the archive's own distribution
        #: when there is a trace, otherwise every pattern equally.
        if availability is not None:
            self._patterns, self._pattern_odds = availability.damage_patterns()
        else:
            codes = np.arange(1, 1 << len(paths.WAVELENGTHS))
            self._patterns = ((codes[:, None] >> np.arange(len(paths.WAVELENGTHS))) & 1
                              ).astype(bool)
            self._pattern_odds = np.full(len(codes), 1.0 / len(codes))
        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."""
        eligible = None
        if self.segments is not None:
            from . import segments as segments_module
            eligible = set(segments_module.targets(self.segments,
                                                   max(self.offsets)))
        found = []
        for satellite in self.satellites:
            for time in sorted(self._by_satellite.get(satellite, {})):
                if eligible is None or time in eligible:
                    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
        if hasattr(shard, "coded"):
            return shard.coded(time)
        frames = shard.frames(time)
        return None if frames is None else encode_for_model(frames)

    def _anchor(self, read, satellite, time, direction, band, trace):
        """Offset of the nearest slot, up to ANCHOR_REACH away, whose frame holds
        `band` usable -- outside any simulated outage, as the production stack
        builder picks anchors from clean times, band by band."""
        available = self._by_satellite.get(satellite, {})
        for step in range(1, ANCHOR_REACH + 1):
            offset = direction * step
            candidate = time + offset * paths.CADENCE
            if candidate not in available:
                continue
            if trace is not None and not trace.usable(satellite, offset)[band]:
                continue
            # A vetted shard lists every slot it saw, including ones stored as
            # an unreadable sentinel; an anchor must actually carry pixels.
            if read(satellite, candidate) is not None:
                return offset
        return None

    def _pattern(self, rng):
        """One loss: which bands a damaged frame loses."""
        return self._patterns[rng.choice(len(self._patterns), p=self._pattern_odds)]

    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
        # Middle-outward: a vetted shard lists slots it stored as unreadable
        # sentinels, so a candidate time is not a promise of pixels on both sides.
        middle = len(candidates) // 2
        order = sorted(range(len(candidates)), key=lambda i: abs(i - middle))
        for index in order[:16]:
            pair_time = candidates[index]
            key = (day, satellite, other, pair_time)
            if key not in self._transfers:
                theirs = counterpart.frames(pair_time)
                ours = local.frames(pair_time) if theirs is not None else None
                if ours is None:
                    continue
                fit = align.fit_photometry(theirs, ours)
                if fit is None:
                    continue    # physically implausible transfer; try another pair
                self._transfers[key] = fit
            return self._transfers[key]
        return None

    def build(self, satellite, time, rng=None):
        """One sample: the stack, its per-frame metadata, and the target.

        `frames` is (S, 6, H, W) in *model* space -- roughly [-1, 1] asinh radiance --
        with zeros wherever a band is absent, and `present` (S, 6) says which bands
        each entry holds.  `target` is the pristine slot and `lost` (6,) the bands
        of it the model must reconstruct; the rest are in the stack, at dt=0 on the
        target's satellite.  `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 band in the stack
        survives -- 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
        cache = {}

        def read(source, when):
            key = (source, when)
            if key not in cache:
                cache[key] = self._read(source, when)
            return cache[key]

        target = read(satellite, time)
        if target is None:
            return None
        day = self._day_of.get(id(self._by_satellite[satellite][time]))

        # This sample's damage: usually a verbatim window of the real archive,
        # centred on a slot where the target satellite lost bands; otherwise a
        # synthetic draw.  A window whose centre sits inside a run longer than
        # MAX_DUAL_GAP in which some band is on neither satellite is redrawn:
        # production declines to fill those, so training there would spend
        # capacity learning to fabricate frames nobody will ever ask for.
        trace = None
        if self.availability is not None and rng.random() < TRACE_PROBABILITY:
            for _ in range(32):
                candidate = self.availability.sample_window(rng, satellite)
                if candidate.dual_gap() <= availability_module.MAX_DUAL_GAP:
                    trace = candidate
                    break
        loss = 0.0
        if trace is not None:
            lost = ~trace.usable(satellite, 0)
        else:
            loss = float(rng.uniform(0.0, self.loss_probability))
            lost = self._pattern(rng)

        layout = list(stack_layout(self.satellites, satellite, self.offsets))
        layout.append((satellite, 0))               # the target's undamaged bands
        anchors = len(layout)
        listed = {offset for source, offset in layout if source == satellite}
        found = set()
        for band in range(len(paths.WAVELENGTHS)):
            for direction in (-1, 1):
                offset = self._anchor(read, satellite, time, direction, band, trace)
                if offset is not None and offset not in listed:
                    found.add(offset)
        layout.extend((satellite, offset) for offset in sorted(found))

        frames, present, states, dts, same, transfers = [], [], [], [], [], []
        for index, (source_satellite, offset) in enumerate(layout):
            when = time + offset * paths.CADENCE
            image = read(source_satellite, when)
            held = np.zeros(len(paths.WAVELENGTHS), dtype=bool)
            if image is not None:
                held[:] = True
                if source_satellite == satellite and offset == 0:
                    held &= ~lost
                elif trace is not None:
                    held &= trace.usable(source_satellite, offset)
                elif index < anchors and rng.random() < loss:
                    held &= ~self._pattern(rng)
            if image is None or not held.any():
                frame = np.zeros(target.shape, dtype=np.float32)
            elif held.all():
                frame = image
            else:
                frame = np.where(held[:, None, None], image, 0.0).astype(np.float32)
            frames.append(frame)
            present.append(held)
            states.append(quality.AVAILABLE if held.any() else quality.MISSING)
            dts.append(float(offset * paths.CADENCE))
            same.append(source_satellite == satellite)
            transfer = None
            if source_satellite != satellite and held.any():
                transfer = self._calibration(day, satellite, source_satellite, time)
            transfers.append(transfer)

        if not any(held.any() for held in present):
            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),
            "present": np.stack(present),
            "states": states,
            "dts": np.array(dts, dtype=np.float32),
            "same_satellite": np.array(same, dtype=bool),
            "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,
            "lost": lost.copy(),
            "slot": (satellite, time),
        }


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"])),
        "present": torch.from_numpy(np.ascontiguousarray(sample["present"])).float(),
        "lost": torch.from_numpy(np.ascontiguousarray(sample["lost"])).float(),
    }