diff --git a/neural_lam/__init__.py b/neural_lam/__init__.py index dc85ce5ca..6f6e404b2 100644 --- a/neural_lam/__init__.py +++ b/neural_lam/__init__.py @@ -9,6 +9,8 @@ import neural_lam.vis # Local +from .forecaster import Forecaster +from .step_predictor import StepPredictor from .weather_dataset import WeatherDataset try: diff --git a/neural_lam/forecaster.py b/neural_lam/forecaster.py new file mode 100644 index 000000000..92c71a56e --- /dev/null +++ b/neural_lam/forecaster.py @@ -0,0 +1,93 @@ +import abc +from typing import Callable, Optional, Tuple + +# Third-party +import torch +import torch.nn as nn + + +class Forecaster(nn.Module, abc.ABC): + """ + Abstract base class for full forecast producers. + + Maps (init_states, forcing, true_states) → (prediction, pred_std) over a + complete forecast window. Sits between ForecasterModule (Lightning + orchestration) and StepPredictor (single-step neural network). + + Responsibilities + ---------------- + * Define the forecasting strategy (AR, direct, ensemble, …). + * Compute the training loss via compute_loss. + * Logging, plotting, and optimizer setup belong in ForecasterModule. + + Shape convention + ---------------- + B : batch size + T_init : conditioning time steps (typically 2) + pred_steps: forecast steps in this batch + N : grid nodes (flat spatial dim) + d_f : state features + d_forcing : forcing features (window already flattened) + """ + + @abc.abstractmethod + def forward( + self, + init_states: torch.Tensor, + forcing_features: torch.Tensor, + true_states: torch.Tensor, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """ + Produce a full forecast. + + Parameters + ---------- + init_states : (B, T_init, N, d_f) — standardized + forcing_features : (B, pred_steps, N, d_forcing) + true_states : (B, pred_steps, N, d_f) — standardized; + used for boundary forcing at each step + + Returns + ------- + prediction : (B, pred_steps, N, d_f) — standardized forecast + pred_std : (B, pred_steps, N, d_f) or (d_f,) — std-dev; + a constant fallback is returned when the predictor + does not produce uncertainty estimates + """ + raise NotImplementedError( + f"{type(self).__name__} must implement forward()." + ) + + def compute_loss( + self, + prediction: torch.Tensor, + target: torch.Tensor, + pred_std: torch.Tensor, + loss_fn: Callable[..., torch.Tensor], + mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Compute a scalar training loss averaged over batch and pred steps. + + Parameters + ---------- + prediction : (B, pred_steps, N, d_f) + target : (B, pred_steps, N, d_f) + pred_std : (B, pred_steps, N, d_f) or (d_f,) + loss_fn : metric function from neural_lam.metrics + mask : (N,) boolean interior-node mask, or None + + Returns + ------- + torch.Tensor — scalar loss + """ + return torch.mean( + loss_fn( + prediction, + target, + pred_std, + mask=mask, + average_grid=True, + sum_vars=True, + ) + ) diff --git a/neural_lam/step_predictor.py b/neural_lam/step_predictor.py new file mode 100644 index 000000000..443e36f19 --- /dev/null +++ b/neural_lam/step_predictor.py @@ -0,0 +1,57 @@ +import abc +from typing import Optional, Tuple + +# Third-party +import torch +import torch.nn as nn + + +class StepPredictor(nn.Module, abc.ABC): + """ + Abstract base class for single-step state predictors. + + Maps (X_{t-1}, X_t, F_t) → X_{t+1}, corresponding to f̂ in Oskarsson + et al. Subclasses implement the neural network (GNN, CNN, ViT, …). + + Responsibilities + ---------------- + * Neural network forward pass only. + * Normalization, boundary masking, and AR unrolling belong in + ARForecaster, not here. + + Shape convention (symbols used throughout) + ------------------------------------------ + B : batch size + N : number of grid nodes (flat spatial dim) + d_f : number of state features + d_forcing: forcing features (window already flattened into last dim) + """ + + # Set to True in subclasses that also output a predicted std-dev. + output_std: bool = False + + @abc.abstractmethod + 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. + + Parameters + ---------- + prev_state : (B, N, d_f) — X_t, standardized + prev_prev_state : (B, N, d_f) — X_{t-1}, standardized + forcing : (B, N, d_forcing) — forcing for this step + + Returns + ------- + new_state : (B, N, d_f) — predicted X_{t+1}, standardized + pred_std : (B, N, d_f) or None — aleatoric std, None if + output_std is False + """ + raise NotImplementedError( + f"{type(self).__name__} must implement forward()." + ) diff --git a/tests/test_imports.py b/tests/test_imports.py index 488deba70..eb20da3f1 100644 --- a/tests/test_imports.py +++ b/tests/test_imports.py @@ -1,9 +1,98 @@ # First-party import neural_lam import neural_lam.vis - +from neural_lam.forecaster import Forecaster +from neural_lam.step_predictor import StepPredictor +import torch.nn as nn +import torch +from neural_lam import metrics def test_import(): assert neural_lam is not None assert neural_lam.vis is not None assert neural_lam.__version__ is not None + + +def test_step_predictor_interface(): + """StepPredictor is importable, is an nn.Module subclass, and cannot + be instantiated directly (it is abstract).""" + + assert issubclass(StepPredictor, nn.Module) + assert StepPredictor.output_std is False + + # Cannot instantiate the abstract base class + try: + StepPredictor() + assert False, "Expected TypeError when instantiating abstract class" + except TypeError: + pass # expected + + +def test_forecaster_interface(): + """Forecaster is importable, is an nn.Module subclass, cannot be + instantiated directly, and exposes the concrete compute_loss method.""" + + assert issubclass(Forecaster, nn.Module) + assert callable(Forecaster.compute_loss) + + # Cannot instantiate the abstract base class + try: + Forecaster() # noqa: F841 + assert False, "Expected TypeError when instantiating abstract class" + except TypeError: + pass # expected + + +def test_forecaster_compute_loss_concrete(): + """compute_loss works end-to-end with a trivial concrete Forecaster.""" + + class _TrivialForecaster(Forecaster): + """Minimal concrete subclass to exercise compute_loss.""" + + def forward(self, init_states, forcing_features, true_states): + # Just return the true states as the prediction (zero loss) + return true_states, torch.ones_like(true_states) + + forecaster = _TrivialForecaster() + + B, T, N, d_f = 2, 3, 10, 4 + prediction = torch.zeros(B, T, N, d_f) + target = torch.zeros(B, T, N, d_f) + pred_std = torch.ones(d_f) # constant per-variable std + + loss = forecaster.compute_loss( + prediction=prediction, + target=target, + pred_std=pred_std, + loss_fn=metrics.mse, + mask=None, + ) + + assert loss.shape == torch.Size([]), "Expected scalar loss" + assert loss.item() == 0.0, "Expected zero loss for identical pred/target" + + +def test_step_predictor_concrete_forward(): + """A minimal concrete StepPredictor can be instantiated and forward + raises NotImplementedError only if the subclass forgets to override it. + A proper override should be callable normally.""" + + class _PassThroughPredictor(StepPredictor): + """Always returns prev_state unchanged (identity predictor).""" + + def forward(self, prev_state, prev_prev_state, forcing): + return prev_state, None + + predictor = _PassThroughPredictor() + assert isinstance(predictor, nn.Module) + assert predictor.output_std is False + + B, N, d_f, d_forcing = 2, 10, 4, 6 + prev_state = torch.zeros(B, N, d_f) + prev_prev_state = torch.zeros(B, N, d_f) + forcing = torch.zeros(B, N, d_forcing) + + new_state, pred_std = predictor(prev_state, prev_prev_state, forcing) + + assert new_state.shape == (B, N, d_f) + assert pred_std is None