neural_lam.train_model ====================== .. py:module:: neural_lam.train_model .. autoapi-nested-parse:: CLI entry point for training Neural-LAM models. Classes ------- .. autoapisummary:: neural_lam.train_model.AdaptiveHelpFormatter neural_lam.train_model.Run Functions --------- .. autoapisummary:: neural_lam.train_model.build_parser neural_lam.train_model.build_predictor neural_lam.train_model.fit neural_lam.train_model.load_forecaster_module_from_checkpoint neural_lam.train_model.main neural_lam.train_model.run Module Contents --------------- .. py:class:: AdaptiveHelpFormatter(prog: str) Bases: :py:obj:`argparse.ArgumentDefaultsHelpFormatter` ``--help`` formatter that scales the column width to the terminal. Pick a help-column width based on the current terminal size. .. py:class:: Run Handle containing output paths for a Neural-LAM run. .. py:method:: example_plots(example_i: int) -> pathlib.Path Frames and GIF of example ``example_i`` (1-based, create_gif). .. py:property:: checkpoint_dir :type: pathlib.Path Directory holding ``min_val_loss.ckpt`` and ``last.ckpt``. .. py:attribute:: checkpoint_path :type: pathlib.Path | None :value: None .. py:property:: plot_dir :type: pathlib.Path Directory evaluation plots are written to (logger ``save_dir``). .. py:property:: rmse_plot :type: pathlib.Path Lead time x variable RMSE heatmap written by test evaluation. .. py:attribute:: run_dir :type: pathlib.Path .. py:function:: build_parser() -> argparse.ArgumentParser Build the argument parser for training and evaluating models. :returns: The parser. :rtype: ArgumentParser .. py:function:: build_predictor(predictor_class: type, args: argparse.Namespace, config: neural_lam.config.NeuralLAMConfig, datastore: neural_lam.datastore.base.BaseDatastore) -> Any 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. .. py:function:: 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 Run training or evaluation using strongly-typed configuration dataclasses. :param model_config: Model architecture and graph configuration. :type model_config: ModelConfig :param train_config: Training hyperparameters and evaluation settings. :type train_config: TrainRunConfig :param data_config: Data loading and forcing parameters. :type data_config: DataConfig :param compute_config: Compute devices, seed, and precision settings. :type compute_config: ComputeConfig :param logging_config: Experiment tracking and logging settings. :type logging_config: LoggingConfig :param config: Loaded Neural-LAM configuration. :type config: NeuralLAMConfig or None, optional :param datastore: Initialized datastore instance. :type datastore: BaseDatastore or None, optional :param datastore_boundary: Initialized boundary datastore, used together with ``config`` and ``datastore``. Loaded from ``config_path`` otherwise. :type datastore_boundary: BaseDatastore or None, optional :param config_path: Path to the configuration file, used if config or datastore is None. :type config_path: str or None, optional :returns: Object containing paths to run outputs and saved checkpoints. :rtype: Run .. py:function:: 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 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. :param ckpt_path: Path to the checkpoint. :type ckpt_path: str :param config: Neural-LAM configuration. :type config: NeuralLAMConfig :param datastore: Interior datastore. :type datastore: BaseDatastore :param datastore_boundary: Boundary datastore, needed to restore the boundary standardization statistics of a checkpoint trained with one. :type datastore_boundary: BaseDatastore, optional :returns: The restored module. :rtype: ForecasterModule .. py:function:: main(input_args: list[str] | None = None) -> None Main function for training and evaluating models. .. py:function:: 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 Run the training or evaluation loop from parsed CLI arguments. :param args: Parsed arguments from :func:`build_parser`. :type args: Namespace :param config: Loaded configuration, loaded from ``args.config_path`` if ``None``. :type config: NeuralLAMConfig or None, optional :param datastore: Interior datastore, used together with ``config``. :type datastore: BaseDatastore or None, optional :param datastore_boundary: Boundary datastore, used together with ``config``. :type datastore_boundary: BaseDatastore or None, optional :returns: Paths to the run outputs and saved checkpoints. :rtype: Run