"""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)