neural_lam.models.step_predictors.base

Base class for step predictors.

Classes

StepPredictor

Abstract base class for step predictors mapping from the two previous

Module Contents

class neural_lam.models.step_predictors.base.StepPredictor(datastore: neural_lam.datastore.BaseDatastore, output_std: bool = False, output_clamping_lower: dict[str, float] | None = None, output_clamping_upper: dict[str, float] | None = None)

Bases: torch.nn.Module, abc.ABC

Abstract base class for step predictors mapping from the two previous time steps plus forcing into a prediction of the next state.

Initialize the StepPredictor.

Parameters:

• datastore (BaseDatastore) – The datastore providing grid metadata and data access.

• output_std (bool, default False) – Whether to output a predicted standard deviation.

• output_clamping_lower (dict, optional) – Lower clamping limits for state variables.

• output_clamping_upper (dict, optional) – Upper clamping limits for state variables.

expand_to_batch(x: torch.Tensor, batch_size: int) → torch.Tensor

Expand a shared node-feature tensor into a batch of copies.

Parameters:

• x (torch.Tensor) – Shape (N, d). Tensor to expand. Dims: N is the number of nodes and d is the feature dimension.

• batch_size (int) – Target batch size B.

Returns:

Shape (B, N, d). Batch-expanded view of x.

Return type:

torch.Tensor

abstract forward(prev_state: torch.Tensor, prev_prev_state: torch.Tensor, forcing: torch.Tensor) → tuple[torch.Tensor, torch.Tensor | None]

Advance the state by one step: (X_t, X_{t-1}, forcing_t) -> X_{t+1}.

Parameters:

• prev_state (torch.Tensor) – Shape (B, num_grid_nodes, num_state_vars). The current state X_t. Dims: B is batch size, num_grid_nodes is the number of spatial nodes, and num_state_vars is the number of state variables.

• prev_prev_state (torch.Tensor) – Shape (B, num_grid_nodes, num_state_vars). The previous state X_{t-1}, used as additional conditioning. Dims: same as prev_state.

• forcing (torch.Tensor) – Shape (B, num_grid_nodes, num_forcing_vars). External forcings for this step (already concatenated past/current/future windows). Dims: B is batch size, num_grid_nodes is the number of spatial nodes, and num_forcing_vars is the forcing feature dimension.

Returns:

• pred_state (torch.Tensor) – Shape (B, num_grid_nodes, num_state_vars). The predicted next state X_{t+1}. Dims: same as prev_state.

• pred_std (torch.Tensor or None) – Shape (B, num_grid_nodes, num_state_vars) when output_std is True, otherwise None. Per-feature predicted standard deviation. Dims: same as prev_state.

get_clamped_new_state(state_delta: torch.Tensor, prev_state: torch.Tensor) → torch.Tensor

Clamp the next-state prediction to its valid feature range.

Instead of the plain residual update X_{t+1} = X_t + delta, the clamped form is X_{t+1} = f(f^{-1}(X_t) + delta) so the model learns to output increments in the domain of the inverse clamping function.

Parameters:

• state_delta (torch.Tensor) – Shape (B, num_grid_nodes, num_state_vars). Raw predicted state increment (network output, already rescaled). Dims: B is batch size, num_grid_nodes is the number of spatial nodes, and num_state_vars is the number of state variables.

• prev_state (torch.Tensor) – Shape (B, num_grid_nodes, num_state_vars). Current state X_t used as the base for the clamped update. Dims: same as state_delta.

Returns:

Shape (B, num_grid_nodes, num_state_vars). Clamped next state. Dims: same as state_delta.

Return type:

torch.Tensor

prepare_clamping_params(datastore: neural_lam.datastore.BaseDatastore) → None

Prepare parameters for clamping predicted values to valid range.

Reads the per-feature lower/upper limits from self._output_clamping_lower and self._output_clamping_upper (set in __init__) and registers the buffers and clamping functions used by get_clamped_new_state.

Parameters:

datastore (BaseDatastore) – The datastore providing variable names.

clamp_lower_idx: torch.Tensor

clamp_lower_upper_idx: torch.Tensor

clamp_upper_idx: torch.Tensor

grid_static_features: torch.Tensor

output_std = False

property predicts_std: bool

Whether this predictor outputs a predicted standard deviation.

Returns:

True if the predictor predicts standard deviation, False otherwise.

Return type:

bool

sigmoid_lower_lims: torch.Tensor

sigmoid_upper_lims: torch.Tensor

softplus_lower_lims: torch.Tensor

softplus_upper_lims: torch.Tensor

state_mean: torch.Tensor

state_std: torch.Tensor

property trainable: bool

Whether the predictor has learnable parameters and can be trained.

Returns:

True if the predictor has trainable parameters, False otherwise.

Return type:

bool