neural_lam.models.ar_model

neural_lam.models.ar_model#

Auto-regressive LightningModule implementations for Neural-LAM.

Module Contents#

class neural_lam.models.ar_model.ARModel(args, config: neural_lam.config.NeuralLAMConfig, datastore: neural_lam.datastore.BaseDatastore)#

Bases: pytorch_lightning.LightningModule

Generic auto-regressive weather model. Abstract class that can be extended.

Parameters:

  • args (argparse.Namespace) – Parsed training arguments controlling rollout length, loss, etc.

  • config (NeuralLAMConfig) – Experiment configuration containing datastore/training settings.

  • datastore (BaseDatastore) – Datastore supplying state/forcing/static arrays.

aggregate_and_plot_metrics(metrics_dict, prefix)#

Aggregate metric tensors and create error-map visualisations.

Parameters:

  • metrics_dict (dict[str, list[torch.Tensor]]) – Mapping from metric name to per-batch tensors of evaluations.

  • prefix (str) – Prefix to use for logger keys.

all_gather_cat(tensor_to_gather)#

Gather tensors across ranks and concatenate along dim-0.

Parameters:

tensor_to_gather (torch.Tensor) –

Tensor distributed across K ranks.

  • Shape: (d1, d2, ...) per rank

Returns:

Concatenated tensor gathered from all ranks.

  • Shape: (K * d1, d2, ...)

Return type:

torch.Tensor

common_step(batch)#

Run a forward pass shared by train/val/test steps.

Parameters:

batch (tuple) –

Tuple of (init_states, target_states, forcing_features, batch_times) produced by WeatherDataset.

  • init_states: (B, 2, num_grid_nodes, d_features)

  • target_states: (B, pred_steps, num_grid_nodes, d_features)

  • forcing_features: (B, pred_steps, num_grid_nodes, d_forcing)

  • batch_times: (B, pred_steps) timestamps

Returns:

(prediction, target_states, pred_std, batch_times).

  • prediction: (B, pred_steps, num_grid_nodes, d_f)

  • target_states: (B, pred_steps, num_grid_nodes, d_f)

  • pred_std: (B, pred_steps, num_grid_nodes, d_f) or (d_f,)

  • batch_times: (B, pred_steps)

Return type:

tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

configure_optimizers()#

Construct the torch.optim.AdamW optimizer for training.

create_metric_log_dict(metric_tensor, prefix, metric_name)#

Assemble logging artefacts for a single metric tensor.

Parameters:

  • metric_tensor (torch.Tensor) –

    Metric values per time step and variable.

    • Shape: (pred_steps, d_f)

  • prefix (str) – Prefix used for logger keys (e.g., "val" or "test").

  • metric_name (str) – Human-readable metric name.

Returns:

Mapping from log keys to figures or scalar tensors.

Return type:

dict[str, object]

static expand_to_batch(x, batch_size)#

Broadcast a tensor by prepending a batch dimension.

Parameters:

  • x (torch.Tensor) – Tensor to expand.

  • batch_size (int) – Batch size to broadcast to.

Returns:

Tensor with a leading batch dimension added via expand.

Return type:

torch.Tensor

on_load_checkpoint(checkpoint)#

Perform any changes to state dict before loading checkpoint

on_test_epoch_end()#

Compute test metrics and make plots at the end of test epoch. Will gather stored tensors and perform plotting and logging on rank 0.

on_validation_epoch_end()#

Compute val metrics at the end of val epoch

plot_examples(batch, n_examples, split, prediction=None)#

Plot the first n_examples forecasts from batch.

Parameters:

  • batch (tuple) – Batch tuple produced by the dataloader.

  • n_examples (int) – Number of forecasts to visualise.

  • split (str) – Dataset split name used for metadata lookups.

  • prediction (torch.Tensor or None, optional) –

    Pre-computed predictions to plot. If None the method runs common_step() to obtain predictions.

    • Shape: (B, pred_steps, num_grid_nodes, d_f)

abstract predict_step(prev_state, prev_prev_state, forcing)#

Advance the state by one step using the prediction model.

Parameters:

  • prev_state (torch.Tensor) –

    Current state X_t.

    • Shape: (B, num_grid_nodes, feature_dim)

  • prev_prev_state (torch.Tensor) –

    Previous state X_{t-1}.

    • Shape: (B, num_grid_nodes, feature_dim)

  • forcing (torch.Tensor) –

    Forcing inputs applied at the prediction step.

    • Shape: (B, num_grid_nodes, forcing_dim)

Returns:

Tuple (new_state, pred_std) describing the next state and optional uncertainty estimate.

Return type:

tuple[torch.Tensor, torch.Tensor | None]

test_step(batch, batch_idx)#

Evaluate batch during testing and log diagnostics.

Parameters:

  • batch (tuple) – Batch sampled from the test dataloader.

  • batch_idx (int) – Index of the current batch.

training_step(batch)#

Execute a single optimization step on batch.

Parameters:

batch (tuple) – Batch sampled from the training dataloader.

unroll_prediction(init_states, forcing_features, true_states)#

Roll out predictions autoregressively over multiple time steps.

Parameters:

  • init_states (torch.Tensor) –

    Initial states providing X_{t-1} and X_t.

    • Shape: (B, 2, num_grid_nodes, d_f)

  • forcing_features (torch.Tensor) –

    Forcing inputs aligned with each rollout step.

    • Shape: (B, pred_steps, num_grid_nodes, d_static_f)

  • true_states (torch.Tensor) –

    Ground-truth states used for boundary replacement.

    • Shape: (B, pred_steps, num_grid_nodes, d_f)

Returns:

Tuple (prediction, pred_std).

  • prediction: (B, pred_steps, num_grid_nodes, d_f)

  • pred_std: (B, pred_steps, num_grid_nodes, d_f) or (d_f,) when a constant per-feature value is used

Return type:

tuple[torch.Tensor, torch.Tensor]

validation_step(batch, batch_idx)#

Evaluate batch during validation.

Parameters:

  • batch (tuple) – Batch sampled from the validation dataloader.

  • batch_idx (int) – Index of the current batch.

args#
feature_weights#
grid_dim#
property interior_mask_bool#

Boolean interior mask identifying non-boundary grid nodes.

Returns:

Boolean mask.

  • Shape: (N,)

Return type:

torch.Tensor

loss#
n_example_pred#
output_std#
plotted_examples = 0#
restore_opt#
spatial_loss_maps: List[Any] = []#
test_metrics: Dict[str, List]#
val_metrics: Dict[str, List]#
Contents