Attention viz

src/ocean_emulators/aggregator/main.py

@@ -5,10 +5,12 @@
 """
 
 import torch
+import torch.nn as nn
 
 from ocean_emulators.aggregator.inference import InferenceEvaluatorAggregator
 from ocean_emulators.aggregator.train import TrainAggregator
 from ocean_emulators.aggregator.validate import ValidateAggregator
+from ocean_emulators.aggregator.validate.attention import AttentionAggregator
 from ocean_emulators.aggregator.validate.map import MapAggregator
 from ocean_emulators.aggregator.validate.reduced import MeanAggregator
 from ocean_emulators.aggregator.validate.snapshot import SnapshotAggregator
@@ -28,6 +30,7 @@ def get_validation_aggregator(
         num_prognostic_channels: int,
         *,
         include_image_aggregators: bool = True,
+        model: nn.Module,
     ) -> ValidateAggregator:
         val_aggregators: dict[str, ValidateSubAggregator] = {
             "reduced": MeanAggregator(area_weights, hist),
@@ -37,6 +40,7 @@ def get_validation_aggregator(
                 {
                     "snapshot": SnapshotAggregator(metadata, hist),
                     "mean_map": MapAggregator(metadata, hist),
+                    "attention": AttentionAggregator(model),
                 }
             )
 

src/ocean_emulators/aggregator/plotting.py

@@ -1,9 +1,11 @@
 import gc
+from typing import Literal
 
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.colors import Colormap
 from matplotlib.figure import Figure
+from wandb.data_types import WBValue
 
 from ocean_emulators.utils.wandb import WandBLogger
 
@@ -133,3 +135,165 @@ def _stitch_data_panels(data: list[list[np.ndarray]], fill_value) -> np.ndarray:
             stitched_data[start_row:end_row, start_col:end_col] = arr
 
     return stitched_data
+
+
+def _downsample_for_display(data: np.ndarray, max_size: int = 512) -> np.ndarray:
+    """Average-pool a 2D array for display if it is too large."""
+    height, width = data.shape
+    if height <= max_size and width <= max_size:
+        return data
+
+    out_height = min(height, max_size)
+    out_width = min(width, max_size)
+    row_edges = np.linspace(0, height, out_height + 1, dtype=int)
+    col_edges = np.linspace(0, width, out_width + 1, dtype=int)
+    pooled = np.empty((out_height, out_width), dtype=np.float32)
+
+    for i in range(out_height):
+        for j in range(out_width):
+            block = data[
+                row_edges[i] : row_edges[i + 1], col_edges[j] : col_edges[j + 1]
+            ]
+            pooled[i, j] = float(block.mean())
+
+    return pooled
+
+
+def plot_attention_map(
+    attn_weights: np.ndarray,
+    axis: Literal["height", "width", "full"],
+    caption: str | None = None,
+) -> WBValue:
+    wandb_logger = WandBLogger.get_instance()
+    display_weights = _downsample_for_display(attn_weights)
+
+    if display_weights.shape != attn_weights.shape:
+        size_note = (
+            f"Displayed as {display_weights.shape[0]}x{display_weights.shape[1]} "
+            f"from original {attn_weights.shape[0]}x{attn_weights.shape[1]}."
+        )
+        caption = f"{caption} {size_note}" if caption else size_note
+
+    fig = Figure(figsize=(6, 5))
+    ax = fig.add_subplot(111)
+    im = ax.imshow(display_weights, cmap="viridis", aspect="auto")
+    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
+
+    labels = {
+        "height": (
+            "Key (latitude index)",
+            "Query (latitude index)",
+            "Height-axis attention",
+        ),
+        "width": (
+            "Key (longitude index)",
+            "Query (longitude index)",
+            "Width-axis attention",
+        ),
+        "full": ("Key token index", "Query token index", "Full 2D attention"),
+    }
+    xlabel, ylabel, title = labels[axis]
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel(ylabel)
+    ax.set_title(title)
+
+    fig.tight_layout()
+    image = wandb_logger.Image(fig, caption=caption)
+    plt.close(fig)
+    gc.collect()
+    return image
+
+
+def plot_attention_receptive_field(
+    height_weights: np.ndarray,
+    width_weights: np.ndarray,
+    query_lat: int,
+    query_lon: int,
+    caption: str | None = None,
+) -> WBValue:
+    wandb_logger = WandBLogger.get_instance()
+
+    h_row = height_weights[query_lat]
+    w_row = width_weights[query_lon]
+    receptive_field = np.outer(h_row, w_row)
+
+    fig = Figure(figsize=(8, 4))
+    ax = fig.add_subplot(111)
+    im = ax.imshow(receptive_field, cmap="inferno", aspect="auto", origin="lower")
+    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
+    ax.set_xlabel("Longitude index")
+    ax.set_ylabel("Latitude index")
+    ax.set_title(f"Receptive field at ({query_lat}, {query_lon})")
+    fig.tight_layout()
+
+    image = wandb_logger.Image(fig, caption=caption)
+    plt.close(fig)
+    gc.collect()
+    return image
+
+
+def plot_full_attention_receptive_field(
+    attn_weights: np.ndarray,
+    grid_shape: tuple[int, int],
+    query_lat: int,
+    query_lon: int,
+    caption: str | None = None,
+) -> WBValue:
+    height, width = grid_shape
+    if not (0 <= query_lat < height and 0 <= query_lon < width):
+        raise IndexError(
+            f"Query ({query_lat}, {query_lon}) is outside the grid shape {grid_shape}."
+        )
+
+    query_index = query_lat * width + query_lon
+    receptive_field = attn_weights[query_index].reshape(height, width)
+
+    wandb_logger = WandBLogger.get_instance()
+    fig = Figure(figsize=(8, 4))
+    ax = fig.add_subplot(111)
+    im = ax.imshow(receptive_field, cmap="inferno", aspect="auto", origin="lower")
+    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
+    ax.set_xlabel("Longitude index")
+    ax.set_ylabel("Latitude index")
+    ax.set_title(f"Full-attention receptive field at ({query_lat}, {query_lon})")
+    fig.tight_layout()
+
+    image = wandb_logger.Image(fig, caption=caption)
+    plt.close(fig)
+    gc.collect()
+    return image
+
+
+def plot_attention_entropy_map(
+    entropy: np.ndarray,
+    grid_shape: tuple[int, int] | None = None,
+    caption: str | None = None,
+) -> WBValue:
+    """Plot per-query attention entropy as a heatmap.
+
+    Args:
+        entropy: Per-query attention entropy. Either a 1D vector or a 2D map.
+        grid_shape: Optional shape used to reshape 1D entropy into a 2D map.
+        caption: Optional caption for the W&B image.
+    """
+    entropy_map = entropy
+    if entropy.ndim == 1:
+        if grid_shape is None:
+            entropy_map = entropy[np.newaxis, :]
+        else:
+            entropy_map = entropy.reshape(grid_shape)
+
+    wandb_logger = WandBLogger.get_instance()
+    fig = Figure(figsize=(8, 4))
+    ax = fig.add_subplot(111)
+    im = ax.imshow(entropy_map, cmap="magma", aspect="auto", origin="lower")
+    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
+    ax.set_xlabel("Query index" if entropy_map.shape[0] == 1 else "Longitude index")
+    ax.set_ylabel("Entropy" if entropy_map.shape[0] == 1 else "Latitude index")
+    ax.set_title("Attention entropy")
+    fig.tight_layout()
+
+    image = wandb_logger.Image(fig, caption=caption)
+    plt.close(fig)
+    gc.collect()
+    return image