neural_lam.models.forecasters.autoregressive

neural_lam.models.forecasters.autoregressive#

Forecaster that uses an auto-regressive strategy to unroll a forecast.

Classes#

ARForecaster

Subclass of Forecaster that uses an auto-regressive strategy to

Module Contents#

class neural_lam.models.forecasters.autoregressive.ARForecaster(predictor: neural_lam.models.step_predictors.base.StepPredictor, datastore: neural_lam.datastore.BaseDatastore)[source]#

Bases: neural_lam.models.forecasters.base.Forecaster

Subclass of Forecaster that uses an auto-regressive strategy to unroll a forecast. Makes use of a StepPredictor at each AR step.

Initialize the ARForecaster.

Parameters:

  • predictor (StepPredictor) – The predictor to use for each step.

  • datastore (BaseDatastore) – The datastore providing grid metadata and boundary masks.

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

Unroll the autoregressive model: at each step i call self.predictor to produce the next state, then overwrite boundary nodes with the true value from boundary_states[:, i].

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 rollout from. Dims: B is batch size, 2 initial time steps (X_{t-1}, X_t), num_grid_nodes is the number of spatial nodes, and num_state_vars is the number of state variables.

  • forcing_features (torch.Tensor) – Shape (B, pred_steps, num_grid_nodes, num_forcing_vars). Forcing features for each predicted step; pred_steps defines the rollout length. Dims: B is batch size, pred_steps is the number of predicted steps, num_grid_nodes is the number of spatial nodes, and num_forcing_vars is the number of forcing variables (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. The interior prediction at step i must not depend on boundary_states[:, i] in any other way. Dims: B is batch size, pred_steps is the number of predicted steps, num_grid_nodes is the number of spatial nodes, and num_state_vars is the state feature dimension.

Returns:

  • prediction (torch.Tensor) – Shape (B, pred_steps, num_grid_nodes, num_state_vars). Stacked per-step forecasts (with boundary overwritten by the true value). Dims: same as boundary_states.

  • pred_std (torch.Tensor or None) – Shape (B, pred_steps, num_grid_nodes, num_state_vars) when the wrapped predictor outputs an std, otherwise None (in which case ForecasterModule substitutes the constant per-variable std). Dims: same as prediction.

boundary_mask: torch.Tensor#
interior_mask: torch.Tensor#
predictor#
property predicts_std: bool#

Whether the forecaster predicts standard deviation.

Returns:

True if the forecaster predicts standard deviation, False otherwise.

Return type:

bool