Refactor code structure and remove redundant code blocks for improved readability and maintainability. Added Sphinx API documentation
4672e76 | """Core FOXES models: released deterministic and coupled single-query Gaussian.""" | |
| from collections import deque | |
| import math | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import pytorch_lightning as pl | |
| from torch.optim.lr_scheduler import CosineAnnealingLR | |
| SXR_LOG_OFFSET = 1e-8 | |
| def validate_bias_penalty_weight(value): | |
| """Accept only finite, nonnegative bias penalty coefficients.""" | |
| if (isinstance(value, bool) or not isinstance(value, (int, float)) | |
| or not math.isfinite(value) or value < 0): | |
| raise ValueError('bias_penalty_weight must be a finite, nonnegative number') | |
| return float(value) | |
| def mean_class_signed_error_penalty(prediction, target, target_raw, weight): | |
| """Average squared signed bias across the GOES classes in a batch.""" | |
| residual = (prediction - target).reshape(-1) | |
| flux = target_raw.reshape(-1) | |
| thresholds = SXRRegressionDynamicLoss.CLASS_THRESHOLDS | |
| class_index = ( | |
| (flux >= thresholds['c']).long() | |
| + (flux >= thresholds['m']).long() | |
| + (flux >= thresholds['x']).long() | |
| ) | |
| sums = residual.new_zeros(4).scatter_add(0, class_index, residual) | |
| counts = residual.new_zeros(4).scatter_add( | |
| 0, class_index, torch.ones_like(residual), | |
| ) | |
| class_means = sums / counts.clamp_min(1) | |
| return weight * class_means.square().sum() / (counts > 0).sum() | |
| def normalize_sxr(unnormalized_values, sxr_norm): | |
| """Convert physical flux using the FOXES ``log10(F + 1e-8)`` convention.""" | |
| log_values = torch.log10(unnormalized_values + SXR_LOG_OFFSET) | |
| normalized = (log_values - float(sxr_norm[0].item())) / float(sxr_norm[1].item()) | |
| return normalized | |
| def unnormalize_sxr(normalized_values, sxr_norm): | |
| """Invert :func:`normalize_sxr` back to physical flux.""" | |
| return ( | |
| 10 ** ( | |
| normalized_values * float(sxr_norm[1].item()) | |
| + float(sxr_norm[0].item()) | |
| ) | |
| - SXR_LOG_OFFSET | |
| ) | |
| class ViTLocal(pl.LightningModule): | |
| """ | |
| Parameters | |
| ---------- | |
| model_kwargs : dict | |
| Forwarded to VisionTransformerLocal (embed_dim, num_heads, mask_mode, ...). | |
| sxr_norm : np.ndarray | |
| (mean, std) used to log-normalize SXR targets. | |
| base_weights : dict, optional | |
| Per-class loss weights (see SXRRegressionDynamicLoss). If None, falls | |
| back to SXRRegressionDynamicLoss's built-in defaults, which were fit to | |
| the originally released dataset — pass real weights (e.g. from | |
| training/train.py's get_four_class_macro_weights) when training on different data. | |
| weight_decay : float | |
| AdamW weight decay. | |
| scheduler_kwargs : dict, optional | |
| Passed to CosineAnnealingLR (T_max, eta_min). ``T_max`` is measured in | |
| epochs because the scheduler is stepped once per epoch. | |
| loss_kwargs : dict, optional | |
| Forwarded to SXRRegressionDynamicLoss (window_size, huber_delta, | |
| adaptive_multipliers) — see that class for defaults. | |
| bias_penalty_weight : float | |
| Coefficient on the mean squared class-specific signed error in | |
| normalized log10 flux. Zero (the default) disables the penalty. | |
| diagnostic_every_n_steps : int | |
| How often (in training steps) to log the adaptive loss's per-class | |
| multipliers to the logger. | |
| """ | |
| DEFAULT_SCHEDULER_KWARGS = {'T_max': 250, 'eta_min': 1e-7} | |
| def __init__(self, model_kwargs, sxr_norm, base_weights=None, weight_decay=1e-5, | |
| scheduler_kwargs=None, loss_kwargs=None, diagnostic_every_n_steps=200, | |
| bias_penalty_weight=0.0): | |
| super().__init__() | |
| self.model_kwargs = model_kwargs | |
| self.lr = model_kwargs.get('learning_rate', model_kwargs.get('lr', 1e-4)) | |
| self.save_hyperparameters() | |
| filtered_kwargs = dict(model_kwargs) | |
| filtered_kwargs.pop('learning_rate', None) | |
| filtered_kwargs.pop('lr', None) | |
| filtered_kwargs.pop('num_classes', None) | |
| self.model = VisionTransformerLocal(**filtered_kwargs) | |
| self.base_weights = base_weights | |
| self.weight_decay = weight_decay | |
| self.scheduler_kwargs = {**self.DEFAULT_SCHEDULER_KWARGS, **(scheduler_kwargs or {})} | |
| self.diagnostic_every_n_steps = diagnostic_every_n_steps | |
| self.adaptive_loss = SXRRegressionDynamicLoss(base_weights=self.base_weights, **(loss_kwargs or {})) | |
| self.bias_penalty_weight = validate_bias_penalty_weight(bias_penalty_weight) | |
| self.huber_delta = self.adaptive_loss.huber_delta | |
| self.sxr_norm = sxr_norm | |
| def forward(self, x, return_attention=False): | |
| return self.model(x, self.sxr_norm, return_attention=return_attention) | |
| def configure_optimizers(self): | |
| # Use AdamW with weight decay for better regularization | |
| optimizer = torch.optim.AdamW( | |
| self.parameters(), | |
| lr=self.lr, | |
| weight_decay=self.weight_decay, | |
| ) | |
| scheduler = CosineAnnealingLR(optimizer, **self.scheduler_kwargs) | |
| return { | |
| 'optimizer': optimizer, | |
| 'lr_scheduler': { | |
| 'scheduler': scheduler, | |
| 'interval': 'epoch', | |
| 'frequency': 1, | |
| 'name': 'learning_rate' | |
| } | |
| } | |
| def _calculate_loss(self, batch, mode="train"): | |
| imgs, sxr = batch | |
| raw_preds, _ = self.model(imgs, self.sxr_norm) | |
| raw_preds_squeezed = torch.squeeze(raw_preds) | |
| if not torch.isfinite(raw_preds_squeezed).all(): | |
| raise RuntimeError( | |
| f"Non-finite model output in {mode} step at global_step={self.global_step}: " | |
| f"raw_preds range=[{raw_preds_squeezed.min().item()}, {raw_preds_squeezed.max().item()}], " | |
| f"imgs range=[{imgs.min().item()}, {imgs.max().item()}], " | |
| f"sxr (normalized target) range=[{sxr.min().item()}, {sxr.max().item()}]. " | |
| "This means the model itself produced NaN/Inf before reaching the loss -- " | |
| "an activation blowup, not a data or loss-weighting issue." | |
| ) | |
| sxr_un = unnormalize_sxr(sxr, self.sxr_norm) | |
| norm_preds_squeezed = normalize_sxr(raw_preds_squeezed, self.sxr_norm) | |
| # Use adaptive rare event loss. Only train batches should feed the | |
| # per-class error history the adaptive multipliers are computed from -- | |
| # otherwise val/test performance silently leaks into training dynamics. | |
| base_loss, _ = self.adaptive_loss.calculate_loss( | |
| norm_preds_squeezed, sxr, sxr_un, update_tracking=(mode == "train") | |
| ) | |
| signed_error = (norm_preds_squeezed - sxr).mean() | |
| bias_penalty = mean_class_signed_error_penalty( | |
| norm_preds_squeezed, sxr, sxr_un, self.bias_penalty_weight, | |
| ) | |
| loss = base_loss + bias_penalty | |
| # Also calculate huber loss for logging | |
| huber_loss = F.huber_loss(norm_preds_squeezed, sxr, delta=self.huber_delta) | |
| mse_loss = F.mse_loss(norm_preds_squeezed, sxr) | |
| mae_loss = F.l1_loss(norm_preds_squeezed, sxr) | |
| rmse_loss = torch.sqrt(mse_loss) | |
| # Log adaptation info | |
| if mode == "train": | |
| # Always log learning rate (every step) | |
| current_lr = self.trainer.optimizers[0].param_groups[0]['lr'] | |
| self.log('train/learning_rate', current_lr, on_step=True, on_epoch=False, | |
| prog_bar=True, logger=True, sync_dist=True) | |
| self.log("train/total_loss", loss, on_step=True, on_epoch=True, | |
| prog_bar=True, logger=True, sync_dist=True) | |
| self.log("train/bias_dex", signed_error * float(self.sxr_norm[1]), | |
| on_step=True, on_epoch=True, logger=True, sync_dist=True) | |
| self.log("train/bias_penalty", bias_penalty, on_step=True, | |
| on_epoch=True, logger=True, sync_dist=True) | |
| self.log("train/huber_loss", huber_loss, on_step=True, on_epoch=True, | |
| prog_bar=True, logger=True, sync_dist=True) | |
| self.log("train/mse_loss", mse_loss, on_step=True, on_epoch=True, | |
| prog_bar=True, logger=True, sync_dist=True) | |
| self.log("train/mae_loss", mae_loss, on_step=True, on_epoch=True, | |
| prog_bar=True, logger=True, sync_dist=True) | |
| self.log("train/rmse_loss", rmse_loss, on_step=True, on_epoch=True, | |
| prog_bar=True, logger=True, sync_dist=True) | |
| # Detailed diagnostics only every N steps | |
| if self.global_step % self.diagnostic_every_n_steps == 0: | |
| multipliers = self.adaptive_loss.get_current_multipliers() | |
| for key, value in multipliers.items(): | |
| self.log(f"adaptive/{key}", value, on_step=True, on_epoch=False, sync_dist=True) | |
| if mode == "val": | |
| # Validation: typically only log epoch aggregates | |
| multipliers = self.adaptive_loss.get_current_multipliers() | |
| for key, value in multipliers.items(): | |
| self.log(f"val/adaptive/{key}", value, on_step=False, on_epoch=True) | |
| self.log("val_total_loss", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True) | |
| self.log("val/bias_dex", signed_error * float(self.sxr_norm[1]), | |
| on_step=False, on_epoch=True, logger=True, sync_dist=True) | |
| self.log("val/bias_penalty", bias_penalty, on_step=False, | |
| on_epoch=True, logger=True, sync_dist=True) | |
| self.log("val_huber_loss", huber_loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, | |
| sync_dist=True) | |
| self.log("val_mse_loss", mse_loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True) | |
| self.log("val_mae_loss", mae_loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True) | |
| self.log("val_rmse_loss", rmse_loss, on_step=False, on_epoch=True, prog_bar=True, logger=True, | |
| sync_dist=True) | |
| return loss | |
| def training_step(self, batch, batch_idx): | |
| return self._calculate_loss(batch, mode="train") | |
| def validation_step(self, batch, batch_idx): | |
| self._calculate_loss(batch, mode="val") | |
| def test_step(self, batch, batch_idx): | |
| self._calculate_loss(batch, mode="test") | |
| class VisionTransformerLocal(nn.Module): | |
| def __init__( | |
| self, | |
| embed_dim, | |
| hidden_dim, | |
| num_channels, | |
| num_heads, | |
| num_layers, | |
| patch_size, | |
| num_patches, | |
| dropout, | |
| mask_mode='inverted', | |
| local_window=9, | |
| ): | |
| """Vision Transformer that outputs flux contributions per patch. | |
| Args: | |
| embed_dim: Dimensionality of the input feature vectors to the Transformer | |
| hidden_dim: Dimensionality of the hidden layer in the feed-forward networks | |
| within the Transformer | |
| num_channels: Number of channels of the input (3 for RGB) | |
| num_heads: Number of heads to use in the Multi-Head Attention block | |
| num_layers: Number of layers to use in the Transformer | |
| patch_size: Number of pixels that the patches have per dimension | |
| num_patches: Maximum number of patches an image can have | |
| dropout: Amount of dropout to apply in the feed-forward network and | |
| on the input encoding | |
| mask_mode: Self-attention masking. 'inverted' (default) reproduces the | |
| original released model exactly; 'local' is true local | |
| attention; 'none' is full global attention. See | |
| InvertedAttentionBlock. | |
| local_window: Side length (in patches) of the local neighbourhood used | |
| by the 'inverted' and 'local' masks. | |
| """ | |
| super().__init__() | |
| self.patch_size = patch_size | |
| # Layers/Networks | |
| self.input_layer = nn.Linear(num_channels * (patch_size ** 2), embed_dim) | |
| self.mask_mode = mask_mode | |
| self.local_window = local_window | |
| self.transformer_blocks = nn.ModuleList([ | |
| InvertedAttentionBlock(embed_dim, hidden_dim, num_heads, num_patches, | |
| dropout=dropout, local_window=local_window, mask_mode=mask_mode) | |
| for _ in range(num_layers) | |
| ]) | |
| self.mlp_head = nn.Sequential(nn.LayerNorm(embed_dim), nn.Linear(embed_dim, 1)) | |
| self.dropout = nn.Dropout(dropout) | |
| # Parameters/Embeddings - using 2D positional encoding for local attention | |
| # No CLS token needed for local attention architecture | |
| self.grid_h = int(math.sqrt(num_patches)) | |
| self.grid_w = int(math.sqrt(num_patches)) | |
| self.pos_embedding_2d = nn.Parameter(torch.randn(1, self.grid_h, self.grid_w, embed_dim)) | |
| def forward(self, x, sxr_norm, return_attention=False): | |
| # Preprocess input | |
| x = img_to_patch(x, self.patch_size) | |
| x = self.input_layer(x) | |
| # Add positional encoding (no CLS token for local attention) | |
| x = self._add_2d_positional_encoding(x) | |
| # Apply Transformer blocks | |
| x = self.dropout(x) | |
| x = x.transpose(0, 1) # [T, B, embed_dim] | |
| attention_weights = [] | |
| for block in self.transformer_blocks: | |
| if return_attention: | |
| x, attn_weights = block(x, return_attention=True) | |
| attention_weights.append(attn_weights) | |
| else: | |
| x = block(x) | |
| patch_embeddings = x.transpose(0, 1) # [B, num_patches, embed_dim] | |
| patch_logits = self.mlp_head(patch_embeddings).squeeze(-1) # normalized log predictions [B, num_patches] | |
| # --- Convert to raw SXR --- | |
| mean, std = sxr_norm # in log10 space | |
| patch_flux_raw = torch.clamp( | |
| 10 ** (patch_logits * std + mean) - SXR_LOG_OFFSET, | |
| min=1e-15, | |
| max=1 | |
| ) | |
| # Sum over patches for raw global flux | |
| global_flux_raw = patch_flux_raw.sum(dim=1, keepdim=True) | |
| # Ensure global flux is never zero (add small epsilon if needed) | |
| global_flux_raw = torch.clamp(global_flux_raw, min=1e-15) | |
| if return_attention: | |
| return global_flux_raw, attention_weights, patch_flux_raw | |
| else: | |
| return global_flux_raw, patch_flux_raw | |
| def _add_2d_positional_encoding(self, x): | |
| """Add learned 2D positional encoding to patch embeddings""" | |
| B, T, embed_dim = x.shape | |
| num_patches = T # All tokens are patches (no CLS token) | |
| # Reshape patches to 2D grid: [B, grid_h, grid_w, embed_dim] | |
| patch_embeddings = x.reshape(B, self.grid_h, self.grid_w, embed_dim) | |
| # Add learned 2D positional encoding | |
| # Broadcasting: [B, grid_h, grid_w, embed_dim] + [1, grid_h, grid_w, embed_dim] | |
| patch_embeddings = patch_embeddings + self.pos_embedding_2d | |
| # Reshape back to sequence format: [B, num_patches, embed_dim] | |
| x = patch_embeddings.reshape(B, num_patches, embed_dim) | |
| return x | |
| def set_mask_mode(self, mask_mode, local_window=None): | |
| """Override the attention mask in every block (e.g. to ablate a loaded | |
| checkpoint). Normal checkpoint loading keeps each block's saved mask.""" | |
| self.mask_mode = mask_mode | |
| if local_window is not None: | |
| self.local_window = local_window | |
| for block in self.transformer_blocks: | |
| block.set_mask_mode(mask_mode, local_window=local_window) | |
| class InvertedAttentionBlock(nn.Module): | |
| """Transformer block whose self-attention can be masked in three ways. | |
| mask_mode (passed through from the model config): | |
| 'inverted' - ORIGINAL FOXES behaviour and the default, so released | |
| checkpoints reproduce exactly. The boolean mask marks LOCAL | |
| pairs as True, and nn.MultiheadAttention treats True as | |
| "blocked", so nearby patches are blocked and every patch | |
| attends only to DISTANT patches. (This is the flipped | |
| local-attention syntax the model was actually trained with.) | |
| 'local' - Correct local attention: block everything OUTSIDE the local | |
| window, so each patch attends only to its neighbours. | |
| 'none' - No mask at all; full global attention. | |
| The mask is registered as a PERSISTENT buffer, so it travels inside the | |
| checkpoint. Loading restores the exact mask a model was trained with -- the | |
| original inverted mask, or any hand-edited mask from past experiments -- | |
| regardless of the mask_mode passed at construction. mask_mode/local_window | |
| only decide the mask for a *fresh* model; to deliberately change the mask of | |
| an already-loaded checkpoint (e.g. for an ablation), call set_mask_mode(). | |
| """ | |
| VALID_MASK_MODES = ('inverted', 'local', 'none') | |
| def __init__(self, embed_dim, hidden_dim, num_heads, num_patches, dropout=0.0, | |
| local_window=9, mask_mode='inverted'): | |
| super().__init__() | |
| if mask_mode not in self.VALID_MASK_MODES: | |
| raise ValueError('mask_mode must be one of %s, got %r' | |
| % (self.VALID_MASK_MODES, mask_mode)) | |
| self.embed_dim = embed_dim | |
| self.num_heads = num_heads | |
| self.local_window = local_window | |
| self.num_patches = num_patches | |
| self.mask_mode = mask_mode | |
| self.layer_norm_1 = nn.LayerNorm(embed_dim) | |
| self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=False) | |
| self.layer_norm_2 = nn.LayerNorm(embed_dim) | |
| self.linear = nn.Sequential( | |
| nn.Linear(embed_dim, hidden_dim), | |
| nn.GELU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(hidden_dim, embed_dim), | |
| nn.Dropout(dropout), | |
| ) | |
| # Persistent: the mask is saved with the weights so every checkpoint | |
| # reproduces exactly what it trained with. None => no masking ('none'). | |
| self.register_buffer('attention_mask', self._build_attention_mask()) | |
| def _build_attention_mask(self): | |
| """Boolean attn_mask for nn.MultiheadAttention (True == position blocked). | |
| """ | |
| if self.mask_mode == 'none': | |
| return None | |
| grid_size = int(math.sqrt(self.num_patches)) | |
| idx = torch.arange(self.num_patches) | |
| rows = (idx // grid_size).to(torch.int16) | |
| cols = (idx % grid_size).to(torch.int16) | |
| r = self.local_window // 2 | |
| # local[i, j] True when patches i and j are within the local window. | |
| local = (((rows[:, None] - rows[None, :]).abs() <= r) & | |
| ((cols[:, None] - cols[None, :]).abs() <= r)) | |
| # 'inverted' blocks the local neighbourhood (original); 'local' blocks its | |
| # complement so attention stays inside the neighbourhood. | |
| return local if self.mask_mode == 'inverted' else ~local | |
| def set_mask_mode(self, mask_mode, local_window=None): | |
| """Rebuild the attention mask in place, overriding whatever is currently | |
| set (including a mask restored from a checkpoint). | |
| Use this to deliberately change a trained model's mask -- e.g. to ablate a | |
| released checkpoint under 'none' or 'local'. Loading a checkpoint normally | |
| keeps its own saved mask; this is the explicit opt-out. | |
| """ | |
| if mask_mode not in self.VALID_MASK_MODES: | |
| raise ValueError('mask_mode must be one of %s, got %r' | |
| % (self.VALID_MASK_MODES, mask_mode)) | |
| self.mask_mode = mask_mode | |
| if local_window is not None: | |
| self.local_window = local_window | |
| mask = self._build_attention_mask() | |
| if mask is not None: | |
| try: | |
| mask = mask.to(self.attention_mask.device) | |
| except AttributeError: # current buffer is None ('none' mode) | |
| mask = mask.to(next(self.parameters()).device) | |
| self.attention_mask = mask | |
| def forward(self, x, return_attention=False): | |
| inp_x = self.layer_norm_1(x) | |
| if return_attention: | |
| # attention_mask is None for mask_mode='none' -> standard full attention. | |
| attn_output, attn_weights = self.attn( | |
| inp_x, inp_x, inp_x, | |
| attn_mask=self.attention_mask, | |
| average_attn_weights=False | |
| ) | |
| x = x + attn_output | |
| x = x + self.linear(self.layer_norm_2(x)) | |
| return x, attn_weights | |
| else: | |
| # Avoid computing and returning the O(T^2) attention weights on | |
| # the normal training/inference path. Attention-map callers use | |
| # the explicit return_attention=True branch above. | |
| attn_output = self.attn( | |
| inp_x, inp_x, inp_x, | |
| attn_mask=self.attention_mask, | |
| need_weights=False, | |
| )[0] | |
| x = x + attn_output | |
| x = x + self.linear(self.layer_norm_2(x)) | |
| return x | |
| def img_to_patch(x, patch_size, flatten_channels=True): | |
| """ | |
| Args: | |
| x: Tensor representing the image of shape [B, H, W, C] | |
| patch_size: Number of pixels per dimension of the patches (integer) | |
| flatten_channels: If True, the patches will be returned in a flattened format | |
| as a feature vector instead of a image grid. | |
| """ | |
| x = x.permute(0, 3, 1, 2) | |
| B, C, H, W = x.shape | |
| x = x.reshape(B, C, H // patch_size, patch_size, W // patch_size, patch_size) | |
| x = x.permute(0, 2, 4, 1, 3, 5) # [B, H', W', C, p_H, p_W] | |
| x = x.flatten(1, 2) # [B, H'*W', C, p_H, p_W] | |
| if flatten_channels: | |
| x = x.flatten(2, 4) # [B, H'*W', C*p_H*p_W] | |
| return x | |
| class SXRRegressionDynamicLoss: | |
| """ | |
| Huber loss with per-class weights (quiet/C/M/X) that adapt based on each | |
| class's recent vs. overall running error — classes that are currently | |
| performing worse than their own history get up-weighted. | |
| Parameters | |
| ---------- | |
| window_size : int | |
| Max length of each class's rolling error history. | |
| base_weights : dict, optional | |
| Static per-class weight, multiplied by the adaptive multiplier below. | |
| If None, falls back to DEFAULT_BASE_WEIGHTS (fit to the originally | |
| released dataset — pass real weights fit to your own data instead, | |
| e.g. via training/train.py's get_four_class_macro_weights). | |
| adaptive_multipliers : dict, optional | |
| Per-class {max_multiplier, min_multiplier, sensitivity, min_samples, | |
| recent_window} controlling how much recent performance can move a | |
| class's weight. Merged over DEFAULT_ADAPTIVE_MULTIPLIERS. | |
| huber_delta : float | |
| Delta for the underlying Huber loss. | |
| """ | |
| # GOES flare-class flux boundaries (W/m^2) — a fixed physical definition, | |
| # not a training knob, so this is intentionally not a constructor param. | |
| CLASS_THRESHOLDS = {'c': 1e-6, 'm': 1e-5, 'x': 1e-4} | |
| DEFAULT_BASE_WEIGHTS = { | |
| 'quiet': 6.643528005464481, | |
| 'c_class': 1.626986450317832, | |
| 'm_class': 4.724573441010383, | |
| 'x_class': 43.13137472283814, | |
| } | |
| DEFAULT_ADAPTIVE_MULTIPLIERS = { | |
| 'quiet': {'max_multiplier': 1.5, 'min_multiplier': 0.6, 'sensitivity': 0.05, 'min_samples': 2500, 'recent_window': 800}, | |
| 'c_class': {'max_multiplier': 2.0, 'min_multiplier': 0.7, 'sensitivity': 0.08, 'min_samples': 2500, 'recent_window': 800}, | |
| 'm_class': {'max_multiplier': 5.0, 'min_multiplier': 0.8, 'sensitivity': 0.1, 'min_samples': 1500, 'recent_window': 500}, | |
| 'x_class': {'max_multiplier': 8.0, 'min_multiplier': 0.8, 'sensitivity': 0.12, 'min_samples': 1000, 'recent_window': 300}, | |
| } | |
| def __init__(self, window_size=15000, base_weights=None, adaptive_multipliers=None, huber_delta=0.3): | |
| self.c_threshold = self.CLASS_THRESHOLDS['c'] | |
| self.m_threshold = self.CLASS_THRESHOLDS['m'] | |
| self.x_threshold = self.CLASS_THRESHOLDS['x'] | |
| self.huber_delta = huber_delta | |
| self.window_size = window_size | |
| self.quiet_errors = deque(maxlen=window_size) | |
| self.c_errors = deque(maxlen=window_size) | |
| self.m_errors = deque(maxlen=window_size) | |
| self.x_errors = deque(maxlen=window_size) | |
| self.base_weights = base_weights if base_weights is not None else dict(self.DEFAULT_BASE_WEIGHTS) | |
| self.multiplier_params = {**self.DEFAULT_ADAPTIVE_MULTIPLIERS, **(adaptive_multipliers or {})} | |
| def calculate_loss(self, preds_norm, sxr_norm, sxr_un, update_tracking=True): | |
| base_loss = F.huber_loss(preds_norm, sxr_norm, delta=self.huber_delta, reduction='none') | |
| weights = self._get_adaptive_weights(sxr_un) | |
| if update_tracking: | |
| self._update_tracking(sxr_un, sxr_norm, preds_norm) | |
| weighted_loss = base_loss * weights | |
| loss = weighted_loss.mean() | |
| return loss, weights | |
| def _class_weight(self, sxrclass, error_history): | |
| mult = self._get_performance_multiplier(error_history, sxrclass) | |
| return self.base_weights[sxrclass] * mult | |
| def _get_adaptive_weights(self, sxr_un): | |
| device = sxr_un.device | |
| quiet_weight = self._class_weight('quiet', self.quiet_errors) | |
| c_weight = self._class_weight('c_class', self.c_errors) | |
| m_weight = self._class_weight('m_class', self.m_errors) | |
| x_weight = self._class_weight('x_class', self.x_errors) | |
| weights = torch.ones_like(sxr_un, device=device) | |
| weights = torch.where(sxr_un < self.c_threshold, quiet_weight, weights) | |
| weights = torch.where((sxr_un >= self.c_threshold) & (sxr_un < self.m_threshold), c_weight, weights) | |
| weights = torch.where((sxr_un >= self.m_threshold) & (sxr_un < self.x_threshold), m_weight, weights) | |
| weights = torch.where(sxr_un >= self.x_threshold, x_weight, weights) | |
| # Normalize so mean weight ~1.0 (optional, helps stability) | |
| mean_weight = torch.mean(weights) | |
| weights = weights / (mean_weight) | |
| return weights | |
| def _get_performance_multiplier(self, error_history, sxrclass): | |
| """Class-dependent performance multiplier: how much recent error deviates | |
| from this class's overall running error, mapped through an exponential | |
| and clipped to [min_multiplier, max_multiplier].""" | |
| params = self.multiplier_params[sxrclass] | |
| if len(error_history) < params['min_samples']: | |
| return 1.0 | |
| recent = np.mean(list(error_history)[-params['recent_window']:]) | |
| overall = np.mean(list(error_history)) | |
| ratio = recent / overall | |
| multiplier = np.exp(params['sensitivity'] * (ratio - 1)) | |
| return np.clip(multiplier, params['min_multiplier'], params['max_multiplier']) | |
| def _update_tracking(self, sxr_un, sxr_norm, preds_norm): | |
| sxr_un_np = sxr_un.detach().cpu().numpy() | |
| error = F.huber_loss(preds_norm, sxr_norm, delta=self.huber_delta, reduction='none') | |
| error = error.detach().cpu().numpy() | |
| quiet_mask = sxr_un_np < self.c_threshold | |
| if quiet_mask.sum() > 0: | |
| self.quiet_errors.append(float(np.mean(error[quiet_mask]))) | |
| c_mask = (sxr_un_np >= self.c_threshold) & (sxr_un_np < self.m_threshold) | |
| if c_mask.sum() > 0: | |
| self.c_errors.append(float(np.mean(error[c_mask]))) | |
| m_mask = (sxr_un_np >= self.m_threshold) & (sxr_un_np < self.x_threshold) | |
| if m_mask.sum() > 0: | |
| self.m_errors.append(float(np.mean(error[m_mask]))) | |
| x_mask = sxr_un_np >= self.x_threshold | |
| if x_mask.sum() > 0: | |
| self.x_errors.append(float(np.mean(error[x_mask]))) | |
| def get_current_multipliers(self): | |
| """Current per-class multipliers/weights for logging — uses the exact | |
| same _get_performance_multiplier params as the loss itself, so what's | |
| logged always matches what was actually applied.""" | |
| quiet_mult = self._get_performance_multiplier(self.quiet_errors, 'quiet') | |
| c_mult = self._get_performance_multiplier(self.c_errors, 'c_class') | |
| m_mult = self._get_performance_multiplier(self.m_errors, 'm_class') | |
| x_mult = self._get_performance_multiplier(self.x_errors, 'x_class') | |
| return { | |
| 'quiet_mult': quiet_mult, | |
| 'c_mult': c_mult, | |
| 'm_mult': m_mult, | |
| 'x_mult': x_mult, | |
| 'quiet_count': len(self.quiet_errors), | |
| 'c_count': len(self.c_errors), | |
| 'm_count': len(self.m_errors), | |
| 'x_count': len(self.x_errors), | |
| 'quiet_error': np.mean(self.quiet_errors) if self.quiet_errors else 0.0, | |
| 'c_error': np.mean(self.c_errors) if self.c_errors else 0.0, | |
| 'm_error': np.mean(self.m_errors) if self.m_errors else 0.0, | |
| 'x_error': np.mean(self.x_errors) if self.x_errors else 0.0, | |
| 'quiet_weight': self.base_weights['quiet'] * quiet_mult, | |
| 'c_weight': self.base_weights['c_class'] * c_mult, | |
| 'm_weight': self.base_weights['m_class'] * m_mult, | |
| 'x_weight': self.base_weights['x_class'] * x_mult, | |
| } | |
| # Coupled single-query uncertainty model. | |
| class CoupledVisionTransformerLocal(VisionTransformerLocal): | |
| """Predict additive patch flux with local attention and one global query. | |
| Local logits scale a positive image-wide baseline spread over the patches. | |
| A zero-initialized global head adjusts each patch using raw AIA tokens and | |
| shared context. The uncertainty head consumes detached local logits. | |
| """ | |
| def __init__(self, *args, uncertainty_enabled=True, | |
| max_abs_log10_patch_multiplier=8.0, | |
| max_abs_log10_global_adjustment=0.3, | |
| global_attention_heads=None, | |
| global_head_hidden_dim=None, | |
| global_attention_dropout=None, | |
| initial_relative_std=1.0, | |
| min_log10_relative_std=-4.0, | |
| max_log10_relative_std=4.0, | |
| uncertainty_hidden_dim=16, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| self.uncertainty_enabled = bool(uncertainty_enabled) | |
| self.max_abs_log10_patch_multiplier = float( | |
| max_abs_log10_patch_multiplier | |
| ) | |
| self.max_abs_log10_global_adjustment = float( | |
| max_abs_log10_global_adjustment | |
| ) | |
| self.initial_relative_std = float(initial_relative_std) | |
| self.min_log10_relative_std = float(min_log10_relative_std) | |
| self.max_log10_relative_std = float(max_log10_relative_std) | |
| self.uncertainty_hidden_dim = int(uncertainty_hidden_dim) | |
| if self.max_abs_log10_patch_multiplier <= 0: | |
| raise ValueError("max_abs_log10_patch_multiplier must be positive") | |
| if self.max_abs_log10_global_adjustment <= 0: | |
| raise ValueError("max_abs_log10_global_adjustment must be positive") | |
| if self.mask_mode != 'local': | |
| raise ValueError( | |
| "The coupled model requires local patch attention so the " | |
| "spatial head cannot import distant activity" | |
| ) | |
| if self.initial_relative_std <= 0: | |
| raise ValueError("initial_relative_std must be positive") | |
| if self.max_log10_relative_std <= self.min_log10_relative_std: | |
| raise ValueError( | |
| "max_log10_relative_std must exceed " | |
| "min_log10_relative_std" | |
| ) | |
| if self.uncertainty_hidden_dim <= 0: | |
| raise ValueError("uncertainty_hidden_dim must be positive") | |
| self.last_log10_global_patch_adjustment = None | |
| embed_dim = self.mlp_head[-1].in_features | |
| local_heads = self.transformer_blocks[0].num_heads | |
| self.global_attention_heads = ( | |
| local_heads if global_attention_heads is None | |
| else global_attention_heads | |
| ) | |
| self.global_head_hidden_dim = ( | |
| embed_dim // 2 if global_head_hidden_dim is None | |
| else global_head_hidden_dim | |
| ) | |
| self.global_attention_dropout = ( | |
| self.dropout.p if global_attention_dropout is None | |
| else global_attention_dropout | |
| ) | |
| if (isinstance(self.global_attention_heads, bool) | |
| or not isinstance(self.global_attention_heads, int) | |
| or self.global_attention_heads <= 0 | |
| or embed_dim % self.global_attention_heads != 0): | |
| raise ValueError( | |
| 'global_attention_heads must be a positive divisor of embed_dim' | |
| ) | |
| if (isinstance(self.global_head_hidden_dim, bool) | |
| or not isinstance(self.global_head_hidden_dim, int) | |
| or self.global_head_hidden_dim <= 0): | |
| raise ValueError('global_head_hidden_dim must be a positive integer') | |
| if (isinstance(self.global_attention_dropout, bool) | |
| or not isinstance(self.global_attention_dropout, (int, float)) | |
| or not 0 <= self.global_attention_dropout < 1): | |
| raise ValueError('global_attention_dropout must be in [0, 1)') | |
| # A summary query sees every original AIA patch. The per-patch | |
| # correction also receives that patch's raw AIA token and its | |
| # locally encoded feature, so remote activity cannot directly | |
| # replace the local evidence at a quiet patch. | |
| self.global_raw_patch_projection = nn.Linear( | |
| self.input_layer.in_features, embed_dim | |
| ) | |
| self.global_query = nn.Parameter(torch.zeros(1, 1, embed_dim)) | |
| self.global_attention = nn.MultiheadAttention( | |
| embed_dim, | |
| self.global_attention_heads, | |
| dropout=self.global_attention_dropout, | |
| batch_first=True, | |
| ) | |
| self.global_patch_head = nn.Sequential( | |
| nn.LayerNorm(3 * embed_dim), | |
| nn.Linear(3 * embed_dim, self.global_head_hidden_dim), | |
| nn.GELU(), | |
| nn.Linear(self.global_head_hidden_dim, 1), | |
| ) | |
| nn.init.zeros_(self.global_patch_head[-1].weight) | |
| nn.init.zeros_(self.global_patch_head[-1].bias) | |
| self.patch_uncertainty_head = None | |
| self.initial_log10_relative_std = None | |
| if self.uncertainty_enabled: | |
| self.initial_log10_relative_std = float( | |
| torch.log10(torch.tensor(self.initial_relative_std)).item() | |
| ) | |
| if not ( | |
| self.min_log10_relative_std | |
| < self.initial_log10_relative_std | |
| < self.max_log10_relative_std | |
| ): | |
| raise ValueError( | |
| "initial log10 relative standard deviation falls outside " | |
| "the configured numerical bounds" | |
| ) | |
| self.patch_uncertainty_head = nn.Sequential( | |
| nn.Linear(1, self.uncertainty_hidden_dim), | |
| nn.SiLU(), | |
| nn.Linear(self.uncertainty_hidden_dim, 1), | |
| ) | |
| nn.init.zeros_(self.patch_uncertainty_head[-1].weight) | |
| nn.init.constant_( | |
| self.patch_uncertainty_head[-1].bias, | |
| self.initial_log10_relative_std, | |
| ) | |
| def _patch_flux_from_logits(self, patch_logits, sxr_norm): | |
| """Convert patch outputs to nonnegative additive raw flux.""" | |
| typical_global_flux = torch.clamp( | |
| 10 ** sxr_norm[0] - SXR_LOG_OFFSET, | |
| min=torch.finfo(torch.float32).tiny, | |
| ) | |
| # Preserve the original normalized-log head scale and default | |
| # initialization, but assign the whole-image training scale across | |
| # patches before applying each patch's multiplicative adjustment. | |
| log10_adjustment = torch.clamp( | |
| patch_logits.float(), | |
| min=-self.max_abs_log10_patch_multiplier, | |
| max=self.max_abs_log10_patch_multiplier, | |
| ) | |
| patch_multiplier = torch.exp( | |
| log10_adjustment * torch.log(torch.tensor( | |
| 10.0, | |
| device=patch_logits.device, | |
| dtype=torch.float32, | |
| )) | |
| ) | |
| return ( | |
| typical_global_flux.float() | |
| / patch_logits.shape[1] | |
| * patch_multiplier | |
| ) | |
| def _encode_patch_embeddings(self, x, return_attention): | |
| """Encode image patches while preserving the public attention output.""" | |
| x = img_to_patch(x, self.patch_size) | |
| x = self.input_layer(x) | |
| x = self._add_2d_positional_encoding(x) | |
| x = self.dropout(x).transpose(0, 1) | |
| attention_weights = [] | |
| for block in self.transformer_blocks: | |
| if return_attention: | |
| x, weights = block(x, return_attention=True) | |
| attention_weights.append(weights) | |
| else: | |
| x = block(x) | |
| return x.transpose(0, 1), attention_weights | |
| def _patch_mean_from_embeddings(self, patch_embeddings, sxr_norm, images): | |
| """Return logits and nonnegative additive patch means.""" | |
| patch_logits = self.mlp_head(patch_embeddings).squeeze(-1) | |
| patch_flux_raw = self._patch_flux_from_logits(patch_logits, sxr_norm) | |
| raw_tokens = self.global_raw_patch_projection( | |
| img_to_patch(images, self.patch_size) | |
| ) | |
| raw_tokens = raw_tokens + self.pos_embedding_2d.reshape( | |
| 1, -1, raw_tokens.shape[-1] | |
| ) | |
| query = self.global_query.expand(patch_embeddings.shape[0], -1, -1) | |
| context, _ = self.global_attention( | |
| query, raw_tokens, raw_tokens, need_weights=False | |
| ) | |
| context_per_patch = context.expand(-1, patch_embeddings.shape[1], -1) | |
| correction_input = torch.cat( | |
| (patch_embeddings, raw_tokens, context_per_patch), dim=-1 | |
| ) | |
| log10_adjustment = ( | |
| self.max_abs_log10_global_adjustment | |
| * torch.tanh(self.global_patch_head(correction_input).squeeze(-1)) | |
| ) | |
| self.last_log10_global_patch_adjustment = log10_adjustment.detach() | |
| patch_flux_raw = patch_flux_raw * torch.pow( | |
| torch.tensor(10.0, device=patch_flux_raw.device), | |
| log10_adjustment.float(), | |
| ) | |
| return patch_logits, patch_flux_raw | |
| def _patch_variance_from_logits(self, patch_logits, patch_flux_raw): | |
| """Predict raw variance contributions from detached patch logits.""" | |
| if not self.uncertainty_enabled: | |
| raise RuntimeError("uncertainty is disabled for this model") | |
| detached_patch_logits = patch_logits.detach().float() | |
| detached_patch_flux = patch_flux_raw.detach().float() | |
| patch_log10_relative_std = self.patch_uncertainty_head( | |
| detached_patch_logits.unsqueeze(-1) | |
| ).squeeze(-1).float() | |
| patch_log10_relative_std = torch.clamp( | |
| patch_log10_relative_std, | |
| min=self.min_log10_relative_std, | |
| max=self.max_log10_relative_std, | |
| ) | |
| relative_std = torch.pow( | |
| torch.tensor( | |
| 10.0, | |
| device=patch_log10_relative_std.device, | |
| dtype=torch.float32, | |
| ), | |
| patch_log10_relative_std, | |
| ) | |
| patch_std_raw = relative_std * detached_patch_flux | |
| return patch_std_raw.square() | |
| def forward(self, x, sxr_norm, return_attention=False): | |
| """Return corrected predictions and additive patch contributions.""" | |
| patch_embeddings, attention_weights = self._encode_patch_embeddings( | |
| x, return_attention | |
| ) | |
| _, sxr_std = sxr_norm | |
| patch_logits, patch_flux_raw = self._patch_mean_from_embeddings( | |
| patch_embeddings, sxr_norm, images=x, | |
| ) | |
| global_flux_raw = patch_flux_raw.sum(dim=1, keepdim=True) | |
| mean_normalized = normalize_sxr(global_flux_raw.squeeze(-1), sxr_norm) | |
| if not self.uncertainty_enabled: | |
| if return_attention: | |
| outputs = ( | |
| global_flux_raw, mean_normalized, | |
| attention_weights, patch_flux_raw, | |
| ) | |
| else: | |
| outputs = global_flux_raw, mean_normalized, patch_flux_raw | |
| return outputs | |
| patch_variance_raw = self._patch_variance_from_logits( | |
| patch_logits, patch_flux_raw | |
| ) | |
| global_variance_raw = patch_variance_raw.sum(dim=1) | |
| # Delta-method conversion for | |
| # y=(log10(raw_flux + offset)-mean)/std. | |
| # The delta-method scale depends numerically on the predicted global | |
| # flux, but its variance gradient must not flow into the mean branch. | |
| # Without this detach, NLL can move an otherwise-correct mean merely to | |
| # shrink or enlarge its normalized variance. | |
| global_flux_for_variance = global_flux_raw.squeeze(-1).float().detach() | |
| log_derivative = 1.0 / ( | |
| torch.log(torch.tensor( | |
| 10.0, device=global_flux_raw.device, dtype=torch.float32 | |
| )) | |
| * sxr_std.float() | |
| * (global_flux_for_variance + SXR_LOG_OFFSET) | |
| ) | |
| variance_normalized = ( | |
| global_variance_raw * log_derivative.square() | |
| ) | |
| if return_attention: | |
| outputs = ( | |
| global_flux_raw, mean_normalized, variance_normalized, | |
| attention_weights, patch_flux_raw, patch_variance_raw, | |
| ) | |
| else: | |
| outputs = ( | |
| global_flux_raw, mean_normalized, variance_normalized, | |
| patch_flux_raw, patch_variance_raw, | |
| ) | |
| return outputs | |
| class CoupledGaussianViTLocal(pl.LightningModule): | |
| """Train the canonical FOXES patch model with Gaussian uncertainty. | |
| The mean/backbone is optimized only by an optionally class-weighted Huber | |
| loss. A separate Gaussian NLL trains relative patch standard deviations | |
| from detached patch logits, so it cannot alter the spatial mean | |
| prediction. | |
| The final corrected-flux Huber loss updates both local and global layers. | |
| The predicted variance is in normalized log10-SXR units. Use | |
| :meth:`predict_interval` to obtain intervals in physical W/m^2 units. | |
| """ | |
| DEFAULT_SCHEDULER_KWARGS = {'T_max': 250, 'eta_min': 1e-7} | |
| predicts_uncertainty = True | |
| uncertainty_kind = 'gaussian' | |
| uncertainty_parameterization = 'logit_conditioned_relative_log10_patch_std' | |
| NETWORK_UNCERTAINTY_KEYS = ( | |
| 'max_abs_log10_patch_multiplier', | |
| 'initial_relative_std', | |
| 'min_log10_relative_std', | |
| 'max_log10_relative_std', | |
| 'uncertainty_hidden_dim', | |
| ) | |
| def __init__(self, model_kwargs, sxr_norm, | |
| weight_decay=1e-5, scheduler_kwargs=None, | |
| uncertainty_kwargs=None): | |
| super().__init__() | |
| self.lr = model_kwargs.get('learning_rate', model_kwargs.get('lr', 1e-4)) | |
| self.save_hyperparameters() | |
| filtered_kwargs = dict(model_kwargs) | |
| filtered_kwargs.pop('learning_rate', None) | |
| filtered_kwargs.pop('lr', None) | |
| filtered_kwargs.pop('num_classes', None) | |
| # This fixed architecture flag is present in the distributed checkpoint. | |
| if filtered_kwargs.pop('global_patch_adjustment', True) is not True: | |
| raise ValueError('The coupled model always applies global patch adjustment') | |
| uncertainty_kwargs = dict(uncertainty_kwargs or {}) | |
| self.uncertainty_enabled = bool( | |
| uncertainty_kwargs.pop('enabled', True) | |
| ) | |
| configured_start_epoch = uncertainty_kwargs.pop('start_epoch', None) | |
| if configured_start_epoch is None: | |
| self.uncertainty_start_epoch = 0 | |
| elif ( | |
| isinstance(configured_start_epoch, bool) | |
| or not isinstance(configured_start_epoch, int) | |
| or configured_start_epoch < 0 | |
| ): | |
| raise ValueError( | |
| "uncertainty.start_epoch must be a nonnegative integer or null" | |
| ) | |
| else: | |
| self.uncertainty_start_epoch = configured_start_epoch | |
| self.predicts_uncertainty = self.uncertainty_enabled | |
| allowed_uncertainty_keys = set(self.NETWORK_UNCERTAINTY_KEYS) | { | |
| 'huber_delta', | |
| 'class_weighting', | |
| 'class_weights', | |
| 'bias_penalty_weight', | |
| } | |
| unknown_keys = set(uncertainty_kwargs) - allowed_uncertainty_keys | |
| if unknown_keys: | |
| raise ValueError( | |
| f"Unknown uncertainty settings: {sorted(unknown_keys)}" | |
| ) | |
| network_uncertainty_kwargs = { | |
| key: uncertainty_kwargs[key] | |
| for key in self.NETWORK_UNCERTAINTY_KEYS | |
| if key in uncertainty_kwargs | |
| } | |
| self.bias_penalty_weight = validate_bias_penalty_weight( | |
| uncertainty_kwargs.get('bias_penalty_weight', 0.0) | |
| ) | |
| self.model = CoupledVisionTransformerLocal( | |
| **filtered_kwargs, | |
| **network_uncertainty_kwargs, | |
| uncertainty_enabled=self.uncertainty_enabled, | |
| ) | |
| self.hparams['uncertainty_parameterization'] = ( | |
| self.uncertainty_parameterization | |
| ) | |
| self.hparams['initial_log10_relative_std'] = ( | |
| self.model.initial_log10_relative_std | |
| ) | |
| self.register_buffer( | |
| 'sxr_norm', torch.as_tensor(sxr_norm, dtype=torch.float32) | |
| ) | |
| self.weight_decay = weight_decay | |
| self.scheduler_kwargs = { | |
| **self.DEFAULT_SCHEDULER_KWARGS, **(scheduler_kwargs or {}) | |
| } | |
| self.huber_delta = float(uncertainty_kwargs.get('huber_delta', 1.0)) | |
| if self.huber_delta <= 0: | |
| raise ValueError("uncertainty.huber_delta must be positive") | |
| self.class_weighting = str(uncertainty_kwargs.get('class_weighting', 'none')) | |
| if self.class_weighting not in { | |
| 'none', 'inverse_frequency', 'sqrt_inverse_frequency', | |
| 'power_inverse_frequency', | |
| }: | |
| raise ValueError( | |
| "uncertainty.class_weighting must be 'none', " | |
| "'inverse_frequency', 'sqrt_inverse_frequency', " | |
| "or 'power_inverse_frequency'" | |
| ) | |
| class_names = ('quiet', 'c_class', 'm_class', 'x_class') | |
| configured_weights = uncertainty_kwargs.get('class_weights') | |
| if self.class_weighting != 'none': | |
| if not isinstance(configured_weights, dict): | |
| raise ValueError( | |
| "class-weighted mean loss requires training-derived " | |
| "uncertainty.class_weights" | |
| ) | |
| if set(configured_weights) != set(class_names): | |
| raise ValueError( | |
| "uncertainty.class_weights must contain exactly " | |
| f"{list(class_names)}" | |
| ) | |
| weights = torch.tensor( | |
| [configured_weights[name] for name in class_names], | |
| dtype=torch.float32, | |
| ) | |
| if not torch.isfinite(weights).all() or (weights <= 0).any(): | |
| raise ValueError( | |
| "uncertainty.class_weights must be finite and positive" | |
| ) | |
| else: | |
| weights = torch.ones(len(class_names), dtype=torch.float32) | |
| # These weights are derived deterministically from the training split. | |
| self.register_buffer( | |
| 'class_weight_values', weights, persistent=False, | |
| ) | |
| def on_load_checkpoint(self, checkpoint): | |
| checkpoint_hparams = checkpoint.get('hyper_parameters', {}) | |
| checkpoint_uncertainty_parameterization = checkpoint_hparams.get( | |
| 'uncertainty_parameterization' | |
| ) | |
| state_dict = checkpoint.get('state_dict', {}) | |
| has_uncertainty_head = any( | |
| 'patch_uncertainty_head' in key for key in state_dict | |
| ) | |
| if ( | |
| has_uncertainty_head | |
| and checkpoint_uncertainty_parameterization | |
| != self.uncertainty_parameterization | |
| ): | |
| raise RuntimeError( | |
| "Checkpoint uncertainty parameterization is missing or " | |
| "incompatible with the logit-conditioned relative-standard-" | |
| "deviation model." | |
| ) | |
| def forward(self, x, return_attention=False): | |
| """Return Gaussian outputs, or deterministic outputs when disabled.""" | |
| outputs = self.model(x, self.sxr_norm, return_attention=return_attention) | |
| if not self.uncertainty_enabled: | |
| if return_attention: | |
| raw, _, attention, patches = outputs | |
| return raw, attention, patches | |
| raw, _, patches = outputs | |
| return raw, patches | |
| if return_attention: | |
| raw, _, variance, attention, patches, patch_variance = outputs | |
| return raw, variance, attention, patches, patch_variance | |
| raw, _, variance, patches, patch_variance = outputs | |
| return raw, variance, patches, patch_variance | |
| def predict_interval(self, x, z=1.96): | |
| """Return point, lower, and upper physical-flux estimates.""" | |
| if not self.uncertainty_enabled: | |
| raise RuntimeError("predict_interval requires uncertainty.enabled=true") | |
| raw, variance_norm, _, _ = self.forward(x, return_attention=False) | |
| mean_norm = normalize_sxr(raw.squeeze(-1), self.sxr_norm) | |
| sigma = torch.sqrt(variance_norm) | |
| lower = unnormalize_sxr(mean_norm - z * sigma, self.sxr_norm) | |
| upper = unnormalize_sxr(mean_norm + z * sigma, self.sxr_norm) | |
| return raw.squeeze(-1), torch.clamp(lower, min=0), torch.clamp(upper, min=0) | |
| def _weighted_huber_loss(self, mean, target, target_raw): | |
| """Return Huber loss, optionally balanced by the four GOES classes.""" | |
| per_sample = F.huber_loss( | |
| mean, target, delta=self.huber_delta, reduction='none', | |
| ) | |
| if self.class_weighting != 'none': | |
| thresholds = SXRRegressionDynamicLoss.CLASS_THRESHOLDS | |
| class_indices = torch.zeros_like(target_raw, dtype=torch.long) | |
| class_indices = torch.where( | |
| target_raw >= thresholds['c'], 1, class_indices, | |
| ) | |
| class_indices = torch.where( | |
| target_raw >= thresholds['m'], 2, class_indices, | |
| ) | |
| class_indices = torch.where( | |
| target_raw >= thresholds['x'], 3, class_indices, | |
| ) | |
| sample_weights = self.class_weight_values[class_indices] | |
| per_sample = per_sample * sample_weights.reshape_as(per_sample) | |
| return per_sample.mean() | |
| def _uncertainty_nll_loss(mean, target, variance): | |
| """Train variance without allowing NLL gradients into the mean.""" | |
| return F.gaussian_nll_loss( | |
| mean.detach(), target, variance, full=True, reduction='mean', | |
| ) | |
| def _uncertainty_active_for_epoch(self, epoch): | |
| return ( | |
| self.uncertainty_enabled | |
| and epoch >= self.uncertainty_start_epoch | |
| ) | |
| def _uncertainty_training_is_active(self): | |
| trainer = getattr(self, '_trainer', None) | |
| epoch = trainer.current_epoch if trainer is not None else 0 | |
| return self._uncertainty_active_for_epoch(epoch) | |
| def _calculate_loss(self, batch, mode): | |
| images, target_norm = batch | |
| outputs = self.model( | |
| images, self.sxr_norm, return_attention=False, | |
| ) | |
| if self.uncertainty_enabled: | |
| ( | |
| raw_prediction, mean_norm, variance_norm, | |
| _, patch_variance_raw, | |
| ) = outputs | |
| else: | |
| raw_prediction, mean_norm, _ = outputs | |
| variance_norm = patch_variance_raw = None | |
| target_norm = target_norm.reshape_as(mean_norm) | |
| target_raw = unnormalize_sxr(target_norm, self.sxr_norm) | |
| mean_loss = self._weighted_huber_loss( | |
| mean_norm, target_norm, target_raw, | |
| ) | |
| signed_error = (mean_norm - target_norm).mean() | |
| bias_penalty = mean_class_signed_error_penalty( | |
| mean_norm, target_norm, target_raw, self.bias_penalty_weight, | |
| ) | |
| mean_objective = mean_loss + bias_penalty | |
| uncertainty_active = self._uncertainty_training_is_active() | |
| uncertainty_loss = ( | |
| self._uncertainty_nll_loss( | |
| mean_norm, target_norm, variance_norm, | |
| ) | |
| if self.uncertainty_enabled else None | |
| ) | |
| loss = ( | |
| mean_objective + uncertainty_loss | |
| if uncertainty_active else mean_objective | |
| ) | |
| mse = F.mse_loss(mean_norm, target_norm) | |
| mae = F.l1_loss(mean_norm, target_norm) | |
| prefix = mode | |
| on_step = mode == 'train' | |
| if self.uncertainty_enabled: | |
| sigma = torch.sqrt(variance_norm) | |
| sigma_dex = sigma * self.sxr_norm[1] | |
| standardized = torch.abs(target_norm - mean_norm) / sigma | |
| coverage_68 = (standardized <= 1.0).float().mean() | |
| coverage_95 = (standardized <= 1.96).float().mean() | |
| self.log(f'{prefix}/nll', uncertainty_loss, on_step=on_step, | |
| on_epoch=True, | |
| prog_bar=True, logger=True, sync_dist=True) | |
| self.log( | |
| f'{prefix}/uncertainty_nll', uncertainty_loss, | |
| on_step=on_step, on_epoch=True, | |
| logger=True, sync_dist=True, | |
| ) | |
| self.log( | |
| f'{prefix}/uncertainty_active', | |
| float(uncertainty_active), on_step=on_step, on_epoch=True, | |
| logger=True, sync_dist=True, | |
| ) | |
| self.log(f'{prefix}/mean_sigma', sigma.mean(), on_step=on_step, | |
| on_epoch=True, logger=True, sync_dist=True) | |
| self.log( | |
| f'{prefix}/mean_sigma_dex', sigma_dex.mean(), | |
| on_step=on_step, on_epoch=True, logger=True, sync_dist=True, | |
| ) | |
| self.log( | |
| f'{prefix}/mean_sqrt_variance_contribution_raw', | |
| torch.sqrt(patch_variance_raw).mean(), | |
| on_step=on_step, on_epoch=True, logger=True, sync_dist=True, | |
| ) | |
| self.log(f'{prefix}/coverage_68', coverage_68, on_step=False, | |
| on_epoch=True, logger=True, sync_dist=True) | |
| self.log(f'{prefix}/coverage_95', coverage_95, on_step=False, | |
| on_epoch=True, logger=True, sync_dist=True) | |
| self.log( | |
| f'{prefix}/mean_loss', mean_loss, on_step=on_step, on_epoch=True, | |
| logger=True, sync_dist=True, | |
| ) | |
| self.log(f'{prefix}/bias_dex', signed_error * self.sxr_norm[1], | |
| on_step=on_step, on_epoch=True, logger=True, sync_dist=True) | |
| self.log(f'{prefix}/bias_penalty', bias_penalty, on_step=on_step, | |
| on_epoch=True, logger=True, sync_dist=True) | |
| self.log(f'{prefix}/total_loss', loss, on_step=on_step, on_epoch=True, | |
| logger=True, sync_dist=True) | |
| self.log(f'{prefix}/mse', mse, on_step=on_step, on_epoch=True, | |
| logger=True, sync_dist=True) | |
| self.log(f'{prefix}/mae', mae, on_step=on_step, on_epoch=True, | |
| logger=True, sync_dist=True) | |
| self.log( | |
| f'{prefix}/global_patch_adjustment_dex', | |
| self.model.last_log10_global_patch_adjustment.abs().mean(), | |
| on_step=on_step, on_epoch=True, logger=True, sync_dist=True, | |
| ) | |
| if mode == 'val': | |
| self.log('val_total_loss', loss, on_step=False, on_epoch=True, | |
| prog_bar=True, logger=True, sync_dist=True) | |
| if mode == 'train': | |
| learning_rate = self.trainer.optimizers[0].param_groups[0]['lr'] | |
| self.log('train/learning_rate', learning_rate, on_step=True, | |
| on_epoch=False, prog_bar=True, logger=True, sync_dist=True) | |
| if mode == 'val': | |
| result = { | |
| 'loss': loss.detach(), | |
| 'mean_norm': mean_norm.detach(), | |
| 'target_norm': target_norm.detach(), | |
| 'target_raw': target_raw.detach(), | |
| 'prediction_raw': raw_prediction.squeeze(-1).detach(), | |
| } | |
| if self.uncertainty_enabled: | |
| result.update({ | |
| 'nll': uncertainty_loss.detach(), | |
| 'variance_norm': variance_norm.detach(), | |
| }) | |
| return result | |
| return loss | |
| def training_step(self, batch, batch_idx): | |
| return self._calculate_loss(batch, 'train') | |
| def validation_step(self, batch, batch_idx): | |
| return self._calculate_loss(batch, 'val') | |
| def test_step(self, batch, batch_idx): | |
| self._calculate_loss(batch, 'test') | |
| def configure_optimizers(self): | |
| trainable = [p for p in self.parameters() if p.requires_grad] | |
| if not trainable: | |
| raise ValueError('No trainable parameters remain') | |
| optimizer = torch.optim.AdamW( | |
| trainable, lr=self.lr, weight_decay=self.weight_decay | |
| ) | |
| scheduler = CosineAnnealingLR(optimizer, **self.scheduler_kwargs) | |
| return { | |
| 'optimizer': optimizer, | |
| 'lr_scheduler': { | |
| 'scheduler': scheduler, 'interval': 'epoch', 'frequency': 1, | |
| 'name': 'learning_rate', | |
| }, | |
| } | |