neural_lam.config ================= .. py:module:: neural_lam.config .. autoapi-nested-parse:: Configuration dataclasses and helpers for Neural-LAM experiments. Exceptions ---------- .. autoapisummary:: neural_lam.config.InvalidConfigError Classes ------- .. autoapisummary:: neural_lam.config.ComputeConfig neural_lam.config.DataConfig neural_lam.config.DatastoreSelection neural_lam.config.LoggingConfig neural_lam.config.ManualStateFeatureWeighting neural_lam.config.ModelConfig neural_lam.config.NeuralLAMConfig neural_lam.config.OutputClamping neural_lam.config.TrainRunConfig neural_lam.config.TrainingConfig neural_lam.config.UniformFeatureWeighting Functions --------- .. autoapisummary:: neural_lam.config.load_config_and_datastore Module Contents --------------- .. py:exception:: InvalidConfigError Bases: :py:obj:`Exception` Raised when the Neural-LAM configuration file is invalid or malformed. Initialize self. See help(type(self)) for accurate signature. .. py:class:: ComputeConfig Compute and accelerator configuration. :ivar seed: Random seed. :vartype seed: int, default 42 :ivar num_nodes: Number of cluster nodes. :vartype num_nodes: int, default 1 :ivar devices: Hardware devices to use. :vartype devices: int or str or list of int or list of str, default ["auto"] :ivar precision: Floating point precision. :vartype precision: str or int, default 32 .. py:attribute:: devices :type: int | str | list[int] | list[str] :value: ['auto'] .. py:attribute:: num_nodes :type: int :value: 1 .. py:attribute:: precision :type: str | int :value: 32 .. py:attribute:: seed :type: int :value: 42 .. py:class:: DataConfig Data loading and forcing parameters. :ivar num_past_forcing_steps: Number of past forcing time steps as input. :vartype num_past_forcing_steps: int, default 1 :ivar num_future_forcing_steps: Number of future forcing time steps as input. :vartype num_future_forcing_steps: int, default 1 :ivar num_past_boundary_steps: Number of past boundary forcing time steps as input. :vartype num_past_boundary_steps: int, default 1 :ivar num_future_boundary_steps: Number of future boundary forcing time steps as input. :vartype num_future_boundary_steps: int, default 1 :ivar num_workers: Number of data loader workers. :vartype num_workers: int, default 4 :ivar load_single_member: Whether to load only ensemble member 0. :vartype load_single_member: bool, default False .. py:attribute:: load_single_member :type: bool :value: False .. py:attribute:: num_future_boundary_steps :type: int :value: 1 .. py:attribute:: num_future_forcing_steps :type: int :value: 1 .. py:attribute:: num_past_boundary_steps :type: int :value: 1 .. py:attribute:: num_past_forcing_steps :type: int :value: 1 .. py:attribute:: num_workers :type: int :value: 4 .. py:class:: DatastoreSelection Configuration for selecting a datastore to use with neural-lam. :ivar kind: The kind of datastore to use, currently `mdp` or `npyfilesmeps` are implemented. :vartype kind: str :ivar config_path: The path to the configuration file for the selected datastore, this is assumed to be relative to the configuration file for neural-lam. :vartype config_path: str :ivar n_boundary_points: Overrides :class:`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. :vartype n_boundary_points: int, optional .. py:attribute:: config_path :type: str .. py:attribute:: kind :type: str .. py:attribute:: n_boundary_points :type: int | None :value: None .. py:class:: LoggingConfig Experiment tracking and artifact logging configuration. :ivar logger: Logger backend name ("wandb" or "mlflow"). :vartype logger: str, default "wandb" :ivar logger_project: Project name for the logger. :vartype logger_project: str, default "neural_lam" :ivar logger_run_name: Custom run name for the logger. :vartype logger_run_name: str or None, default None :ivar runs_root: Root directory for run artifacts and checkpoints. :vartype runs_root: str, default "runs" :ivar wandb_id: W&B run ID for resuming experiments. :vartype wandb_id: str or None, default None .. py:attribute:: logger :type: str :value: 'wandb' .. py:attribute:: logger_project :type: str :value: 'neural_lam' .. py:attribute:: logger_run_name :type: str | None :value: None .. py:attribute:: runs_root :type: str :value: 'runs' .. py:attribute:: wandb_id :type: str | None :value: None .. py:class:: ManualStateFeatureWeighting Configuration for weighting the state features in the loss function where the weights are manually specified. :ivar weights: Manual weights for the state features. :vartype weights: dict[str, float] .. py:attribute:: weights :type: dict[str, float] .. py:class:: ModelConfig Configuration for neural-lam model architecture. :ivar model: Model architecture name. :vartype model: str, default "graph_lam" :ivar graph: Name of the graph to load from the datastore. :vartype graph: str, default "multiscale" :ivar hidden_dim: Dimensionality of hidden representations. :vartype hidden_dim: int, default 64 :ivar hidden_layers: Number of hidden layers in MLPs. :vartype hidden_layers: int, default 1 :ivar processor_layers: Number of GNN layers in processor GNN. :vartype processor_layers: int, default 4 :ivar mesh_aggr: Aggregation method for m2m GNN layers ("sum" or "mean"). :vartype mesh_aggr: str, default "sum" :ivar output_std: Whether models should output standard deviation per feature. :vartype output_std: bool, default False :ivar g2m_gnn_type: GNN layer type for grid-to-mesh encoding. :vartype g2m_gnn_type: str, default "InteractionNet" :ivar m2g_gnn_type: GNN layer type for mesh-to-grid decoding. :vartype m2g_gnn_type: str, default "InteractionNet" :ivar mesh_up_gnn_type: GNN layer type for upward mesh message passing. :vartype mesh_up_gnn_type: str, default "InteractionNet" :ivar mesh_down_gnn_type: GNN layer type for downward mesh message passing. :vartype mesh_down_gnn_type: str, default "InteractionNet" .. py:attribute:: g2m_gnn_type :type: str :value: 'InteractionNet' .. py:attribute:: graph :type: str :value: 'multiscale' .. py:attribute:: hidden_dim :type: int :value: 64 .. py:attribute:: hidden_layers :type: int :value: 1 .. py:attribute:: m2g_gnn_type :type: str :value: 'InteractionNet' .. py:attribute:: mesh_aggr :type: str :value: 'sum' .. py:attribute:: mesh_down_gnn_type :type: str :value: 'InteractionNet' .. py:attribute:: mesh_up_gnn_type :type: str :value: 'InteractionNet' .. py:attribute:: model :type: str :value: 'graph_lam' .. py:attribute:: output_std :type: bool :value: False .. py:attribute:: processor_layers :type: int :value: 4 .. py:class:: NeuralLAMConfig Bases: :py:obj:`dataclass_wizard.JSONWizard`, :py:obj:`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``. :ivar datastores: 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). :func:`load_config_and_datastore` requires exactly one interior and at most one boundary. :vartype datastores: dict[str, DatastoreSelection] :ivar training: Configuration for training the model, including loss function and feature-weighting strategy. Defaults to ``TrainingConfig()``. :vartype training: TrainingConfig .. py:attribute:: datastores :type: dict[str, DatastoreSelection] .. py:attribute:: training :type: TrainingConfig .. py:class:: OutputClamping Configuration for clamping the output of the model. :ivar lower: The minimum value to clamp each output feature to. :vartype lower: dict[str, float] :ivar upper: The maximum value to clamp each output feature to. :vartype upper: dict[str, float] .. py:attribute:: lower :type: dict[str, float] .. py:attribute:: upper :type: dict[str, float] .. py:class:: TrainRunConfig Hyperparameters and settings for a training or evaluation run. :ivar epochs: Number of training epochs. :vartype epochs: int, default 200 :ivar batch_size: Batch size. :vartype batch_size: int, default 4 :ivar ar_steps_train: Autoregressive rollout steps during training. :vartype ar_steps_train: int, default 1 :ivar ar_steps_eval: Autoregressive rollout steps during evaluation. :vartype ar_steps_eval: int, default 10 :ivar loss: Loss function to use ("mse" or "wmse"). :vartype loss: str, default "wmse" :ivar lr: Learning rate. :vartype lr: float, default 1e-3 :ivar val_interval: Validation epoch frequency. :vartype val_interval: int, default 1 :ivar num_sanity_val_steps: Number of validation batches to run prior to training. :vartype num_sanity_val_steps: int, default 2 :ivar val_steps_to_log: Rollout steps to log during validation. :vartype val_steps_to_log: list of int, default [1, 2, 3, 5, 10] :ivar train_steps_to_log: Rollout steps to log during training. :vartype train_steps_to_log: list of int, default [] :ivar metrics_watch: List of watched metrics. :vartype metrics_watch: list of str, default [] :ivar var_leads_metrics_watch: Mapping from variable index to list of watched rollout steps. :vartype var_leads_metrics_watch: dict of {int: list of int}, default {} :ivar load: Checkpoint path to load. :vartype load: str or None, default None :ivar restore_opt: Whether to restore optimizer state from checkpoint. :vartype restore_opt: bool, default False :ivar eval: Evaluation split name ("val", "test") or None if training. :vartype eval: str or None, default None :ivar n_example_pred: Number of example predictions to plot during testing. :vartype n_example_pred: int, default 1 :ivar create_gif: Whether to create animation GIFs during evaluation. :vartype create_gif: bool, default False .. py:attribute:: ar_steps_eval :type: int :value: 10 .. py:attribute:: ar_steps_train :type: int :value: 1 .. py:attribute:: batch_size :type: int :value: 4 .. py:attribute:: create_gif :type: bool :value: False .. py:attribute:: epochs :type: int :value: 200 .. py:attribute:: eval :type: str | None :value: None .. py:attribute:: load :type: str | None :value: None .. py:attribute:: loss :type: str :value: 'wmse' .. py:attribute:: lr :type: float :value: 0.001 .. py:attribute:: metrics_watch :type: list[str] :value: [] .. py:attribute:: n_example_pred :type: int :value: 1 .. py:attribute:: num_sanity_val_steps :type: int :value: 2 .. py:attribute:: restore_opt :type: bool :value: False .. py:attribute:: train_steps_to_log :type: list[int] :value: [] .. py:attribute:: val_interval :type: int :value: 1 .. py:attribute:: val_steps_to_log :type: list[int] :value: [1, 2, 3, 5, 10] .. py:attribute:: var_leads_metrics_watch :type: dict[int, list[int]] .. py:class:: TrainingConfig Configuration related to training neural-lam :ivar 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). :ivar output_clamping: Per-feature lower / upper clamping bounds applied to the model output. Defaults to an empty ``OutputClamping`` (no clamping). :vartype output_clamping: OutputClamping .. py:attribute:: output_clamping :type: OutputClamping .. py:attribute:: state_feature_weighting :type: ManualStateFeatureWeighting | UniformFeatureWeighting .. py:class:: UniformFeatureWeighting Configuration for weighting the state features in the loss function where all state features are weighted equally. .. py:function:: 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. :param config_path: Path to the Neural-LAM configuration file. :type config_path: str :returns: The configuration, the interior datastore (the one with `state` data), and the boundary datastore (the one without, or None). :rtype: tuple[NeuralLAMConfig, datastore, datastore_boundary] :raises InvalidConfigError: If not exactly one datastore provides `state` data, or if more than one omits it.