neural_lam.models.forecasters.base

neural_lam.models.forecasters.base#

Base class for forecasters.

Classes#

Forecaster

Generic forecaster capable of mapping from a set of initial states,

Module Contents#

class neural_lam.models.forecasters.base.Forecaster(*args: Any, **kwargs: Any)[source]#

Bases: torch.nn.Module, abc.ABC

Generic forecaster capable of mapping from a set of initial states, forcing and forces and previous states into a full forecast of the requested length.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

abstract forward(init_states: torch.Tensor, forcing_features: torch.Tensor, boundary_states: torch.Tensor) → tuple[torch.Tensor, torch.Tensor | None][source]#

Produce a forecast of length pred_steps from two initial states, the per-step forcing features, and the per-step true boundary states.

Parameters:

  • init_states (torch.Tensor) – Shape (B, 2, num_grid_nodes, num_state_vars). The two initial states [X_{t-1}, X_t] used to start the forecast from. Dims: B is batch size, 2 is the time index ([X_{t-1}, X_t]), num_grid_nodes is the number of spatial nodes, and num_state_vars is the state feature dimension.

  • forcing_features (torch.Tensor) – Shape (B, pred_steps, num_grid_nodes, num_forcing_vars). External forcings provided at each predicted step. Dims: B is batch size, pred_steps is the autoregressive rollout length, num_grid_nodes is the number of spatial nodes, and num_forcing_vars is the forcing feature dimension (already concatenated past/current/future windows).

  • boundary_states (torch.Tensor) – Shape (B, pred_steps, num_grid_nodes, num_state_vars). True state values used ONLY to overwrite boundary nodes at each AR step; interior predictions must not depend on boundary_states in any other way. Dims: B is batch size, pred_steps is the rollout length, num_grid_nodes is the number of spatial nodes, and num_state_vars is the state feature dimension. This is a temporary mechanism that mirrors the pre-refactor ARModel behavior; it will be replaced by a dedicated boundary-forcing input in #138 (training on interior + boundary datastore), at which point this parameter will be removed.

Returns:

  • prediction (torch.Tensor) – Shape (B, pred_steps, num_grid_nodes, num_state_vars). Forecast of state at each predicted step. Dims: same as boundary_states.

  • pred_std (torch.Tensor or None) – Shape (B, pred_steps, num_grid_nodes, num_state_vars) when predicts_std is True, otherwise None. Per-feature predicted standard deviation; when None, the constant per-variable std is substituted upstream by ForecasterModule. Dims: same as prediction.

property predicts_std: bool#
Abstractmethod:

Whether this forecaster outputs a predicted standard deviation.

Returns:

True if the forecaster predicts standard deviation, False otherwise.

Return type:

bool