neural_lam.train_model#

CLI entry point for training Neural-LAM models.

Classes#

`AdaptiveHelpFormatter`	`--help` formatter that scales the column width to the terminal.
`Run`	Handle containing output paths for a Neural-LAM run.

Functions#

`build_parser`(→ argparse.ArgumentParser)	Build the argument parser for training and evaluating models.
`build_predictor`(→ Any)	Instantiate a step predictor with the GNN kwargs its family accepts.
`fit`(→ Run)	Run training or evaluation using strongly-typed configuration dataclasses.
`load_forecaster_module_from_checkpoint`(...)	Reconstruct a ForecasterModule from a checkpoint without requiring the
`main`(→ None)	Main function for training and evaluating models.
`run`(→ Run)	Run the training or evaluation loop from parsed CLI arguments.

Module Contents#

class neural_lam.train_model.AdaptiveHelpFormatter(prog: str)[source]#

Bases: argparse.ArgumentDefaultsHelpFormatter

--help formatter that scales the column width to the terminal.

Pick a help-column width based on the current terminal size.

class neural_lam.train_model.Run[source]#

Handle containing output paths for a Neural-LAM run.

example_plots(example_i: int) → pathlib.Path[source]#: Frames and GIF of example example_i (1-based, create_gif).

property checkpoint_dir: pathlib.Path#: Directory holding min_val_loss.ckpt and last.ckpt.

checkpoint_path: pathlib.Path | None = None#

property plot_dir: pathlib.Path#: Directory evaluation plots are written to (logger save_dir).

property rmse_plot: pathlib.Path#: Lead time x variable RMSE heatmap written by test evaluation.

run_dir: pathlib.Path#

neural_lam.train_model.build_parser() → argparse.ArgumentParser[source]#

Build the argument parser for training and evaluating models.

Returns:: The parser.
Return type:: ArgumentParser

neural_lam.train_model.build_predictor(predictor_class: type, args: argparse.Namespace, config: neural_lam.config.NeuralLAMConfig, datastore: neural_lam.datastore.base.BaseDatastore) → Any[source]#

Instantiate a step predictor with the GNN kwargs its family accepts.

Hierarchical GNN kwargs are only passed to BaseHiGraphModel subclasses, gating on the class hierarchy so that future hierarchical models are covered without maintaining a model-name list. GNN type arguments fall back to InteractionNet for checkpoints saved before those CLI flags existed.

neural_lam.train_model.fit(model_config: neural_lam.config.ModelConfig, train_config: neural_lam.config.TrainRunConfig, data_config: neural_lam.config.DataConfig, compute_config: neural_lam.config.ComputeConfig, logging_config: neural_lam.config.LoggingConfig, *, config: neural_lam.config.NeuralLAMConfig | None = None, datastore: neural_lam.datastore.base.BaseDatastore | None = None, datastore_boundary: neural_lam.datastore.base.BaseDatastore | None = None, config_path: str | None = None) → Run[source]#

Run training or evaluation using strongly-typed configuration dataclasses.

Parameters:

model_config (ModelConfig) – Model architecture and graph configuration.
train_config (TrainRunConfig) – Training hyperparameters and evaluation settings.
data_config (DataConfig) – Data loading and forcing parameters.
compute_config (ComputeConfig) – Compute devices, seed, and precision settings.
logging_config (LoggingConfig) – Experiment tracking and logging settings.
config (NeuralLAMConfig or None, optional) – Loaded Neural-LAM configuration.
datastore (BaseDatastore or None, optional) – Initialized datastore instance.
datastore_boundary (BaseDatastore or None, optional) – Initialized boundary datastore, used together with config and datastore. Loaded from config_path otherwise.
config_path (str or None, optional) – Path to the configuration file, used if config or datastore is None.

Returns:

Object containing paths to run outputs and saved checkpoints.

Return type:

Run

neural_lam.train_model.load_forecaster_module_from_checkpoint(ckpt_path: str, config: neural_lam.config.NeuralLAMConfig, datastore: neural_lam.datastore.base.BaseDatastore, datastore_boundary: neural_lam.datastore.base.BaseDatastore | None = None) → neural_lam.models.ForecasterModule[source]#

Reconstruct a ForecasterModule from a checkpoint without requiring the caller to know the original architecture kwargs.

The checkpoint must have been saved with args in hyper_parameters (i.e. created via train_model.main), so that model class and architecture kwargs can be recovered automatically.

Parameters:

ckpt_path (str) – Path to the checkpoint.
config (NeuralLAMConfig) – Neural-LAM configuration.
datastore (BaseDatastore) – Interior datastore.
datastore_boundary (BaseDatastore, optional) – Boundary datastore, needed to restore the boundary standardization statistics of a checkpoint trained with one.

Returns:

The restored module.

Return type:

ForecasterModule

neural_lam.train_model.main(input_args: list[str] | None = None) → None[source]#: Main function for training and evaluating models.

neural_lam.train_model.run(args: argparse.Namespace, config: neural_lam.config.NeuralLAMConfig | None = None, datastore: neural_lam.datastore.base.BaseDatastore | None = None, datastore_boundary: neural_lam.datastore.base.BaseDatastore | None = None) → Run[source]#

Run the training or evaluation loop from parsed CLI arguments.

Parameters:

args (Namespace) – Parsed arguments from build_parser().
config (NeuralLAMConfig or None, optional) – Loaded configuration, loaded from args.config_path if None.
datastore (BaseDatastore or None, optional) – Interior datastore, used together with config.
datastore_boundary (BaseDatastore or None, optional) – Boundary datastore, used together with config.

Returns:

Paths to the run outputs and saved checkpoints.

Return type:

Run