neural_lam.models.module ======================== .. py:module:: neural_lam.models.module .. autoapi-nested-parse:: Lightning module handling training, validation and testing loops. Attributes ---------- .. autoapisummary:: neural_lam.models.module.Batch Classes ------- .. autoapisummary:: neural_lam.models.module.ForecasterModule Module Contents --------------- .. py:class:: 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: :py:obj:`pytorch_lightning.LightningModule` Lightning module handling training, validation and testing loops. Wraps a Forecaster instance which performs the actual prediction. Initialize the ForecasterModule. :param forecaster: The forecaster model to use for predictions. :type forecaster: Forecaster :param config: Configuration object for the neural LAM model. :type config: NeuralLAMConfig :param datastore: Datastore providing grid metadata and data access. :type datastore: BaseDatastore :param datastore_boundary: Boundary forcing datastore, used for its standardization statistics. :type datastore_boundary: BaseDatastore, optional :param loss: The loss function to use. :type loss: str, default "wmse" :param lr: Learning rate for the optimizer. :type lr: float, default 1e-3 :param restore_opt: Whether to restore optimizer state from checkpoint. :type restore_opt: bool, default False :param n_example_pred: Number of example predictions to plot during testing. :type n_example_pred: int, default 1 :param create_gif: Whether to create GIFs of example predictions. :type create_gif: bool, default False :param val_steps_to_log: Specific rollout steps to log during validation/testing. :type val_steps_to_log: list of int, optional :param metrics_watch: List of metrics to watch and log specifically. :type metrics_watch: list of str, optional :param var_leads_metrics_watch: Mapping from variable index to a list of rollout steps to log individually for the configured metrics. :type var_leads_metrics_watch: dict of {int: list of int}, optional :param args: 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. :type args: argparse.Namespace, optional .. py:method:: aggregate_and_plot_metrics(metrics_dict: dict[str, list[torch.Tensor]], prefix: str) -> None Aggregate metrics from all GPUs and plot them. :param metrics_dict: Dictionary of metrics lists. :type metrics_dict: dict :param prefix: Prefix for the log names. :type prefix: str .. py:method:: all_gather_cat(tensor_to_gather: torch.Tensor) -> torch.Tensor Gather tensors from all GPUs and concatenate them. :param tensor_to_gather: The tensor to gather from all processes. :type tensor_to_gather: torch.Tensor :returns: The concatenated tensor from all processes. :rtype: torch.Tensor .. py:method:: common_step(batch: Batch) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor] Perform a common prediction step for training, validation, and testing. :param batch: The batch of data containing initial states, target states, forcing features, boundary forcing features, and batch times. :type batch: tuple :returns: A tuple containing prediction, target states, predicted standard deviation, and batch times. :rtype: tuple .. py:method:: configure_optimizers() -> torch.optim.Optimizer Configure the optimizers and learning rate schedulers. :returns: The configured optimizer. :rtype: torch.optim.Optimizer .. py:method:: create_metric_log_dict(metric_tensor: torch.Tensor, prefix: str, metric_name: str) -> dict[str, Any] Create a dictionary of metrics to log. :param metric_tensor: The metric values. :type metric_tensor: torch.Tensor :param prefix: Prefix for the log names (e.g., 'val', 'test'). :type prefix: str :param metric_name: The name of the metric. :type metric_name: str :returns: Dictionary of logged metrics and figures. :rtype: dict .. py:method:: 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. :param batch: ``(init_states, target_states, forcing, boundary, target_times)`` as returned by the dataloader. :type batch: Batch :param dataloader_idx: Index of the dataloader the batch comes from. :type dataloader_idx: int :returns: The batch with state, forcing and boundary standardized. :rtype: Batch .. py:method:: on_load_checkpoint(checkpoint: dict[str, Any]) -> None Perform actions when loading a checkpoint. Handles backward compatibility for older checkpoints. :param checkpoint: The loaded checkpoint dictionary. :type checkpoint: dict .. py:method:: on_test_epoch_end() -> None Perform actions at the end of the test epoch. Aggregates and plots test metrics and spatial loss maps. .. py:method:: on_validation_epoch_end() -> None Perform actions at the end of the validation epoch. Aggregates and plots validation metrics. .. py:method:: plot_examples(batch: Batch, n_examples: int, split: str, prediction: torch.Tensor) -> None Plot example predictions. :param batch: The batch of data. :type batch: tuple :param n_examples: Number of examples to plot. :type n_examples: int :param split: The data split. :type split: str :param prediction: The model predictions. :type prediction: torch.Tensor .. py:method:: test_step(batch: Batch, batch_idx: int) -> None Perform a single test step. :param batch: The batch of data. :type batch: tuple :param batch_idx: The index of the batch. :type batch_idx: int .. py:method:: training_step(batch: Batch) -> torch.Tensor Perform a single training step. :param batch: The batch of data. :type batch: tuple :returns: The computed loss for the training step. :rtype: torch.Tensor .. py:method:: validation_step(batch: Batch, batch_idx: int) -> None Perform a single validation step. :param batch: The batch of data. :type batch: tuple :param batch_idx: The index of the batch. :type batch_idx: int .. py:attribute:: boundary_mean :type: torch.Tensor | None .. py:attribute:: boundary_mean_tiled :type: torch.Tensor | None .. py:attribute:: boundary_std :type: torch.Tensor | None .. py:attribute:: boundary_std_tiled :type: torch.Tensor | None .. py:attribute:: create_gif :value: False .. py:attribute:: datastore .. py:attribute:: datastore_boundary :value: None .. py:attribute:: forcing_mean :type: torch.Tensor | None .. py:attribute:: forcing_mean_tiled :type: torch.Tensor | None .. py:attribute:: forcing_std :type: torch.Tensor | None .. py:attribute:: forcing_std_tiled :type: torch.Tensor | None .. py:attribute:: forecaster .. py:attribute:: interior_mask_bool :type: torch.Tensor .. py:attribute:: loss .. py:attribute:: matched_metrics :type: set .. py:attribute:: n_example_pred :value: 1 .. py:attribute:: per_var_std :type: torch.Tensor | None .. py:attribute:: plotted_examples :value: 0 .. py:attribute:: restore_opt :value: False .. py:attribute:: spatial_loss_maps :type: list[Any] :value: [] .. py:attribute:: state_mean :type: torch.Tensor .. py:attribute:: state_std :type: torch.Tensor .. py:attribute:: test_metrics :type: dict[str, list] .. py:attribute:: val_metrics :type: dict[str, list] .. py:data:: Batch