""" PyTorch Lightning callbacks for visualizing training progress: predicted vs. true SXR flux, and Vision Transformer attention maps, logged to Weights & Biases. Used by train.py — not meant to be run standalone. """ import wandb from pytorch_lightning import Callback import matplotlib.pyplot as plt import numpy as np import torch from forecasting.model import unnormalize_sxr class ImagePredictionLogger_SXR(Callback): """ PyTorch Lightning callback for logging AIA input images and corresponding true vs predicted Soft X-Ray (SXR) flux values to Weights & Biases (wandb). This helps monitor model performance across validation epochs by comparing predicted vs. ground-truth flare intensities. """ def __init__(self, data_samples, sxr_norm): """ Initialize callback with validation samples and normalization parameters. Parameters ---------- data_samples : list List of validation samples (AIA image, SXR target pairs). sxr_norm : np.ndarray Normalization statistics used to unnormalize predicted flux values. """ super().__init__() self.data_samples = data_samples self.sxr_norm = sxr_norm def on_validation_epoch_end(self, trainer, pl_module): """ Log scatter plots comparing predicted and true SXR flux values at the end of each validation epoch. """ true_sxr = [] pred_sxr = [] for aia, target in self.data_samples: aia = aia.to(pl_module.device).unsqueeze(0) # forward() always returns a tuple (global_flux_raw, ...); we only need the flux. pred, *_ = pl_module(aia, return_attention=False) pred_sxr.append(pred.item()) true_sxr.append(target.item()) true_unorm = unnormalize_sxr(np.array(true_sxr, dtype=np.float32), self.sxr_norm) pred_unnorm = unnormalize_sxr(np.array(pred_sxr, dtype=np.float32), self.sxr_norm) fig1 = self.plot_aia_sxr(true_unorm, pred_unnorm) trainer.logger.experiment.log({"Soft X-ray flux plots": wandb.Image(fig1)}) plt.close(fig1) # Flare-class range: A-class starts at 1e-8 W/m^2, X10 is 10x the 1e-4 X-class threshold. AXIS_MIN = 1e-8 AXIS_MAX = 1e-3 def plot_aia_sxr(self, val_sxr, pred_sxr): """Log-log parity plot: predicted vs. true SXR flux, with a 1:1 reference line.""" fig, ax = plt.subplots(1, 1, figsize=(4, 4)) ax.plot([self.AXIS_MIN, self.AXIS_MAX], [self.AXIS_MIN, self.AXIS_MAX], color='gray', linestyle='--', linewidth=1, label='Perfect prediction') ax.scatter(val_sxr, pred_sxr, color='blue', alpha=0.7,s=10, label='Predictions') ax.set_xscale('log') ax.set_yscale('log') ax.set_xlim(self.AXIS_MIN, self.AXIS_MAX) ax.set_ylim(self.AXIS_MIN, self.AXIS_MAX) ax.set_xlabel("True SXR flux [W/m$^2$]") ax.set_ylabel("Predicted SXR flux [W/m$^2$]") ax.legend() fig.tight_layout() return fig class AttentionMapCallback(Callback): """ PyTorch Lightning callback for visualizing transformer attention maps during validation epochs. Supports CLS-token-based and local patch attention visualization. """ def __init__(self, log_every_n_epochs=1, num_samples=4, patch_size=8, use_local_attention=False): """ Initialize callback. Parameters ---------- log_every_n_epochs : int Frequency of logging attention maps. num_samples : int Number of samples to visualize per epoch. patch_size : int Patch size used in the Vision Transformer. use_local_attention : bool If True, visualize local attention patterns instead of CLS attention. """ super().__init__() self.patch_size = patch_size self.log_every_n_epochs = log_every_n_epochs self.num_samples = num_samples self.use_local_attention = use_local_attention def on_validation_epoch_end(self, trainer, pl_module): """Trigger visualization of attention maps at the end of validation epochs.""" if trainer.current_epoch % self.log_every_n_epochs == 0: self._visualize_attention(trainer, pl_module) def _visualize_attention(self, trainer, pl_module): """Generate and log attention maps from the model's attention weights.""" val_dataloader = trainer.val_dataloaders if val_dataloader is None: return pl_module.eval() with torch.no_grad(): batch = next(iter(val_dataloader)) imgs, labels = batch imgs = imgs[:self.num_samples].to(pl_module.device) # forward(return_attention=True) -> (global_flux_raw, attention_weights, patch_flux_raw) _, attention_weights, patch_flux_raw = pl_module(imgs, return_attention=True) for sample_idx in range(min(self.num_samples, imgs.size(0))): fig = self._plot_attention_map( imgs[sample_idx], attention_weights, sample_idx, trainer.current_epoch, patch_size=self.patch_size, patch_flux=patch_flux_raw[sample_idx] if patch_flux_raw is not None else None ) trainer.logger.experiment.log({"Attention plots": wandb.Image(fig)}) plt.close(fig) def _plot_attention_map(self, image, attention_weights, sample_idx, epoch, patch_size, patch_flux=None): """Plot and return a visualization of the attention heatmaps for a single image.""" img_np = image.cpu().numpy() if len(img_np.shape) == 3 and img_np.shape[0] in [1, 3]: img_np = np.transpose(img_np, (1, 2, 0)) H, W = img_np.shape[:2] grid_h, grid_w = H // patch_size, W // patch_size last_layer_attention = attention_weights[-1] sample_attention = last_layer_attention[sample_idx] avg_attention = sample_attention.mean(dim=0) if self.use_local_attention: # Spatial center of the grid, not the middle of the flattened sequence # (those only coincide when grid_w == 1) — row-major flatten: idx = row*grid_w + col. center_patch_idx = (grid_h // 2) * grid_w + (grid_w // 2) center_attention = avg_attention[center_patch_idx, :].cpu() avg_attention_map = avg_attention.mean(dim=0).cpu() attention_map = avg_attention_map.reshape(grid_h, grid_w) center_map = center_attention.reshape(grid_h, grid_w) else: cls_attention = avg_attention[0, 1:].cpu() attention_map = cls_attention.reshape(grid_h, grid_w) center_map = None if len(img_np[0, 0, :]) >= 6: rgb_channels = [0, 2, 4] img_display = np.stack([(img_np[:, :, i] + 1) / 2 for i in rgb_channels], axis=2) img_display = np.clip(img_display, 0, 1) else: img_display = (img_np[:, :, 0] + 1) / 2 img_display = np.stack([img_display] * 3, axis=2) fig, axes = plt.subplots(2, 3, figsize=(15, 10)) fig.suptitle(f'Attention Visualization - Epoch {epoch}, Sample {sample_idx}', fontsize=16) axes[0, 0].imshow(img_display) axes[0, 0].set_title('Original Image') axes[0, 0].axis('off') im1 = axes[0, 1].imshow(attention_map, cmap='hot', interpolation='nearest') axes[0, 1].set_title('Attention Map') axes[0, 1].axis('off') plt.colorbar(im1, ax=axes[0, 1]) axes[0, 2].imshow(img_display) axes[0, 2].imshow(attention_map, cmap='hot', alpha=0.6, interpolation='nearest') axes[0, 2].set_title('Attention Overlay') axes[0, 2].axis('off') if center_map is not None: im2 = axes[1, 0].imshow(center_map, cmap='hot', interpolation='nearest') axes[1, 0].set_title('Center Patch Attention') axes[1, 0].axis('off') plt.colorbar(im2, ax=axes[1, 0]) else: axes[1, 0].text(0.5, 0.5, 'Center attention\nnot available', ha='center', va='center', transform=axes[1, 0].transAxes) axes[1, 0].set_title('Center Patch Attention') axes[1, 0].axis('off') if patch_flux is not None: patch_flux_np = patch_flux.cpu().numpy().reshape(grid_h, grid_w) im3 = axes[1, 1].imshow(patch_flux_np, cmap='viridis', interpolation='nearest') axes[1, 1].set_title('Patch Flux') axes[1, 1].axis('off') plt.colorbar(im3, ax=axes[1, 1]) else: axes[1, 1].text(0.5, 0.5, 'Patch flux\nnot available', ha='center', va='center', transform=axes[1, 1].transAxes) axes[1, 1].set_title('Patch Flux') axes[1, 1].axis('off') axes[1, 2].hist(attention_map.flatten(), bins=50, alpha=0.7) axes[1, 2].set_title('Attention Distribution') axes[1, 2].set_xlabel('Attention Weight') axes[1, 2].set_ylabel('Frequency') plt.tight_layout() return fig