noaa-goes-visualization/suvi/model.py

"""The learned filler: choose, per pixel, which aligned observation to trust.

The stack this model receives has already been aligned by :mod:`suvi.align`: every
frame is rotation-warped to the target instant and cross-satellite frames are on the
target instrument's radiance scale.  Everything with a closed form -- solar rotation,
the photometric transfer -- was applied there, deterministically.  What remains is
exactly what has no closed form:

* **Per-pixel trust.**  Which candidate is right varies spatially: on-disc, a warped
  temporal neighbour is sharp and co-rotating; above the limb the corona does not
  co-rotate and the counterpart's simultaneous view wins; inside a torn or partly
  corrupted frame, half the pixels are good and half are not.  The weight head emits
  per-pixel, per-band blend logits over the stack, and the softmax over candidates
  does the selection.
* **A bounded photometric polish.**  The daily gain fit is a global affine, and the
  measured residual varies 28-61% with radius and drifts between fits.  The polish
  head emits a per-candidate log-gain field at working resolution, upsampled and
  bounded to ``exp(+/-1.0)`` -- wide and local enough to track a compact
  brightening a candidate shows dimly, yet still multiplicative on real pixels:
  it can only amplify structure a candidate carries, never draw its own.
* **A bounded residual flow.**  Solar rotation is warped out deterministically,
  but plasma also *moves*: loops rise after a flare, jets and dimming fronts
  propagate, and no re-weighting or re-scaling of misplaced pixels can express
  displacement.  The flow head emits one displacement field per candidate
  (shared across bands -- all six watch the same plasma), bounded by
  :data:`MAX_FLOW_FRACTION` scaled with each frame's dt, applied by grid
  sampling before the blend.  A flow-warped real frame is still real pixels.
* **Trend candidates.**  A convex blend of observations can interpolate between
  its endpoints but never represent an extremum inside a gap -- a flare that rises
  and peaks while both spacecraft are dark would be clipped to the brightest thing
  any candidate shows.  So the stack is augmented (in :meth:`StackFiller.forward`,
  identically at training and inference) with two *derived* candidates: per-pixel
  linear extrapolations to the target instant from the two nearest valid
  target-satellite frames on each side.  Where the data itself shows brightness
  rising, the extrapolated frame overshoots the real ones, and the blend may
  choose it; reach is capped at :data:`TREND_MAX_REACH` times the pair's own
  baseline so a short-baseline slope cannot be amplified across a long gap.

Why so little machine learning
------------------------------
The previous architecture predicted flow, affine, gain, offset, blend, and a free
residual from eight zero-initialised heads, and three training runs failed -- the
last collapsed to an input-independent output because an unbounded additive head is
the cheapest way to satisfy a badly conditioned loss.  Here no head can produce an
image on its own: the output is always a convex combination of real aligned
observations, times a gain pinned near one.  An untrained model *is* the prior-
weighted blend of its candidates, and every hand-written filler is one softmax
saturation away, so training starts from a sane policy and can only refine it.

Cross-frame reasoning -- "trust the counterpart *because* the temporal neighbours
disagree" -- needs no attention: each frame's conditioning carries an **agreement**
feature, its RMS deviation from the stack's per-pixel median, computed identically
at training and inference.  A frozen, torn or mis-gained frame announces itself
there, whether or not any detector flagged it.
"""

import math

import torch
import torch.nn as nn
import torch.nn.functional as F

#: Resolution the network reasons at.  Blend and polish are smooth fields, so they
#: are predicted here and upsampled to the frames' native resolution -- 640 for the
#: training shards, 1280 on the bench -- where they are applied to the real pixels.
#: Sharpness comes from the observations, not the network.
WORK_SIZE = 160
#: Bands per frame, in suvi.paths.WAVELENGTHS order.
BANDS = 6
#: Half-width of the polish in log-gain: gains live in [exp(-1), exp(1)], about
#: [0.37, 2.72].  The original +/-0.3 covered instrument-transfer drift but clipped
#: real solar dynamics: an EUV kernel brightens far more than 35% between slots, so
#: a flare visible dimly in a candidate was structurally unreachable.  Still a
#: multiplicative bound on real pixels -- not a free residual.
POLISH_RANGE = 1.0
#: How far a trend candidate may extrapolate, in multiples of its source pair's own
#: time baseline.  Slope error grows linearly with reach, so unbounded reach would
#: let a 1-slot baseline extrapolate pure noise across a 300-slot gap.
TREND_MAX_REACH = 4.0
#: Local cost-volume radius for the `corr_flow` variant, in working-resolution
#: pixels: (2r+1)^2 correlation channels between each candidate and the stack
#: median.  Turns the flow head from a motion *prior* (one frame plus dt) into
#: measured correspondence, the PWC-Net/RAFT primitive at screening scale.
CORR_RADIUS = 3
#: Token grid for the `attention` variant: features pooled to this many cells
#: per side, attention across the stack axis per cell.  Coarse on purpose --
#: cross-frame *judgment* is low-frequency; sharpness still comes from pixels.
ATTN_GRID = 10
ATTN_HEADS = 4

#: Ceiling of the flow head's displacement, as a fraction of frame width -- the
#: resolution-independent unit, since training frames are 640^2 and bench frames
#: 1280^2 over the same field of view.  The effective bound scales with each
#: frame's dt encoding (motion needs and uncertainty both grow with time
#: distance), reaching this ceiling only at the encoding's saturation: ~1.2% of
#: width for a 1-slot neighbour (~15 native px on the bench, generous for
#: 50-100 km/s loop motion over 4 minutes), ~9% for a 300-slot anchor.  A
#: simultaneous counterpart (dt=0) gets no flow at all: a real observation of the
#: right instant never needs motion correction.
MAX_FLOW_FRACTION = 0.10

#: Per-frame conditioning, one column each:
#:   available / missing     state one-hot
#:   target_satellite        the frame is from the target's satellite
#:   dt_sign, dt_log         dt encoding (sign, log1p|dt| scaled; ~1.0 at 300 slots)
#:   agree_mean, agree_max   RMS deviation from the stack median, band mean and max
#:   derived                 a trend candidate built by the model, not a real frame
#: The agreement columns are written by the model itself, from the aligned stack --
#: so they exist identically at training and inference, unlike detector scores,
#: which were always zero in training and populated on the bench.
#:
#: There is no `suspect` state (v9).  A band judged bad is withheld everywhere --
#: training, bench and render -- so it reads as absent; v8 and earlier carried a
#: suspect column, and their checkpoints do not load under this layout.
COND_LAYOUT = ("available", "missing", "target_satellite", "dt_sign", "dt_log",
               "agree_mean", "agree_max", "derived")
COND_DIM = len(COND_LAYOUT)
AVAILABLE, MISSING, TARGET_SATELLITE, DT_SIGN, DT_LOG, AGREE_MEAN, AGREE_MAX, DERIVED = (
    range(COND_DIM))

#: Frames the encoder, and the full-resolution blend, take at a time during inference.
ENCODER_CHUNK = 4
BLEND_CHUNK = 4

#: Blend prior: fixed logits added to the weight head's output, so the untrained
#: model starts at a sane hand-written policy instead of the uniform stack mean.
#: With candidates aligned and gain-matched there is no cross-satellite penalty any
#: more -- the counterpart arrives on the right scale -- leaving recency, and the
#: derived-candidate penalty below.
PRIOR_TEMPORAL_DECAY = 1.5
#: Derived (trend) candidates start disfavoured: extrapolation is noisier than
#: observation, so the untrained model remains the same sane observation blend it
#: always was, and training has to *earn* trust in the trends.
PRIOR_DERIVED_PENALTY = 2.0


def frame_conditioning(state, is_target_satellite, dt_seconds):
    """Build one frame's static conditioning vector.  See :data:`COND_DIM`."""
    vector = torch.zeros(COND_DIM)
    vector[{"available": AVAILABLE, "missing": MISSING}[state]] = 1.0
    vector[TARGET_SATELLITE] = 1.0 if is_target_satellite else 0.0
    dt = float(dt_seconds)
    vector[DT_SIGN] = math.copysign(1.0, dt) if dt else 0.0
    vector[DT_LOG] = math.log1p(abs(dt) / 240.0) / 6.0
    return vector


def dt_from_conditioning(condition):
    """Invert the dt encoding, exactly.  (..., COND_DIM) -> seconds.

    The trend builder needs real time deltas, and re-deriving them from the
    conditioning (rather than threading a second dts tensor through every call
    site) keeps forward()'s signature -- and every caller -- unchanged."""
    return condition[..., DT_SIGN] * torch.expm1(condition[..., DT_LOG] * 6.0) * 240.0


def blend_prior(condition):
    """Fixed blend logits from each frame's conditioning.  (B, S) -> (B, S)."""
    distance = condition[..., DT_LOG] * 6.0
    return (-PRIOR_TEMPORAL_DECAY * distance
            - PRIOR_DERIVED_PENALTY * condition[..., DERIVED])


# ------------------------------------------------------------------------ backbone


class FiLM(nn.Module):
    """Per-frame feature modulation from that frame's conditioning vector.

    The stack is processed as a batch of frames, so each frame's `dt`, availability
    and agreement have to reach its own features and no other's.  Scale-and-shift is
    the cheapest thing that does that.
    """

    def __init__(self, channels, cond_dim=COND_DIM):
        super().__init__()
        self.to_scale_shift = nn.Sequential(
            nn.Linear(cond_dim, channels * 2), nn.SiLU(),
            nn.Linear(channels * 2, channels * 2),
        )
        nn.init.zeros_(self.to_scale_shift[-1].weight)
        nn.init.zeros_(self.to_scale_shift[-1].bias)

    def forward(self, features, condition):
        scale, shift = self.to_scale_shift(condition).chunk(2, dim=1)
        return features * (1 + scale[..., None, None]) + shift[..., None, None]


class Block(nn.Module):
    def __init__(self, inputs, outputs, cond_dim=COND_DIM):
        super().__init__()
        self.first = nn.Conv2d(inputs, outputs, 3, padding=1)
        self.second = nn.Conv2d(outputs, outputs, 3, padding=1)
        self.norm = nn.GroupNorm(8, outputs)
        self.film = FiLM(outputs, cond_dim)
        self.skip = (
            nn.Identity() if inputs == outputs else nn.Conv2d(inputs, outputs, 1)
        )

    def forward(self, x, condition):
        h = F.silu(self.first(x))
        h = self.norm(self.second(h))
        return F.silu(self.film(h, condition) + self.skip(x))


class Encoder(nn.Module):
    """Shared per-frame encoder: every frame in the stack goes through one set of
    weights, and a frame's role -- which satellite, how far in time, how damaged --
    reaches its features through conditioning rather than through separate branches.
    """

    def __init__(self, in_channels, base=32, depth=3, cond_dim=COND_DIM):
        super().__init__()
        widths = [base * min(2**i, 8) for i in range(depth + 1)]
        self.stem = Block(in_channels, widths[0], cond_dim)
        self.down = nn.ModuleList(
            [Block(widths[i], widths[i + 1], cond_dim) for i in range(depth)]
        )
        self.up = nn.ModuleList(
            [Block(widths[i + 1] + widths[i], widths[i], cond_dim)
             for i in reversed(range(depth))]
        )
        self.width = widths[0]

    def forward(self, x, condition):
        h = self.stem(x, condition)
        skips = []
        for block in self.down:
            skips.append(h)
            h = block(F.avg_pool2d(h, 2), condition)
        for block, skip in zip(self.up, reversed(skips)):
            h = F.interpolate(h, size=skip.shape[-2:], mode="nearest")
            h = block(torch.cat([h, skip], dim=1), condition)
        return h


# --------------------------------------------------------------------------- model


#: Blend one band at a time at inference (no gradients): identical output, a sixth
#: the full-resolution memory.  A switch only so tests can compare the two paths.
BAND_AT_A_TIME = True


class StackFiller(nn.Module):
    """Reconstruct one frame as a per-pixel convex combination of aligned candidates.

    Forward takes:
      sources   (B, S, BANDS, H, W)  the **aligned** stack from suvi.align, asinh
                space, zeros where a frame is missing.  H is 640 in training and
                1280 on the bench; the network downsamples internally.
      condition (B, S, COND_DIM)     static per-frame conditioning; the agreement
                columns are overwritten here from the stack itself.

    Returns the reconstruction at the resolution of `sources`.
    """

    def __init__(self, base=32, depth=3, relative=False, corr_flow=False,
                 attention=False):
        super().__init__()
        #: Screened cross-frame pathways, one per experiment arm.  Every one is
        #: identity at zero-init, so an untrained variant is still the prior
        #: blend and the collapse mode stays structurally unreachable.
        self.relative = relative
        self.corr_flow = corr_flow
        # +1 input channel: per-frame validity, so the encoder can tell a genuinely
        # dark frame (an eclipse) from one that is merely absent.  The `relative`
        # variant adds each frame's per-pixel deviation from the stack median:
        # spatial disagreement as an input map, not a 2-scalar bottleneck.
        in_channels = BANDS + 1 + (BANDS if relative else 0)
        self.encoder = Encoder(in_channels, base, depth, COND_DIM)
        self.weight = nn.Conv2d(self.encoder.width, BANDS, 3, padding=1)
        self.polish = nn.Conv2d(self.encoder.width, BANDS, 1)
        # One displacement field per candidate, shared across bands: the flow is a
        # physical movement of plasma, which all six bands watch together.  The
        # `corr_flow` variant feeds the head a local cost volume against the
        # stack median: correspondence, not just a motion prior.
        flow_in = self.encoder.width + ((2 * CORR_RADIUS + 1) ** 2 if corr_flow else 0)
        self.flow = nn.Conv2d(flow_in, 2, 3, padding=1)
        # The `attention` variant: learned mixing across the stack axis on a
        # coarse token grid, gated by a zero-init scalar so it starts silent.
        self.attend = (nn.MultiheadAttention(self.encoder.width, ATTN_HEADS,
                                             batch_first=True)
                       if attention else None)
        self.attn_gate = nn.Parameter(torch.zeros(1)) if attention else None
        for head in (self.weight, self.polish, self.flow):
            nn.init.zeros_(head.weight)
            nn.init.zeros_(head.bias)

    def forward(self, sources, condition, present=None):
        """`present` (B, S, BANDS), 1 where a frame holds that band, lets a frame
        missing some bands still contribute the rest: a missing band takes no blend
        weight in that band, and the encoder sees the stack's per-band median there
        instead of empty pixels.  None means every frame holds every band."""
        if present is None:
            present = sources.new_ones(sources.shape[:2] + (BANDS,))
        present = present.to(sources.dtype)
        trends, condition, present = self._trends(sources, condition, present)
        batch, stack = condition.shape[:2]
        size = sources.shape[-2:]
        valid = condition[..., AVAILABLE]

        def candidates(bands):
            """The real frames then the trends for a slice of bands, at full
            resolution, (B*stack, n, H, W).  Joined here, a slice at a time, not
            once for the whole stack: that copy is ~1 GB at 1280px, and inference
            blends one band at a time."""
            joined = torch.cat([sources[:, :, bands], trends[:, :, bands]], dim=1)
            return joined.reshape(batch * stack, -1, *size)

        def small(part):
            """(B, n, BANDS, H, W) -> the working resolution.  Per frame, so the real
            frames and the trends resampled apart equal the joined stack resampled."""
            count = part.shape[1]
            if count == 0:
                return part.new_zeros(batch, 0, BANDS, WORK_SIZE, WORK_SIZE)
            return F.interpolate(
                part.reshape(batch * count, BANDS, *size),
                size=(WORK_SIZE, WORK_SIZE), mode="bilinear", align_corners=False,
            ).view(batch, count, BANDS, WORK_SIZE, WORK_SIZE)

        if size[-1] != WORK_SIZE:
            work = torch.cat([small(sources), small(trends)], dim=1).reshape(
                batch * stack, BANDS, WORK_SIZE, WORK_SIZE)
        else:
            work = candidates(slice(None))

        median = self._stack_median(work, valid, batch, stack, present)
        held = present.reshape(batch * stack, BANDS, 1, 1)
        if bool((held < 0.5).any()):
            filler = torch.nan_to_num(median, nan=0.0).to(work.dtype).expand(
                batch, stack, BANDS, *work.shape[-2:]
            ).reshape(batch * stack, BANDS, *work.shape[-2:])
            work = torch.where(held > 0.5, work, filler)
        condition = self._with_agreement(work, condition, valid, batch, stack, median,
                                         present)
        mask = valid.reshape(batch * stack, 1, 1, 1).expand(-1, 1, *work.shape[-2:])
        inputs = [work, mask]
        flat_median = None
        if self.relative or self.corr_flow:
            flat_median = torch.nan_to_num(median, nan=0.0).to(work.dtype).expand(
                batch, stack, BANDS, *work.shape[-2:]
            ).reshape(batch * stack, BANDS, *work.shape[-2:])
        if self.relative:
            inputs.append((work - flat_median)
                          * valid.reshape(batch * stack, 1, 1, 1))
        inputs = torch.cat(inputs, dim=1)
        flat_condition = condition.reshape(batch * stack, -1)
        if torch.is_grad_enabled():
            features = self.encoder(inputs, flat_condition)
        else:
            # Every layer is per frame (convolutions, GroupNorm, FiLM), so a few
            # frames at a time is the same computation; all of them at once held
            # ~0.9 GB of activations for a 22-frame stack.
            features = torch.cat([
                self.encoder(inputs[start:start + ENCODER_CHUNK],
                             flat_condition[start:start + ENCODER_CHUNK])
                for start in range(0, inputs.shape[0], ENCODER_CHUNK)])
        if self.attend is not None:
            features = features + self.attn_gate * self._attended(features, batch, stack)

        # The heads run once, at working resolution.  Everything after them is
        # per-pixel and per-band at full resolution.
        weight = self.weight(features)
        polish, polish_range = self.polish(features), POLISH_RANGE
        # Residual motion: rotation was warped out deterministically in suvi.align;
        # whatever moved besides -- rising loops, jets, expanding dimmings -- is
        # displaced here, per candidate, before the blend sees it.  Zero-init makes
        # the untrained flow the identity, and a flow-warped real frame is still
        # real pixels, so the no-synthesis contract survives.
        flow_input = features
        if self.corr_flow:
            flow_input = torch.cat([features, _correlation(work, flat_median)], dim=1)
        coarse_flow = self.flow(flow_input)
        allowance = (MAX_FLOW_FRACTION
                     * condition[..., DT_LOG].reshape(batch * stack, 1, 1, 1))

        def grid_for(rows=slice(None)):
            """The sampling grid of the frames in `rows` (of B*stack), full size."""
            flow = F.interpolate(coarse_flow[rows], size=size, mode="bilinear",
                                 align_corners=False)
            if torch.is_grad_enabled():
                return _flow_grid(allowance[rows] * torch.tanh(flow))
            return _flow_grid(flow.tanh_().mul_(allowance[rows]), inplace=True)

        prior = blend_prior(condition).reshape(batch * stack, 1, 1, 1)
        # Per band: a frame with no pixels, or without this band, takes no weight.
        absent = (valid.reshape(batch * stack, 1, 1, 1) < 0.5) | (held < 0.5)

        def logits_for(rows, bands):
            """Full-size blend logits of the frames in `rows`, absent ones -inf."""
            logits = F.interpolate(weight[rows, bands], size=size, mode="bilinear",
                                   align_corners=False)
            return logits.add_(prior[rows]).masked_fill_(absent[rows, bands],
                                                         float("-inf"))

        def blend_chunked(bands):
            """:func:`blend` without gradients, a few frames at a time.

            Every full-resolution intermediate -- logits, flow grid, warped frames,
            gain -- is (frames, n_bands, H, W), ~6.5 MB per frame per band at 1280px;
            the whole stack at once held ~1 GB of them.  The softmax is streamed:
            one pass for the per-pixel maximum, one for the sum, one to blend.  The
            same weights up to rounding."""
            real = sources.shape[1]
            count = len(range(BANDS)[bands])
            out = sources.new_zeros(batch, count, *size)
            for item in range(batch):
                chunks = [slice(item * stack + lo, item * stack + min(lo + BLEND_CHUNK,
                                                                      stack))
                          for lo in range(0, stack, BLEND_CHUNK)]
                peak = torch.full_like(out[item], float("-inf"))
                for rows in chunks:
                    torch.maximum(peak, logits_for(rows, bands).amax(dim=0), out=peak)
                total = torch.zeros_like(out[item])
                for rows in chunks:
                    total += logits_for(rows, bands).sub_(peak).exp_().sum(dim=0)
                for rows in chunks:
                    lo, hi = rows.start - item * stack, rows.stop - item * stack
                    frames = torch.cat([sources[item, lo:min(hi, real), bands],
                                        trends[item, max(lo - real, 0):max(hi - real, 0),
                                               bands]])
                    frames = _warp(frames, grid_for(rows))
                    gain = F.interpolate(polish[rows, bands], size=size,
                                         mode="bilinear", align_corners=False)
                    frames.mul_(gain.tanh_().mul_(polish_range).exp_())
                    del gain
                    weights = logits_for(rows, bands).sub_(peak).exp_().div_(total)
                    frames.mul_(weights.nan_to_num_(nan=0.0))
                    out[item] += frames.sum(dim=0)
            return out

        def blend(bands):
            """The convex blend for a slice of bands, (B, n_bands, H, W)."""
            if not torch.is_grad_enabled():
                return blend_chunked(bands)
            logits = F.interpolate(weight[:, bands], size=size,
                                   mode="bilinear", align_corners=False)
            gain = torch.exp(polish_range * torch.tanh(F.interpolate(
                polish[:, bands], size=size, mode="bilinear", align_corners=False)))
            frames = _warp(candidates(bands), grid_for())
            logits = logits + prior
            # A frame with no pixels in a band must not win weight there, however
            # confident the head is.
            logits = logits.masked_fill(absent[:, bands], float("-inf"))
            count = logits.shape[1]
            weights = torch.softmax(logits.view(batch, stack, count, *size), dim=1)
            weights = torch.nan_to_num(weights, nan=0.0)
            adjusted = (gain * frames).view(batch, stack, count, *size)
            return (adjusted * weights).sum(dim=1)

        if torch.is_grad_enabled() or not BAND_AT_A_TIME:
            return blend(slice(None))
        # Inference: a band at a time.  The bands blend independently, so the result
        # is the same, but the full-resolution intermediates are a sixth the size.
        # All six at once took a 24-candidate stack at 1280px from 1.9 GB to a
        # 9.4 GB peak.
        out = sources.new_empty(batch, BANDS, *size)
        for band in range(BANDS):
            out[:, band:band + 1] = blend(slice(band, band + 1))
        return out

    def _trends(self, sources, condition, present=None):
        """The two trend candidates: per-pixel linear extrapolations to t=0.

        Returns (frames (B, 2, BANDS, H, W), condition and present for the stack with
        the trends appended).  The frames come back on their own, not appended to
        `sources`, because that copy is the whole stack at full resolution.

        For each side (past, future), take the two nearest *valid, target-satellite*
        frames holding all six bands and continue their per-pixel line to the
        target instant.  Where the
        data shows brightness rising into the gap, the extrapolated frame overshoots
        every real candidate -- which is exactly the reach a convex blend of
        observations lacks; the blend can then select the trend where it helps.
        Derived frames are conditioned with the nearest source's dt (recency prior
        applies) and the derived flag, and run through the same agreement statistics
        as everything else, so a wild extrapolation announces itself there.
        """
        batch, stack = sources.shape[:2]
        if present is None:
            present = sources.new_ones(batch, stack, BANDS)
        if stack < 2:                            # nothing to draw a line through
            return sources.new_zeros(batch, 0, *sources.shape[2:]), condition, present
        dts = dt_from_conditioning(condition)
        # A frame missing a band has no line to continue in it.
        valid = (condition[..., AVAILABLE] > 0.5) & (present > 0.5).all(dim=-1)
        same = condition[..., TARGET_SATELLITE] > 0.5
        rows = torch.arange(batch, device=sources.device)

        extra_frames, extra_condition = [], []
        for future in (False, True):
            side = (dts > 0) if future else (dts < 0)
            eligible = valid & same & side
            rank = torch.where(eligible, dts.abs(),
                               torch.full_like(dts, float("inf")))
            order = rank.argsort(dim=1)
            near, far = order[:, 0], order[:, 1]
            ok = eligible[rows, near] & eligible[rows, far]

            dt1, dt2 = dts[rows, near], dts[rows, far]
            f1, f2 = sources[rows, near], sources[rows, far]
            # f(t) = f1 + (t - dt1) * (f1 - f2)/(dt1 - dt2), evaluated at t = 0 and
            # capped: reach is at most TREND_MAX_REACH times the pair's baseline.
            reach = (-dt1 / (dt1 - dt2)).nan_to_num(0.0).clamp(0.0, TREND_MAX_REACH)
            frame = (f1 + reach[:, None, None, None] * (f1 - f2)).clamp(-1.0, 1.0)
            frame = frame * ok[:, None, None, None]

            extra = torch.zeros(batch, COND_DIM, device=condition.device,
                                dtype=condition.dtype)
            extra[:, AVAILABLE] = ok        # available...
            extra[:, MISSING] = ~ok         # ...or missing, if a side lacks two frames
            extra[:, TARGET_SATELLITE] = 1.0
            extra[:, DT_SIGN] = torch.where(ok, torch.sign(dt1), torch.zeros_like(dt1))
            extra[:, DT_LOG] = torch.where(
                ok, torch.log1p(dt1.abs() / 240.0) / 6.0, torch.zeros_like(dt1))
            extra[:, DERIVED] = ok.to(condition.dtype)
            extra_frames.append(frame)
            extra_condition.append(extra)

        condition = torch.cat([condition, torch.stack(extra_condition, dim=1)], dim=1)
        present = torch.cat([present, present.new_ones(batch, 2, BANDS)], dim=1)
        return torch.stack(extra_frames, dim=1), condition, present

    def _stack_median(self, work, valid, batch, stack, present=None):
        """The per-pixel median over valid frames holding each band,
        (B, 1, BANDS, h, w).

        NaN where no valid frame covers a pixel -- deliberately, so the agreement
        statistic ignores such pixels; consumers that need numbers nan_to_num it
        themselves.
        """
        with torch.no_grad():
            grouped = work.detach().float().view(batch, stack, BANDS,
                                                 *work.shape[-2:])
            keep = valid.reshape(batch, stack, 1, 1, 1) > 0.5
            if present is not None:
                keep = keep & (present.reshape(batch, stack, BANDS, 1, 1) > 0.5)
            hidden = torch.where(keep, grouped, torch.full_like(grouped, float("nan")))
            return hidden.nanmedian(dim=1, keepdim=True).values

    def _with_agreement(self, work, condition, valid, batch, stack, median=None,
                        present=None):
        """Fill the agreement columns of the conditioning from the stack itself.

        Each frame's RMS deviation from the stack's per-pixel median, over the valid
        frames only, and over the bands each frame holds.  A frozen, torn, or mis-gained frame stands out here whether or
        not anything flagged it -- this is what replaces the detector scores of the
        previous design.  float32 throughout: the deviations are ~0.01-0.1 in asinh
        units, below bf16's comfort.
        """
        with torch.no_grad():
            grouped = work.detach().float().view(batch, stack, BANDS, *work.shape[-2:])
            if median is None:
                median = self._stack_median(work, valid, batch, stack, present)
            deviation = torch.sqrt(((grouped - median) ** 2).mean(dim=(-2, -1)))
            deviation = torch.nan_to_num(deviation, nan=0.0) * valid[..., None]
            held = (torch.ones_like(deviation) if present is None
                    else (present > 0.5).to(deviation.dtype))
            deviation = deviation * held
            condition = condition.clone()
            condition[..., AGREE_MEAN] = (deviation.sum(dim=-1)
                                 / held.sum(dim=-1).clamp(min=1.0)).clamp(0.0, 3.0)
            condition[..., AGREE_MAX] = deviation.amax(dim=-1).clamp(0.0, 3.0)
        return condition

    def _attended(self, features, batch, stack):
        """Cross-frame mixing on a coarse token grid, one token per (frame, cell)."""
        grid = F.adaptive_avg_pool2d(features, ATTN_GRID)
        channels = grid.shape[1]
        tokens = grid.view(batch, stack, channels, -1).permute(0, 3, 1, 2)
        tokens = tokens.reshape(batch * ATTN_GRID * ATTN_GRID, stack, channels)
        mixed, _ = self.attend(tokens, tokens, tokens, need_weights=False)
        mixed = mixed.reshape(batch, ATTN_GRID * ATTN_GRID, stack, channels)
        mixed = mixed.permute(0, 2, 3, 1).reshape(
            batch * stack, channels, ATTN_GRID, ATTN_GRID)
        return F.interpolate(mixed, size=features.shape[-2:], mode="bilinear",
                             align_corners=False)


def _displace(frames, flow):
    """Warp frames by a flow field given as a fraction of frame width.

    `flow` is (N, 2, H, W) -- x then y displacement, resolution-independent units.
    float32 throughout: grid coordinates near +/-1 lose whole pixels in bf16.
    Border padding so a constant frame stays constant, which keeps the bounded-
    combination invariant intact whatever the flow says.
    """
    return _warp(frames, _flow_grid(flow))


def _flow_grid(flow, inplace=False):
    """The (N, H, W, 2) sampling grid for a flow field, as `grid_sample` takes it.
    `inplace` builds it in `flow`'s own storage (float32 only, no autograd)."""
    _, _, h, w = flow.shape
    ys, xs = torch.meshgrid(
        torch.linspace(-1.0, 1.0, h, device=flow.device),
        torch.linspace(-1.0, 1.0, w, device=flow.device),
        indexing="ij",
    )
    base = torch.stack([xs, ys])                        # (2, H, W), grid units
    if inplace and flow.dtype == torch.float32:
        return flow.mul_(2.0).add_(base).permute(0, 2, 3, 1)
    return (base + 2.0 * flow.float()).permute(0, 2, 3, 1)


def _warp(frames, grid):
    """Sample frames on a grid from :func:`_flow_grid`; float32 throughout."""
    warped = F.grid_sample(frames.float(), grid, mode="bilinear",
                           padding_mode="border", align_corners=True)
    return warped.to(frames.dtype)


def _correlation(work, median, radius=CORR_RADIUS):
    """Local cost volume between each frame and the stack median, band-means.

    (N, (2r+1)^2, h, w): the classic correspondence primitive -- at each pixel,
    how well does this candidate match the consensus when shifted by (dy, dx)?
    Normalised by the frame's own RMS so bright frames do not dominate.
    """
    a = work.mean(dim=1, keepdim=True)
    b = median.mean(dim=1, keepdim=True)
    scale = a.square().mean(dim=(-2, -1), keepdim=True).sqrt() + 1e-6
    a, b = a / scale, b / scale
    padded = F.pad(b, (radius,) * 4)
    height, width = a.shape[-2:]
    volumes = [a * padded[..., dy:dy + height, dx:dx + width]
               for dy in range(2 * radius + 1) for dx in range(2 * radius + 1)]
    return torch.cat(volumes, dim=1)


def build(base=32, depth=3, relative=False, corr_flow=False, attention=False):
    return StackFiller(base=base, depth=depth, relative=relative,
                       corr_flow=corr_flow, attention=attention)