diff --git a/CHANGELOG.md b/CHANGELOG.md index 300b8559..939b7761 100644 --- a/CHANGELOG.md +++ b/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`, diff --git a/neural_lam/__init__.py b/neural_lam/__init__.py index b1d01939..623e5a48 100644 --- a/neural_lam/__init__.py +++ b/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 diff --git a/neural_lam/cnn_layers.py b/neural_lam/cnn_layers.py new file mode 100644 index 00000000..dfff2fe9 --- /dev/null +++ b/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) diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py index cb87d76d..28e769dc 100644 --- a/neural_lam/models/__init__.py +++ b/neural_lam/models/__init__.py @@ -5,6 +5,7 @@ from .forecasters.base import Forecaster from .module import ForecasterModule from .step_predictors.base import StepPredictor +from .step_predictors.cnn import CNNPredictor from .step_predictors.graph.base import BaseGraphModel from .step_predictors.graph.graph_lam import GraphLAM from .step_predictors.graph.hi_lam import HiLAM @@ -12,6 +13,7 @@ from .step_predictors.graph.hierarchical import BaseHiGraphModel MODELS = { + "cnn_predictor": CNNPredictor, "graph_lam": GraphLAM, "hi_lam": HiLAM, "hi_lam_parallel": HiLAMParallel, diff --git a/neural_lam/models/step_predictors/__init__.py b/neural_lam/models/step_predictors/__init__.py index 4ddc8b98..e1147d85 100644 --- a/neural_lam/models/step_predictors/__init__.py +++ b/neural_lam/models/step_predictors/__init__.py @@ -4,3 +4,4 @@ # Local from .base import StepPredictor +from .cnn import CNNPredictor diff --git a/neural_lam/models/step_predictors/base.py b/neural_lam/models/step_predictors/base.py index daf8b0e6..efcef45a 100644 --- a/neural_lam/models/step_predictors/base.py +++ b/neural_lam/models/step_predictors/base.py @@ -84,6 +84,22 @@ def __init__( torch.tensor(da_state_stats.state_std.values, dtype=torch.float32), persistent=False, ) + self.register_buffer( + "diff_mean", + torch.tensor( + da_state_stats.state_diff_mean_standardized.values, + dtype=torch.float32, + ), + persistent=False, + ) + self.register_buffer( + "diff_std", + torch.tensor( + da_state_stats.state_diff_std_standardized.values, + dtype=torch.float32, + ), + persistent=False, + ) self.output_std = bool(output_std) if self.output_std: @@ -91,7 +107,7 @@ def __init__( else: self.grid_output_dim = num_state_vars - (self.num_grid_nodes, _) = self.grid_static_features.shape + self.num_grid_nodes, _ = self.grid_static_features.shape @property def predicts_std(self) -> bool: diff --git a/neural_lam/models/step_predictors/cnn/__init__.py b/neural_lam/models/step_predictors/cnn/__init__.py new file mode 100644 index 00000000..53b26ef9 --- /dev/null +++ b/neural_lam/models/step_predictors/cnn/__init__.py @@ -0,0 +1,4 @@ +"""CNN-based step predictors for regular-grid datastores.""" + +# Local +from .cnn_predictor import CNNPredictor diff --git a/neural_lam/models/step_predictors/cnn/cnn_predictor.py b/neural_lam/models/step_predictors/cnn/cnn_predictor.py new file mode 100644 index 00000000..9e13e81f --- /dev/null +++ b/neural_lam/models/step_predictors/cnn/cnn_predictor.py @@ -0,0 +1,235 @@ +"""Regular-grid CNN predictor using a ResHRRR-inspired backbone. + +The backbone follows the residual CNN, Squeeze-and-Excitation, and FiLM +conditioning ideas from HRRRCast: https://arxiv.org/abs/2507.05658. +""" + +# Standard library +from typing import Dict, Optional + +# Third-party +import torch + +# Local +from ....cnn_layers import ( + ResHRRRBackbone, + grid_to_node, + node_to_grid, + validate_padding_mode_for_grid, +) +from ....datastore.base import BaseDatastore, BaseRegularGridDatastore +from ..base import StepPredictor + + +class CNNPredictor(StepPredictor): + """ + ResHRRR-style CNN step predictor for regular-grid datastores. + + The predictor keeps the Neural-LAM ``StepPredictor`` interface in node + space, but runs the CNN backbone on regular-grid tensors. + """ + + def __init__( + self, + datastore: BaseDatastore, + cnn_channels: int = 128, + cnn_blocks: int = 8, + cnn_kernel_size: int = 3, + cnn_se_reduction: int = 16, + cnn_film: bool = False, + cnn_padding_mode: str = "zeros", + num_past_forcing_steps: int = 1, + num_future_forcing_steps: int = 1, + output_std: bool = False, + output_clamping_lower: Optional[Dict[str, float]] = None, + output_clamping_upper: Optional[Dict[str, float]] = None, + ): + """ + Initialize a CNN predictor for regular-grid datastores. + + Parameters + ---------- + datastore : BaseDatastore + Datastore providing state, forcing, static features, and grid + shape metadata. + cnn_channels : int + Number of hidden channels in the CNN backbone. + cnn_blocks : int + Number of residual CNN blocks. + cnn_kernel_size : int + Odd convolution kernel size used by residual blocks. + cnn_se_reduction : int + Channel reduction ratio for Squeeze-and-Excitation blocks. + cnn_film : bool + Whether to condition residual blocks with global forcing FiLM. + cnn_padding_mode : str + Padding mode passed to convolution layers. + num_past_forcing_steps : int + Number of past forcing steps in the forcing window. + num_future_forcing_steps : int + Number of future forcing steps in the forcing window. + output_std : bool + Whether to predict a positive per-feature standard deviation. + output_clamping_lower : dict[str, float], optional + Lower physical output clamps by state feature name. + output_clamping_upper : dict[str, float], optional + Upper physical output clamps by state feature name. + """ + if not isinstance(datastore, BaseRegularGridDatastore): + raise TypeError("CNNPredictor requires a BaseRegularGridDatastore") + + super().__init__( + datastore=datastore, + output_std=output_std, + output_clamping_lower=output_clamping_lower, + output_clamping_upper=output_clamping_upper, + ) + + self.grid_shape = datastore.grid_shape_state + expected_nodes = self.grid_shape.x * self.grid_shape.y + if expected_nodes != self.num_grid_nodes: + raise ValueError( + "datastore grid_shape_state does not match num_grid_points: " + f"got {expected_nodes}, expected {self.num_grid_nodes}" + ) + + self.num_state_vars = datastore.get_num_data_vars(category="state") + num_forcing_vars = datastore.get_num_data_vars(category="forcing") + forcing_window_steps = ( + num_past_forcing_steps + num_future_forcing_steps + 1 + ) + self.forcing_input_dim = num_forcing_vars * forcing_window_steps + + grid_static_dim = self.grid_static_features.shape[1] + self.grid_input_dim = ( + 2 * self.num_state_vars + self.forcing_input_dim + grid_static_dim + ) + context_dim = self.forcing_input_dim if cnn_film else None + validate_padding_mode_for_grid( + grid_shape=self.grid_shape, + kernel_size=cnn_kernel_size, + padding_mode=cnn_padding_mode, + ) + + self.backbone = ResHRRRBackbone( + input_channels=self.grid_input_dim, + output_channels=self.grid_output_dim, + hidden_channels=cnn_channels, + num_blocks=cnn_blocks, + kernel_size=cnn_kernel_size, + reduction=cnn_se_reduction, + context_dim=context_dim, + padding_mode=cnn_padding_mode, + ) + + self.cnn_film = cnn_film + self.prepare_clamping_params(datastore) + + def forward( + self, + prev_state: torch.Tensor, + prev_prev_state: torch.Tensor, + forcing: torch.Tensor, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """ + Predict the next state from two previous states and forcing. + + Parameters + ---------- + prev_state : torch.Tensor + Current state tensor with shape ``(B, N, d_state)``. + prev_prev_state : torch.Tensor + Previous state tensor with shape ``(B, N, d_state)``. + forcing : torch.Tensor + Windowed forcing tensor with shape ``(B, N, d_forcing)``. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor | None] + Predicted next state and optional predicted standard deviation. + """ + self._validate_inputs(prev_state, prev_prev_state, forcing) + batch_size = prev_state.shape[0] + + grid_features = torch.cat( + ( + prev_state, + prev_prev_state, + forcing, + self.expand_to_batch(self.grid_static_features, batch_size), + ), + dim=-1, + ) + grid_features = node_to_grid(grid_features, self.grid_shape) + + context = ( + self._get_global_forcing_context(forcing) if self.cnn_film else None + ) + net_output = self.backbone(grid_features, context=context) + net_output = grid_to_node(net_output, self.grid_shape) + + if self.output_std: + pred_delta_mean, pred_std_raw = net_output.chunk(2, dim=-1) + pred_std = torch.nn.functional.softplus(pred_std_raw) + else: + pred_delta_mean = net_output + pred_std = None + + rescaled_delta_mean = pred_delta_mean * self.diff_std + self.diff_mean + new_state = self.get_clamped_new_state(rescaled_delta_mean, prev_state) + + return new_state, pred_std + + def _get_global_forcing_context( + self, + forcing: torch.Tensor, + ) -> torch.Tensor: + """ + Average node-wise forcing into the global context used by FiLM. + + Spatially varying forcing is still included in ``grid_features`` before + the CNN backbone; this context only controls global channel modulation. + """ + return forcing.mean(dim=1) + + def _validate_inputs( + self, + prev_state: torch.Tensor, + prev_prev_state: torch.Tensor, + forcing: torch.Tensor, + ) -> None: + """Validate state and forcing tensors before grid conversion.""" + if prev_state.ndim != 3: + raise ValueError( + "prev_state must have shape (B, N, d_state), " + f"got {tuple(prev_state.shape)}" + ) + if prev_prev_state.shape != prev_state.shape: + raise ValueError( + "prev_prev_state must have the same shape as prev_state" + ) + if forcing.ndim != 3: + raise ValueError( + "forcing must have shape (B, N, d_forcing), " + f"got {tuple(forcing.shape)}" + ) + if forcing.shape[:2] != prev_state.shape[:2]: + raise ValueError( + "forcing must have the same batch and node dimensions as " + "prev_state" + ) + if prev_state.shape[1] != self.num_grid_nodes: + raise ValueError( + f"state node dimension must be {self.num_grid_nodes}, " + f"got {prev_state.shape[1]}" + ) + if prev_state.shape[2] != self.num_state_vars: + raise ValueError( + f"state feature dimension must be {self.num_state_vars}, " + f"got {prev_state.shape[2]}" + ) + if forcing.shape[2] != self.forcing_input_dim: + raise ValueError( + f"forcing feature dimension must be {self.forcing_input_dim}, " + f"got {forcing.shape[2]}" + ) diff --git a/neural_lam/models/step_predictors/graph/base.py b/neural_lam/models/step_predictors/graph/base.py index 1a747971..2ff1c2f2 100644 --- a/neural_lam/models/step_predictors/graph/base.py +++ b/neural_lam/models/step_predictors/graph/base.py @@ -72,25 +72,6 @@ def __init__( self.g2m_gnn_type = g2m_gnn_type self.m2g_gnn_type = m2g_gnn_type - # Retrieve difference statistics for rescaling in forward pass - da_state_stats = datastore.get_standardization_dataarray("state") - self.register_buffer( - "diff_mean", - torch.tensor( - da_state_stats.state_diff_mean_standardized.values, - dtype=torch.float32, - ), - persistent=False, - ) - self.register_buffer( - "diff_std", - torch.tensor( - da_state_stats.state_diff_std_standardized.values, - dtype=torch.float32, - ), - persistent=False, - ) - # Store architecture hyperparameters for subclass use self.hidden_dim = hidden_dim self.hidden_layers = hidden_layers diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index f98065c4..e2e9aed6 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -22,6 +22,63 @@ from .models import MODELS, ARForecaster, ForecasterModule from .weather_dataset import WeatherDataModule +CNN_PADDING_MODES = ("zeros", "reflect", "replicate", "circular") + + +def get_predictor_kwargs(args, config): + """Return architecture kwargs for the selected predictor.""" + clamping_kwargs = { + "output_clamping_lower": config.training.output_clamping.lower, + "output_clamping_upper": config.training.output_clamping.upper, + } + + common_kwargs = { + "num_past_forcing_steps": args.num_past_forcing_steps, + "num_future_forcing_steps": args.num_future_forcing_steps, + "output_std": args.output_std, + **clamping_kwargs, + } + + if args.model == "cnn_predictor": + return { + **common_kwargs, + "cnn_channels": args.cnn_channels, + "cnn_blocks": args.cnn_blocks, + "cnn_kernel_size": args.cnn_kernel_size, + "cnn_se_reduction": args.cnn_se_reduction, + "cnn_film": args.cnn_film, + "cnn_padding_mode": args.cnn_padding_mode, + } + + graph_kwargs = { + **common_kwargs, + "graph_name": args.graph, + "hidden_dim": args.hidden_dim, + "hidden_layers": args.hidden_layers, + "processor_layers": args.processor_layers, + "mesh_aggr": args.mesh_aggr, + "g2m_gnn_type": args.g2m_gnn_type, + "m2g_gnn_type": args.m2g_gnn_type, + } + if args.model in {"hi_lam", "hi_lam_parallel"}: + graph_kwargs.update( + { + "mesh_up_gnn_type": args.mesh_up_gnn_type, + "mesh_down_gnn_type": args.mesh_down_gnn_type, + } + ) + + return graph_kwargs + + +def build_predictor(args, config, datastore): + """Build the selected predictor from parsed CLI or checkpoint args.""" + predictor_class = MODELS[args.model] + return predictor_class( + datastore=datastore, + **get_predictor_kwargs(args, config), + ) + class AdaptiveHelpFormatter(ArgumentDefaultsHelpFormatter): """``--help`` formatter that scales the column width to the terminal.""" @@ -49,20 +106,7 @@ def load_forecaster_module_from_checkpoint(ckpt_path, config, datastore): """ ckpt = torch.load(ckpt_path, weights_only=False) args = ckpt["hyper_parameters"]["args"] - predictor_class = MODELS[args.model] - predictor = predictor_class( - datastore=datastore, - graph_name=args.graph, - hidden_dim=args.hidden_dim, - hidden_layers=args.hidden_layers, - processor_layers=args.processor_layers, - mesh_aggr=args.mesh_aggr, - num_past_forcing_steps=args.num_past_forcing_steps, - num_future_forcing_steps=args.num_future_forcing_steps, - output_std=args.output_std, - output_clamping_lower=config.training.output_clamping.lower, - output_clamping_upper=config.training.output_clamping.upper, - ) + predictor = build_predictor(args, config, datastore) forecaster = ARForecaster(predictor, datastore) return ForecasterModule.load_from_checkpoint( ckpt_path, @@ -205,6 +249,42 @@ def main(input_args=None): help="GNN type for downward mesh message passing in " "hierarchical models", ) + parser.add_argument( + "--cnn_channels", + type=int, + default=128, + help="Number of hidden channels in CNN predictor layers", + ) + parser.add_argument( + "--cnn_blocks", + type=int, + default=8, + help="Number of residual blocks in CNN predictor", + ) + parser.add_argument( + "--cnn_kernel_size", + type=int, + default=3, + help="Convolution kernel size in CNN predictor", + ) + parser.add_argument( + "--cnn_se_reduction", + type=int, + default=16, + help="Squeeze-and-Excitation reduction ratio in CNN predictor", + ) + parser.add_argument( + "--cnn_film", + action="store_true", + help="Use global forcing FiLM conditioning in CNN predictor", + ) + parser.add_argument( + "--cnn_padding_mode", + type=str, + default="zeros", + choices=CNN_PADDING_MODES, + help="Padding mode for CNN predictor convolutions", + ) # Training options train_group = parser.add_argument_group("Training Options") @@ -439,24 +519,7 @@ def main(input_args=None): # Build predictor and forecaster externally, then inject into # ForecasterModule - predictor_class = MODELS[args.model] - predictor = predictor_class( - datastore=datastore, - graph_name=args.graph, - hidden_dim=args.hidden_dim, - hidden_layers=args.hidden_layers, - processor_layers=args.processor_layers, - mesh_aggr=args.mesh_aggr, - num_past_forcing_steps=args.num_past_forcing_steps, - num_future_forcing_steps=args.num_future_forcing_steps, - output_std=args.output_std, - output_clamping_lower=config.training.output_clamping.lower, - output_clamping_upper=config.training.output_clamping.upper, - g2m_gnn_type=args.g2m_gnn_type, - m2g_gnn_type=args.m2g_gnn_type, - mesh_up_gnn_type=args.mesh_up_gnn_type, - mesh_down_gnn_type=args.mesh_down_gnn_type, - ) + predictor = build_predictor(args, config, datastore) forecaster = ARForecaster(predictor, datastore) model = ForecasterModule( diff --git a/tests/test_cnn_layers.py b/tests/test_cnn_layers.py new file mode 100644 index 00000000..5a3feec3 --- /dev/null +++ b/tests/test_cnn_layers.py @@ -0,0 +1,251 @@ +# Third-party +import pytest +import torch + +# First-party +from neural_lam.cnn_layers import ( + FiLM2d, + GridToNode, + NodeToGrid, + ResHRRRBackbone, + ResHRRRBlock, + SqueezeExcitation2d, + grid_to_node, + node_to_grid, + validate_padding_mode_for_grid, +) +from neural_lam.datastore.base import CartesianGridShape + + +def test_node_to_grid_matches_datastore_flattening_order(): + batch_size = 2 + grid_shape = CartesianGridShape(x=3, y=4) + num_channels = 5 + node_features = torch.arange( + batch_size * grid_shape.x * grid_shape.y * num_channels, + dtype=torch.float32, + ).reshape(batch_size, grid_shape.x * grid_shape.y, num_channels) + + grid_features = node_to_grid(node_features, grid_shape) + + expected = node_features.reshape( + batch_size, grid_shape.x, grid_shape.y, num_channels + ).permute(0, 3, 1, 2) + assert torch.equal(grid_features, expected) + + +def test_grid_to_node_round_trip(): + grid_shape = (3, 4) + node_features = torch.randn(2, 12, 5) + + round_trip = grid_to_node(node_to_grid(node_features, grid_shape)) + + assert torch.equal(round_trip, node_features) + + +def test_grid_transform_layers_round_trip(): + grid_shape = CartesianGridShape(x=2, y=3) + node_features = torch.randn(4, 6, 7) + + grid_features = NodeToGrid(grid_shape)(node_features) + round_trip = GridToNode(grid_shape)(grid_features) + + assert torch.equal(round_trip, node_features) + + +def test_node_to_grid_rejects_wrong_node_count(): + node_features = torch.randn(2, 11, 5) + + with pytest.raises(ValueError, match="node dimension"): + node_to_grid(node_features, (3, 4)) + + +def test_node_to_grid_rejects_non_positive_grid_shape(): + node_features = torch.randn(2, 0, 5) + + with pytest.raises(ValueError, match="positive"): + node_to_grid(node_features, (0, 4)) + + with pytest.raises(ValueError, match="positive"): + node_to_grid(node_features, CartesianGridShape(x=0, y=4)) + + +def test_grid_to_node_rejects_wrong_grid_shape(): + grid_features = torch.randn(2, 5, 3, 4) + + with pytest.raises(ValueError, match="spatial dimensions"): + grid_to_node(grid_features, (4, 3)) + + +def test_validate_padding_mode_for_grid_rejects_unsafe_reflect_padding(): + with pytest.raises(ValueError, match="reflect padding"): + validate_padding_mode_for_grid( + (1, 4), kernel_size=3, padding_mode="reflect" + ) + + +def test_validate_padding_mode_for_grid_rejects_unsafe_circular_padding(): + with pytest.raises(ValueError, match="circular padding"): + validate_padding_mode_for_grid( + (4, 4), + kernel_size=11, + padding_mode="circular", + ) + + +def test_validate_padding_mode_for_grid_accepts_safe_padding_modes(): + validate_padding_mode_for_grid((1, 1), kernel_size=3, padding_mode="zeros") + validate_padding_mode_for_grid( + (1, 1), kernel_size=3, padding_mode="replicate" + ) + validate_padding_mode_for_grid( + (2, 2), kernel_size=3, padding_mode="reflect" + ) + validate_padding_mode_for_grid( + (1, 1), kernel_size=3, padding_mode="circular" + ) + + +def test_squeeze_excitation_preserves_shape(): + x = torch.randn(2, 8, 5, 4) + + y = SqueezeExcitation2d(channels=8, reduction=4)(x) + + assert y.shape == x.shape + + +def test_squeeze_excitation_rejects_invalid_args(): + with pytest.raises(ValueError, match="channels"): + SqueezeExcitation2d(channels=0) + + with pytest.raises(ValueError, match="reduction"): + SqueezeExcitation2d(channels=8, reduction=0) + + +def test_film2d_preserves_shape(): + x = torch.randn(2, 8, 5, 4) + context = torch.randn(2, 3) + + y = FiLM2d(context_dim=3, channels=8)(x, context) + + assert y.shape == x.shape + + +def test_film2d_applies_channel_affine_conditioning(): + x = torch.ones(2, 2, 3, 4) + context = torch.ones(2, 3) + film = FiLM2d(context_dim=3, channels=2) + film.proj.weight.data.zero_() + film.proj.bias.data = torch.tensor([1.0, 2.0, 3.0, 4.0]) + + y = film(x, context) + + expected = torch.empty_like(x) + expected[:, 0] = 5.0 + expected[:, 1] = 7.0 + assert torch.equal(y, expected) + + +def test_film2d_rejects_shape_mismatch(): + x = torch.ones(2, 2, 3, 4) + film = FiLM2d(context_dim=3, channels=2) + + with pytest.raises(ValueError, match="context must have shape"): + film(x, torch.ones(2, 3, 1)) + + with pytest.raises(ValueError, match="same batch size"): + film(x, torch.ones(3, 3)) + + with pytest.raises(ValueError, match="feature dimension"): + film(x, torch.ones(2, 4)) + + with pytest.raises(ValueError, match="channel dimension"): + film(torch.ones(2, 3, 3, 4), torch.ones(2, 3)) + + +def test_reshrrr_block_preserves_shape_without_context(): + x = torch.randn(2, 8, 5, 4) + + y = ResHRRRBlock(channels=8, reduction=4)(x) + + assert y.shape == x.shape + + +def test_reshrrr_block_preserves_shape_with_context(): + x = torch.randn(2, 8, 5, 4) + context = torch.randn(2, 3) + + y = ResHRRRBlock(channels=8, reduction=4, context_dim=3)(x, context=context) + + assert y.shape == x.shape + + +def test_reshrrr_block_requires_context_when_film_enabled(): + x = torch.randn(2, 8, 5, 4) + block = ResHRRRBlock(channels=8, reduction=4, context_dim=3) + + with pytest.raises(ValueError, match="context is required"): + block(x) + + +def test_reshrrr_block_rejects_even_kernel_size(): + with pytest.raises(ValueError, match="kernel_size"): + ResHRRRBlock(channels=8, kernel_size=2) + + +def test_reshrrr_block_rejects_non_positive_args(): + with pytest.raises(ValueError, match="channels"): + ResHRRRBlock(channels=0) + + with pytest.raises(ValueError, match="kernel_size"): + ResHRRRBlock(channels=8, kernel_size=0) + + +def test_reshrrr_backbone_returns_output_channels_without_context(): + x = torch.randn(2, 5, 6, 4) + backbone = ResHRRRBackbone( + input_channels=5, + output_channels=3, + hidden_channels=8, + num_blocks=2, + reduction=4, + ) + + y = backbone(x) + + assert y.shape == (2, 3, 6, 4) + + +def test_reshrrr_backbone_returns_output_channels_with_context(): + x = torch.randn(2, 5, 6, 4) + context = torch.randn(2, 3) + backbone = ResHRRRBackbone( + input_channels=5, + output_channels=3, + hidden_channels=8, + num_blocks=2, + reduction=4, + context_dim=3, + ) + + y = backbone(x, context=context) + + assert y.shape == (2, 3, 6, 4) + + +def test_reshrrr_backbone_rejects_invalid_args(): + with pytest.raises(ValueError, match="input_channels"): + ResHRRRBackbone(input_channels=0, output_channels=3) + + with pytest.raises(ValueError, match="output_channels"): + ResHRRRBackbone(input_channels=5, output_channels=0) + + with pytest.raises(ValueError, match="hidden_channels"): + ResHRRRBackbone( + input_channels=5, + output_channels=3, + hidden_channels=0, + ) + + with pytest.raises(ValueError, match="num_blocks"): + ResHRRRBackbone(input_channels=5, output_channels=3, num_blocks=0) diff --git a/tests/test_cnn_predictor.py b/tests/test_cnn_predictor.py new file mode 100644 index 00000000..e16637f7 --- /dev/null +++ b/tests/test_cnn_predictor.py @@ -0,0 +1,173 @@ +# Third-party +import pytest +import torch + +# First-party +from neural_lam.models import ARForecaster, CNNPredictor +from tests.dummy_datastore import DummyDatastore + + +class NoStaticDummyDatastore(DummyDatastore): + """Dummy datastore variant without static features.""" + + def get_dataarray(self, category, split, standardize=False): + if category == "static": + return None + return super().get_dataarray(category, split, standardize=standardize) + + +def _make_predictor(datastore, output_std=False, cnn_film=False): + return CNNPredictor( + datastore=datastore, + cnn_channels=8, + cnn_blocks=2, + cnn_se_reduction=4, + cnn_film=cnn_film, + num_past_forcing_steps=1, + num_future_forcing_steps=1, + output_std=output_std, + ) + + +def _make_inputs(datastore, batch_size=2, pred_steps=1): + num_grid_nodes = datastore.num_grid_points + d_state = datastore.get_num_data_vars(category="state") + d_forcing = datastore.get_num_data_vars(category="forcing") * 3 + + prev_state = torch.randn(batch_size, num_grid_nodes, d_state) + prev_prev_state = torch.randn(batch_size, num_grid_nodes, d_state) + forcing = torch.randn(batch_size, num_grid_nodes, d_forcing) + + init_states = torch.stack((prev_prev_state, prev_state), dim=1) + forcing_features = torch.randn( + batch_size, + pred_steps, + num_grid_nodes, + d_forcing, + ) + boundary_states = torch.randn( + batch_size, + pred_steps, + num_grid_nodes, + d_state, + ) + + return ( + prev_state, + prev_prev_state, + forcing, + init_states, + forcing_features, + boundary_states, + ) + + +def test_cnn_predictor_forward_shape_without_std(): + datastore = DummyDatastore(n_grid_points=16) + predictor = _make_predictor(datastore, output_std=False) + prev_state, prev_prev_state, forcing, *_ = _make_inputs(datastore) + + prediction, pred_std = predictor(prev_state, prev_prev_state, forcing) + + assert prediction.shape == prev_state.shape + assert pred_std is None + + +def test_cnn_predictor_forward_shape_with_std(): + datastore = DummyDatastore(n_grid_points=16) + predictor = _make_predictor(datastore, output_std=True) + prev_state, prev_prev_state, forcing, *_ = _make_inputs(datastore) + + prediction, pred_std = predictor(prev_state, prev_prev_state, forcing) + + assert prediction.shape == prev_state.shape + assert pred_std.shape == prev_state.shape + assert torch.all(pred_std > 0) + + +def test_cnn_predictor_forward_shape_with_film(): + datastore = DummyDatastore(n_grid_points=16) + predictor = _make_predictor(datastore, cnn_film=True) + prev_state, prev_prev_state, forcing, *_ = _make_inputs(datastore) + + prediction, pred_std = predictor(prev_state, prev_prev_state, forcing) + + assert prediction.shape == prev_state.shape + assert pred_std is None + + +def test_cnn_predictor_global_film_context_averages_nodes(): + datastore = DummyDatastore(n_grid_points=16) + predictor = _make_predictor(datastore, cnn_film=True) + _, _, forcing, *_ = _make_inputs(datastore) + + context = predictor._get_global_forcing_context(forcing) + + assert torch.equal(context, forcing.mean(dim=1)) + + +def test_cnn_predictor_no_static_features(): + datastore = NoStaticDummyDatastore(n_grid_points=16) + predictor = _make_predictor(datastore) + prev_state, prev_prev_state, forcing, *_ = _make_inputs(datastore) + + prediction, pred_std = predictor(prev_state, prev_prev_state, forcing) + + assert predictor.grid_static_features.shape == ( + datastore.num_grid_points, + 0, + ) + assert prediction.shape == prev_state.shape + assert pred_std is None + + +def test_cnn_predictor_ar_forecaster_rollout_shape(): + datastore = DummyDatastore(n_grid_points=16) + predictor = _make_predictor(datastore) + forecaster = ARForecaster(predictor, datastore) + _, _, _, init_states, forcing_features, boundary_states = _make_inputs( + datastore, + pred_steps=3, + ) + + prediction, pred_std = forecaster( + init_states, + forcing_features, + boundary_states, + ) + + assert prediction.shape == boundary_states.shape + assert pred_std is None + + +def test_cnn_predictor_rejects_non_regular_datastore(): + with pytest.raises(TypeError, match="BaseRegularGridDatastore"): + CNNPredictor(datastore=object()) + + +def test_cnn_predictor_rejects_unsafe_reflect_padding(): + datastore = DummyDatastore(n_grid_points=1) + + with pytest.raises(ValueError, match="reflect padding"): + CNNPredictor( + datastore=datastore, + cnn_padding_mode="reflect", + ) + + +def test_cnn_predictor_rejects_wrong_forcing_size(): + datastore = DummyDatastore(n_grid_points=16) + predictor = _make_predictor(datastore) + prev_state, prev_prev_state, forcing, *_ = _make_inputs(datastore) + + with pytest.raises(ValueError, match="forcing feature dimension"): + predictor(prev_state, prev_prev_state, forcing[..., :-1]) + + +def test_cnn_predictor_rejects_wrong_state_size(): + datastore = DummyDatastore(n_grid_points=16) + predictor = _make_predictor(datastore) + prev_state, prev_prev_state, forcing, *_ = _make_inputs(datastore) + + with pytest.raises(ValueError, match="state feature dimension"): + predictor(prev_state[..., :-1], prev_prev_state[..., :-1], forcing) diff --git a/tests/test_train_model_cnn.py b/tests/test_train_model_cnn.py new file mode 100644 index 00000000..379b3f51 --- /dev/null +++ b/tests/test_train_model_cnn.py @@ -0,0 +1,239 @@ +# Standard library +from argparse import Namespace +from types import SimpleNamespace + +# Third-party +import pytest +import pytorch_lightning as pl +import torch + +# First-party +from neural_lam import config as nlconfig +from neural_lam.models import ( + MODELS, + ARForecaster, + CNNPredictor, + ForecasterModule, +) +from neural_lam.train_model import ( + build_predictor, + get_predictor_kwargs, + load_forecaster_module_from_checkpoint, + main, +) +from neural_lam.weather_dataset import WeatherDataModule +from tests.dummy_datastore import DummyDatastore + + +def _config(): + return SimpleNamespace( + training=SimpleNamespace( + output_clamping=SimpleNamespace(lower={}, upper={}) + ) + ) + + +def _neural_lam_config(): + return nlconfig.NeuralLAMConfig( + datastore=nlconfig.DatastoreSelection( + kind="mdp", + config_path=".", + ) + ) + + +def _cnn_args(**overrides): + args = Namespace( + model="cnn_predictor", + num_past_forcing_steps=1, + num_future_forcing_steps=1, + output_std=False, + cnn_channels=8, + cnn_blocks=2, + cnn_kernel_size=3, + cnn_se_reduction=4, + cnn_film=False, + cnn_padding_mode="zeros", + loss="mse", + lr=1e-3, + restore_opt=False, + n_example_pred=1, + create_gif=False, + val_steps_to_log=[1], + metrics_watch=[], + var_leads_metrics_watch={}, + ) + for key, value in overrides.items(): + setattr(args, key, value) + return args + + +def _make_cnn_forecaster_module(datastore, config, args): + predictor = build_predictor(args, config, datastore) + forecaster = ARForecaster(predictor, datastore) + return ForecasterModule( + forecaster=forecaster, + config=config, + datastore=datastore, + loss=args.loss, + lr=args.lr, + restore_opt=args.restore_opt, + n_example_pred=args.n_example_pred, + val_steps_to_log=args.val_steps_to_log, + metrics_watch=args.metrics_watch, + var_leads_metrics_watch=args.var_leads_metrics_watch, + args=args, + ) + + +def test_cnn_predictor_registered_in_models(): + assert MODELS["cnn_predictor"] is CNNPredictor + + +def test_cnn_predictor_kwargs_do_not_include_graph_options(): + kwargs = get_predictor_kwargs(_cnn_args(cnn_film=True), _config()) + + assert kwargs["cnn_channels"] == 8 + assert kwargs["cnn_blocks"] == 2 + assert kwargs["cnn_film"] is True + assert "graph_name" not in kwargs + assert "hidden_dim" not in kwargs + assert "g2m_gnn_type" not in kwargs + + +def test_build_predictor_constructs_cnn_predictor(): + datastore = DummyDatastore(n_grid_points=16) + predictor = build_predictor( + _cnn_args(cnn_padding_mode="reflect"), + _config(), + datastore, + ) + + assert isinstance(predictor, CNNPredictor) + assert predictor.backbone.blocks[0].conv1.padding_mode == "reflect" + + +def test_train_model_cli_constructs_cnn_predictor(monkeypatch): + datastore = DummyDatastore(n_grid_points=16) + captured = {} + + def capture_forecaster_module_init(_self, **kwargs): + captured["predictor"] = kwargs["forecaster"].predictor + raise SystemExit(0) + + monkeypatch.setattr( + "neural_lam.train_model.load_config_and_datastore", + lambda config_path: (_config(), datastore), + ) + monkeypatch.setattr( + "neural_lam.train_model.WeatherDataModule", + lambda **kwargs: None, + ) + monkeypatch.setattr( + "neural_lam.train_model.ForecasterModule.__init__", + capture_forecaster_module_init, + ) + + with pytest.raises(SystemExit): + main( + [ + "--config_path", + "dummy.yaml", + "--model", + "cnn_predictor", + "--cnn_channels", + "6", + "--cnn_blocks", + "1", + "--cnn_se_reduction", + "2", + "--cnn_film", + "--cnn_padding_mode", + "reflect", + "--ar_steps_eval", + "1", + "--val_steps_to_log", + "1", + ] + ) + + predictor = captured["predictor"] + assert isinstance(predictor, CNNPredictor) + assert predictor.cnn_film is True + assert predictor.backbone.blocks[0].conv1.padding_mode == "reflect" + + +def test_load_cnn_predictor_from_checkpoint(tmp_path): + datastore = DummyDatastore(n_grid_points=16) + config = _neural_lam_config() + args = _cnn_args(cnn_padding_mode="reflect") + model = _make_cnn_forecaster_module(datastore, config, args) + + ckpt_path = tmp_path / "cnn_predictor.ckpt" + trainer = pl.Trainer( + max_epochs=1, + accelerator="cpu", + logger=False, + enable_checkpointing=False, + ) + trainer.strategy.connect(model) + trainer.save_checkpoint(ckpt_path) + + loaded_model = load_forecaster_module_from_checkpoint( + ckpt_path, + config, + datastore, + ) + + assert isinstance(loaded_model.forecaster.predictor, CNNPredictor) + assert ( + loaded_model.forecaster.predictor.backbone.blocks[0].conv1.padding_mode + == "reflect" + ) + + batch_size = 2 + num_grid_nodes = datastore.num_grid_points + d_state = datastore.get_num_data_vars(category="state") + d_forcing = datastore.get_num_data_vars(category="forcing") * 3 + init_states = torch.randn(batch_size, 2, num_grid_nodes, d_state) + forcing = torch.randn(batch_size, 1, num_grid_nodes, d_forcing) + boundary = torch.randn(batch_size, 1, num_grid_nodes, d_state) + + with torch.no_grad(): + prediction_before, _ = model.forecaster(init_states, forcing, boundary) + prediction_after, _ = loaded_model.forecaster( + init_states, + forcing, + boundary, + ) + + assert torch.allclose(prediction_before, prediction_after) + + +def test_cnn_predictor_training_smoke(): + datastore = DummyDatastore(n_grid_points=16, n_timesteps=6) + config = _neural_lam_config() + args = _cnn_args(cnn_channels=4, cnn_blocks=1, cnn_se_reduction=2) + model = _make_cnn_forecaster_module(datastore, config, args) + data_module = WeatherDataModule( + datastore=datastore, + ar_steps_train=1, + ar_steps_eval=1, + batch_size=2, + num_workers=0, + num_past_forcing_steps=args.num_past_forcing_steps, + num_future_forcing_steps=args.num_future_forcing_steps, + ) + trainer = pl.Trainer( + max_epochs=1, + accelerator="cpu", + devices=1, + logger=False, + enable_checkpointing=False, + enable_model_summary=False, + limit_train_batches=1, + limit_val_batches=1, + num_sanity_val_steps=0, + ) + + trainer.fit(model=model, datamodule=data_module)