634 lines
33 KiB
Python
634 lines
33 KiB
Python
"""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)
|