neural_lam.models.module

Lightning module handling training, validation and testing loops.

Attributes

Batch

Classes

ForecasterModule

Lightning module handling training, validation and testing loops.

Module Contents

class neural_lam.models.module.ForecasterModule(forecaster: neural_lam.models.forecasters.base.Forecaster, config: neural_lam.config.NeuralLAMConfig, datastore: neural_lam.datastore.base.BaseDatastore, datastore_boundary: neural_lam.datastore.base.BaseDatastore | None = None, loss: str = 'wmse', lr: float = 0.001, restore_opt: bool = False, n_example_pred: int = 1, create_gif: bool = False, val_steps_to_log: list[int] | None = None, train_steps_to_log: list[int] | None = None, metrics_watch: list[str] | None = None, var_leads_metrics_watch: dict[int, list[int]] | None = None, args: Any | None = None)

Bases: pytorch_lightning.LightningModule

Lightning module handling training, validation and testing loops. Wraps a Forecaster instance which performs the actual prediction.

Initialize the ForecasterModule.

Parameters:

• forecaster (Forecaster) – The forecaster model to use for predictions.

• config (NeuralLAMConfig) – Configuration object for the neural LAM model.

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

• datastore_boundary (BaseDatastore, optional) – Boundary forcing datastore, used for its standardization statistics.

• loss (str, default "wmse") – The loss function to use.

• lr (float, default 1e-3) – Learning rate for the optimizer.

• restore_opt (bool, default False) – Whether to restore optimizer state from checkpoint.

• n_example_pred (int, default 1) – Number of example predictions to plot during testing.

• create_gif (bool, default False) – Whether to create GIFs of example predictions.

• val_steps_to_log (list of int, optional) – Specific rollout steps to log during validation/testing.

• metrics_watch (list of str, optional) – List of metrics to watch and log specifically.

• var_leads_metrics_watch (dict of {int: list of int}, optional) – Mapping from variable index to a list of rollout steps to log individually for the configured metrics.

• args (argparse.Namespace, optional) – Pre-refactor ARModel checkpoint hyperparameters. When provided, attributes on args take precedence over the corresponding explicit kwargs (loss, lr, restore_opt, n_example_pred, create_gif, val_steps_to_log, metrics_watch, var_leads_metrics_watch) so legacy checkpoints round-trip through load_from_checkpoint correctly.

aggregate_and_plot_metrics(metrics_dict: dict[str, list[torch.Tensor]], prefix: str) → None

Aggregate metrics from all GPUs and plot them.

Parameters:

• metrics_dict (dict) – Dictionary of metrics lists.

• prefix (str) – Prefix for the log names.

all_gather_cat(tensor_to_gather: torch.Tensor) → torch.Tensor

Gather tensors from all GPUs and concatenate them.

Parameters:

tensor_to_gather (torch.Tensor) – The tensor to gather from all processes.

Returns:

The concatenated tensor from all processes.

Return type:

torch.Tensor

common_step(batch: Batch) → tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor]

Perform a common prediction step for training, validation, and testing.

Parameters:

batch (tuple) – The batch of data containing initial states, target states, forcing features, boundary forcing features, and batch times.

Returns:

A tuple containing prediction, target states, predicted standard deviation, and batch times.

Return type:

tuple

configure_optimizers() → torch.optim.Optimizer

Configure the optimizers and learning rate schedulers.

Returns:

The configured optimizer.

Return type:

torch.optim.Optimizer

create_metric_log_dict(metric_tensor: torch.Tensor, prefix: str, metric_name: str) → dict[str, Any]

Create a dictionary of metrics to log.

Parameters:

• metric_tensor (torch.Tensor) – The metric values.

• prefix (str) – Prefix for the log names (e.g., ‘val’, ‘test’).

• metric_name (str) – The name of the metric.

Returns:

Dictionary of logged metrics and figures.

Return type:

dict

on_after_batch_transfer(batch: Batch, dataloader_idx: int) → Batch

Standardize a batch on-device after transfer to the accelerator.

Lightning calls this for every train/val/test/predict batch. WeatherDataset returns unstandardized state, forcing and boundary; all three are normalized here so the work runs on the accelerator. Forcing and boundary share the same feature-major (feature, window) stacking, so per-feature mean/std are tiled once and cached.

Parameters:

• batch (Batch) – (init_states, target_states, forcing, boundary, target_times) as returned by the dataloader.

• dataloader_idx (int) – Index of the dataloader the batch comes from.

Returns:

The batch with state, forcing and boundary standardized.

Return type:

Batch

on_load_checkpoint(checkpoint: dict[str, Any]) → None

Perform actions when loading a checkpoint. Handles backward compatibility for older checkpoints.

Parameters:

checkpoint (dict) – The loaded checkpoint dictionary.

on_test_epoch_end() → None

Perform actions at the end of the test epoch. Aggregates and plots test metrics and spatial loss maps.

on_validation_epoch_end() → None

Perform actions at the end of the validation epoch. Aggregates and plots validation metrics.

plot_examples(batch: Batch, n_examples: int, split: str, prediction: torch.Tensor) → None

Plot example predictions.

Parameters:

• batch (tuple) – The batch of data.

• n_examples (int) – Number of examples to plot.

• split (str) – The data split.

• prediction (torch.Tensor) – The model predictions.

test_step(batch: Batch, batch_idx: int) → None

Perform a single test step.

Parameters:

• batch (tuple) – The batch of data.

• batch_idx (int) – The index of the batch.

training_step(batch: Batch) → torch.Tensor

Perform a single training step.

Parameters:

batch (tuple) – The batch of data.

Returns:

The computed loss for the training step.

Return type:

torch.Tensor

validation_step(batch: Batch, batch_idx: int) → None

Perform a single validation step.

Parameters:

• batch (tuple) – The batch of data.

• batch_idx (int) – The index of the batch.

boundary_mean: torch.Tensor | None

boundary_mean_tiled: torch.Tensor | None

boundary_std: torch.Tensor | None

boundary_std_tiled: torch.Tensor | None

create_gif = False

datastore

datastore_boundary = None

forcing_mean: torch.Tensor | None

forcing_mean_tiled: torch.Tensor | None

forcing_std: torch.Tensor | None

forcing_std_tiled: torch.Tensor | None

forecaster

interior_mask_bool: torch.Tensor

loss

matched_metrics: set

n_example_pred = 1

per_var_std: torch.Tensor | None

plotted_examples = 0

restore_opt = False

spatial_loss_maps: list[Any] = []

state_mean: torch.Tensor

state_std: torch.Tensor

test_metrics: dict[str, list]

val_metrics: dict[str, list]

neural_lam.models.module.Batch