257 lines
12 KiB
Python
257 lines
12 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 20x20, upsampled and bounded to
|
||
|
|
``exp(+/-0.3)`` -- enough to fix a level, a radial profile or a flare's broad
|
||
|
|
brightening, and structurally incapable of synthesising an image.
|
||
|
|
|
||
|
|
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
|
||
|
|
#: Resolution of the polish head's log-gain field. Deliberately coarse: the measured
|
||
|
|
#: cross-satellite gain residual varies over hundreds of pixels, and a field this
|
||
|
|
#: smooth cannot carry image content.
|
||
|
|
POLISH_GRID = 20
|
||
|
|
#: Half-width of the polish in log-gain: gains live in [exp(-0.3), exp(0.3)], about
|
||
|
|
#: [0.74, 1.35]. Covers the measured drift and radial residual with margin, and is
|
||
|
|
#: the structural bound that keeps the head from becoming a free residual.
|
||
|
|
POLISH_RANGE = 0.3
|
||
|
|
|
||
|
|
#: Per-frame conditioning:
|
||
|
|
#: 3 state one-hot (available / missing / suspect)
|
||
|
|
#: 1 satellite is the target's
|
||
|
|
#: 2 dt encoding (sign, log1p|dt| scaled; ~1.0 at a 300-slot gap)
|
||
|
|
#: 2 agreement (RMS deviation from the stack median: band mean, band max)
|
||
|
|
#: The last two 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.
|
||
|
|
COND_DIM = 8
|
||
|
|
|
||
|
|
#: 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 two terms:
|
||
|
|
#: recency, and distrust of frames known to be damaged.
|
||
|
|
PRIOR_TEMPORAL_DECAY = 1.5
|
||
|
|
PRIOR_SUSPECT_PENALTY = 4.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": 0, "missing": 1, "suspect": 2}[state]] = 1.0
|
||
|
|
vector[3] = 1.0 if is_target_satellite else 0.0
|
||
|
|
dt = float(dt_seconds)
|
||
|
|
vector[4] = math.copysign(1.0, dt) if dt else 0.0
|
||
|
|
vector[5] = math.log1p(abs(dt) / 240.0) / 6.0
|
||
|
|
return vector
|
||
|
|
|
||
|
|
|
||
|
|
def blend_prior(condition):
|
||
|
|
"""Fixed blend logits from each frame's conditioning. (B, S) -> (B, S)."""
|
||
|
|
distance = condition[..., 5] * 6.0
|
||
|
|
suspect = condition[..., 2]
|
||
|
|
return -PRIOR_TEMPORAL_DECAY * distance - PRIOR_SUSPECT_PENALTY * suspect
|
||
|
|
|
||
|
|
|
||
|
|
# ------------------------------------------------------------------------ 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
|
||
|
|
|
||
|
|
|
||
|
|
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, cond_dim=COND_DIM):
|
||
|
|
super().__init__()
|
||
|
|
# +1 input channel: per-frame validity, so the encoder can tell a genuinely
|
||
|
|
# dark frame (an eclipse) from one that is merely absent.
|
||
|
|
self.encoder = Encoder(BANDS + 1, base, depth, cond_dim)
|
||
|
|
self.weight = nn.Conv2d(self.encoder.width, BANDS, 3, padding=1)
|
||
|
|
self.polish = nn.Conv2d(self.encoder.width, BANDS, 1)
|
||
|
|
for head in (self.weight, self.polish):
|
||
|
|
nn.init.zeros_(head.weight)
|
||
|
|
nn.init.zeros_(head.bias)
|
||
|
|
|
||
|
|
def forward(self, sources, condition):
|
||
|
|
batch, stack = sources.shape[:2]
|
||
|
|
size = sources.shape[-2:]
|
||
|
|
valid = (condition[..., 0] + condition[..., 2]).clamp(0, 1)
|
||
|
|
|
||
|
|
flat = sources.reshape(batch * stack, BANDS, *size)
|
||
|
|
work = flat
|
||
|
|
if size[-1] != WORK_SIZE:
|
||
|
|
work = F.interpolate(flat, size=(WORK_SIZE, WORK_SIZE),
|
||
|
|
mode="bilinear", align_corners=False)
|
||
|
|
|
||
|
|
condition = self._with_agreement(work, condition, valid, batch, stack)
|
||
|
|
mask = valid.reshape(batch * stack, 1, 1, 1).expand(-1, 1, *work.shape[-2:])
|
||
|
|
features = self.encoder(torch.cat([work, mask], dim=1),
|
||
|
|
condition.reshape(batch * stack, -1))
|
||
|
|
|
||
|
|
logits = F.interpolate(self.weight(features), size=size,
|
||
|
|
mode="bilinear", align_corners=False)
|
||
|
|
polish = F.interpolate(self.polish(F.adaptive_avg_pool2d(features, POLISH_GRID)),
|
||
|
|
size=size, mode="bilinear", align_corners=False)
|
||
|
|
gain = torch.exp(POLISH_RANGE * torch.tanh(polish))
|
||
|
|
|
||
|
|
logits = logits + blend_prior(condition).reshape(batch * stack, 1, 1, 1)
|
||
|
|
# A frame with no pixels must not win weight, however confident the head is.
|
||
|
|
logits = logits.masked_fill(
|
||
|
|
valid.reshape(batch * stack, 1, 1, 1) < 0.5, float("-inf")
|
||
|
|
)
|
||
|
|
weights = torch.softmax(logits.view(batch, stack, BANDS, *size), dim=1)
|
||
|
|
weights = torch.nan_to_num(weights, nan=0.0)
|
||
|
|
|
||
|
|
adjusted = (gain * flat).view(batch, stack, BANDS, *size)
|
||
|
|
return (adjusted * weights).sum(dim=1)
|
||
|
|
|
||
|
|
def _with_agreement(self, work, condition, valid, batch, stack):
|
||
|
|
"""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. A frozen, torn, or mis-gained frame stands out here whether or
|
||
|
|
not anything flagged it -- this is what replaces both cross-frame attention
|
||
|
|
and 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:])
|
||
|
|
hidden = torch.where(
|
||
|
|
valid.reshape(batch, stack, 1, 1, 1) > 0.5, grouped,
|
||
|
|
torch.full_like(grouped, float("nan")),
|
||
|
|
)
|
||
|
|
median = hidden.nanmedian(dim=1, keepdim=True).values
|
||
|
|
deviation = torch.sqrt(((grouped - median) ** 2).mean(dim=(-2, -1)))
|
||
|
|
deviation = torch.nan_to_num(deviation, nan=0.0) * valid[..., None]
|
||
|
|
condition = condition.clone()
|
||
|
|
condition[..., 6] = deviation.mean(dim=-1).clamp(0.0, 3.0)
|
||
|
|
condition[..., 7] = deviation.amax(dim=-1).clamp(0.0, 3.0)
|
||
|
|
return condition
|
||
|
|
|
||
|
|
|
||
|
|
def build(base=32, depth=3):
|
||
|
|
return StackFiller(base=base, depth=depth)
|