Add CNN predictor

CHANGELOG.md

@@ -9,6 +9,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Add an initial regular-grid `CNNPredictor`, CLI model construction support,
+  reusable CNN grid transforms, FiLM conditioning, Squeeze-and-Excitation, and
+  ResHRRR-style residual backbone layers with validation coverage.
+
 - Add `PropagationNet` GNN layer that incentivises directional message
   propagation from sender to receiver nodes, and expose it alongside
   `InteractionNet` through four new CLI arguments (`--g2m_gnn_type`,

neural_lam/__init__.py

@@ -4,6 +4,7 @@
 import importlib.metadata
 
 # First-party
+import neural_lam.cnn_layers
 import neural_lam.gnn_layers
 import neural_lam.metrics
 import neural_lam.models

neural_lam/cnn_layers.py

@@ -0,0 +1,375 @@
+"""CNN layers and node/grid transforms for regular-grid predictors.
+
+The residual backbone is inspired by the ResHRRR model in HRRRCast:
+https://arxiv.org/abs/2507.05658.
+"""
+
+# Third-party
+import torch
+from torch import nn
+
+# Local
+from .datastore.base import CartesianGridShape
+
+GridShape = CartesianGridShape | tuple[int, int]
+PADDING_MODES = {"zeros", "reflect", "replicate", "circular"}
+
+
+def _grid_shape_xy(grid_shape: GridShape) -> tuple[int, int]:
+    """
+    Return ``(x, y)`` dimensions from a datastore or tuple grid shape.
+    """
+    if isinstance(grid_shape, CartesianGridShape):
+        grid_x, grid_y = grid_shape.x, grid_shape.y
+    else:
+        if len(grid_shape) != 2:
+            raise ValueError("grid_shape must contain exactly two dimensions")
+        grid_x, grid_y = int(grid_shape[0]), int(grid_shape[1])
+
+    if grid_x <= 0 or grid_y <= 0:
+        raise ValueError("grid_shape dimensions must be positive")
+
+    return int(grid_x), int(grid_y)
+
+
+def node_to_grid(
+    node_features: torch.Tensor,
+    grid_shape: GridShape,
+) -> torch.Tensor:
+    """
+    Convert Neural-LAM node tensors to CNN grid tensors.
+
+    Parameters
+    ----------
+    node_features : torch.Tensor
+        Tensor of shape ``(B, N, C)``.
+    grid_shape : CartesianGridShape or tuple[int, int]
+        Regular grid shape as ``(x, y)``.
+
+    Returns
+    -------
+    torch.Tensor
+        Tensor of shape ``(B, C, x, y)``.
+    """
+    if node_features.ndim != 3:
+        raise ValueError(
+            "node_features must have shape (B, N, C), "
+            f"got {tuple(node_features.shape)}"
+        )
+
+    grid_x, grid_y = _grid_shape_xy(grid_shape)
+    batch_size, num_nodes, num_channels = node_features.shape
+    expected_nodes = grid_x * grid_y
+
+    if num_nodes != expected_nodes:
+        raise ValueError(
+            "node_features node dimension does not match grid_shape: "
+            f"got {num_nodes}, expected {expected_nodes}"
+        )
+
+    return (
+        node_features.reshape(batch_size, grid_x, grid_y, num_channels)
+        .permute(0, 3, 1, 2)
+        .contiguous()
+    )
+
+
+def grid_to_node(
+    grid_features: torch.Tensor,
+    grid_shape: GridShape | None = None,
+) -> torch.Tensor:
+    """
+    Convert CNN grid tensors to Neural-LAM node tensors.
+
+    Parameters
+    ----------
+    grid_features : torch.Tensor
+        Tensor of shape ``(B, C, x, y)``.
+    grid_shape : CartesianGridShape or tuple[int, int], optional
+        Expected regular grid shape as ``(x, y)``.
+
+    Returns
+    -------
+    torch.Tensor
+        Tensor of shape ``(B, x * y, C)``.
+    """
+    if grid_features.ndim != 4:
+        raise ValueError(
+            "grid_features must have shape (B, C, x, y), "
+            f"got {tuple(grid_features.shape)}"
+        )
+
+    batch_size, num_channels, grid_x, grid_y = grid_features.shape
+
+    if grid_shape is not None:
+        expected_x, expected_y = _grid_shape_xy(grid_shape)
+        if (grid_x, grid_y) != (expected_x, expected_y):
+            raise ValueError(
+                "grid_features spatial dimensions do not match grid_shape: "
+                f"got {(grid_x, grid_y)}, expected {(expected_x, expected_y)}"
+            )
+
+    return (
+        grid_features.permute(0, 2, 3, 1)
+        .reshape(batch_size, grid_x * grid_y, num_channels)
+        .contiguous()
+    )
+
+
+def validate_padding_mode_for_grid(
+    grid_shape: GridShape,
+    kernel_size: int,
+    padding_mode: str,
+) -> None:
+    """
+    Validate that a CNN padding mode is safe for a fixed regular grid.
+
+    PyTorch reflection padding requires each spatial padding width to be
+    smaller than the corresponding input dimension. Circular padding cannot
+    wrap around an axis more than once.
+    """
+    if kernel_size <= 0 or kernel_size % 2 == 0:
+        raise ValueError("kernel_size must be a positive odd integer")
+    if padding_mode not in PADDING_MODES:
+        raise ValueError(
+            f"padding_mode must be one of {sorted(PADDING_MODES)}, "
+            f"got {padding_mode}"
+        )
+
+    padding = kernel_size // 2
+    grid_x, grid_y = _grid_shape_xy(grid_shape)
+    min_grid_dim = min(grid_x, grid_y)
+
+    if padding_mode == "reflect" and padding >= min_grid_dim:
+        raise ValueError(
+            "reflect padding requires kernel_size // 2 to be smaller than "
+            "each grid dimension: "
+            f"got padding={padding}, grid_shape={(grid_x, grid_y)}"
+        )
+    if padding_mode == "circular" and padding > min_grid_dim:
+        raise ValueError(
+            "circular padding requires kernel_size // 2 to be no larger than "
+            "each grid dimension: "
+            f"got padding={padding}, grid_shape={(grid_x, grid_y)}"
+        )
+
+
+class NodeToGrid(nn.Module):
+    """Layer wrapper for ``node_to_grid``."""
+
+    def __init__(self, grid_shape: GridShape):
+        """Initialize the transform with a fixed regular grid shape."""
+        super().__init__()
+        self.grid_shape = _grid_shape_xy(grid_shape)
+
+    def forward(self, node_features: torch.Tensor) -> torch.Tensor:
+        """Convert node features from ``(B, N, C)`` to ``(B, C, x, y)``."""
+        return node_to_grid(node_features, self.grid_shape)
+
+
+class GridToNode(nn.Module):
+    """Layer wrapper for ``grid_to_node``."""
+
+    def __init__(self, grid_shape: GridShape | None = None):
+        """Initialize the transform with an optional expected grid shape."""
+        super().__init__()
+        self.grid_shape = (
+            None if grid_shape is None else _grid_shape_xy(grid_shape)
+        )
+
+    def forward(self, grid_features: torch.Tensor) -> torch.Tensor:
+        """Convert grid features from ``(B, C, x, y)`` to ``(B, N, C)``."""
+        return grid_to_node(grid_features, self.grid_shape)
+
+
+class SqueezeExcitation2d(nn.Module):
+    """Channel attention for 2D feature maps."""
+
+    def __init__(self, channels: int, reduction: int = 16):
+        """Initialize channel attention with a bottleneck reduction ratio."""
+        super().__init__()
+        if channels <= 0:
+            raise ValueError("channels must be positive")
+        if reduction <= 0:
+            raise ValueError("reduction must be positive")
+
+        hidden_channels = max(1, channels // reduction)
+        self.net = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(channels, hidden_channels, kernel_size=1),
+            nn.SiLU(),
+            nn.Conv2d(hidden_channels, channels, kernel_size=1),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply channel-wise gates to a 2D feature map."""
+        return x * self.net(x)
+
+
+class FiLM2d(nn.Module):
+    """Feature-wise affine conditioning for 2D feature maps."""
+
+    def __init__(self, context_dim: int, channels: int):
+        """Initialize a context projection for per-channel affine factors."""
+        super().__init__()
+        if context_dim <= 0:
+            raise ValueError("context_dim must be positive")
+        if channels <= 0:
+            raise ValueError("channels must be positive")
+
+        self.context_dim = context_dim
+        self.channels = channels
+        self.proj = nn.Linear(context_dim, 2 * channels)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor,
+    ) -> torch.Tensor:
+        """Apply context-conditioned affine modulation to feature maps."""
+        if x.ndim != 4:
+            raise ValueError(f"x must have shape (B, C, H, W), got {x.shape}")
+        if context.ndim != 2:
+            raise ValueError(
+                "context must have shape (B, context_dim), "
+                f"got {tuple(context.shape)}"
+            )
+        if x.shape[0] != context.shape[0]:
+            raise ValueError("x and context must have the same batch size")
+        if x.shape[1] != self.channels:
+            raise ValueError(
+                f"x channel dimension must be {self.channels}, "
+                f"got {x.shape[1]}"
+            )
+        if context.shape[1] != self.context_dim:
+            raise ValueError(
+                f"context feature dimension must be {self.context_dim}, "
+                f"got {context.shape[1]}"
+            )
+
+        gamma, beta = self.proj(context).chunk(2, dim=-1)
+        gamma = gamma[:, :, None, None]
+        beta = beta[:, :, None, None]
+        return x * (1 + gamma) + beta
+
+
+class ResHRRRBlock(nn.Module):
+    """Residual CNN block with optional SE and FiLM conditioning."""
+
+    def __init__(
+        self,
+        channels: int,
+        kernel_size: int = 3,
+        reduction: int = 16,
+        context_dim: int | None = None,
+        padding_mode: str = "zeros",
+    ):
+        """Initialize a shape-preserving residual convolution block."""
+        super().__init__()
+        if channels <= 0:
+            raise ValueError("channels must be positive")
+        if kernel_size <= 0 or kernel_size % 2 == 0:
+            raise ValueError("kernel_size must be a positive odd integer")
+
+        padding = kernel_size // 2
+        self.conv1 = nn.Conv2d(
+            channels,
+            channels,
+            kernel_size=kernel_size,
+            padding=padding,
+            padding_mode=padding_mode,
+        )
+        self.act = nn.SiLU()
+        self.conv2 = nn.Conv2d(
+            channels,
+            channels,
+            kernel_size=kernel_size,
+            padding=padding,
+            padding_mode=padding_mode,
+        )
+        self.se = SqueezeExcitation2d(channels, reduction=reduction)
+        self.film = (
+            None if context_dim is None else FiLM2d(context_dim, channels)
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Run the residual block with optional FiLM conditioning."""
+        residual = x
+        x = self.conv1(x)
+        x = self.act(x)
+        x = self.conv2(x)
+        x = self.se(x)
+
+        if self.film is not None:
+            if context is None:
+                raise ValueError("context is required when FiLM is enabled")
+            x = self.film(x, context)
+
+        return self.act(x + residual)
+
+
+class ResHRRRBackbone(nn.Module):
+    """ResHRRR-style CNN backbone."""
+
+    def __init__(
+        self,
+        input_channels: int,
+        output_channels: int,
+        hidden_channels: int = 128,
+        num_blocks: int = 8,
+        kernel_size: int = 3,
+        reduction: int = 16,
+        context_dim: int | None = None,
+        padding_mode: str = "zeros",
+    ):
+        """Initialize input projection, residual blocks, and output head."""
+        super().__init__()
+        if input_channels <= 0:
+            raise ValueError("input_channels must be positive")
+        if output_channels <= 0:
+            raise ValueError("output_channels must be positive")
+        if hidden_channels <= 0:
+            raise ValueError("hidden_channels must be positive")
+        if num_blocks <= 0:
+            raise ValueError("num_blocks must be positive")
+
+        self.input_proj = nn.Conv2d(
+            input_channels,
+            hidden_channels,
+            kernel_size=1,
+        )
+        self.blocks = nn.ModuleList(
+            [
+                ResHRRRBlock(
+                    channels=hidden_channels,
+                    kernel_size=kernel_size,
+                    reduction=reduction,
+                    context_dim=context_dim,
+                    padding_mode=padding_mode,
+                )
+                for _ in range(num_blocks)
+            ]
+        )
+        self.output_proj = nn.Conv2d(
+            hidden_channels,
+            output_channels,
+            kernel_size=1,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Map input grid features to output grid features."""
+        x = self.input_proj(x)
+
+        for block in self.blocks:
+            x = block(x, context=context)
+
+        return self.output_proj(x)