add support for Step Predictor and Forecaster abstract Interfaces.

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:

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,
+            )
+        )

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()."
+        )

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