neural_lam.config

Configuration dataclasses and helpers for Neural-LAM experiments.

Exceptions

InvalidConfigError

Raised when the Neural-LAM configuration file is invalid or malformed.

Classes

ComputeConfig	Compute and accelerator configuration.
DataConfig	Data loading and forcing parameters.
DatastoreSelection	Configuration for selecting a datastore to use with neural-lam.
LoggingConfig	Experiment tracking and artifact logging configuration.
ManualStateFeatureWeighting	Configuration for weighting the state features in the loss function where
ModelConfig	Configuration for neural-lam model architecture.
NeuralLAMConfig	Configuration for the Neural-LAM model and training pipeline.
OutputClamping	Configuration for clamping the output of the model.
TrainRunConfig	Hyperparameters and settings for a training or evaluation run.
TrainingConfig	Configuration related to training neural-lam
UniformFeatureWeighting	Configuration for weighting the state features in the loss function where

Functions

load_config_and_datastore(→ tuple[NeuralLAMConfig, ...)

Load the neural-lam configuration and the datastores specified in the

Module Contents

exception neural_lam.config.InvalidConfigError

Bases: Exception

Raised when the Neural-LAM configuration file is invalid or malformed.

Initialize self. See help(type(self)) for accurate signature.

class neural_lam.config.ComputeConfig

Compute and accelerator configuration.

Variables:

• seed (int, default 42) – Random seed.

• num_nodes (int, default 1) – Number of cluster nodes.

• devices (int or str or list of int or list of str, default ["auto"]) – Hardware devices to use.

• precision (str or int, default 32) – Floating point precision.

devices: int | str | list[int] | list[str] = ['auto']

num_nodes: int = 1

precision: str | int = 32

seed: int = 42

class neural_lam.config.DataConfig

Data loading and forcing parameters.

Variables:

• num_past_forcing_steps (int, default 1) – Number of past forcing time steps as input.

• num_future_forcing_steps (int, default 1) – Number of future forcing time steps as input.

• num_past_boundary_steps (int, default 1) – Number of past boundary forcing time steps as input.

• num_future_boundary_steps (int, default 1) – Number of future boundary forcing time steps as input.

• num_workers (int, default 4) – Number of data loader workers.

• load_single_member (bool, default False) – Whether to load only ensemble member 0.

load_single_member: bool = False

num_future_boundary_steps: int = 1

num_future_forcing_steps: int = 1

num_past_boundary_steps: int = 1

num_past_forcing_steps: int = 1

num_workers: int = 4

class neural_lam.config.DatastoreSelection

Configuration for selecting a datastore to use with neural-lam.

Variables:

• kind (str) – The kind of datastore to use, currently mdp or npyfilesmeps are implemented.

• config_path (str) – The path to the configuration file for the selected datastore, this is assumed to be relative to the configuration file for neural-lam.

• n_boundary_points (int, optional) – Overrides MDPDatastore’s default boundary-mask width (30), only valid for kind: mdp. Set it to 0 when the boundary forcing comes from a separate datastore.

config_path: str

kind: str

n_boundary_points: int | None = None

class neural_lam.config.LoggingConfig

Experiment tracking and artifact logging configuration.

Variables:

• logger (str, default "wandb") – Logger backend name (“wandb” or “mlflow”).

• logger_project (str, default "neural_lam") – Project name for the logger.

• logger_run_name (str or None, default None) – Custom run name for the logger.

• runs_root (str, default "runs") – Root directory for run artifacts and checkpoints.

• wandb_id (str or None, default None) – W&B run ID for resuming experiments.

logger: str = 'wandb'

logger_project: str = 'neural_lam'

logger_run_name: str | None = None

runs_root: str = 'runs'

wandb_id: str | None = None

class neural_lam.config.ManualStateFeatureWeighting

Configuration for weighting the state features in the loss function where the weights are manually specified.

Variables:

weights (dict[str, float]) – Manual weights for the state features.

weights: dict[str, float]

class neural_lam.config.ModelConfig

Configuration for neural-lam model architecture.

Variables:

• model (str, default "graph_lam") – Model architecture name.

• graph (str, default "multiscale") – Name of the graph to load from the datastore.

• hidden_dim (int, default 64) – Dimensionality of hidden representations.

• hidden_layers (int, default 1) – Number of hidden layers in MLPs.

• processor_layers (int, default 4) – Number of GNN layers in processor GNN.

• mesh_aggr (str, default "sum") – Aggregation method for m2m GNN layers (“sum” or “mean”).

• output_std (bool, default False) – Whether models should output standard deviation per feature.

• g2m_gnn_type (str, default "InteractionNet") – GNN layer type for grid-to-mesh encoding.

• m2g_gnn_type (str, default "InteractionNet") – GNN layer type for mesh-to-grid decoding.

• mesh_up_gnn_type (str, default "InteractionNet") – GNN layer type for upward mesh message passing.

• mesh_down_gnn_type (str, default "InteractionNet") – GNN layer type for downward mesh message passing.

g2m_gnn_type: str = 'InteractionNet'

graph: str = 'multiscale'

hidden_dim: int = 64

hidden_layers: int = 1

m2g_gnn_type: str = 'InteractionNet'

mesh_aggr: str = 'sum'

mesh_down_gnn_type: str = 'InteractionNet'

mesh_up_gnn_type: str = 'InteractionNet'

model: str = 'graph_lam'

output_std: bool = False

processor_layers: int = 4

class neural_lam.config.NeuralLAMConfig

Bases: dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard

Configuration for the Neural-LAM model and training pipeline.

Loads and stores all settings needed to run Neural-LAM, including datastore selection and training hyperparameters. Serialisation and deserialisation from YAML/JSON is handled via dataclass_wizard.

Variables:

• datastores (dict[str, DatastoreSelection]) – Mapping from a user-chosen datastore name to its selection config. A datastore’s role follows from the categories it provides: one with state data is the interior (input and output), one without is input-only (boundary forcing). load_config_and_datastore() requires exactly one interior and at most one boundary.

• training (TrainingConfig) – Configuration for training the model, including loss function and feature-weighting strategy. Defaults to TrainingConfig().

datastores: dict[str, DatastoreSelection]

training: TrainingConfig

class neural_lam.config.OutputClamping

Configuration for clamping the output of the model.

Variables:

• lower (dict[str, float]) – The minimum value to clamp each output feature to.

• upper (dict[str, float]) – The maximum value to clamp each output feature to.

lower: dict[str, float]

upper: dict[str, float]

class neural_lam.config.TrainRunConfig

Hyperparameters and settings for a training or evaluation run.

Variables:

• epochs (int, default 200) – Number of training epochs.

• batch_size (int, default 4) – Batch size.

• ar_steps_train (int, default 1) – Autoregressive rollout steps during training.

• ar_steps_eval (int, default 10) – Autoregressive rollout steps during evaluation.

• loss (str, default "wmse") – Loss function to use (“mse” or “wmse”).

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

• val_interval (int, default 1) – Validation epoch frequency.

• num_sanity_val_steps (int, default 2) – Number of validation batches to run prior to training.

• val_steps_to_log (list of int, default [1, 2, 3, 5, 10]) – Rollout steps to log during validation.

• train_steps_to_log (list of int, default []) – Rollout steps to log during training.

• metrics_watch (list of str, default []) – List of watched metrics.

• var_leads_metrics_watch (dict of {int: list of int}, default {}) – Mapping from variable index to list of watched rollout steps.

• load (str or None, default None) – Checkpoint path to load.

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

• eval (str or None, default None) – Evaluation split name (“val”, “test”) or None if training.

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

• create_gif (bool, default False) – Whether to create animation GIFs during evaluation.

ar_steps_eval: int = 10

ar_steps_train: int = 1

batch_size: int = 4

create_gif: bool = False

epochs: int = 200

eval: str | None = None

load: str | None = None

loss: str = 'wmse'

lr: float = 0.001

metrics_watch: list[str] = []

n_example_pred: int = 1

num_sanity_val_steps: int = 2

restore_opt: bool = False

train_steps_to_log: list[int] = []

val_interval: int = 1

val_steps_to_log: list[int] = [1, 2, 3, 5, 10]

var_leads_metrics_watch: dict[int, list[int]]

class neural_lam.config.TrainingConfig

Configuration related to training neural-lam

Variables:

• state_feature_weighting – ManualStateFeatureWeighting | UniformFeatureWeighting The method to use for weighting the state features in the loss function. Defaults to uniform weighting (UniformFeatureWeighting, i.e. all features are weighted equally).

• output_clamping (OutputClamping) – Per-feature lower / upper clamping bounds applied to the model output. Defaults to an empty OutputClamping (no clamping).

output_clamping: OutputClamping

state_feature_weighting: ManualStateFeatureWeighting | UniformFeatureWeighting

class neural_lam.config.UniformFeatureWeighting

Configuration for weighting the state features in the loss function where all state features are weighted equally.

neural_lam.config.load_config_and_datastore(config_path: str) → tuple[NeuralLAMConfig, neural_lam.datastore.MDPDatastore | neural_lam.datastore.NpyFilesDatastoreMEPS, neural_lam.datastore.MDPDatastore | neural_lam.datastore.NpyFilesDatastoreMEPS | None]

Load the neural-lam configuration and the datastores specified in the configuration.

Parameters:

config_path (str) – Path to the Neural-LAM configuration file.

Returns:

The configuration, the interior datastore (the one with state data), and the boundary datastore (the one without, or None).

Return type:

tuple[NeuralLAMConfig, datastore, datastore_boundary]

Raises:

InvalidConfigError – If not exactly one datastore provides state data, or if more than one omits it.