From 2fe4079d7a1b2579f9d9e78b41ec4c6a79239886 Mon Sep 17 00:00:00 2001
From: Mohit-Lakra <mohitlakra2007777@gmail.com>
Date: Wed, 25 Feb 2026 00:37:01 +0530
Subject: [PATCH 01/16] docs : add NumPy-style docstrings to core model
 components

---
 neural_lam/__init__.py           |  1 +
 neural_lam/config.py             | 14 ++++--
 neural_lam/datastore/__init__.py |  1 +
 neural_lam/interaction_net.py    | 46 +++++++++++-------
 neural_lam/metrics.py            | 80 +++++++++++++++++++++-----------
 neural_lam/models/__init__.py    |  1 +
 6 files changed, 95 insertions(+), 48 deletions(-)

diff --git a/neural_lam/__init__.py b/neural_lam/__init__.py
index dc85ce5ca..b434957fe 100644
--- a/neural_lam/__init__.py
+++ b/neural_lam/__init__.py
@@ -1,3 +1,4 @@
+"""Neural-LAM: graph-based neural weather prediction models."""
 # Standard library
 import importlib.metadata
 
diff --git a/neural_lam/config.py b/neural_lam/config.py
index f4195ec36..7b00cb3b2 100644
--- a/neural_lam/config.py
+++ b/neural_lam/config.py
@@ -107,15 +107,20 @@ class TrainingConfig:
 @dataclasses.dataclass
 class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard):
     """
-    Dataclass for Neural-LAM configuration. This class is used to load and
-    store the configuration for using Neural-LAM.
+    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``.
 
     Attributes
     ----------
     datastore : DatastoreSelection
-        The configuration for the datastore to use.
+        Configuration specifying which datastore backend to use and its
+        associated settings.
     training : TrainingConfig
-        The configuration for training the model.
+        Configuration for training the model, including loss function and
+        feature-weighting strategy. Defaults to ``TrainingConfig()``.
     """
 
     datastore: DatastoreSelection
@@ -150,6 +155,7 @@ class _(dataclass_wizard.JSONWizard.Meta):
 
 
 class InvalidConfigError(Exception):
+    """Raised when the Neural-LAM configuration file is invalid or malformed."""
     pass
 
 
diff --git a/neural_lam/datastore/__init__.py b/neural_lam/datastore/__init__.py
index dead77135..29876a601 100644
--- a/neural_lam/datastore/__init__.py
+++ b/neural_lam/datastore/__init__.py
@@ -1,3 +1,4 @@
+"""Datastore backends for loading and serving weather model data."""
 # Local
 from .base import BaseDatastore  # noqa
 from .mdp import MDPDatastore  # noqa
diff --git a/neural_lam/interaction_net.py b/neural_lam/interaction_net.py
index 2f45b03fa..715b94f56 100644
--- a/neural_lam/interaction_net.py
+++ b/neural_lam/interaction_net.py
@@ -27,23 +27,37 @@ def __init__(
         aggr_chunk_sizes=None,
         aggr="sum",
     ):
+
         """
-        Create a new InteractionNet
-
-        edge_index: (2,M), Edges in pyg format
-        input_dim: Dimensionality of input representations,
-            for both nodes and edges
-        update_edges: If new edge representations should be computed
-            and returned
-        hidden_layers: Number of hidden layers in MLPs
-        hidden_dim: Dimensionality of hidden layers, if None then same
-            as input_dim
-        edge_chunk_sizes: List of chunks sizes to split edge representation
-            into and use separate MLPs for (None = no chunking, same MLP)
-        aggr_chunk_sizes: List of chunks sizes to split aggregated node
-            representation into and use separate MLPs for
-            (None = no chunking, same MLP)
-        aggr: Message aggregation method (sum/mean)
+        Initialise an InteractionNet message-passing layer.
+
+        Parameters
+        ----------
+        edge_index : torch.Tensor
+            Edge connectivity tensor of shape ``(2, M)`` in PyG format,
+            where ``M`` is the number of edges.
+        input_dim : int
+            Dimensionality of both node and edge input representations.
+        update_edges : bool, optional
+            If ``True``, compute and return updated edge representations in
+            addition to node representations. Default is ``True``.
+        hidden_layers : int, optional
+            Number of hidden layers in each MLP. Default is ``1``.
+        hidden_dim : int or None, optional
+            Width of hidden layers. If ``None``, defaults to ``input_dim``.
+        edge_chunk_sizes : list of int or None, optional
+            Chunk sizes for splitting edge representations across separate
+            MLPs. ``None`` uses a single shared MLP.
+        aggr_chunk_sizes : list of int or None, optional
+            Chunk sizes for splitting aggregated node representations across
+            separate MLPs. ``None`` uses a single shared MLP.
+        aggr : {'sum', 'mean'}, optional
+            Message aggregation method. Default is ``'sum'``.
+
+        Raises
+        ------
+        AssertionError
+            If ``aggr`` is not one of ``'sum'`` or ``'mean'``.
         """
         assert aggr in ("sum", "mean"), f"Unknown aggregation method: {aggr}"
         super().__init__(aggr=aggr)
diff --git a/neural_lam/metrics.py b/neural_lam/metrics.py
index 7db2cca6d..a078b9cc5 100644
--- a/neural_lam/metrics.py
+++ b/neural_lam/metrics.py
@@ -20,19 +20,30 @@ def get_metric(metric_name):
 
 def mask_and_reduce_metric(metric_entry_vals, mask, average_grid, sum_vars):
     """
-    Masks and (optionally) reduces entry-wise metric values
-
-    (...,) is any number of batch dimensions, potentially different
-        but broadcastable
-    metric_entry_vals: (..., N, d_state), prediction
-    mask: (N,), boolean mask describing which grid nodes to use in metric
-    average_grid: boolean, if grid dimension -2 should be reduced (mean over N)
-    sum_vars: boolean, if variable dimension -1 should be reduced (sum
-        over d_state)
-
-    Returns:
-    metric_val: One of (...,), (..., d_state), (..., N), (..., N, d_state),
-    depending on reduction arguments.
+    Apply a spatial mask and optionally reduce a per-entry metric tensor.
+
+    Parameters
+    ----------
+    metric_entry_vals : torch.Tensor
+        Entry-wise metric values of shape ``(..., N, d_state)``, where
+        ``...`` denotes any number of broadcastable batch dimensions.
+    mask : torch.Tensor or None
+        Boolean mask of shape ``(N,)`` selecting which grid nodes to include.
+        Pass ``None`` to use all nodes.
+    average_grid : bool
+        If ``True``, reduce the grid dimension ``N`` by taking the mean,
+        producing shape ``(..., d_state)``.
+    sum_vars : bool
+        If ``True``, reduce the variable dimension ``d_state`` by summing,
+        producing shape ``(..., N)`` or ``(...,)`` depending on
+        ``average_grid``.
+
+    Returns
+    -------
+    torch.Tensor
+        Reduced metric tensor. Shape is one of ``(...,)``,
+        ``(..., d_state)``, ``(..., N)``, or ``(..., N, d_state)``
+        depending on the reduction arguments.
     """
     # Only keep grid nodes in mask
     if mask is not None:
@@ -55,21 +66,34 @@ def mask_and_reduce_metric(metric_entry_vals, mask, average_grid, sum_vars):
 
 def wmse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     """
-    Weighted Mean Squared Error
-
-    (...,) is any number of batch dimensions, potentially different
-        but broadcastable
-    pred: (..., N, d_state), prediction
-    target: (..., N, d_state), target
-    pred_std: (..., N, d_state) or (d_state,), predicted std.-dev.
-    mask: (N,), boolean mask describing which grid nodes to use in metric
-    average_grid: boolean, if grid dimension -2 should be reduced (mean over N)
-    sum_vars: boolean, if variable dimension -1 should be reduced (sum
-        over d_state)
-
-    Returns:
-    metric_val: One of (...,), (..., d_state), (..., N), (..., N, d_state),
-    depending on reduction arguments.
+    Compute the Weighted Mean Squared Error (wMSE).
+
+    Scales the squared error at each grid node and variable by the inverse
+    variance ``1 / pred_std**2``, then applies masking and reduction via
+    :func:`mask_and_reduce_metric`.
+
+    Parameters
+    ----------
+    pred : torch.Tensor
+        Model predictions of shape ``(..., N, d_state)``.
+    target : torch.Tensor
+        Ground-truth values of shape ``(..., N, d_state)``.
+    pred_std : torch.Tensor
+        Predicted standard deviation of shape ``(..., N, d_state)`` or
+        ``(d_state,)`` used as the per-entry weighting.
+    mask : torch.Tensor or None, optional
+        Boolean mask of shape ``(N,)`` selecting grid nodes. Default is
+        ``None`` (all nodes used).
+    average_grid : bool, optional
+        If ``True``, average over the grid dimension. Default is ``True``.
+    sum_vars : bool, optional
+        If ``True``, sum over the variable dimension. Default is ``True``.
+
+    Returns
+    -------
+    torch.Tensor
+        Weighted MSE, with shape determined by ``average_grid`` and
+        ``sum_vars`` (see :func:`mask_and_reduce_metric`).
     """
     entry_mse = torch.nn.functional.mse_loss(
         pred, target, reduction="none"
diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py
index f65387ab6..45170c789 100644
--- a/neural_lam/models/__init__.py
+++ b/neural_lam/models/__init__.py
@@ -1,3 +1,4 @@
+"""Neural-LAM model architectures including GraphLAM, HiLAM, and variants."""
 # Local
 from .base_graph_model import BaseGraphModel
 from .base_hi_graph_model import BaseHiGraphModel

From fc3bde7c5169d3add75c31f80f6ef006a2e84c5f Mon Sep 17 00:00:00 2001
From: Mohit-Lakra <mohitlakra2007777@gmail.com>
Date: Fri, 27 Feb 2026 08:10:55 +0530
Subject: [PATCH 02/16] Datastore docstring Updated

---
 neural_lam/datastore/__init__.py              | 20 +++++++
 neural_lam/datastore/base.py                  | 14 +++--
 neural_lam/datastore/mdp.py                   |  2 +
 neural_lam/datastore/npyfilesmeps/__init__.py |  2 +
 .../compute_standardization_stats.py          | 57 ++++++++++++++++---
 neural_lam/datastore/npyfilesmeps/config.py   |  2 +
 neural_lam/datastore/npyfilesmeps/store.py    | 31 ++++++++--
 neural_lam/datastore/plot_example.py          |  3 +
 8 files changed, 112 insertions(+), 19 deletions(-)

diff --git a/neural_lam/datastore/__init__.py b/neural_lam/datastore/__init__.py
index 29876a601..77c909a7e 100644
--- a/neural_lam/datastore/__init__.py
+++ b/neural_lam/datastore/__init__.py
@@ -16,6 +16,26 @@
 
 
 def init_datastore(datastore_kind, config_path):
+    """
+    Instantiate a datastore based on its short-name identifier.
+
+    Parameters
+    ----------
+    datastore_kind : str
+        Key corresponding to one of :data:`DATASTORES`.
+    config_path : str | pathlib.Path
+        Path to the datastore-specific configuration file.
+
+    Returns
+    -------
+    BaseDatastore
+        Concrete datastore instance configured for ``config_path``.
+
+    Raises
+    ------
+    NotImplementedError
+        If ``datastore_kind`` is not registered.
+    """
     DatastoreClass = DATASTORES.get(datastore_kind)
 
     if DatastoreClass is None:
diff --git a/neural_lam/datastore/base.py b/neural_lam/datastore/base.py
index fc096595c..316ea6dc6 100644
--- a/neural_lam/datastore/base.py
+++ b/neural_lam/datastore/base.py
@@ -1,3 +1,5 @@
+"""Abstract base classes describing Neural-LAM datastore APIs."""
+
 # Standard library
 import abc
 import collections
@@ -85,8 +87,10 @@ def config(self) -> collections.abc.Mapping:
     def step_length(self) -> timedelta:
         """The step length of the dataset as a time interval.
 
-        Returns:
-            timedelta: The step length as a datetime.timedelta object.
+        Returns
+        -------
+        datetime.timedelta
+            The step length as a ``datetime.timedelta`` object.
 
         """
         pass
@@ -366,8 +370,10 @@ def state_feature_weights_values(self) -> List[float]:
         the loss function for each state variable (e.g. via the standard
         deviation of the 1-step differences of the state variables).
 
-        Returns:
-            List[float]: The weights for each state feature.
+        Returns
+        -------
+        List[float]
+            The weights for each state feature.
         """
         pass
 
diff --git a/neural_lam/datastore/mdp.py b/neural_lam/datastore/mdp.py
index ad5150118..2ad1f2c8c 100644
--- a/neural_lam/datastore/mdp.py
+++ b/neural_lam/datastore/mdp.py
@@ -1,3 +1,5 @@
+"""Datastore implementation wrapping ``mllam-data-prep`` outputs."""
+
 # Standard library
 import copy
 import warnings
diff --git a/neural_lam/datastore/npyfilesmeps/__init__.py b/neural_lam/datastore/npyfilesmeps/__init__.py
index 397a5075a..27cc356dc 100644
--- a/neural_lam/datastore/npyfilesmeps/__init__.py
+++ b/neural_lam/datastore/npyfilesmeps/__init__.py
@@ -1,2 +1,4 @@
+"""MEPS-specific datastore exposing numpy-based datasets."""
+
 # Local
 from .store import NpyFilesDatastoreMEPS  # noqa
diff --git a/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py b/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py
index 813d7b8e0..7c6790312 100644
--- a/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py
+++ b/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py
@@ -1,3 +1,5 @@
+"""Utilities for computing MEPS datastore standardization statistics."""
+
 # Standard library
 import os
 import subprocess
@@ -18,7 +20,19 @@
 
 
 class PaddedWeatherDataset(torch.utils.data.Dataset):
+    """Wrap :class:`WeatherDataset` to pad samples for distributed runners."""
+
     def __init__(self, base_dataset, world_size, batch_size):
+        """
+        Parameters
+        ----------
+        base_dataset : WeatherDataset
+            Dataset to pad.
+        world_size : int
+            Total number of distributed ranks participating.
+        batch_size : int
+            Per-rank batch size.
+        """
         super().__init__()
         self.base_dataset = base_dataset
         self.world_size = world_size
@@ -33,6 +47,7 @@ def __init__(self, base_dataset, world_size, batch_size):
         )
 
     def __getitem__(self, idx):
+        """Return an item, repeating the final sample for padded indices."""
         return self.base_dataset[
             (
                 self.original_indices[-1]
@@ -42,12 +57,15 @@ def __getitem__(self, idx):
         ]
 
     def __len__(self):
+        """Return the padded dataset length."""
         return self.total_samples + self.padded_samples
 
     def get_original_indices(self):
+        """Return indices of the non-padded samples."""
         return self.original_indices
 
     def get_original_window_indices(self, step_length):
+        """Return index mapping for sub-sampled windows at ``step_length``."""
         step_int, _ = get_integer_time(step_length.total_seconds())
         return [
             i // step_int for i in range(len(self.original_indices) * step_int)
@@ -55,10 +73,12 @@ def get_original_window_indices(self, step_length):
 
 
 def get_rank():
+    """Return the rank inferred from SLURM or default to 0."""
     return int(os.environ.get("SLURM_PROCID", 0))
 
 
 def get_world_size():
+    """Return the world size inferred from SLURM or default to 1."""
     return int(os.environ.get("SLURM_NTASKS", 1))
 
 
@@ -97,6 +117,24 @@ def setup(rank, world_size):  # pylint: disable=redefined-outer-name
 def save_stats(
     static_dir_path, means, squares, flux_means, flux_squares, filename_prefix
 ):
+    """
+    Aggregate running statistics and persist them to ``static_dir_path``.
+
+    Parameters
+    ----------
+    static_dir_path : str or pathlib.Path
+        Directory where ``*.pt`` files should be written.
+    means : Sequence[torch.Tensor]
+        Batch-wise means with shape ``(N_batch, d_features)``.
+    squares : Sequence[torch.Tensor]
+        Batch-wise second moments with shape ``(N_batch, d_features)``.
+    flux_means : Sequence[torch.Tensor]
+        Optional flux means of shape ``(N_batch,)``.
+    flux_squares : Sequence[torch.Tensor]
+        Optional flux second moments of shape ``(N_batch,)``.
+    filename_prefix : str
+        Prefix (e.g., ``"parameter"`` or ``"diff"``) for saved tensors.
+    """
     means = (
         torch.stack(means) if len(means) > 1 else means[0]
     )  # (N_batch, d_features,)
@@ -139,20 +177,20 @@ def main(
     datastore_config_path, batch_size, step_length, n_workers, distributed
 ):
     """
-    Pre-compute parameter weights to be used in loss function
+    Pre-compute and persist standardization statistics from the datastore.
 
-    Arguments
-    ---------
-    datastore_config_path : str
-        Path to datastore config file
+    Parameters
+    ----------
+    datastore_config_path : str or pathlib.Path
+        Path to the MEPS datastore configuration file.
     batch_size : int
-        Batch size when iterating over the dataset
+        Batch size used while iterating through the dataset.
     step_length : datetime.timedelta
-        Step length to consider single time step
+        Temporal sampling interval for the difference statistics.
     n_workers : int
-        Number of workers in data loader
+        Number of dataloader workers.
     distributed : bool
-        Run the script in distributed
+        If ``True``, run using torch.distributed with SLURM settings.
     """
 
     rank = get_rank()
@@ -378,6 +416,7 @@ def main(
 
 
 def cli():
+    """Parse CLI arguments and trigger :func:`main`."""
     parser = ArgumentParser(description="Training arguments")
     parser.add_argument(
         "--datastore_config_path",
diff --git a/neural_lam/datastore/npyfilesmeps/config.py b/neural_lam/datastore/npyfilesmeps/config.py
index 1d36b6fe0..b5f015b6b 100644
--- a/neural_lam/datastore/npyfilesmeps/config.py
+++ b/neural_lam/datastore/npyfilesmeps/config.py
@@ -1,3 +1,5 @@
+"""Dataclasses describing the MEPS numpy-file datastore configuration."""
+
 # Standard library
 from dataclasses import dataclass, field
 from datetime import timedelta
diff --git a/neural_lam/datastore/npyfilesmeps/store.py b/neural_lam/datastore/npyfilesmeps/store.py
index 26214e30c..96c287542 100644
--- a/neural_lam/datastore/npyfilesmeps/store.py
+++ b/neural_lam/datastore/npyfilesmeps/store.py
@@ -35,6 +35,7 @@
 
 
 def _load_np(fp, add_feature_dim, feature_dim_mask=None):
+    """Load an ``.npy`` file and optionally expand/mask the feature axis."""
     arr = np.load(fp)
     if add_feature_dim:
         arr = arr[..., np.newaxis]
@@ -44,25 +45,25 @@ def _load_np(fp, add_feature_dim, feature_dim_mask=None):
 
 
 class NpyFilesDatastoreMEPS(BaseRegularGridDatastore):
-    __doc__ = f"""
+    """
     Represents a dataset stored as numpy files on disk. The dataset is assumed
     to be stored in a directory structure where each sample is stored in a
-    separate file. The file-name format is assumed to be
-    '{STATE_FILENAME_FORMAT}'
+    separate file. The file-name format is assumed to be described by
+    ``STATE_FILENAME_FORMAT``.
 
     The MEPS dataset is organised into three splits: train, val, and test. Each
     split has a set of files which are:
 
-    - `{STATE_FILENAME_FORMAT}`:
+    - ``STATE_FILENAME_FORMAT``:
         The state variables for a forecast started at `analysis_time` with
         member id `member_id`. The dimensions of the array are
         `[forecast_timestep, y, x, feature]`.
 
-    - `{TOA_SW_DOWN_FLUX_FILENAME_FORMAT}`:
+    - ``TOA_SW_DOWN_FLUX_FILENAME_FORMAT``:
         The top-of-atmosphere downwelling shortwave flux at `time`. The
         dimensions of the array are `[forecast_timestep, y, x]`.
 
-    - `{OPEN_WATER_FILENAME_FORMAT}`:
+    - ``OPEN_WATER_FILENAME_FORMAT``:
         The open water fraction at `time`. The dimensions of the array are
         `[y, x]`.
 
@@ -550,6 +551,19 @@ def _get_analysis_times(self, split) -> List[np.datetime64]:
         return times
 
     def _calc_datetime_forcing_features(self, da_time: xr.DataArray):
+        """
+        Compute sinusoidal encodings of hour-of-day and day-of-year.
+
+        Parameters
+        ----------
+        da_time : xr.DataArray
+            Time coordinate with dimension ``time``.
+
+        Returns
+        -------
+        xr.DataArray
+            Normalized sine/cosine features with dims ``("feature",)``.
+        """
         da_hour_angle = da_time.dt.hour / 12 * np.pi
         da_year_angle = da_time.dt.dayofyear / 365 * 2 * np.pi
 
@@ -573,6 +587,7 @@ def _calc_datetime_forcing_features(self, da_time: xr.DataArray):
         return da_datetime_forcing
 
     def get_vars_units(self, category: str) -> List[str]:
+        """Return unit strings for the variables in ``category``."""
         if category == "state":
             return self.config.dataset.var_units
         elif category == "forcing":
@@ -590,6 +605,7 @@ def get_vars_units(self, category: str) -> List[str]:
             raise NotImplementedError(f"Category {category} not supported")
 
     def get_vars_names(self, category: str) -> List[str]:
+        """Return canonical short names for the variables in ``category``."""
         if category == "state":
             return self.config.dataset.var_names
         elif category == "forcing":
@@ -609,6 +625,7 @@ def get_vars_names(self, category: str) -> List[str]:
             raise NotImplementedError(f"Category {category} not supported")
 
     def get_vars_long_names(self, category: str) -> List[str]:
+        """Return descriptive names for the variables in ``category``."""
         if category == "state":
             return self.config.dataset.var_longnames
         else:
@@ -616,6 +633,7 @@ def get_vars_long_names(self, category: str) -> List[str]:
             return self.get_vars_names(category=category)
 
     def get_num_data_vars(self, category: str) -> int:
+        """Return the number of variables available in ``category``."""
         return len(self.get_vars_names(category=category))
 
     def get_xy(self, category: str, stacked: bool) -> np.ndarray:
@@ -731,6 +749,7 @@ def get_standardization_dataarray(self, category: str) -> xr.Dataset:
         """
 
         def load_pickled_tensor(fn):
+            """Load a serialized tensor from ``static`` and convert to numpy."""
             return torch.load(
                 self.root_path / "static" / fn, weights_only=True
             ).numpy()
diff --git a/neural_lam/datastore/plot_example.py b/neural_lam/datastore/plot_example.py
index 4f61ac7e7..912478db4 100644
--- a/neural_lam/datastore/plot_example.py
+++ b/neural_lam/datastore/plot_example.py
@@ -1,3 +1,5 @@
+"""CLI helper to plot slices from datastores for manual inspection."""
+
 # Third-party
 import matplotlib.pyplot as plt
 
@@ -93,6 +95,7 @@ def plot_example_from_datastore(
     import argparse
 
     def _parse_dict(arg_str):
+        """Parse ``key=value`` CLI arguments into typed dictionary entries."""
         key, value = arg_str.split("=")
         for op in [int, float]:
             try:

From 19cbe51d1e493d8311049435e95a1a44026ce3b7 Mon Sep 17 00:00:00 2001
From: Mohit-Lakra <mohitlakra2007777@gmail.com>
Date: Fri, 27 Feb 2026 08:53:32 +0530
Subject: [PATCH 03/16] Datastore docstring Updated

---
 neural_lam/datastore/__init__.py           | 1 +
 neural_lam/datastore/npyfilesmeps/store.py | 1 +
 2 files changed, 2 insertions(+)

diff --git a/neural_lam/datastore/__init__.py b/neural_lam/datastore/__init__.py
index 77c909a7e..e2117217f 100644
--- a/neural_lam/datastore/__init__.py
+++ b/neural_lam/datastore/__init__.py
@@ -1,4 +1,5 @@
 """Datastore backends for loading and serving weather model data."""
+
 # Local
 from .base import BaseDatastore  # noqa
 from .mdp import MDPDatastore  # noqa
diff --git a/neural_lam/datastore/npyfilesmeps/store.py b/neural_lam/datastore/npyfilesmeps/store.py
index 96c287542..9e84d959a 100644
--- a/neural_lam/datastore/npyfilesmeps/store.py
+++ b/neural_lam/datastore/npyfilesmeps/store.py
@@ -141,6 +141,7 @@ class NpyFilesDatastoreMEPS(BaseRegularGridDatastore):
     d_features = 8
     d_forcing = 1
     """
+
     SHORT_NAME = "npyfilesmeps"
 
     is_ensemble = True

From b026969008cfd207175b650798718c9729fcbf3a Mon Sep 17 00:00:00 2001
From: Mohit-Lakra <mohitlakra2007777@gmail.com>
Date: Fri, 27 Feb 2026 08:54:17 +0530
Subject: [PATCH 04/16] Models Docstring updated

---
 neural_lam/models/__init__.py            |   1 +
 neural_lam/models/ar_model.py            | 222 ++++++++++++++++++-----
 neural_lam/models/base_graph_model.py    | 125 ++++++++++---
 neural_lam/models/base_hi_graph_model.py |  76 ++++++--
 neural_lam/models/graph_lam.py           |  40 +++-
 neural_lam/models/hi_lam.py              |  91 ++++++++--
 neural_lam/models/hi_lam_parallel.py     |  37 +++-
 7 files changed, 478 insertions(+), 114 deletions(-)

diff --git a/neural_lam/models/__init__.py b/neural_lam/models/__init__.py
index 45170c789..924203f0b 100644
--- a/neural_lam/models/__init__.py
+++ b/neural_lam/models/__init__.py
@@ -1,4 +1,5 @@
 """Neural-LAM model architectures including GraphLAM, HiLAM, and variants."""
+
 # Local
 from .base_graph_model import BaseGraphModel
 from .base_hi_graph_model import BaseHiGraphModel
diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py
index 10fe64190..76e2b81e4 100644
--- a/neural_lam/models/ar_model.py
+++ b/neural_lam/models/ar_model.py
@@ -1,3 +1,5 @@
+"""Auto-regressive LightningModule implementations for Neural-LAM."""
+
 # Standard library
 import os
 import warnings
@@ -36,6 +38,16 @@ def __init__(
         config: NeuralLAMConfig,
         datastore: BaseDatastore,
     ):
+        """
+        Parameters
+        ----------
+        args : argparse.Namespace
+            Parsed training arguments controlling rollout length, loss, etc.
+        config : NeuralLAMConfig
+            Experiment configuration containing datastore/training settings.
+        datastore : BaseDatastore
+            Datastore supplying state/forcing/static arrays.
+        """
         super().__init__()
         self.save_hyperparameters(ignore=["datastore"])
         self.args = args
@@ -176,9 +188,9 @@ def _create_dataarray_from_tensor(
         Parameters
         ----------
         tensor : torch.Tensor
-            The tensor to convert to a `xr.DataArray` with dimensions [time,
-            grid_index, feature]. The tensor will be copied to the CPU if it is
-            not already there.
+            Tensor to convert back to an ``xr.DataArray``.
+
+            * **Shape**: ``(time, grid_index, feature)``
         time : torch.Tensor
             The time index or indices for the data, given as tensor representing
             epoch time in nanoseconds. The tensor will be
@@ -199,6 +211,7 @@ def _create_dataarray_from_tensor(
         return da
 
     def configure_optimizers(self):
+        """Construct the :class:`torch.optim.AdamW` optimizer for training."""
         opt = torch.optim.AdamW(
             self.parameters(), lr=self.args.lr, betas=(0.9, 0.95)
         )
@@ -207,32 +220,90 @@ def configure_optimizers(self):
     @property
     def interior_mask_bool(self):
         """
-        Get the interior mask as a boolean (N,) mask.
+        Boolean interior mask identifying non-boundary grid nodes.
+
+        Returns
+        -------
+        torch.Tensor
+            Boolean mask.
+
+            * **Shape**: ``(N,)``
         """
         return self.interior_mask[:, 0].to(torch.bool)
 
     @staticmethod
     def expand_to_batch(x, batch_size):
         """
-        Expand tensor with initial batch dimension
+        Broadcast a tensor by prepending a batch dimension.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Tensor to expand.
+        batch_size : int
+            Batch size to broadcast to.
+
+        Returns
+        -------
+        torch.Tensor
+            Tensor with a leading batch dimension added via ``expand``.
         """
         return x.unsqueeze(0).expand(batch_size, -1, -1)
 
     def predict_step(self, prev_state, prev_prev_state, forcing):
         """
-        Step state one step ahead using prediction model, X_{t-1}, X_t -> X_t+1
-        prev_state: (B, num_grid_nodes, feature_dim), X_t prev_prev_state: (B,
-        num_grid_nodes, feature_dim), X_{t-1} forcing: (B, num_grid_nodes,
-        forcing_dim)
+        Advance the state by one step using the prediction model.
+
+        Parameters
+        ----------
+        prev_state : torch.Tensor
+            Current state ``X_t``.
+
+            * **Shape**: ``(B, num_grid_nodes, feature_dim)``
+        prev_prev_state : torch.Tensor
+            Previous state ``X_{t-1}``.
+
+            * **Shape**: ``(B, num_grid_nodes, feature_dim)``
+        forcing : torch.Tensor
+            Forcing inputs applied at the prediction step.
+
+            * **Shape**: ``(B, num_grid_nodes, forcing_dim)``
+
+        Returns
+        -------
+        tuple[torch.Tensor, torch.Tensor | None]
+            Tuple ``(new_state, pred_std)`` describing the next state and
+            optional uncertainty estimate.
         """
         raise NotImplementedError("No prediction step implemented")
 
     def unroll_prediction(self, init_states, forcing_features, true_states):
         """
-        Roll out prediction taking multiple autoregressive steps with model
-        init_states: (B, 2, num_grid_nodes, d_f) forcing_features: (B,
-        pred_steps, num_grid_nodes, d_static_f) true_states: (B, pred_steps,
-        num_grid_nodes, d_f)
+        Roll out predictions autoregressively over multiple time steps.
+
+        Parameters
+        ----------
+        init_states : torch.Tensor
+            Initial states providing ``X_{t-1}`` and ``X_t``.
+
+            * **Shape**: ``(B, 2, num_grid_nodes, d_f)``
+        forcing_features : torch.Tensor
+            Forcing inputs aligned with each rollout step.
+
+            * **Shape**: ``(B, pred_steps, num_grid_nodes, d_static_f)``
+        true_states : torch.Tensor
+            Ground-truth states used for boundary replacement.
+
+            * **Shape**: ``(B, pred_steps, num_grid_nodes, d_f)``
+
+        Returns
+        -------
+        tuple[torch.Tensor, torch.Tensor]
+            Tuple ``(prediction, pred_std)``.
+
+            * **prediction**: ``(B, pred_steps, num_grid_nodes, d_f)``
+            * **pred_std**: ``(B, pred_steps, num_grid_nodes, d_f)`` or
+              ``(d_f,)`` when a constant per-feature value is used
         """
         prev_prev_state = init_states[:, 0]
         prev_state = init_states[:, 1]
@@ -278,11 +349,30 @@ def unroll_prediction(self, init_states, forcing_features, true_states):
 
     def common_step(self, batch):
         """
-        Predict on single batch batch consists of: init_states: (B, 2,
-        num_grid_nodes, d_features) target_states: (B, pred_steps,
-        num_grid_nodes, d_features) forcing_features: (B, pred_steps,
-        num_grid_nodes, d_forcing),
-            where index 0 corresponds to index 1 of init_states
+        Run a forward pass shared by train/val/test steps.
+
+        Parameters
+        ----------
+        batch : tuple
+            Tuple of ``(init_states, target_states, forcing_features,
+            batch_times)`` produced by :class:`WeatherDataset`.
+
+            * **init_states**: ``(B, 2, num_grid_nodes, d_features)``
+            * **target_states**: ``(B, pred_steps, num_grid_nodes, d_features)``
+            * **forcing_features**: ``(B, pred_steps, num_grid_nodes,
+              d_forcing)``
+            * **batch_times**: ``(B, pred_steps)`` timestamps
+
+        Returns
+        -------
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
+            ``(prediction, target_states, pred_std, batch_times)``.
+
+            * **prediction**: ``(B, pred_steps, num_grid_nodes, d_f)``
+            * **target_states**: ``(B, pred_steps, num_grid_nodes, d_f)``
+            * **pred_std**: ``(B, pred_steps, num_grid_nodes, d_f)`` or
+              ``(d_f,)``
+            * **batch_times**: ``(B, pred_steps)``
         """
         (init_states, target_states, forcing_features, batch_times) = batch
 
@@ -295,8 +385,12 @@ def common_step(self, batch):
         return prediction, target_states, pred_std, batch_times
 
     def training_step(self, batch):
-        """
-        Train on single batch
+        """Execute a single optimization step on ``batch``.
+
+        Parameters
+        ----------
+        batch : tuple
+            Batch sampled from the training dataloader.
         """
         prediction, target, pred_std, _ = self.common_step(batch)
 
@@ -320,20 +414,35 @@ def training_step(self, batch):
 
     def all_gather_cat(self, tensor_to_gather):
         """
-        Gather tensors across all ranks, and concatenate across dim. 0 (instead
-        of stacking in new dim. 0)
+        Gather tensors across ranks and concatenate along dim-0.
+
+        Parameters
+        ----------
+        tensor_to_gather : torch.Tensor
+            Tensor distributed across ``K`` ranks.
 
-        tensor_to_gather: (d1, d2, ...), distributed over K ranks
+            * **Shape**: ``(d1, d2, ...)`` per rank
 
-        returns: (K*d1, d2, ...)
+        Returns
+        -------
+        torch.Tensor
+            Concatenated tensor gathered from all ranks.
+
+            * **Shape**: ``(K * d1, d2, ...)``
         """
         return self.all_gather(tensor_to_gather).flatten(0, 1)
 
     # newer lightning versions requires batch_idx argument, even if unused
     # pylint: disable-next=unused-argument
     def validation_step(self, batch, batch_idx):
-        """
-        Run validation on single batch
+        """Evaluate ``batch`` during validation.
+
+        Parameters
+        ----------
+        batch : tuple
+            Batch sampled from the validation dataloader.
+        batch_idx : int
+            Index of the current batch.
         """
         prediction, target, pred_std, _ = self.common_step(batch)
 
@@ -383,8 +492,14 @@ def on_validation_epoch_end(self):
 
     # pylint: disable-next=unused-argument
     def test_step(self, batch, batch_idx):
-        """
-        Run test on single batch
+        """Evaluate ``batch`` during testing and log diagnostics.
+
+        Parameters
+        ----------
+        batch : tuple
+            Batch sampled from the test dataloader.
+        batch_idx : int
+            Index of the current batch.
         """
         # TODO Here batch_times can be used for plotting routines
         prediction, target, pred_std, batch_times = self.common_step(batch)
@@ -465,12 +580,21 @@ def test_step(self, batch, batch_idx):
 
     def plot_examples(self, batch, n_examples, split, prediction=None):
         """
-        Plot the first n_examples forecasts from batch
+        Plot the first ``n_examples`` forecasts from ``batch``.
 
-        batch: batch with data to plot corresponding forecasts for n_examples:
-        number of forecasts to plot prediction: (B, pred_steps, num_grid_nodes,
-        d_f), existing prediction.
-            Generate if None.
+        Parameters
+        ----------
+        batch : tuple
+            Batch tuple produced by the dataloader.
+        n_examples : int
+            Number of forecasts to visualise.
+        split : str
+            Dataset split name used for metadata lookups.
+        prediction : torch.Tensor or None, optional
+            Pre-computed predictions to plot. If ``None`` the method runs
+            :meth:`common_step` to obtain predictions.
+
+            * **Shape**: ``(B, pred_steps, num_grid_nodes, d_f)``
         """
         if prediction is None:
             prediction, target, _, _ = self.common_step(batch)
@@ -592,14 +716,23 @@ def plot_examples(self, batch, n_examples, split, prediction=None):
 
     def create_metric_log_dict(self, metric_tensor, prefix, metric_name):
         """
-        Put together a dict with everything to log for one metric. Also saves
-        plots as pdf and csv if using test prefix.
+        Assemble logging artefacts for a single metric tensor.
+
+        Parameters
+        ----------
+        metric_tensor : torch.Tensor
+            Metric values per time step and variable.
 
-        metric_tensor: (pred_steps, d_f), metric values per time and variable
-        prefix: string, prefix to use for logging metric_name: string, name of
-        the metric
+            * **Shape**: ``(pred_steps, d_f)``
+        prefix : str
+            Prefix used for logger keys (e.g., ``"val"`` or ``"test"``).
+        metric_name : str
+            Human-readable metric name.
 
-        Return: log_dict: dict with everything to log for given metric
+        Returns
+        -------
+        dict[str, object]
+            Mapping from log keys to figures or scalar tensors.
         """
         log_dict = {}
         metric_fig = vis.plot_error_map(
@@ -634,11 +767,14 @@ def create_metric_log_dict(self, metric_tensor, prefix, metric_name):
 
     def aggregate_and_plot_metrics(self, metrics_dict, prefix):
         """
-        Aggregate and create error map plots for all metrics in metrics_dict
+        Aggregate metric tensors and create error-map visualisations.
 
-        metrics_dict: dictionary with metric_names and list of tensors
-            with step-evals.
-        prefix: string, prefix to use for logging
+        Parameters
+        ----------
+        metrics_dict : dict[str, list[torch.Tensor]]
+            Mapping from metric name to per-batch tensors of evaluations.
+        prefix : str
+            Prefix to use for logger keys.
         """
         log_dict = {}
         for metric_name, metric_val_list in metrics_dict.items():
diff --git a/neural_lam/models/base_graph_model.py b/neural_lam/models/base_graph_model.py
index 35b1ab126..eaac26c2c 100644
--- a/neural_lam/models/base_graph_model.py
+++ b/neural_lam/models/base_graph_model.py
@@ -1,3 +1,5 @@
+"""Base classes for Neural-LAM graph models."""
+
 # Third-party
 import torch
 
@@ -16,6 +18,18 @@ class BaseGraphModel(ARModel):
     """
 
     def __init__(self, args, config: NeuralLAMConfig, datastore: BaseDatastore):
+        """
+        Initialize the graph-model scaffolding shared by concrete variants.
+
+        Parameters
+        ----------
+        args : argparse.Namespace
+            Training/runtime arguments describing graph paths and widths.
+        config : NeuralLAMConfig
+            Experiment configuration for clamping and weighting.
+        datastore : BaseDatastore
+            Datastore providing static features and metadata (e.g. graph path).
+        """
         super().__init__(args, config=config, datastore=datastore)
 
         # Load graph with static features
@@ -86,7 +100,14 @@ def prepare_clamping_params(
         self, config: NeuralLAMConfig, datastore: BaseDatastore
     ):
         """
-        Prepare parameters for clamping predicted values to valid range
+        Prepare per-feature parameters for clamping model outputs.
+
+        Parameters
+        ----------
+        config : NeuralLAMConfig
+            Model and training configuration containing clamping settings.
+        datastore : BaseDatastore
+            Datastore that provides the ordering of state variables.
         """
 
         # Read configs
@@ -219,18 +240,30 @@ def prepare_clamping_params(
 
     def get_clamped_new_state(self, state_delta, prev_state):
         """
-        Clamp prediction to valid range supplied in config
-        Returns the clamped new state after adding delta to original state
-
-        Instead of the new state being computed as
-        $X_{t+1} = X_t + \\delta = X_t + model(\\{X_t,X_{t-1},...\\}, forcing)$
-        The clamped values will be
-        $f(f^{-1}(X_t) + model(\\{X_t, X_{t-1},... \\}, forcing))$
-        Which means the model will learn to output values in the range of the
-        inverse clamping function
-
-        state_delta: (B, num_grid_nodes, feature_dim)
-        prev_state: (B, num_grid_nodes, feature_dim)
+        Clamp predicted deltas and add them to the previous state.
+
+        The clamped values follow
+        ``f(f^{-1}(X_t) + model({X_t, X_{t-1}, ...}, forcing))`` so that the
+        model learns to emit outputs in the range of the inverse clamping
+        function.
+
+        Parameters
+        ----------
+        state_delta : torch.Tensor
+            Predicted change to apply to the previous state.
+
+            * **Shape**: ``(B, num_grid_nodes, feature_dim)``
+        prev_state : torch.Tensor
+            Previous state ``X_t``.
+
+            * **Shape**: ``(B, num_grid_nodes, feature_dim)``
+
+        Returns
+        -------
+        torch.Tensor
+            Clamped next state ``X_{t+1}``.
+
+            * **Shape**: ``(B, num_grid_nodes, feature_dim)``
         """
 
         # Assign new state, but overwrite clamped values of each type later
@@ -267,34 +300,76 @@ def get_clamped_new_state(self, state_delta, prev_state):
 
     def get_num_mesh(self):
         """
-        Compute number of mesh nodes from loaded features,
-        and number of mesh nodes that should be ignored in encoding/decoding
+        Compute mesh node counts used for encoding and decoding.
+
+        Returns
+        -------
+        tuple[int, int]
+            Total number of mesh nodes and the number that should be ignored
+            during encoding/decoding.
         """
         raise NotImplementedError("get_num_mesh not implemented")
 
     def embedd_mesh_nodes(self):
         """
-        Embed static mesh features
-        Returns tensor of shape (num_mesh_nodes, d_h)
+        Embed static mesh features for downstream processing.
+
+        Returns
+        -------
+        torch.Tensor
+            Embedded mesh node representations.
+
+            * **Shape**: ``(num_mesh_nodes, d_h)``
         """
         raise NotImplementedError("embedd_mesh_nodes not implemented")
 
     def process_step(self, mesh_rep):
         """
-        Process step of embedd-process-decode framework
-        Processes the representation on the mesh, possible in multiple steps
+        Run the processor portion of the encode-process-decode framework.
+
+        Parameters
+        ----------
+        mesh_rep : torch.Tensor
+            Mesh node representations prior to the processor.
+
+            * **Shape**: ``(B, num_mesh_nodes, d_h)``
 
-        mesh_rep: has shape (B, num_mesh_nodes, d_h)
-        Returns mesh_rep: (B, num_mesh_nodes, d_h)
+        Returns
+        -------
+        torch.Tensor
+            Updated mesh representations after processing.
+
+            * **Shape**: ``(B, num_mesh_nodes, d_h)``
         """
         raise NotImplementedError("process_step not implemented")
 
     def predict_step(self, prev_state, prev_prev_state, forcing):
         """
-        Step state one step ahead using prediction model, X_{t-1}, X_t -> X_t+1
-        prev_state: (B, num_grid_nodes, feature_dim), X_t
-        prev_prev_state: (B, num_grid_nodes, feature_dim), X_{t-1}
-        forcing: (B, num_grid_nodes, forcing_dim)
+        Advance the state by one step using the prediction model.
+
+        Parameters
+        ----------
+        prev_state : torch.Tensor
+            Current state ``X_t``.
+
+            * **Shape**: ``(B, num_grid_nodes, feature_dim)``
+        prev_prev_state : torch.Tensor
+            Previous state ``X_{t-1}``.
+
+            * **Shape**: ``(B, num_grid_nodes, feature_dim)``
+        forcing : torch.Tensor
+            Forcing inputs applied at the prediction step.
+
+            * **Shape**: ``(B, num_grid_nodes, forcing_dim)``
+
+        Returns
+        -------
+        tuple[torch.Tensor, torch.Tensor | None]
+            Tuple ``(new_state, pred_std)`` where ``pred_std`` is ``None`` when
+            the model does not emit uncertainty estimates.
+
+            * **Shape**: ``(B, num_grid_nodes, feature_dim)`` for ``new_state``
+              and ``(B, num_grid_nodes, d_f)`` for ``pred_std`` when present.
         """
         batch_size = prev_state.shape[0]
 
diff --git a/neural_lam/models/base_hi_graph_model.py b/neural_lam/models/base_hi_graph_model.py
index 882dbf4da..5f1534fd9 100644
--- a/neural_lam/models/base_hi_graph_model.py
+++ b/neural_lam/models/base_hi_graph_model.py
@@ -1,3 +1,5 @@
+"""Base implementations for hierarchical (multi-level) graph models."""
+
 # Third-party
 from torch import nn
 
@@ -15,6 +17,7 @@ class BaseHiGraphModel(BaseGraphModel):
     """
 
     def __init__(self, args, config: NeuralLAMConfig, datastore: BaseDatastore):
+        """Extend :class:`BaseGraphModel` with hierarchical mesh structures."""
         super().__init__(args, config=config, datastore=datastore)
 
         # Track number of nodes, edges on each level
@@ -103,8 +106,13 @@ def __init__(self, args, config: NeuralLAMConfig, datastore: BaseDatastore):
 
     def get_num_mesh(self):
         """
-        Compute number of mesh nodes from loaded features,
-        and number of mesh nodes that should be ignored in encoding/decoding
+        Compute mesh node counts used for encoding and decoding.
+
+        Returns
+        -------
+        tuple[int, int]
+            Total number of mesh nodes and the number to ignore during
+            encoding/decoding.
         """
         num_mesh_nodes = sum(
             node_feat.shape[0] for node_feat in self.mesh_static_features
@@ -116,20 +124,34 @@ def get_num_mesh(self):
 
     def embedd_mesh_nodes(self):
         """
-        Embed static mesh features
-        This embeds only bottom level, rest is done at beginning of
-        processing step
-        Returns tensor of shape (num_mesh_nodes[0], d_h)
+        Embed static mesh features for the bottom level of the hierarchy.
+
+        Returns
+        -------
+        torch.Tensor
+            Embedded representations for the base-level mesh nodes.
+
+            * **Shape**: ``(num_mesh_nodes[0], d_h)``
         """
         return self.mesh_embedders[0](self.mesh_static_features[0])
 
     def process_step(self, mesh_rep):
         """
-        Process step of embedd-process-decode framework
-        Processes the representation on the mesh, possible in multiple steps
+        Run the processor portion of the hierarchical encode-process-decode.
+
+        Parameters
+        ----------
+        mesh_rep : torch.Tensor
+            Base-level mesh representations prior to the processor.
+
+            * **Shape**: ``(B, num_mesh_nodes, d_h)``
 
-        mesh_rep: has shape (B, num_mesh_nodes, d_h)
-        Returns mesh_rep: (B, num_mesh_nodes, d_h)
+        Returns
+        -------
+        torch.Tensor
+            Updated base-level mesh representations.
+
+            * **Shape**: ``(B, num_mesh_nodes, d_h)``
         """
         batch_size = mesh_rep.shape[0]
 
@@ -222,16 +244,34 @@ def hi_processor_step(
         self, mesh_rep_levels, mesh_same_rep, mesh_up_rep, mesh_down_rep
     ):
         """
-        Internal processor step of hierarchical graph models.
-        Between mesh init and read out.
+        Internal processor step executed between mesh init and read-out.
+
+        Parameters
+        ----------
+        mesh_rep_levels : list[torch.Tensor]
+            Mesh representations for each level.
+
+            * **Shape**: ``(B, num_mesh_nodes[l], d_h)``
+        mesh_same_rep : list[torch.Tensor]
+            Same-level edge representations per level.
 
-        Each input is list with representations, each with shape
+            * **Shape**: ``(B, M_same[l], d_h)``
+        mesh_up_rep : list[torch.Tensor]
+            Edge representations from level ``l`` to ``l+1``.
 
-        mesh_rep_levels: (B, num_mesh_nodes[l], d_h)
-        mesh_same_rep: (B, M_same[l], d_h)
-        mesh_up_rep: (B, M_up[l -> l+1], d_h)
-        mesh_down_rep: (B, M_down[l <- l+1], d_h)
+            * **Shape**: ``(B, M_up[l -> l+1], d_h)``
+        mesh_down_rep : list[torch.Tensor]
+            Edge representations from level ``l+1`` down to ``l``.
 
-        Returns same lists
+            * **Shape**: ``(B, M_down[l <- l+1], d_h)``
+
+        Returns
+        -------
+        tuple[
+            list[torch.Tensor], list[torch.Tensor], list[torch.Tensor],
+            list[torch.Tensor]
+        ]
+            Updated representations for (mesh, same-level, up edges, down edges)
+            in that order.
         """
         raise NotImplementedError("hi_process_step not implemented")
diff --git a/neural_lam/models/graph_lam.py b/neural_lam/models/graph_lam.py
index 0a5b6b574..156fb950c 100644
--- a/neural_lam/models/graph_lam.py
+++ b/neural_lam/models/graph_lam.py
@@ -1,3 +1,5 @@
+"""GraphLAM: the non-hierarchical Neural-LAM architecture."""
+
 # Third-party
 import torch_geometric as pyg
 
@@ -18,6 +20,7 @@ class GraphLAM(BaseGraphModel):
     """
 
     def __init__(self, args, config: NeuralLAMConfig, datastore: BaseDatastore):
+        """Initialize the non-hierarchical GraphLAM variant."""
         super().__init__(args, config=config, datastore=datastore)
 
         assert (
@@ -58,25 +61,46 @@ def __init__(self, args, config: NeuralLAMConfig, datastore: BaseDatastore):
 
     def get_num_mesh(self):
         """
-        Compute number of mesh nodes from loaded features,
-        and number of mesh nodes that should be ignored in encoding/decoding
+        Compute mesh node counts used for encoding and decoding.
+
+        Returns
+        -------
+        tuple[int, int]
+            Total number of mesh nodes and the number to ignore during
+            encoding/decoding.
         """
         return self.mesh_static_features.shape[0], 0
 
     def embedd_mesh_nodes(self):
         """
-        Embed static mesh features
-        Returns tensor of shape (N_mesh, d_h)
+        Embed static mesh features.
+
+        Returns
+        -------
+        torch.Tensor
+            Embedded mesh node representations.
+
+            * **Shape**: ``(N_mesh, d_h)``
         """
         return self.mesh_embedder(self.mesh_static_features)  # (N_mesh, d_h)
 
     def process_step(self, mesh_rep):
         """
-        Process step of embedd-process-decode framework
-        Processes the representation on the mesh, possible in multiple steps
+        Run the processor portion of the encode-process-decode framework.
+
+        Parameters
+        ----------
+        mesh_rep : torch.Tensor
+            Mesh node representations before processing.
+
+            * **Shape**: ``(B, N_mesh, d_h)``
+
+        Returns
+        -------
+        torch.Tensor
+            Updated mesh representations.
 
-        mesh_rep: has shape (B, N_mesh, d_h)
-        Returns mesh_rep: (B, N_mesh, d_h)
+            * **Shape**: ``(B, N_mesh, d_h)``
         """
         # Embed m2m here first
         batch_size = mesh_rep.shape[0]
diff --git a/neural_lam/models/hi_lam.py b/neural_lam/models/hi_lam.py
index c340c95da..ab9384e27 100644
--- a/neural_lam/models/hi_lam.py
+++ b/neural_lam/models/hi_lam.py
@@ -1,3 +1,5 @@
+"""Sequential up/down hierarchical Neural-LAM model (Hi-LAM)."""
+
 # Third-party
 from torch import nn
 
@@ -16,6 +18,7 @@ class HiLAM(BaseHiGraphModel):
     """
 
     def __init__(self, args, config: NeuralLAMConfig, datastore: BaseDatastore):
+        """Initialize the sequential up/down hierarchical processor."""
         super().__init__(args, config=config, datastore=datastore)
 
         # Make down GNNs, both for down edges and same level
@@ -88,8 +91,31 @@ def mesh_down_step(
         same_gnns,
     ):
         """
-        Run down-part of vertical processing, sequentially alternating between
-        processing using down edges and same-level edges.
+        Run the downward half of the hierarchical processing sweep.
+
+        Parameters
+        ----------
+        mesh_rep_levels : list[torch.Tensor]
+            Mesh representations for each level.
+
+            * **Shape**: ``(B, N_mesh[l], d_h)``
+        mesh_same_rep : list[torch.Tensor]
+            Same-level edge representations.
+
+            * **Shape**: ``(B, M_same[l], d_h)``
+        mesh_down_rep : list[torch.Tensor]
+            Downward edge representations.
+
+            * **Shape**: ``(B, M_down[l], d_h)``
+        down_gnns : Sequence[InteractionNet]
+            Message-passing networks applied to downward edges.
+        same_gnns : Sequence[InteractionNet]
+            Message-passing networks for same-level processing.
+
+        Returns
+        -------
+        tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]
+            Updated ``(mesh_rep_levels, mesh_same_rep, mesh_down_rep)``.
         """
         # Run same level processing on level L
         mesh_rep_levels[-1], mesh_same_rep[-1] = same_gnns[-1](
@@ -127,8 +153,31 @@ def mesh_up_step(
         self, mesh_rep_levels, mesh_same_rep, mesh_up_rep, up_gnns, same_gnns
     ):
         """
-        Run up-part of vertical processing, sequentially alternating between
-        processing using up edges and same-level edges.
+        Run the upward half of the hierarchical processing sweep.
+
+        Parameters
+        ----------
+        mesh_rep_levels : list[torch.Tensor]
+            Mesh representations for each level.
+
+            * **Shape**: ``(B, N_mesh[l], d_h)``
+        mesh_same_rep : list[torch.Tensor]
+            Same-level edge representations.
+
+            * **Shape**: ``(B, M_same[l], d_h)``
+        mesh_up_rep : list[torch.Tensor]
+            Upward edge representations.
+
+            * **Shape**: ``(B, M_up[l], d_h)``
+        up_gnns : Sequence[InteractionNet]
+            Message-passing networks applied to upward edges.
+        same_gnns : Sequence[InteractionNet]
+            Message-passing networks for same-level processing.
+
+        Returns
+        -------
+        tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]
+            Updated ``(mesh_rep_levels, mesh_same_rep, mesh_up_rep)``.
         """
 
         # Run same level processing on level 0
@@ -166,17 +215,35 @@ def hi_processor_step(
         self, mesh_rep_levels, mesh_same_rep, mesh_up_rep, mesh_down_rep
     ):
         """
-        Internal processor step of hierarchical graph models.
-        Between mesh init and read out.
+        Execute one full processor iteration (down + up sweeps).
+
+        Parameters
+        ----------
+        mesh_rep_levels : list[torch.Tensor]
+            Mesh representations for each level.
+
+            * **Shape**: ``(B, N_mesh[l], d_h)``
+        mesh_same_rep : list[torch.Tensor]
+            Same-level edge representations.
+
+            * **Shape**: ``(B, M_same[l], d_h)``
+        mesh_up_rep : list[torch.Tensor]
+            Upward edge representations.
 
-        Each input is list with representations, each with shape
+            * **Shape**: ``(B, M_up[l], d_h)``
+        mesh_down_rep : list[torch.Tensor]
+            Downward edge representations.
 
-        mesh_rep_levels: (B, N_mesh[l], d_h)
-        mesh_same_rep: (B, M_same[l], d_h)
-        mesh_up_rep: (B, M_up[l -> l+1], d_h)
-        mesh_down_rep: (B, M_down[l <- l+1], d_h)
+            * **Shape**: ``(B, M_down[l], d_h)``
 
-        Returns same lists
+        Returns
+        -------
+        tuple[
+            list[torch.Tensor], list[torch.Tensor], list[torch.Tensor],
+            list[torch.Tensor]
+        ]
+            Updated representations ``(mesh_rep_levels, mesh_same_rep,
+            mesh_up_rep, mesh_down_rep)`` after both sweeps.
         """
         for down_gnns, down_same_gnns, up_gnns, up_same_gnns in zip(
             self.mesh_down_gnns,
diff --git a/neural_lam/models/hi_lam_parallel.py b/neural_lam/models/hi_lam_parallel.py
index a0a84d293..b839df6e6 100644
--- a/neural_lam/models/hi_lam_parallel.py
+++ b/neural_lam/models/hi_lam_parallel.py
@@ -1,3 +1,5 @@
+"""Parallel message-passing variant of the Hi-LAM architecture."""
+
 # Third-party
 import torch
 import torch_geometric as pyg
@@ -19,6 +21,7 @@ class HiLAMParallel(BaseHiGraphModel):
     """
 
     def __init__(self, args, config: NeuralLAMConfig, datastore: BaseDatastore):
+        """Initialize the parallel hierarchical message-passing processor."""
         super().__init__(args, config=config, datastore=datastore)
 
         # Processor GNNs
@@ -56,17 +59,35 @@ def hi_processor_step(
         self, mesh_rep_levels, mesh_same_rep, mesh_up_rep, mesh_down_rep
     ):
         """
-        Internal processor step of hierarchical graph models.
-        Between mesh init and read out.
+        Internal processor step executed between mesh init and read-out.
+
+        Parameters
+        ----------
+        mesh_rep_levels : list[torch.Tensor]
+            Mesh representations for each level.
+
+            * **Shape**: ``(B, N_mesh[l], d_h)``
+        mesh_same_rep : list[torch.Tensor]
+            Same-level edge representations.
 
-        Each input is list with representations, each with shape
+            * **Shape**: ``(B, M_same[l], d_h)``
+        mesh_up_rep : list[torch.Tensor]
+            Upward edge representations.
 
-        mesh_rep_levels: (B, N_mesh[l], d_h)
-        mesh_same_rep: (B, M_same[l], d_h)
-        mesh_up_rep: (B, M_up[l -> l+1], d_h)
-        mesh_down_rep: (B, M_down[l <- l+1], d_h)
+            * **Shape**: ``(B, M_up[l], d_h)``
+        mesh_down_rep : list[torch.Tensor]
+            Downward edge representations.
 
-        Returns same lists
+            * **Shape**: ``(B, M_down[l], d_h)``
+
+        Returns
+        -------
+        tuple[
+            list[torch.Tensor], list[torch.Tensor], list[torch.Tensor],
+            list[torch.Tensor]
+        ]
+            Updated representations ``(mesh_rep_levels, mesh_same_rep,
+            mesh_up_rep, mesh_down_rep)`` after the parallel pass.
         """
 
         # First join all node and edge representations to single tensors

From 73a70e44f673a2ad5aa44d50142616168d26b529 Mon Sep 17 00:00:00 2001
From: Mohit-Lakra <mohitlakra2007777@gmail.com>
Date: Fri, 27 Feb 2026 09:03:21 +0530
Subject: [PATCH 05/16] Docstring Updated

---
 neural_lam/__init__.py        |   1 +
 neural_lam/config.py          |  11 ++
 neural_lam/create_graph.py    |  48 ++++++
 neural_lam/custom_loggers.py  |  35 +++-
 neural_lam/interaction_net.py |  94 +++++++----
 neural_lam/loss_weighting.py  |   2 +
 neural_lam/metrics.py         | 294 ++++++++++++++++++++++------------
 neural_lam/plot_graph.py      |   2 +
 neural_lam/train_model.py     |   2 +
 neural_lam/utils.py           | 156 ++++++++++++------
 neural_lam/vis.py             |  70 ++++++--
 neural_lam/weather_dataset.py |  63 +++++++-
 12 files changed, 588 insertions(+), 190 deletions(-)

diff --git a/neural_lam/__init__.py b/neural_lam/__init__.py
index b434957fe..c42f8354f 100644
--- a/neural_lam/__init__.py
+++ b/neural_lam/__init__.py
@@ -1,4 +1,5 @@
 """Neural-LAM: graph-based neural weather prediction models."""
+
 # Standard library
 import importlib.metadata
 
diff --git a/neural_lam/config.py b/neural_lam/config.py
index 7b00cb3b2..5d449f545 100644
--- a/neural_lam/config.py
+++ b/neural_lam/config.py
@@ -1,3 +1,5 @@
+"""Configuration dataclasses and helpers for Neural-LAM experiments."""
+
 # Standard library
 import dataclasses
 from pathlib import Path
@@ -33,6 +35,14 @@ class DatastoreSelection:
     kind: str
 
     def __post_init__(self):
+        """
+        Validate that the selected datastore kind is implemented.
+
+        Raises
+        ------
+        ValueError
+            If the provided ``kind`` is not part of :data:`DATASTORES`.
+        """
         if self.kind not in DATASTORES:
             raise ValueError(f"Datastore kind {self.kind} is not implemented")
 
@@ -156,6 +166,7 @@ class _(dataclass_wizard.JSONWizard.Meta):
 
 class InvalidConfigError(Exception):
     """Raised when the Neural-LAM configuration file is invalid or malformed."""
+
     pass
 
 
diff --git a/neural_lam/create_graph.py b/neural_lam/create_graph.py
index e0d81ead5..35b132211 100644
--- a/neural_lam/create_graph.py
+++ b/neural_lam/create_graph.py
@@ -1,3 +1,5 @@
+"""Graph construction utilities for Neural-LAM meshes and grids."""
+
 # Standard library
 import os
 from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
@@ -19,6 +21,21 @@
 
 
 def plot_graph(graph, title=None):
+    """
+    Render a PyTorch Geometric graph using stored node coordinates.
+
+    Parameters
+    ----------
+    graph : torch_geometric.data.Data
+        Graph containing ``edge_index`` and ``pos`` attributes.
+    title : str or None, optional
+        Optional subplot title.
+
+    Returns
+    -------
+    tuple[matplotlib.figure.Figure, matplotlib.axes.Axes]
+        Figure and axis handles for further customization.
+    """
     fig, axis = plt.subplots(figsize=(8, 8), dpi=200)  # W,H
     edge_index = graph.edge_index
     pos = graph.pos
@@ -70,6 +87,7 @@ def plot_graph(graph, title=None):
 
 
 def sort_nodes_internally(nx_graph):
+    """Return a copy of ``nx_graph`` with deterministically ordered nodes."""
     # For some reason the networkx .nodes() return list can not be sorted,
     # but this is the ordering used by pyg when converting.
     # This function fixes this.
@@ -80,6 +98,7 @@ def sort_nodes_internally(nx_graph):
 
 
 def save_edges(graph, name, base_path):
+    """Persist edge indices/features for a PyG graph under ``base_path``."""
     torch.save(
         graph.edge_index, os.path.join(base_path, f"{name}_edge_index.pt")
     )
@@ -90,6 +109,7 @@ def save_edges(graph, name, base_path):
 
 
 def save_edges_list(graphs, name, base_path):
+    """Persist edge indices/features for a list of graphs."""
     torch.save(
         [graph.edge_index for graph in graphs],
         os.path.join(base_path, f"{name}_edge_index.pt"),
@@ -104,12 +124,14 @@ def save_edges_list(graphs, name, base_path):
 
 
 def from_networkx_with_start_index(nx_graph, start_index):
+    """Convert a NetworkX graph to PyG and offset node indices."""
     pyg_graph = from_networkx(nx_graph)
     pyg_graph.edge_index += start_index
     return pyg_graph
 
 
 def mk_2d_graph(xy, nx, ny):
+    """Create a diagonal 2-D grid graph over the ``xy`` positions."""
     xm, xM = np.amin(xy[:, :, 0][:, 0]), np.amax(xy[:, :, 0][:, 0])
     ym, yM = np.amin(xy[:, :, 1][0, :]), np.amax(xy[:, :, 1][0, :])
 
@@ -149,6 +171,7 @@ def mk_2d_graph(xy, nx, ny):
 
 
 def prepend_node_index(graph, new_index):
+    """Relabel each node by prepending ``new_index`` to its tuple identifier."""
     # Relabel node indices in graph, insert (graph_level, i, j)
     ijk = [tuple((new_index,) + x) for x in graph.nodes]
     to_mapping = dict(zip(graph.nodes, ijk))
@@ -543,6 +566,22 @@ def create_graph_from_datastore(
     hierarchical: bool = False,
     create_plot: bool = False,
 ):
+    """
+    Generate graph components for ``datastore`` and persist them on disk.
+
+    Parameters
+    ----------
+    datastore : BaseRegularGridDatastore
+        Datastore providing ``get_xy`` for state nodes.
+    output_root_path : str
+        Directory where the resulting ``*.pt`` graph files are stored.
+    n_max_levels : int or None, optional
+        Optional limit of hierarchical mesh levels to build.
+    hierarchical : bool, optional
+        If ``True``, create multi-level hierarchical graphs. Default ``False``.
+    create_plot : bool, optional
+        If ``True``, display matplotlib previews of the generated graphs.
+    """
     if isinstance(datastore, BaseRegularGridDatastore):
         xy = datastore.get_xy(category="state", stacked=False)
     else:
@@ -560,6 +599,15 @@ def create_graph_from_datastore(
 
 
 def cli(input_args=None):
+    """
+    Parse CLI arguments and call :func:`create_graph_from_datastore`.
+
+    Parameters
+    ----------
+    input_args : list[str] or None, optional
+        Argument list forwarded to :class:`argparse.ArgumentParser`. When
+        ``None``, ``sys.argv`` is used.
+    """
     parser = ArgumentParser(
         description="Graph generation for neural-lam",
         formatter_class=ArgumentDefaultsHelpFormatter,
diff --git a/neural_lam/custom_loggers.py b/neural_lam/custom_loggers.py
index 635f515ed..795618603 100644
--- a/neural_lam/custom_loggers.py
+++ b/neural_lam/custom_loggers.py
@@ -1,3 +1,5 @@
+"""Custom logging utilities (e.g., MLFlow wrappers) used in Neural-LAM."""
+
 # Standard library
 import sys
 
@@ -16,6 +18,18 @@ class CustomMLFlowLogger(pl.loggers.MLFlowLogger):
     """
 
     def __init__(self, experiment_name, tracking_uri, run_name):
+        """
+        Initialize the logger and start an MLflow run.
+
+        Parameters
+        ----------
+        experiment_name : str
+            Target MLflow experiment.
+        tracking_uri : str
+            MLflow tracking server URI.
+        run_name : str
+            Human-readable run name stored as ``mlflow.runName``.
+        """
         super().__init__(
             experiment_name=experiment_name, tracking_uri=tracking_uri
         )
@@ -39,14 +53,21 @@ def save_dir(self):
 
     def log_image(self, key, images, step=None):
         """
-        Log a matplotlib figure as an image to MLFlow
+        Log one or more Matplotlib figures as images in MLflow.
+
+        Parameters
+        ----------
+        key : str
+            Identifier under which to log the image.
+        images : Sequence[matplotlib.figure.Figure]
+            Figures to export; only the first element is logged.
+        step : int or None, optional
+            Optional training step index appended to ``key``.
 
-        key: str
-            Key to log the image under
-        images: list
-            List of matplotlib figures to log
-        step: Union[int, None]
-            Step to log the image under. If None, logs under the key directly
+        Raises
+        ------
+        SystemExit
+            If AWS credentials for the MLflow artifact store are missing.
         """
         # Third-party
         from botocore.exceptions import NoCredentialsError
diff --git a/neural_lam/interaction_net.py b/neural_lam/interaction_net.py
index 715b94f56..61d247aa4 100644
--- a/neural_lam/interaction_net.py
+++ b/neural_lam/interaction_net.py
@@ -1,3 +1,5 @@
+"""Interaction Network layers and helper modules used by Neural-LAM."""
+
 # Third-party
 import torch
 import torch_geometric as pyg
@@ -27,15 +29,15 @@ def __init__(
         aggr_chunk_sizes=None,
         aggr="sum",
     ):
-
         """
         Initialise an InteractionNet message-passing layer.
 
         Parameters
         ----------
         edge_index : torch.Tensor
-            Edge connectivity tensor of shape ``(2, M)`` in PyG format,
-            where ``M`` is the number of edges.
+            Edge connectivity tensor in PyG format.
+
+            * **Shape**: ``(2, M)`` where ``M`` is the number of edges.
         input_dim : int
             Dimensionality of both node and edge input representations.
         update_edges : bool, optional
@@ -45,19 +47,19 @@ def __init__(
             Number of hidden layers in each MLP. Default is ``1``.
         hidden_dim : int or None, optional
             Width of hidden layers. If ``None``, defaults to ``input_dim``.
-        edge_chunk_sizes : list of int or None, optional
+        edge_chunk_sizes : list[int] or None, optional
             Chunk sizes for splitting edge representations across separate
             MLPs. ``None`` uses a single shared MLP.
-        aggr_chunk_sizes : list of int or None, optional
+        aggr_chunk_sizes : list[int] or None, optional
             Chunk sizes for splitting aggregated node representations across
             separate MLPs. ``None`` uses a single shared MLP.
-        aggr : {'sum', 'mean'}, optional
-            Message aggregation method. Default is ``'sum'``.
+        aggr : {"sum", "mean"}, optional
+            Message aggregation method. Default is ``"sum"``.
 
         Raises
         ------
         AssertionError
-            If ``aggr`` is not one of ``'sum'`` or ``'mean'``.
+            If ``aggr`` is not one of ``"sum"`` or ``"mean"``.
         """
         assert aggr in ("sum", "mean"), f"Unknown aggregation method: {aggr}"
         super().__init__(aggr=aggr)
@@ -99,17 +101,33 @@ def __init__(
 
     def forward(self, send_rep, rec_rep, edge_rep):
         """
-        Apply interaction network to update the representations of receiver
-        nodes, and optionally the edge representations.
+        Update receiver (and optionally edge) representations via message
+        passing.
+
+        Parameters
+        ----------
+        send_rep : torch.Tensor
+            Vector representations of sender nodes.
+
+            * **Shape**: ``(N_send, d_h)``
+        rec_rep : torch.Tensor
+            Vector representations of receiver nodes.
 
-        send_rep: (N_send, d_h), vector representations of sender nodes
-        rec_rep: (N_rec, d_h), vector representations of receiver nodes
-        edge_rep: (M, d_h), vector representations of edges used
+            * **Shape**: ``(N_rec, d_h)``
+        edge_rep : torch.Tensor
+            Edge representations used during message passing.
 
-        Returns:
-        rec_rep: (N_rec, d_h), updated vector representations of receiver nodes
-        (optionally) edge_rep: (M, d_h), updated vector representations
-            of edges
+            * **Shape**: ``(M, d_h)``
+
+        Returns
+        -------
+        torch.Tensor or tuple[torch.Tensor, torch.Tensor]
+            Updated receiver representations. If ``self.update_edges`` is
+            ``True``, the tuple ``(rec_rep, edge_rep)`` containing the updated
+            receiver and edge representations is returned.
+
+            * **Shape**: ``(N_rec, d_h)`` for receivers and ``(M, d_h)`` for
+              edges.
         """
         # Always concatenate to [rec_nodes, send_nodes] for propagation,
         # but only aggregate to rec_nodes
@@ -129,18 +147,12 @@ def forward(self, send_rep, rec_rep, edge_rep):
         return rec_rep
 
     def message(self, x_j, x_i, edge_attr):
-        """
-        Compute messages from node j to node i.
-        """
+        """Compute messages from node ``j`` to ``i`` using edge features."""
         return self.edge_mlp(torch.cat((edge_attr, x_j, x_i), dim=-1))
 
     # pylint: disable-next=signature-differs
     def aggregate(self, inputs, index, ptr, dim_size):
-        """
-        Overridden aggregation function to:
-        * return both aggregated and original messages,
-        * only aggregate to number of receiver nodes.
-        """
+        """Aggregate messages while also returning the per-edge values."""
         aggr = super().aggregate(inputs, index, ptr, self.num_rec)
         return aggr, inputs
 
@@ -153,6 +165,21 @@ class SplitMLPs(nn.Module):
     """
 
     def __init__(self, mlps, chunk_sizes):
+        """
+        Create a module that dispatches chunks of the input to separate MLPs.
+
+        Parameters
+        ----------
+        mlps : Iterable[nn.Module]
+            Sequence of MLPs to apply to each chunk.
+        chunk_sizes : Sequence[int]
+            Sizes used when splitting the input along ``dim=-2``.
+
+        Raises
+        ------
+        AssertionError
+            If the number of ``mlps`` and ``chunk_sizes`` differ.
+        """
         super().__init__()
         assert len(mlps) == len(
             chunk_sizes
@@ -163,12 +190,21 @@ def __init__(self, mlps, chunk_sizes):
 
     def forward(self, x):
         """
-        Chunk up input and feed through MLPs
+        Chunk up input tensor and feed each slice through its MLP.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor to split and process.
+
+            * **Shape**: ``(..., N, d)`` where ``N = sum(chunk_sizes)``.
 
-        x: (..., N, d), where N = sum(chunk_sizes)
+        Returns
+        -------
+        torch.Tensor
+            Concatenated MLP outputs assembled along the chunk dimension.
 
-        Returns:
-        joined_output: (..., N, d), concatenated results from the MLPs
+            * **Shape**: ``(..., N, d)``
         """
         chunks = torch.split(x, self.chunk_sizes, dim=-2)
         chunk_outputs = [
diff --git a/neural_lam/loss_weighting.py b/neural_lam/loss_weighting.py
index c842b2023..3335fb6d7 100644
--- a/neural_lam/loss_weighting.py
+++ b/neural_lam/loss_weighting.py
@@ -1,3 +1,5 @@
+"""Utility functions for configuring state-feature loss weighting."""
+
 # Local
 from .config import (
     ManualStateFeatureWeighting,
diff --git a/neural_lam/metrics.py b/neural_lam/metrics.py
index a078b9cc5..2d2eeb941 100644
--- a/neural_lam/metrics.py
+++ b/neural_lam/metrics.py
@@ -1,15 +1,27 @@
+"""Evaluation metrics shared across training and validation routines."""
+
 # Third-party
 import torch
 
 
 def get_metric(metric_name):
     """
-    Get a defined metric with given name
+    Retrieve a registered metric function by name.
 
-    metric_name: str, name of the metric
+    Parameters
+    ----------
+    metric_name : str
+        Name of the metric to load (case-insensitive).
 
-    Returns:
-    metric: function implementing the metric
+    Returns
+    -------
+    callable
+        Metric function implementing the requested metric.
+
+    Raises
+    ------
+    AssertionError
+        If ``metric_name`` is not part of :data:`DEFINED_METRICS`.
     """
     metric_name_lower = metric_name.lower()
     assert (
@@ -25,25 +37,29 @@ def mask_and_reduce_metric(metric_entry_vals, mask, average_grid, sum_vars):
     Parameters
     ----------
     metric_entry_vals : torch.Tensor
-        Entry-wise metric values of shape ``(..., N, d_state)``, where
-        ``...`` denotes any number of broadcastable batch dimensions.
+        Entry-wise metric values.
+
+        * **Shape**: ``(..., N, d_state)`` where ``...`` are broadcastable
+          leading dimensions.
     mask : torch.Tensor or None
-        Boolean mask of shape ``(N,)`` selecting which grid nodes to include.
-        Pass ``None`` to use all nodes.
+        Boolean mask selecting which grid nodes to include. Pass ``None`` to
+        use all nodes.
+
+        * **Shape**: ``(N,)``
     average_grid : bool
         If ``True``, reduce the grid dimension ``N`` by taking the mean,
-        producing shape ``(..., d_state)``.
+        producing ``(..., d_state)``.
     sum_vars : bool
         If ``True``, reduce the variable dimension ``d_state`` by summing,
-        producing shape ``(..., N)`` or ``(...,)`` depending on
-        ``average_grid``.
+        producing ``(..., N)`` or ``(...,)`` depending on ``average_grid``.
 
     Returns
     -------
     torch.Tensor
-        Reduced metric tensor. Shape is one of ``(...,)``,
-        ``(..., d_state)``, ``(..., N)``, or ``(..., N, d_state)``
-        depending on the reduction arguments.
+        Reduced metric tensor.
+
+        * **Shape**: one of ``(...,)``, ``(..., d_state)``, ``(..., N)``, or
+          ``(..., N, d_state)`` depending on the reduction flags.
     """
     # Only keep grid nodes in mask
     if mask is not None:
@@ -75,15 +91,21 @@ def wmse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     Parameters
     ----------
     pred : torch.Tensor
-        Model predictions of shape ``(..., N, d_state)``.
+        Model predictions.
+
+        * **Shape**: ``(..., N, d_state)``
     target : torch.Tensor
-        Ground-truth values of shape ``(..., N, d_state)``.
+        Ground-truth values.
+
+        * **Shape**: ``(..., N, d_state)``
     pred_std : torch.Tensor
-        Predicted standard deviation of shape ``(..., N, d_state)`` or
-        ``(d_state,)`` used as the per-entry weighting.
+        Predicted standard deviation used as the per-entry weighting.
+
+        * **Shape**: ``(..., N, d_state)`` or ``(d_state,)``
     mask : torch.Tensor or None, optional
-        Boolean mask of shape ``(N,)`` selecting grid nodes. Default is
-        ``None`` (all nodes used).
+        Boolean mask selecting grid nodes. Default is ``None`` (all nodes).
+
+        * **Shape**: ``(N,)``
     average_grid : bool, optional
         If ``True``, average over the grid dimension. Default is ``True``.
     sum_vars : bool, optional
@@ -92,8 +114,10 @@ def wmse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     Returns
     -------
     torch.Tensor
-        Weighted MSE, with shape determined by ``average_grid`` and
-        ``sum_vars`` (see :func:`mask_and_reduce_metric`).
+        Weighted MSE after masking and reduction (see
+        :func:`mask_and_reduce_metric`).
+
+        * **Shape**: determined by ``average_grid`` and ``sum_vars``.
     """
     entry_mse = torch.nn.functional.mse_loss(
         pred, target, reduction="none"
@@ -110,21 +134,35 @@ def wmse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
 
 def mse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     """
-    (Unweighted) Mean Squared Error
-
-    (...,) is any number of batch dimensions, potentially different
-        but broadcastable
-    pred: (..., N, d_state), prediction
-    target: (..., N, d_state), target
-    pred_std: (..., N, d_state) or (d_state,), predicted std.-dev.
-    mask: (N,), boolean mask describing which grid nodes to use in metric
-    average_grid: boolean, if grid dimension -2 should be reduced (mean over N)
-    sum_vars: boolean, if variable dimension -1 should be reduced (sum
-        over d_state)
-
-    Returns:
-    metric_val: One of (...,), (..., d_state), (..., N), (..., N, d_state),
-    depending on reduction arguments.
+    Compute the unweighted Mean Squared Error (MSE).
+
+    Parameters
+    ----------
+    pred : torch.Tensor
+        Model predictions.
+
+        * **Shape**: ``(..., N, d_state)``
+    target : torch.Tensor
+        Ground-truth values.
+
+        * **Shape**: ``(..., N, d_state)``
+    pred_std : torch.Tensor
+        Unused argument for API parity with :func:`wmse`.
+
+        * **Shape**: ``(..., N, d_state)`` or ``(d_state,)``
+    mask : torch.Tensor or None, optional
+        Boolean mask selecting grid nodes. Default is ``None`` (all nodes).
+
+        * **Shape**: ``(N,)``
+    average_grid : bool, optional
+        If ``True``, average over the grid dimension. Default is ``True``.
+    sum_vars : bool, optional
+        If ``True``, sum over the variable dimension. Default is ``True``.
+
+    Returns
+    -------
+    torch.Tensor
+        MSE with shape determined by ``average_grid`` and ``sum_vars``.
     """
     # Replace pred_std with constant ones
     return wmse(
@@ -134,21 +172,36 @@ def mse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
 
 def wmae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     """
-    Weighted Mean Absolute Error
-
-    (...,) is any number of batch dimensions, potentially different
-        but broadcastable
-    pred: (..., N, d_state), prediction
-    target: (..., N, d_state), target
-    pred_std: (..., N, d_state) or (d_state,), predicted std.-dev.
-    mask: (N,), boolean mask describing which grid nodes to use in metric
-    average_grid: boolean, if grid dimension -2 should be reduced (mean over N)
-    sum_vars: boolean, if variable dimension -1 should be reduced (sum
-        over d_state)
-
-    Returns:
-    metric_val: One of (...,), (..., d_state), (..., N), (..., N, d_state),
-    depending on reduction arguments.
+    Compute the Weighted Mean Absolute Error (wMAE).
+
+    Parameters
+    ----------
+    pred : torch.Tensor
+        Model predictions.
+
+        * **Shape**: ``(..., N, d_state)``
+    target : torch.Tensor
+        Ground-truth values.
+
+        * **Shape**: ``(..., N, d_state)``
+    pred_std : torch.Tensor
+        Predicted standard deviation used as the per-entry weighting.
+
+        * **Shape**: ``(..., N, d_state)`` or ``(d_state,)``
+    mask : torch.Tensor or None, optional
+        Boolean mask selecting grid nodes. Default is ``None`` (all nodes).
+
+        * **Shape**: ``(N,)``
+    average_grid : bool, optional
+        If ``True``, average over the grid dimension. Default is ``True``.
+    sum_vars : bool, optional
+        If ``True``, sum over the variable dimension. Default is ``True``.
+
+    Returns
+    -------
+    torch.Tensor
+        Weighted MAE with shape determined by ``average_grid`` and
+        ``sum_vars``.
     """
     entry_mae = torch.nn.functional.l1_loss(
         pred, target, reduction="none"
@@ -165,21 +218,35 @@ def wmae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
 
 def mae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     """
-    (Unweighted) Mean Absolute Error
-
-    (...,) is any number of batch dimensions, potentially different
-        but broadcastable
-    pred: (..., N, d_state), prediction
-    target: (..., N, d_state), target
-    pred_std: (..., N, d_state) or (d_state,), predicted std.-dev.
-    mask: (N,), boolean mask describing which grid nodes to use in metric
-    average_grid: boolean, if grid dimension -2 should be reduced (mean over N)
-    sum_vars: boolean, if variable dimension -1 should be reduced (sum
-        over d_state)
-
-    Returns:
-    metric_val: One of (...,), (..., d_state), (..., N), (..., N, d_state),
-    depending on reduction arguments.
+    Compute the unweighted Mean Absolute Error (MAE).
+
+    Parameters
+    ----------
+    pred : torch.Tensor
+        Model predictions.
+
+        * **Shape**: ``(..., N, d_state)``
+    target : torch.Tensor
+        Ground-truth values.
+
+        * **Shape**: ``(..., N, d_state)``
+    pred_std : torch.Tensor
+        Unused argument for compatibility with :func:`wmae`.
+
+        * **Shape**: ``(..., N, d_state)`` or ``(d_state,)``
+    mask : torch.Tensor or None, optional
+        Boolean mask selecting grid nodes. Default is ``None`` (all nodes).
+
+        * **Shape**: ``(N,)``
+    average_grid : bool, optional
+        If ``True``, average over the grid dimension. Default is ``True``.
+    sum_vars : bool, optional
+        If ``True``, sum over the variable dimension. Default is ``True``.
+
+    Returns
+    -------
+    torch.Tensor
+        MAE with shape determined by ``average_grid`` and ``sum_vars``.
     """
     # Replace pred_std with constant ones
     return wmae(
@@ -189,21 +256,36 @@ def mae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
 
 def nll(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     """
-    Negative Log Likelihood loss, for isotropic Gaussian likelihood
-
-    (...,) is any number of batch dimensions, potentially different
-        but broadcastable
-    pred: (..., N, d_state), prediction
-    target: (..., N, d_state), target
-    pred_std: (..., N, d_state) or (d_state,), predicted std.-dev.
-    mask: (N,), boolean mask describing which grid nodes to use in metric
-    average_grid: boolean, if grid dimension -2 should be reduced (mean over N)
-    sum_vars: boolean, if variable dimension -1 should be reduced (sum
-        over d_state)
-
-    Returns:
-    metric_val: One of (...,), (..., d_state), (..., N), (..., N, d_state),
-    depending on reduction arguments.
+    Compute the Negative Log Likelihood for an isotropic Gaussian likelihood.
+
+    Parameters
+    ----------
+    pred : torch.Tensor
+        Distribution mean predictions.
+
+        * **Shape**: ``(..., N, d_state)``
+    target : torch.Tensor
+        Ground-truth values.
+
+        * **Shape**: ``(..., N, d_state)``
+    pred_std : torch.Tensor
+        Predicted standard deviation parameter of the Gaussian.
+
+        * **Shape**: ``(..., N, d_state)`` or ``(d_state,)``
+    mask : torch.Tensor or None, optional
+        Boolean mask selecting grid nodes. Default is ``None`` (all nodes).
+
+        * **Shape**: ``(N,)``
+    average_grid : bool, optional
+        If ``True``, average over the grid dimension. Default is ``True``.
+    sum_vars : bool, optional
+        If ``True``, sum over the variable dimension. Default is ``True``.
+
+    Returns
+    -------
+    torch.Tensor
+        Negative log-likelihood with shape determined by ``average_grid`` and
+        ``sum_vars``.
     """
     # Broadcast pred_std if shaped (d_state,), done internally in Normal class
     dist = torch.distributions.Normal(pred, pred_std)  # (..., N, d_state)
@@ -218,22 +300,38 @@ def crps_gauss(
     pred, target, pred_std, mask=None, average_grid=True, sum_vars=True
 ):
     """
-    (Negative) Continuous Ranked Probability Score (CRPS)
-    Closed-form expression based on Gaussian predictive distribution
-
-    (...,) is any number of batch dimensions, potentially different
-            but broadcastable
-    pred: (..., N, d_state), prediction
-    target: (..., N, d_state), target
-    pred_std: (..., N, d_state) or (d_state,), predicted std.-dev.
-    mask: (N,), boolean mask describing which grid nodes to use in metric
-    average_grid: boolean, if grid dimension -2 should be reduced (mean over N)
-    sum_vars: boolean, if variable dimension -1 should be reduced (sum
-        over d_state)
-
-    Returns:
-    metric_val: One of (...,), (..., d_state), (..., N), (..., N, d_state),
-    depending on reduction arguments.
+    Compute the (negative) Continuous Ranked Probability Score (CRPS).
+
+    A closed-form expression for a Gaussian predictive distribution is used.
+
+    Parameters
+    ----------
+    pred : torch.Tensor
+        Distribution mean predictions.
+
+        * **Shape**: ``(..., N, d_state)``
+    target : torch.Tensor
+        Ground-truth values.
+
+        * **Shape**: ``(..., N, d_state)``
+    pred_std : torch.Tensor
+        Predicted standard deviation parameter of the Gaussian.
+
+        * **Shape**: ``(..., N, d_state)`` or ``(d_state,)``
+    mask : torch.Tensor or None, optional
+        Boolean mask selecting grid nodes. Default is ``None`` (all nodes).
+
+        * **Shape**: ``(N,)``
+    average_grid : bool, optional
+        If ``True``, average over the grid dimension. Default is ``True``.
+    sum_vars : bool, optional
+        If ``True``, sum over the variable dimension. Default is ``True``.
+
+    Returns
+    -------
+    torch.Tensor
+        Negative CRPS values with shape determined by ``average_grid`` and
+        ``sum_vars``.
     """
     std_normal = torch.distributions.Normal(
         torch.zeros((), device=pred.device), torch.ones((), device=pred.device)
diff --git a/neural_lam/plot_graph.py b/neural_lam/plot_graph.py
index 39b8639f3..22985a8fa 100644
--- a/neural_lam/plot_graph.py
+++ b/neural_lam/plot_graph.py
@@ -1,3 +1,5 @@
+"""Command-line utility for plotting saved Neural-LAM graphs."""
+
 # Standard library
 import os
 from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py
index 7c7a4eefe..30f882ea7 100644
--- a/neural_lam/train_model.py
+++ b/neural_lam/train_model.py
@@ -1,3 +1,5 @@
+"""CLI entry point for training Neural-LAM models."""
+
 # Standard library
 import json
 import random
diff --git a/neural_lam/utils.py b/neural_lam/utils.py
index e3030c620..f61a5cc23 100644
--- a/neural_lam/utils.py
+++ b/neural_lam/utils.py
@@ -1,3 +1,5 @@
+"""Utility helpers shared across Neural-LAM training and evaluation."""
+
 # Standard library
 import os
 import shutil
@@ -29,18 +31,31 @@ class BufferList(nn.Module):
     """
 
     def __init__(self, buffer_tensors, persistent=True):
+        """
+        Register a collection of tensors as buffers inside a module.
+
+        Parameters
+        ----------
+        buffer_tensors : Sequence[torch.Tensor]
+            Buffers to register in the order they should be indexed.
+        persistent : bool, optional
+            If ``True``, buffers are saved in checkpoints. Default ``True``.
+        """
         super().__init__()
         self.n_buffers = len(buffer_tensors)
         for buffer_i, tensor in enumerate(buffer_tensors):
             self.register_buffer(f"b{buffer_i}", tensor, persistent=persistent)
 
     def __getitem__(self, key):
+        """Return the buffer at ``key`` (0-indexed)."""
         return getattr(self, f"b{key}")
 
     def __len__(self):
+        """Return the number of registered buffers."""
         return self.n_buffers
 
     def __iter__(self):
+        """Iterate over the registered buffers in ascending index order."""
         return (self[i] for i in range(len(self)))
 
 
@@ -90,6 +105,10 @@ def load_graph(graph_dir_path, device="cpu"):
     """
 
     def loads_file(fn):
+        """Load ``torch.load`` data from ``graph_dir_path``.
+
+        Applies ``map_location`` so tensors land on the requested device.
+        """
         return torch.load(
             os.path.join(graph_dir_path, fn),
             map_location=device,
@@ -201,13 +220,20 @@ def loads_file(fn):
 
 def make_mlp(blueprint, layer_norm=True):
     """
-    Create MLP from list blueprint, with
-    input dimensionality: blueprint[0]
-    output dimensionality: blueprint[-1] and
-    hidden layers of dimensions: blueprint[1], ..., blueprint[-2]
+    Construct a multilayer perceptron from a blueprint of layer widths.
+
+    Parameters
+    ----------
+    blueprint : list[int]
+        Sequence of layer dimensions where ``blueprint[0]`` is the input size
+        and ``blueprint[-1]`` is the output size.
+    layer_norm : bool, optional
+        If ``True``, append a ``LayerNorm`` to the output as in GraphCast.
 
-    if layer_norm is True, includes a LayerNorm layer at
-    the output (as used in GraphCast)
+    Returns
+    -------
+    torch.nn.Sequential
+        Sequential module implementing the specified MLP.
     """
     hidden_layers = len(blueprint) - 2
     assert hidden_layers >= 0, "Invalid MLP blueprint"
@@ -228,7 +254,12 @@ def make_mlp(blueprint, layer_norm=True):
 @cache
 def has_working_latex():
     """
-    Check if LaTeX is available or its toolchain
+    Check whether a LaTeX toolchain is available on the system.
+
+    Returns
+    -------
+    bool
+        ``True`` if ``latex`` and the required auxiliary tools are callable.
     """
     # If latex/toolchain is not available, some visualizations might not render
     # correctly, but will at least not raise an error. Alternatively, use
@@ -290,8 +321,17 @@ def has_working_latex():
 
 def fractional_plot_bundle(fraction):
     """
-    Get the tueplots bundle, but with figure width as a fraction of
-    the page width.
+    Return a ``tueplots`` bundle scaled to a fraction of the page width.
+
+    Parameters
+    ----------
+    fraction : float
+        Denominator applied to the default NeurIPS figure width.
+
+    Returns
+    -------
+    dict
+        Matplotlib rcParams bundle with updated ``figure.figsize``.
     """
 
     usetex = has_working_latex()
@@ -307,13 +347,20 @@ def fractional_plot_bundle(fraction):
 
 @rank_zero_only
 def rank_zero_print(*args, **kwargs):
-    """Print only from rank 0 process"""
+    """Print arguments only from the rank-zero process in distributed runs."""
     print(*args, **kwargs)
 
 
 def init_training_logger_metrics(training_logger, val_steps):
     """
-    Set up logger metrics to track
+    Configure validation metric aggregation for the active training logger.
+
+    Parameters
+    ----------
+    training_logger : pytorch_lightning.loggers.Logger
+        Logger instance used during training.
+    val_steps : Iterable[int]
+        Autoregressive rollout lengths to log as separate metrics.
     """
     experiment = training_logger.experiment
     if isinstance(training_logger, WandbLogger):
@@ -332,15 +379,14 @@ def init_training_logger_metrics(training_logger, val_steps):
 @rank_zero_only
 def setup_training_logger(datastore, args, run_name):
     """
+    Instantiate the configured experiment logger.
 
     Parameters
     ----------
-    datastore : Datastore
-        Datastore object.
-
+    datastore : BaseDatastore
+        Datastore providing metadata for logging configuration.
     args : argparse.Namespace
-        Arguments from command line.
-
+        Parsed training arguments controlling the logger backend.
     run_name : str
         Name of the run.
 
@@ -376,13 +422,21 @@ def setup_training_logger(datastore, args, run_name):
 
 def inverse_softplus(x, beta=1, threshold=20):
     """
-    Inverse of torch.nn.functional.softplus
+    Approximate the inverse of :func:`torch.nn.functional.softplus`.
 
-    Input is clamped to approximately positive values of x, and the function is
-    linear for inputs above x*beta for numerical stability.
+    Parameters
+    ----------
+    x : torch.Tensor
+        Input tensor whose softplus inverse should be computed.
+    beta : float, optional
+        Softplus ``beta`` parameter that controls the sharpness. Default ``1``.
+    threshold : float, optional
+        Threshold applied to the input for numerical stability. Default ``20``.
 
-    Note that this torch.clamp will make gradients 0, but this is not a
-    problem as values of x that are this close to 0 have gradients of 0 anyhow.
+    Returns
+    -------
+    torch.Tensor
+        Tensor containing the inverse-softplus values.
     """
     x_clamped = torch.clamp(
         x, min=torch.log(torch.tensor(1e-6 + 1)) / beta, max=threshold / beta
@@ -399,12 +453,18 @@ def inverse_softplus(x, beta=1, threshold=20):
 
 def inverse_sigmoid(x):
     """
-    Inverse of torch.sigmoid
+    Compute the logit (inverse sigmoid) while clamping to ``(0, 1)``.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Input tensor assumed to contain logits after a sigmoid.
 
-    Sigmoid output takes values in [0,1], this makes sure input is just within
-    this interval.
-    Note that this torch.clamp will make gradients 0, but this is not a problem
-    as values of x that are this close to 0 or 1 have gradients of 0 anyhow.
+    Returns
+    -------
+    torch.Tensor
+        Tensor containing ``log(x / (1 - x))`` after clamping away from the
+        saturation limits.
     """
     x_clamped = torch.clamp(x, min=1e-6, max=1 - 1e-6)
     return torch.log(x_clamped / (1 - x_clamped))
@@ -412,26 +472,30 @@ def inverse_sigmoid(x):
 
 def get_integer_time(tdelta) -> tuple[int, str]:
     """
-    Get the largest time unit that can represent the given timedelta as an
-    integer.
-
-    Returns:
-        int: The integer value of the timedelta in the largest time unit, or
-                1 if no such unit exists.
-        str: The time unit as a string ('weeks', 'days', 'hours', 'minutes',
-                'seconds', 'milliseconds', 'microseconds'). If no unit can
-                represent the timedelta as an integer, returns 'unknown'.
-
-    Examples:
-        >>> from datetime import timedelta
-        >>> get_integer_time(timedelta(days=14))
-        (2, 'weeks')
-        >>> get_integer_time(timedelta(hours=5))
-        (5, 'hours')
-        >>> get_integer_time(timedelta(milliseconds=1000))
-        (1, 'seconds')
-        >>> get_integer_time(timedelta(days=0.001))
-        (1, 'unknown')
+    Express a :class:`datetime.timedelta` as an integer number of time units.
+
+    Parameters
+    ----------
+    tdelta : datetime.timedelta
+        Time interval to convert.
+
+    Returns
+    -------
+    tuple[int, str]
+        Integer value and the corresponding unit (e.g. ``"hours"``). If no
+        unit yields an integer count, ``(1, "unknown")`` is returned.
+
+    Examples
+    --------
+    >>> from datetime import timedelta
+    >>> get_integer_time(timedelta(days=14))
+    (2, 'weeks')
+    >>> get_integer_time(timedelta(hours=5))
+    (5, 'hours')
+    >>> get_integer_time(timedelta(milliseconds=1000))
+    (1, 'seconds')
+    >>> get_integer_time(timedelta(days=0.001))
+    (1, 'unknown')
     """
     total_seconds = tdelta.total_seconds()
 
diff --git a/neural_lam/vis.py b/neural_lam/vis.py
index 3db4365a4..cdb58f3e2 100644
--- a/neural_lam/vis.py
+++ b/neural_lam/vis.py
@@ -1,3 +1,5 @@
+"""Visualization helpers for analysing Neural-LAM predictions and errors."""
+
 # Third-party
 import matplotlib
 import matplotlib.pyplot as plt
@@ -12,9 +14,23 @@
 @matplotlib.rc_context(utils.fractional_plot_bundle(1))
 def plot_error_map(errors, datastore: BaseRegularGridDatastore, title=None):
     """
-    Plot a heatmap of errors of different variables at different
-    predictions horizons
-    errors: (pred_steps, d_f)
+    Plot a heatmap of per-variable errors across prediction horizons.
+
+    Parameters
+    ----------
+    errors : torch.Tensor
+        Error values for each horizon and feature.
+
+        * **Shape**: ``(pred_steps, d_f)``
+    datastore : BaseRegularGridDatastore
+        Datastore providing metadata for labels and units.
+    title : str or None, optional
+        Optional plot title.
+
+    Returns
+    -------
+    matplotlib.figure.Figure
+        Figure handle containing the rendered heatmap.
     """
     errors_np = errors.T.cpu().numpy()  # (d_f, pred_steps)
     d_f, pred_steps = errors_np.shape
@@ -75,10 +91,29 @@ def plot_prediction(
     vrange=None,
 ):
     """
-    Plot example prediction and grond truth.
-
-    Each has shape (N_grid,)
-
+    Plot a prediction alongside the corresponding ground truth field.
+
+    Parameters
+    ----------
+    datastore : BaseRegularGridDatastore
+        Datastore providing coordinate metadata and projection details.
+    da_prediction : xarray.DataArray
+        Predicted field flattened over the grid.
+
+        * **Shape**: ``(N_grid,)``
+    da_target : xarray.DataArray
+        Ground-truth field flattened over the grid.
+
+        * **Shape**: ``(N_grid,)``
+    title : str or None, optional
+        Optional figure title.
+    vrange : tuple[float, float] or None, optional
+        Explicit value range ``(vmin, vmax)`` for the color scale.
+
+    Returns
+    -------
+    matplotlib.figure.Figure
+        Figure handle containing the two subplots.
     """
     # Get common scale for values
     if vrange is None:
@@ -131,8 +166,25 @@ def plot_spatial_error(
     error, datastore: BaseRegularGridDatastore, title=None, vrange=None
 ):
     """
-    Plot errors over spatial map
-    Error and obs_mask has shape (N_grid,)
+    Plot spatial error magnitudes on the datastore grid.
+
+    Parameters
+    ----------
+    error : torch.Tensor
+        Error magnitudes on the flattened grid.
+
+        * **Shape**: ``(N_grid,)``
+    datastore : BaseRegularGridDatastore
+        Datastore providing coordinate metadata and boundary masks.
+    title : str or None, optional
+        Optional figure title.
+    vrange : tuple[float, float] or None, optional
+        Explicit value range ``(vmin, vmax)`` for the color scale.
+
+    Returns
+    -------
+    matplotlib.figure.Figure
+        Figure handle containing the plotted map.
     """
     # Get common scale for values
     if vrange is None:
diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py
index 8a9cc8253..fcb02ded0 100644
--- a/neural_lam/weather_dataset.py
+++ b/neural_lam/weather_dataset.py
@@ -1,3 +1,5 @@
+"""Dataset helpers wrapping Neural-LAM datastores for PyTorch Lightning."""
+
 # Standard library
 import datetime
 import warnings
@@ -49,6 +51,26 @@ def __init__(
         num_future_forcing_steps: int = 1,
         standardize: bool = True,
     ):
+        """
+        Parameters
+        ----------
+        datastore : BaseDatastore
+            Datastore providing access to state/forcing/static arrays.
+        split : str, optional
+            Data split (``"train"``, ``"val"``, or ``"test"``).
+            Default ``"train"``.
+        ar_steps : int, optional
+            Number of autoregressive steps per training sample. Default ``3``.
+        num_past_forcing_steps : int, optional
+            Past forcing window length ``i`` so that ``[t-i, ..., t]`` forcings
+            are concatenated. Default ``1``.
+        num_future_forcing_steps : int, optional
+            Future forcing window length ``j`` so that ``[t, ..., t+j]``
+            forcings are available. Default ``1``.
+        standardize : bool, optional
+            If ``True``, normalize state/forcing arrays via datastore stats.
+        """
+
         super().__init__()
 
         self.split = split
@@ -116,6 +138,14 @@ def __init__(
                 self.da_forcing_std = self.ds_forcing_stats.forcing_std
 
     def __len__(self):
+        """
+        Return the number of autoregressive training samples available.
+
+        Returns
+        -------
+        int
+            Number of (init, target) pairs derivable from the datastore.
+        """
         if self.datastore.is_forecast:
             # for now we simply create a single sample for each analysis time
             # and then take the first (2 + ar_steps) forecast times. In
@@ -544,7 +574,7 @@ def create_dataarray_from_tensor(
         """
 
         def _is_listlike(obj):
-            # match list, tuple, numpy array
+            """Return ``True`` for list/tuple/ndarray-like containers."""
             return hasattr(obj, "__iter__") and not isinstance(obj, str)
 
         add_time_as_dim = False
@@ -616,6 +646,28 @@ def __init__(
         num_workers: int = 16,
         eval_split: str = "test",
     ):
+        """
+        Parameters
+        ----------
+        datastore : BaseDatastore
+            Datastore used for all splits.
+        ar_steps_train : int, optional
+            Number of AR steps for training batches. Default ``3``.
+        ar_steps_eval : int, optional
+            Number of AR steps for validation/test batches. Default ``25``.
+        standardize : bool, optional
+            If ``True``, datasets are returned standardized. Default ``True``.
+        num_past_forcing_steps : int, optional
+            Number of past forcing steps to include. Default ``1``.
+        num_future_forcing_steps : int, optional
+            Number of future forcing steps to include. Default ``1``.
+        batch_size : int, optional
+            Mini-batch size for dataloaders. Default ``4``.
+        num_workers : int, optional
+            Number of background workers per dataloader. Default ``16``.
+        eval_split : str, optional
+            Dataset split to use for ``test_dataloader``. Default ``"test"``.
+        """
         super().__init__()
         self._datastore = datastore
         self.num_past_forcing_steps = num_past_forcing_steps
@@ -636,6 +688,15 @@ def __init__(
             self.multiprocessing_context = "spawn"
 
     def setup(self, stage=None):
+        """
+        Instantiate datasets for the requested trainer stage.
+
+        Parameters
+        ----------
+        stage : str or None, optional
+            Trainer stage identifier (``"fit"``/``"test"``/``None``). When
+            ``None``, both train and evaluation datasets are created.
+        """
         if stage == "fit" or stage is None:
             self.train_dataset = WeatherDataset(
                 datastore=self._datastore,

From dff0bc1ce5ca63a1cf4d607a4af56c96a97b8f69 Mon Sep 17 00:00:00 2001
From: Mohit-Lakra <mohitlakra2007777@gmail.com>
Date: Tue, 3 Mar 2026 23:32:42 +0530
Subject: [PATCH 06/16] Precommit Updated for 100 percent docs coverage check

---
 .pre-commit-config.yaml | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index fe2ff11ae..45bb18444 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -36,9 +36,16 @@ repos:
       - id: flake8
         description: Check Python code for correctness, consistency and adherence to best practices
         additional_dependencies: [Flake8-pyproject]
+  - repo: https://github.com/econchick/interrogate
+    rev: 1.7.0
+    hooks:
+      - id: interrogate
+        description: Ensure documentation coverage stays perfect
+        pass_filenames: false
+        args: ["--fail-under=100", "neural_lam"]
   - repo: https://github.com/pre-commit/mirrors-mypy
     rev: v1.19.0
     hooks:
-    -   id: mypy
+      - id: mypy
         additional_dependencies: [types-PyYAML, types-Pillow, types-tqdm]
         description: Check for type errors

From 2c1017e86088419f9a73a2df2f4a0c010393e0d1 Mon Sep 17 00:00:00 2001
From: Mohit-Lakra <mohitlakra2007777@gmail.com>
Date: Tue, 17 Mar 2026 08:25:33 +0530
Subject: [PATCH 07/16] docs: address reviewer feedback on PR #252 docstrings

---
 .../compute_standardization_stats.py          |   6 +-
 neural_lam/interaction_net.py                 |  14 +-
 neural_lam/metrics.py                         | 138 +++++++++++-------
 neural_lam/models/ar_model.py                 |  94 +++++++-----
 neural_lam/models/base_graph_model.py         |  66 +++++----
 neural_lam/models/base_hi_graph_model.py      |  39 +++--
 neural_lam/utils.py                           |  32 +++-
 neural_lam/vis.py                             |   8 +-
 neural_lam/weather_dataset.py                 |  29 +---
 9 files changed, 248 insertions(+), 178 deletions(-)

diff --git a/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py b/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py
index 7c6790312..ddc9d2687 100644
--- a/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py
+++ b/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py
@@ -129,9 +129,11 @@ def save_stats(
     squares : Sequence[torch.Tensor]
         Batch-wise second moments with shape ``(N_batch, d_features)``.
     flux_means : Sequence[torch.Tensor]
-        Optional flux means of shape ``(N_batch,)``.
+        Flux means of shape ``(N_batch,)``; pass an empty sequence to skip
+        saving.
     flux_squares : Sequence[torch.Tensor]
-        Optional flux second moments of shape ``(N_batch,)``.
+        Flux second moments of shape ``(N_batch,)``; pass an empty sequence to
+        skip saving.
     filename_prefix : str
         Prefix (e.g., ``"parameter"`` or ``"diff"``) for saved tensors.
     """
diff --git a/neural_lam/interaction_net.py b/neural_lam/interaction_net.py
index 61d247aa4..27048fde1 100644
--- a/neural_lam/interaction_net.py
+++ b/neural_lam/interaction_net.py
@@ -37,7 +37,7 @@ def __init__(
         edge_index : torch.Tensor
             Edge connectivity tensor in PyG format.
 
-            * **Shape**: ``(2, M)`` where ``M`` is the number of edges.
+            * **Shape**: ``(2, num_edges)``
         input_dim : int
             Dimensionality of both node and edge input representations.
         update_edges : bool, optional
@@ -109,15 +109,15 @@ def forward(self, send_rep, rec_rep, edge_rep):
         send_rep : torch.Tensor
             Vector representations of sender nodes.
 
-            * **Shape**: ``(N_send, d_h)``
+            * **Shape**: ``(num_send, input_dim)``
         rec_rep : torch.Tensor
             Vector representations of receiver nodes.
 
-            * **Shape**: ``(N_rec, d_h)``
+            * **Shape**: ``(num_rec, input_dim)``
         edge_rep : torch.Tensor
             Edge representations used during message passing.
 
-            * **Shape**: ``(M, d_h)``
+            * **Shape**: ``(num_edges, input_dim)``
 
         Returns
         -------
@@ -126,8 +126,8 @@ def forward(self, send_rep, rec_rep, edge_rep):
             ``True``, the tuple ``(rec_rep, edge_rep)`` containing the updated
             receiver and edge representations is returned.
 
-            * **Shape**: ``(N_rec, d_h)`` for receivers and ``(M, d_h)`` for
-              edges.
+            * **Shape**: ``(num_rec, hidden_dim)`` for receivers and
+              ``(num_edges, hidden_dim)`` for edges.
         """
         # Always concatenate to [rec_nodes, send_nodes] for propagation,
         # but only aggregate to rec_nodes
@@ -147,7 +147,7 @@ def forward(self, send_rep, rec_rep, edge_rep):
         return rec_rep
 
     def message(self, x_j, x_i, edge_attr):
-        """Compute messages from node ``j`` to ``i`` using edge features."""
+        """Compute messages from node ``j`` to ``i``."""
         return self.edge_mlp(torch.cat((edge_attr, x_j, x_i), dim=-1))
 
     # pylint: disable-next=signature-differs
diff --git a/neural_lam/metrics.py b/neural_lam/metrics.py
index 2d2eeb941..fc5586f16 100644
--- a/neural_lam/metrics.py
+++ b/neural_lam/metrics.py
@@ -39,43 +39,45 @@ def mask_and_reduce_metric(metric_entry_vals, mask, average_grid, sum_vars):
     metric_entry_vals : torch.Tensor
         Entry-wise metric values.
 
-        * **Shape**: ``(..., N, d_state)`` where ``...`` are broadcastable
-          leading dimensions.
+        * **Shape**: ``(..., num_grid_nodes, num_variables)`` where leading
+          dimensions are broadcastable.
     mask : torch.Tensor or None
         Boolean mask selecting which grid nodes to include. Pass ``None`` to
         use all nodes.
 
-        * **Shape**: ``(N,)``
+        * **Shape**: ``(num_grid_nodes,)``
     average_grid : bool
-        If ``True``, reduce the grid dimension ``N`` by taking the mean,
-        producing ``(..., d_state)``.
+        If ``True``, reduce ``num_grid_nodes`` by taking the mean,
+        producing ``(..., num_variables)``.
     sum_vars : bool
-        If ``True``, reduce the variable dimension ``d_state`` by summing,
-        producing ``(..., N)`` or ``(...,)`` depending on ``average_grid``.
+        If ``True``, reduce the variable dimension ``num_variables`` by
+        summing, producing ``(..., num_grid_nodes)`` or ``(...,)`` depending on
+        ``average_grid``.
 
     Returns
     -------
     torch.Tensor
         Reduced metric tensor.
 
-        * **Shape**: one of ``(...,)``, ``(..., d_state)``, ``(..., N)``, or
-          ``(..., N, d_state)`` depending on the reduction flags.
+        * **Shape**: one of ``(...,)``, ``(..., num_variables)``,
+          ``(..., num_grid_nodes)``, or ``(..., num_grid_nodes, num_variables)``
+          depending on the reduction flags.
     """
     # Only keep grid nodes in mask
     if mask is not None:
         metric_entry_vals = metric_entry_vals[
             ..., mask, :
-        ]  # (..., N', d_state)
+        ]  # (..., num_selected_nodes, num_variables)
 
     # Optionally reduce last two dimensions
     if average_grid:  # Reduce grid first
         metric_entry_vals = torch.mean(
             metric_entry_vals, dim=-2
-        )  # (..., d_state)
+        )  # (..., num_variables)
     if sum_vars:  # Reduce vars second
         metric_entry_vals = torch.sum(
             metric_entry_vals, dim=-1
-        )  # (..., N) or (...,)
+        )  # (..., num_grid_nodes) or (...,)
 
     return metric_entry_vals
 
@@ -93,19 +95,20 @@ def wmse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     pred : torch.Tensor
         Model predictions.
 
-        * **Shape**: ``(..., N, d_state)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)``
     target : torch.Tensor
         Ground-truth values.
 
-        * **Shape**: ``(..., N, d_state)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)``
     pred_std : torch.Tensor
         Predicted standard deviation used as the per-entry weighting.
 
-        * **Shape**: ``(..., N, d_state)`` or ``(d_state,)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)`` or
+          ``(num_variables,)``
     mask : torch.Tensor or None, optional
         Boolean mask selecting grid nodes. Default is ``None`` (all nodes).
 
-        * **Shape**: ``(N,)``
+        * **Shape**: ``(num_grid_nodes,)``
     average_grid : bool, optional
         If ``True``, average over the grid dimension. Default is ``True``.
     sum_vars : bool, optional
@@ -121,8 +124,10 @@ def wmse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     """
     entry_mse = torch.nn.functional.mse_loss(
         pred, target, reduction="none"
-    )  # (..., N, d_state)
-    entry_mse_weighted = entry_mse / (pred_std**2)  # (..., N, d_state)
+    )  # (..., num_grid_nodes, num_variables)
+    entry_mse_weighted = entry_mse / (
+        pred_std**2
+    )  # (..., num_grid_nodes, num_variables)
 
     return mask_and_reduce_metric(
         entry_mse_weighted,
@@ -141,19 +146,20 @@ def mse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     pred : torch.Tensor
         Model predictions.
 
-        * **Shape**: ``(..., N, d_state)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)``
     target : torch.Tensor
         Ground-truth values.
 
-        * **Shape**: ``(..., N, d_state)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)``
     pred_std : torch.Tensor
         Unused argument for API parity with :func:`wmse`.
 
-        * **Shape**: ``(..., N, d_state)`` or ``(d_state,)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)`` or
+          ``(num_variables,)``
     mask : torch.Tensor or None, optional
         Boolean mask selecting grid nodes. Default is ``None`` (all nodes).
 
-        * **Shape**: ``(N,)``
+        * **Shape**: ``(num_grid_nodes,)``
     average_grid : bool, optional
         If ``True``, average over the grid dimension. Default is ``True``.
     sum_vars : bool, optional
@@ -162,7 +168,10 @@ def mse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     Returns
     -------
     torch.Tensor
-        MSE with shape determined by ``average_grid`` and ``sum_vars``.
+        MSE after masking and reduction (see
+        :func:`mask_and_reduce_metric`).
+
+        * **Shape**: determined by ``average_grid`` and ``sum_vars``.
     """
     # Replace pred_std with constant ones
     return wmse(
@@ -174,24 +183,29 @@ def wmae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     """
     Compute the Weighted Mean Absolute Error (wMAE).
 
+    Scales the absolute error at each grid node and variable by the inverse
+    standard deviation ``1 / pred_std``, then applies masking and reduction via
+    :func:`mask_and_reduce_metric`.
+
     Parameters
     ----------
     pred : torch.Tensor
         Model predictions.
 
-        * **Shape**: ``(..., N, d_state)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)``
     target : torch.Tensor
         Ground-truth values.
 
-        * **Shape**: ``(..., N, d_state)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)``
     pred_std : torch.Tensor
         Predicted standard deviation used as the per-entry weighting.
 
-        * **Shape**: ``(..., N, d_state)`` or ``(d_state,)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)`` or
+          ``(num_variables,)``
     mask : torch.Tensor or None, optional
         Boolean mask selecting grid nodes. Default is ``None`` (all nodes).
 
-        * **Shape**: ``(N,)``
+        * **Shape**: ``(num_grid_nodes,)``
     average_grid : bool, optional
         If ``True``, average over the grid dimension. Default is ``True``.
     sum_vars : bool, optional
@@ -200,13 +214,17 @@ def wmae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     Returns
     -------
     torch.Tensor
-        Weighted MAE with shape determined by ``average_grid`` and
-        ``sum_vars``.
+        Weighted MAE after masking and reduction (see
+        :func:`mask_and_reduce_metric`).
+
+        * **Shape**: determined by ``average_grid`` and ``sum_vars``.
     """
     entry_mae = torch.nn.functional.l1_loss(
         pred, target, reduction="none"
-    )  # (..., N, d_state)
-    entry_mae_weighted = entry_mae / pred_std  # (..., N, d_state)
+    )  # (..., num_grid_nodes, num_variables)
+    entry_mae_weighted = (
+        entry_mae / pred_std
+    )  # (..., num_grid_nodes, num_variables)
 
     return mask_and_reduce_metric(
         entry_mae_weighted,
@@ -225,19 +243,20 @@ def mae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     pred : torch.Tensor
         Model predictions.
 
-        * **Shape**: ``(..., N, d_state)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)``
     target : torch.Tensor
         Ground-truth values.
 
-        * **Shape**: ``(..., N, d_state)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)``
     pred_std : torch.Tensor
         Unused argument for compatibility with :func:`wmae`.
 
-        * **Shape**: ``(..., N, d_state)`` or ``(d_state,)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)`` or
+          ``(num_variables,)``
     mask : torch.Tensor or None, optional
         Boolean mask selecting grid nodes. Default is ``None`` (all nodes).
 
-        * **Shape**: ``(N,)``
+        * **Shape**: ``(num_grid_nodes,)``
     average_grid : bool, optional
         If ``True``, average over the grid dimension. Default is ``True``.
     sum_vars : bool, optional
@@ -246,7 +265,10 @@ def mae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     Returns
     -------
     torch.Tensor
-        MAE with shape determined by ``average_grid`` and ``sum_vars``.
+        MAE after masking and reduction (see
+        :func:`mask_and_reduce_metric`).
+
+        * **Shape**: determined by ``average_grid`` and ``sum_vars``.
     """
     # Replace pred_std with constant ones
     return wmae(
@@ -263,19 +285,20 @@ def nll(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     pred : torch.Tensor
         Distribution mean predictions.
 
-        * **Shape**: ``(..., N, d_state)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)``
     target : torch.Tensor
         Ground-truth values.
 
-        * **Shape**: ``(..., N, d_state)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)``
     pred_std : torch.Tensor
         Predicted standard deviation parameter of the Gaussian.
 
-        * **Shape**: ``(..., N, d_state)`` or ``(d_state,)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)`` or
+          ``(num_variables,)``
     mask : torch.Tensor or None, optional
         Boolean mask selecting grid nodes. Default is ``None`` (all nodes).
 
-        * **Shape**: ``(N,)``
+        * **Shape**: ``(num_grid_nodes,)``
     average_grid : bool, optional
         If ``True``, average over the grid dimension. Default is ``True``.
     sum_vars : bool, optional
@@ -284,12 +307,16 @@ def nll(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     Returns
     -------
     torch.Tensor
-        Negative log-likelihood with shape determined by ``average_grid`` and
-        ``sum_vars``.
+        Negative log-likelihood after masking and reduction (see
+        :func:`mask_and_reduce_metric`).
+
+        * **Shape**: determined by ``average_grid`` and ``sum_vars``.
     """
-    # Broadcast pred_std if shaped (d_state,), done internally in Normal class
-    dist = torch.distributions.Normal(pred, pred_std)  # (..., N, d_state)
-    entry_nll = -dist.log_prob(target)  # (..., N, d_state)
+    # Broadcast pred_std if shaped (num_variables,) via distribution internals
+    dist = torch.distributions.Normal(
+        pred, pred_std
+    )  # (..., num_grid_nodes, num_variables)
+    entry_nll = -dist.log_prob(target)  # (..., num_grid_nodes, num_variables)
 
     return mask_and_reduce_metric(
         entry_nll, mask=mask, average_grid=average_grid, sum_vars=sum_vars
@@ -309,19 +336,20 @@ def crps_gauss(
     pred : torch.Tensor
         Distribution mean predictions.
 
-        * **Shape**: ``(..., N, d_state)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)``
     target : torch.Tensor
         Ground-truth values.
 
-        * **Shape**: ``(..., N, d_state)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)``
     pred_std : torch.Tensor
         Predicted standard deviation parameter of the Gaussian.
 
-        * **Shape**: ``(..., N, d_state)`` or ``(d_state,)``
+        * **Shape**: ``(..., num_grid_nodes, num_variables)`` or
+          ``(num_variables,)``
     mask : torch.Tensor or None, optional
         Boolean mask selecting grid nodes. Default is ``None`` (all nodes).
 
-        * **Shape**: ``(N,)``
+        * **Shape**: ``(num_grid_nodes,)``
     average_grid : bool, optional
         If ``True``, average over the grid dimension. Default is ``True``.
     sum_vars : bool, optional
@@ -330,19 +358,23 @@ def crps_gauss(
     Returns
     -------
     torch.Tensor
-        Negative CRPS values with shape determined by ``average_grid`` and
-        ``sum_vars``.
+        Negative CRPS values after masking and reduction (see
+        :func:`mask_and_reduce_metric`).
+
+        * **Shape**: determined by ``average_grid`` and ``sum_vars``.
     """
     std_normal = torch.distributions.Normal(
         torch.zeros((), device=pred.device), torch.ones((), device=pred.device)
     )
-    target_standard = (target - pred) / pred_std  # (..., N, d_state)
+    target_standard = (
+        target - pred
+    ) / pred_std  # (..., num_grid_nodes, num_variables)
 
     entry_crps = -pred_std * (
         torch.pi ** (-0.5)
         - 2 * torch.exp(std_normal.log_prob(target_standard))
         - target_standard * (2 * std_normal.cdf(target_standard) - 1)
-    )  # (..., N, d_state)
+    )  # (..., num_grid_nodes, num_variables)
 
     return mask_and_reduce_metric(
         entry_crps, mask=mask, average_grid=average_grid, sum_vars=sum_vars
diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py
index 76e2b81e4..fd19ecc11 100644
--- a/neural_lam/models/ar_model.py
+++ b/neural_lam/models/ar_model.py
@@ -46,7 +46,8 @@ def __init__(
         config : NeuralLAMConfig
             Experiment configuration containing datastore/training settings.
         datastore : BaseDatastore
-            Datastore supplying state/forcing/static arrays.
+            Datastore supplying data and information about dataset and forecast
+            region.
         """
         super().__init__()
         self.save_hyperparameters(ignore=["datastore"])
@@ -188,7 +189,8 @@ def _create_dataarray_from_tensor(
         Parameters
         ----------
         tensor : torch.Tensor
-            Tensor to convert back to an ``xr.DataArray``.
+            Tensor to convert back to an ``xr.DataArray``. The tensor will be
+            copied to CPU memory before conversion.
 
             * **Shape**: ``(time, grid_index, feature)``
         time : torch.Tensor
@@ -252,28 +254,33 @@ def expand_to_batch(x, batch_size):
 
     def predict_step(self, prev_state, prev_prev_state, forcing):
         """
-        Advance the state by one step using the prediction model.
+        Advance the state by one step using the prediction model, as
+        ``X_{t-2}, X_{t-1} -> X_t``.
 
         Parameters
         ----------
         prev_state : torch.Tensor
             Current state ``X_t``.
 
-            * **Shape**: ``(B, num_grid_nodes, feature_dim)``
+            * **Shape**: ``(B, num_grid_nodes, num_state_vars)``
         prev_prev_state : torch.Tensor
             Previous state ``X_{t-1}``.
 
-            * **Shape**: ``(B, num_grid_nodes, feature_dim)``
+            * **Shape**: ``(B, num_grid_nodes, num_state_vars)``
         forcing : torch.Tensor
             Forcing inputs applied at the prediction step.
 
-            * **Shape**: ``(B, num_grid_nodes, forcing_dim)``
+            * **Shape**: ``(B, num_grid_nodes, num_forcing_vars)``
 
         Returns
         -------
         tuple[torch.Tensor, torch.Tensor | None]
             Tuple ``(new_state, pred_std)`` describing the next state and
             optional uncertainty estimate.
+
+            * **new_state**: ``(B, num_grid_nodes, num_state_vars)``
+            * **pred_std**: ``(B, num_grid_nodes, num_state_vars)`` or
+              ``(num_state_vars,)`` when using constant per-feature values
         """
         raise NotImplementedError("No prediction step implemented")
 
@@ -286,24 +293,26 @@ def unroll_prediction(self, init_states, forcing_features, true_states):
         init_states : torch.Tensor
             Initial states providing ``X_{t-1}`` and ``X_t``.
 
-            * **Shape**: ``(B, 2, num_grid_nodes, d_f)``
+            * **Shape**: ``(B, 2, num_grid_nodes, num_state_vars)``
         forcing_features : torch.Tensor
             Forcing inputs aligned with each rollout step.
 
-            * **Shape**: ``(B, pred_steps, num_grid_nodes, d_static_f)``
+            * **Shape**: ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``
         true_states : torch.Tensor
             Ground-truth states used for boundary replacement.
 
-            * **Shape**: ``(B, pred_steps, num_grid_nodes, d_f)``
+            * **Shape**: ``(B, pred_steps, num_grid_nodes, num_state_vars)``
 
         Returns
         -------
         tuple[torch.Tensor, torch.Tensor]
             Tuple ``(prediction, pred_std)``.
 
-            * **prediction**: ``(B, pred_steps, num_grid_nodes, d_f)``
-            * **pred_std**: ``(B, pred_steps, num_grid_nodes, d_f)`` or
-              ``(d_f,)`` when a constant per-feature value is used
+            * **prediction**: ``(B, pred_steps, num_grid_nodes,
+              num_state_vars)``
+            * **pred_std**: ``(B, pred_steps, num_grid_nodes, num_state_vars)``
+              or ``(num_state_vars,)`` when a constant per-feature value is
+              used
         """
         prev_prev_state = init_states[:, 0]
         prev_state = init_states[:, 1]
@@ -318,8 +327,8 @@ def unroll_prediction(self, init_states, forcing_features, true_states):
             pred_state, pred_std = self.predict_step(
                 prev_state, prev_prev_state, forcing
             )
-            # state: (B, num_grid_nodes, d_f) pred_std: (B, num_grid_nodes,
-            # d_f) or None
+            # state: (B, num_grid_nodes, num_state_vars)
+            # pred_std: (B, num_grid_nodes, num_state_vars) or None
 
             # Overwrite border with true state
             new_state = (
@@ -337,13 +346,13 @@ def unroll_prediction(self, init_states, forcing_features, true_states):
 
         prediction = torch.stack(
             prediction_list, dim=1
-        )  # (B, pred_steps, num_grid_nodes, d_f)
+        )  # (B, pred_steps, num_grid_nodes, num_state_vars)
         if self.output_std:
             pred_std = torch.stack(
                 pred_std_list, dim=1
-            )  # (B, pred_steps, num_grid_nodes, d_f)
+            )  # (B, pred_steps, num_grid_nodes, num_state_vars)
         else:
-            pred_std = self.per_var_std  # (d_f,)
+            pred_std = self.per_var_std  # (num_state_vars,)
 
         return prediction, pred_std
 
@@ -357,10 +366,11 @@ def common_step(self, batch):
             Tuple of ``(init_states, target_states, forcing_features,
             batch_times)`` produced by :class:`WeatherDataset`.
 
-            * **init_states**: ``(B, 2, num_grid_nodes, d_features)``
-            * **target_states**: ``(B, pred_steps, num_grid_nodes, d_features)``
+            * **init_states**: ``(B, 2, num_grid_nodes, num_state_vars)``
+            * **target_states**: ``(B, pred_steps, num_grid_nodes,
+              num_state_vars)``
             * **forcing_features**: ``(B, pred_steps, num_grid_nodes,
-              d_forcing)``
+              num_forcing_vars)``
             * **batch_times**: ``(B, pred_steps)`` timestamps
 
         Returns
@@ -368,19 +378,22 @@ def common_step(self, batch):
         tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
             ``(prediction, target_states, pred_std, batch_times)``.
 
-            * **prediction**: ``(B, pred_steps, num_grid_nodes, d_f)``
-            * **target_states**: ``(B, pred_steps, num_grid_nodes, d_f)``
-            * **pred_std**: ``(B, pred_steps, num_grid_nodes, d_f)`` or
-              ``(d_f,)``
+            * **prediction**: ``(B, pred_steps, num_grid_nodes,
+              num_state_vars)``
+            * **target_states**: ``(B, pred_steps, num_grid_nodes,
+              num_state_vars)``
+            * **pred_std**: ``(B, pred_steps, num_grid_nodes, num_state_vars)``
+              or ``(num_state_vars,)``
             * **batch_times**: ``(B, pred_steps)``
         """
         (init_states, target_states, forcing_features, batch_times) = batch
 
         prediction, pred_std = self.unroll_prediction(
             init_states, forcing_features, target_states
-        )  # (B, pred_steps, num_grid_nodes, d_f)
-        # prediction: (B, pred_steps, num_grid_nodes, d_f) pred_std: (B,
-        # pred_steps, num_grid_nodes, d_f) or (d_f,)
+        )  # (B, pred_steps, num_grid_nodes, num_state_vars)
+        # prediction: (B, pred_steps, num_grid_nodes, num_state_vars)
+        # pred_std: (B, pred_steps, num_grid_nodes, num_state_vars) or
+        # (num_state_vars,)
 
         return prediction, target_states, pred_std, batch_times
 
@@ -476,7 +489,7 @@ def validation_step(self, batch, batch_idx):
             pred_std,
             mask=self.interior_mask_bool,
             sum_vars=False,
-        )  # (B, pred_steps, d_f)
+        )  # (B, pred_steps, num_state_vars)
         self.val_metrics["mse"].append(entry_mses)
 
     def on_validation_epoch_end(self):
@@ -503,8 +516,9 @@ def test_step(self, batch, batch_idx):
         """
         # TODO Here batch_times can be used for plotting routines
         prediction, target, pred_std, batch_times = self.common_step(batch)
-        # prediction: (B, pred_steps, num_grid_nodes, d_f) pred_std: (B,
-        # pred_steps, num_grid_nodes, d_f) or (d_f,)
+        # prediction: (B, pred_steps, num_grid_nodes, num_state_vars)
+        # pred_std: (B, pred_steps, num_grid_nodes, num_state_vars) or
+        # (num_state_vars,)
 
         time_step_loss = torch.mean(
             self.loss(
@@ -540,14 +554,14 @@ def test_step(self, batch, batch_idx):
                 pred_std,
                 mask=self.interior_mask_bool,
                 sum_vars=False,
-            )  # (B, pred_steps, d_f)
+            )  # (B, pred_steps, num_state_vars)
             self.test_metrics[metric_name].append(batch_metric_vals)
 
         if self.output_std:
             # Store output std. per variable, spatially averaged
             mean_pred_std = torch.mean(
                 pred_std[..., self.interior_mask_bool, :], dim=-2
-            )  # (B, pred_steps, d_f)
+            )  # (B, pred_steps, num_state_vars)
             self.test_metrics["output_std"].append(mean_pred_std)
 
         # Save per-sample spatial loss for specific times
@@ -594,7 +608,7 @@ def plot_examples(self, batch, n_examples, split, prediction=None):
             Pre-computed predictions to plot. If ``None`` the method runs
             :meth:`common_step` to obtain predictions.
 
-            * **Shape**: ``(B, pred_steps, num_grid_nodes, d_f)``
+            * **Shape**: ``(B, pred_steps, num_grid_nodes, num_state_vars)``
         """
         if prediction is None:
             prediction, target, _, _ = self.common_step(batch)
@@ -612,7 +626,7 @@ def plot_examples(self, batch, n_examples, split, prediction=None):
             target_rescaled[:n_examples],
             time[:n_examples],
         ):
-            # Each slice is (pred_steps, num_grid_nodes, d_f)
+            # Each slice is (pred_steps, num_grid_nodes, num_state_vars)
             self.plotted_examples += 1  # Increment already here
 
             da_prediction = self._create_dataarray_from_tensor(
@@ -635,7 +649,7 @@ def plot_examples(self, batch, n_examples, split, prediction=None):
                 )
                 .cpu()
                 .numpy()
-            )  # (d_f,)
+            )  # (num_state_vars,)
             var_vmax = (
                 torch.maximum(
                     pred_slice.flatten(0, 1).max(dim=0)[0],
@@ -643,7 +657,7 @@ def plot_examples(self, batch, n_examples, split, prediction=None):
                 )
                 .cpu()
                 .numpy()
-            )  # (d_f,)
+            )  # (num_state_vars,)
             var_vranges = list(zip(var_vmin, var_vmax))
 
             # Iterate over prediction horizon time steps
@@ -723,7 +737,7 @@ def create_metric_log_dict(self, metric_tensor, prefix, metric_name):
         metric_tensor : torch.Tensor
             Metric values per time step and variable.
 
-            * **Shape**: ``(pred_steps, d_f)``
+            * **Shape**: ``(pred_steps, num_state_vars)``
         prefix : str
             Prefix used for logger keys (e.g., ``"val"`` or ``"test"``).
         metric_name : str
@@ -780,11 +794,11 @@ def aggregate_and_plot_metrics(self, metrics_dict, prefix):
         for metric_name, metric_val_list in metrics_dict.items():
             metric_tensor = self.all_gather_cat(
                 torch.cat(metric_val_list, dim=0)
-            )  # (N_eval, pred_steps, d_f)
+            )  # (N_eval, pred_steps, num_state_vars)
 
             if self.trainer.is_global_zero:
                 metric_tensor_averaged = torch.mean(metric_tensor, dim=0)
-                # (pred_steps, d_f)
+                # (pred_steps, num_state_vars)
 
                 # Take square root after all averaging to change MSE to RMSE
                 if "mse" in metric_name:
@@ -793,7 +807,7 @@ def aggregate_and_plot_metrics(self, metrics_dict, prefix):
 
                 # NOTE: we here assume rescaling for all metrics is linear
                 metric_rescaled = metric_tensor_averaged * self.state_std
-                # (pred_steps, d_f)
+                # (pred_steps, num_state_vars)
                 log_dict.update(
                     self.create_metric_log_dict(
                         metric_rescaled, prefix, metric_name
diff --git a/neural_lam/models/base_graph_model.py b/neural_lam/models/base_graph_model.py
index eaac26c2c..2655da0f9 100644
--- a/neural_lam/models/base_graph_model.py
+++ b/neural_lam/models/base_graph_model.py
@@ -1,4 +1,4 @@
-"""Base classes for Neural-LAM graph models."""
+"""Base class for Neural-LAM graph models."""
 
 # Third-party
 import torch
@@ -239,31 +239,42 @@ def prepare_clamping_params(
         )
 
     def get_clamped_new_state(self, state_delta, prev_state):
-        """
+        r"""
         Clamp predicted deltas and add them to the previous state.
 
         The clamped values follow
-        ``f(f^{-1}(X_t) + model({X_t, X_{t-1}, ...}, forcing))`` so that the
-        model learns to emit outputs in the range of the inverse clamping
-        function.
+        ``X_{t+1} = f(f^{-1}(X_t) + \Delta X_t)`` where ``\Delta X_t`` is the
+        model output. ``f(·)`` applies the appropriate clamp per variable,
+        while ``f^{-1}(·)`` makes the inverse transformation so the network
+        predicts in the unclamped space. The element-wise clamp is implemented
+        as
+
+        * ``f(z) = a + (b - a) * \sigma(z)`` for variables with bounds
+          ``(a, b)`` (``\sigma`` is the logistic function)
+        * ``f(z) = a + \operatorname{softplus}(z)`` for lower-bounded variables
+        * ``f(z) = b - \operatorname{softplus}(-z)`` for upper-bounded variables
+
+        which ensures ``X_{t+1}`` respects the physical limits encoded in the
+        datastore configuration while keeping the model output centered in a
+        numerically stable range.
 
         Parameters
         ----------
         state_delta : torch.Tensor
             Predicted change to apply to the previous state.
 
-            * **Shape**: ``(B, num_grid_nodes, feature_dim)``
+            * **Shape**: ``(B, num_grid_nodes, num_state_vars)``
         prev_state : torch.Tensor
             Previous state ``X_t``.
 
-            * **Shape**: ``(B, num_grid_nodes, feature_dim)``
+            * **Shape**: ``(B, num_grid_nodes, num_state_vars)``
 
         Returns
         -------
         torch.Tensor
             Clamped next state ``X_{t+1}``.
 
-            * **Shape**: ``(B, num_grid_nodes, feature_dim)``
+            * **Shape**: ``(B, num_grid_nodes, num_state_vars)``
         """
 
         # Assign new state, but overwrite clamped values of each type later
@@ -319,7 +330,7 @@ def embedd_mesh_nodes(self):
         torch.Tensor
             Embedded mesh node representations.
 
-            * **Shape**: ``(num_mesh_nodes, d_h)``
+            * **Shape**: ``(num_mesh_nodes, hidden_dim)``
         """
         raise NotImplementedError("embedd_mesh_nodes not implemented")
 
@@ -332,14 +343,14 @@ def process_step(self, mesh_rep):
         mesh_rep : torch.Tensor
             Mesh node representations prior to the processor.
 
-            * **Shape**: ``(B, num_mesh_nodes, d_h)``
+            * **Shape**: ``(B, num_mesh_nodes, hidden_dim)``
 
         Returns
         -------
         torch.Tensor
             Updated mesh representations after processing.
 
-            * **Shape**: ``(B, num_mesh_nodes, d_h)``
+            * **Shape**: ``(B, num_mesh_nodes, hidden_dim)``
         """
         raise NotImplementedError("process_step not implemented")
 
@@ -352,24 +363,25 @@ def predict_step(self, prev_state, prev_prev_state, forcing):
         prev_state : torch.Tensor
             Current state ``X_t``.
 
-            * **Shape**: ``(B, num_grid_nodes, feature_dim)``
+            * **Shape**: ``(B, num_grid_nodes, num_state_vars)``
         prev_prev_state : torch.Tensor
             Previous state ``X_{t-1}``.
 
-            * **Shape**: ``(B, num_grid_nodes, feature_dim)``
+            * **Shape**: ``(B, num_grid_nodes, num_state_vars)``
         forcing : torch.Tensor
             Forcing inputs applied at the prediction step.
 
-            * **Shape**: ``(B, num_grid_nodes, forcing_dim)``
+            * **Shape**: ``(B, num_grid_nodes, num_forcing_vars)``
 
         Returns
         -------
         tuple[torch.Tensor, torch.Tensor | None]
-            Tuple ``(new_state, pred_std)`` where ``pred_std`` is ``None`` when
-            the model does not emit uncertainty estimates.
+            Tuple ``(new_state, pred_std)`` where ``pred_std`` is ``None``
+            when the model does not emit uncertainty estimates.
 
-            * **Shape**: ``(B, num_grid_nodes, feature_dim)`` for ``new_state``
-              and ``(B, num_grid_nodes, d_f)`` for ``pred_std`` when present.
+            * **Shape**: ``(B, num_grid_nodes, num_state_vars)`` for
+              ``new_state`` and ``(B, num_grid_nodes, num_state_vars)`` for
+              ``pred_std`` when present.
         """
         batch_size = prev_state.shape[0]
 
@@ -385,25 +397,27 @@ def predict_step(self, prev_state, prev_prev_state, forcing):
         )
 
         # Embed all features
-        grid_emb = self.grid_embedder(grid_features)  # (B, num_grid_nodes, d_h)
-        g2m_emb = self.g2m_embedder(self.g2m_features)  # (M_g2m, d_h)
-        m2g_emb = self.m2g_embedder(self.m2g_features)  # (M_m2g, d_h)
+        grid_emb = self.grid_embedder(
+            grid_features
+        )  # (B, num_grid_nodes, hidden_dim)
+        g2m_emb = self.g2m_embedder(self.g2m_features)  # (M_g2m, hidden_dim)
+        m2g_emb = self.m2g_embedder(self.m2g_features)  # (M_m2g, hidden_dim)
         mesh_emb = self.embedd_mesh_nodes()
 
         # Map from grid to mesh
         mesh_emb_expanded = self.expand_to_batch(
             mesh_emb, batch_size
-        )  # (B, num_mesh_nodes, d_h)
+        )  # (B, num_mesh_nodes, hidden_dim)
         g2m_emb_expanded = self.expand_to_batch(g2m_emb, batch_size)
 
         # This also splits representation into grid and mesh
         mesh_rep = self.g2m_gnn(
             grid_emb, mesh_emb_expanded, g2m_emb_expanded
-        )  # (B, num_mesh_nodes, d_h)
+        )  # (B, num_mesh_nodes, hidden_dim)
         # Also MLP with residual for grid representation
         grid_rep = grid_emb + self.encoding_grid_mlp(
             grid_emb
-        )  # (B, num_grid_nodes, d_h)
+        )  # (B, num_grid_nodes, hidden_dim)
 
         # Run processor step
         mesh_rep = self.process_step(mesh_rep)
@@ -412,7 +426,7 @@ def predict_step(self, prev_state, prev_prev_state, forcing):
         m2g_emb_expanded = self.expand_to_batch(m2g_emb, batch_size)
         grid_rep = self.m2g_gnn(
             mesh_rep, grid_rep, m2g_emb_expanded
-        )  # (B, num_grid_nodes, d_h)
+        )  # (B, num_grid_nodes, hidden_dim)
 
         # Map to output dimension, only for grid
         net_output = self.output_map(
@@ -422,7 +436,7 @@ def predict_step(self, prev_state, prev_prev_state, forcing):
         if self.output_std:
             pred_delta_mean, pred_std_raw = net_output.chunk(
                 2, dim=-1
-            )  # both (B, num_grid_nodes, d_f)
+            )  # both (B, num_grid_nodes, num_state_vars)
             # NOTE: The predicted std. is not scaled in any way here
             # linter for some reason does not think softplus is callable
             # pylint: disable-next=not-callable
diff --git a/neural_lam/models/base_hi_graph_model.py b/neural_lam/models/base_hi_graph_model.py
index 5f1534fd9..c8e4511d2 100644
--- a/neural_lam/models/base_hi_graph_model.py
+++ b/neural_lam/models/base_hi_graph_model.py
@@ -131,7 +131,7 @@ def embedd_mesh_nodes(self):
         torch.Tensor
             Embedded representations for the base-level mesh nodes.
 
-            * **Shape**: ``(num_mesh_nodes[0], d_h)``
+            * **Shape**: ``(num_mesh_nodes[0], hidden_dim)``
         """
         return self.mesh_embedders[0](self.mesh_static_features[0])
 
@@ -144,20 +144,21 @@ def process_step(self, mesh_rep):
         mesh_rep : torch.Tensor
             Base-level mesh representations prior to the processor.
 
-            * **Shape**: ``(B, num_mesh_nodes, d_h)``
+            * **Shape**: ``(B, num_mesh_nodes[0], hidden_dim)`` (only the
+              bottom-level nodes are present at this point)
 
         Returns
         -------
         torch.Tensor
             Updated base-level mesh representations.
 
-            * **Shape**: ``(B, num_mesh_nodes, d_h)``
+            * **Shape**: ``(B, num_mesh_nodes[0], hidden_dim)``
         """
         batch_size = mesh_rep.shape[0]
 
         # EMBED REMAINING MESH NODES (levels >= 1) -
         # Create list of mesh node representations for each level,
-        # each of size (B, num_mesh_nodes[l], d_h)
+        # each of size (B, num_mesh_nodes[l], hidden_dim)
         mesh_rep_levels = [mesh_rep] + [
             self.expand_to_batch(emb(node_static_features), batch_size)
             for emb, node_static_features in zip(
@@ -193,10 +194,10 @@ def process_step(self, mesh_rep):
             # Extract representations
             send_node_rep = mesh_rep_levels[
                 level_l - 1
-            ]  # (B, num_mesh_nodes[l-1], d_h)
+            ]  # (B, num_mesh_nodes[l-1], hidden_dim)
             rec_node_rep = mesh_rep_levels[
                 level_l
-            ]  # (B, num_mesh_nodes[l], d_h)
+            ]  # (B, num_mesh_nodes[l], hidden_dim)
             edge_rep = mesh_up_rep[level_l - 1]
 
             # Apply GNN
@@ -206,9 +207,11 @@ def process_step(self, mesh_rep):
 
             # Update node and edge vectors in lists
             mesh_rep_levels[level_l] = (
-                new_node_rep  # (B, num_mesh_nodes[l], d_h)
+                new_node_rep  # (B, num_mesh_nodes[l], hidden_dim)
+            )
+            mesh_up_rep[level_l - 1] = (
+                new_edge_rep  # (B, M_up[l-1], hidden_dim)
             )
-            mesh_up_rep[level_l - 1] = new_edge_rep  # (B, M_up[l-1], d_h)
 
         # - PROCESSOR -
         mesh_rep_levels, _, _, mesh_down_rep = self.hi_processor_step(
@@ -223,10 +226,10 @@ def process_step(self, mesh_rep):
             # Extract representations
             send_node_rep = mesh_rep_levels[
                 level_l + 1
-            ]  # (B, num_mesh_nodes[l+1], d_h)
+            ]  # (B, num_mesh_nodes[l+1], hidden_dim)
             rec_node_rep = mesh_rep_levels[
                 level_l
-            ]  # (B, num_mesh_nodes[l], d_h)
+            ]  # (B, num_mesh_nodes[l], hidden_dim)
             edge_rep = mesh_down_rep[level_l]
 
             # Apply GNN
@@ -234,11 +237,11 @@ def process_step(self, mesh_rep):
 
             # Update node and edge vectors in lists
             mesh_rep_levels[level_l] = (
-                new_node_rep  # (B, num_mesh_nodes[l], d_h)
+                new_node_rep  # (B, num_mesh_nodes[l], hidden_dim)
             )
 
         # Return only bottom level representation
-        return mesh_rep_levels[0]  # (B, num_mesh_nodes[0], d_h)
+        return mesh_rep_levels[0]  # (B, num_mesh_nodes[0], hidden_dim)
 
     def hi_processor_step(
         self, mesh_rep_levels, mesh_same_rep, mesh_up_rep, mesh_down_rep
@@ -251,19 +254,21 @@ def hi_processor_step(
         mesh_rep_levels : list[torch.Tensor]
             Mesh representations for each level.
 
-            * **Shape**: ``(B, num_mesh_nodes[l], d_h)``
+            * Each element ``l`` has shape ``(B, num_mesh_nodes[l],
+              hidden_dim)``.
         mesh_same_rep : list[torch.Tensor]
             Same-level edge representations per level.
 
-            * **Shape**: ``(B, M_same[l], d_h)``
+            * Each element ``l`` has shape ``(B, M_same[l], hidden_dim)``.
         mesh_up_rep : list[torch.Tensor]
             Edge representations from level ``l`` to ``l+1``.
 
-            * **Shape**: ``(B, M_up[l -> l+1], d_h)``
+            * Each element ``l`` has shape ``(B, M_up[l -> l+1], hidden_dim)``.
         mesh_down_rep : list[torch.Tensor]
             Edge representations from level ``l+1`` down to ``l``.
 
-            * **Shape**: ``(B, M_down[l <- l+1], d_h)``
+            * Each element ``l`` has shape ``(B, M_down[l <- l+1],
+              hidden_dim)``.
 
         Returns
         -------
@@ -273,5 +278,7 @@ def hi_processor_step(
         ]
             Updated representations for (mesh, same-level, up edges, down edges)
             in that order.
+
+            * Each list preserves the element-wise shapes described above.
         """
         raise NotImplementedError("hi_process_step not implemented")
diff --git a/neural_lam/utils.py b/neural_lam/utils.py
index f61a5cc23..eddc7256a 100644
--- a/neural_lam/utils.py
+++ b/neural_lam/utils.py
@@ -225,8 +225,9 @@ def make_mlp(blueprint, layer_norm=True):
     Parameters
     ----------
     blueprint : list[int]
-        Sequence of layer dimensions where ``blueprint[0]`` is the input size
-        and ``blueprint[-1]`` is the output size.
+        Sequence of layer dimensions where ``blueprint[0]`` is the input size,
+        ``blueprint[-1]`` is the output size, and the intermediate entries
+        specify the hidden widths.
     layer_norm : bool, optional
         If ``True``, append a ``LayerNorm`` to the output as in GraphCast.
 
@@ -422,7 +423,11 @@ def setup_training_logger(datastore, args, run_name):
 
 def inverse_softplus(x, beta=1, threshold=20):
     """
-    Approximate the inverse of :func:`torch.nn.functional.softplus`.
+    Inverse of :func:`torch.nn.functional.softplus`.
+
+    Input is clamped to approximately positive values of ``x`` for numerical
+    stability; everything above ``threshold / beta`` is treated as linear and
+    exactly matches the softplus inverse within numerical precision.
 
     Parameters
     ----------
@@ -437,6 +442,11 @@ def inverse_softplus(x, beta=1, threshold=20):
     -------
     torch.Tensor
         Tensor containing the inverse-softplus values.
+
+    Notes
+    -----
+    ``torch.clamp`` will zero the gradients near the bounds, but values this
+    close to zero or ``threshold / beta`` already have negligible gradients.
     """
     x_clamped = torch.clamp(
         x, min=torch.log(torch.tensor(1e-6 + 1)) / beta, max=threshold / beta
@@ -453,7 +463,10 @@ def inverse_softplus(x, beta=1, threshold=20):
 
 def inverse_sigmoid(x):
     """
-    Compute the logit (inverse sigmoid) while clamping to ``(0, 1)``.
+    Inverse of ``torch.sigmoid`` with clamping for numerical stability.
+
+    Sigmoid output takes values in ``[0, 1]``; we clamp the input slightly
+    within that open interval before applying ``log(x / (1 - x))``.
 
     Parameters
     ----------
@@ -465,6 +478,11 @@ def inverse_sigmoid(x):
     torch.Tensor
         Tensor containing ``log(x / (1 - x))`` after clamping away from the
         saturation limits.
+
+    Notes
+    -----
+    ``torch.clamp`` zeroes gradients for values at the bounds, but values this
+    close to 0 or 1 already have negligible gradients.
     """
     x_clamped = torch.clamp(x, min=1e-6, max=1 - 1e-6)
     return torch.log(x_clamped / (1 - x_clamped))
@@ -482,8 +500,10 @@ def get_integer_time(tdelta) -> tuple[int, str]:
     Returns
     -------
     tuple[int, str]
-        Integer value and the corresponding unit (e.g. ``"hours"``). If no
-        unit yields an integer count, ``(1, "unknown")`` is returned.
+        Integer value and the corresponding unit (``"weeks"``, ``"days"``,
+        ``"hours"``, ``"minutes"``, ``"seconds"``, ``"milliseconds"``, or
+        ``"microseconds"``). If no unit yields an integer count,
+        ``(1, "unknown")`` is returned.
 
     Examples
     --------
diff --git a/neural_lam/vis.py b/neural_lam/vis.py
index cdb58f3e2..456d0b925 100644
--- a/neural_lam/vis.py
+++ b/neural_lam/vis.py
@@ -98,13 +98,13 @@ def plot_prediction(
     datastore : BaseRegularGridDatastore
         Datastore providing coordinate metadata and projection details.
     da_prediction : xarray.DataArray
-        Predicted field flattened over the grid.
+        Predicted field.
 
-        * **Shape**: ``(N_grid,)``
+        * **Shape**: ``(num_grid_nodes,)``
     da_target : xarray.DataArray
-        Ground-truth field flattened over the grid.
+        Ground-truth field.
 
-        * **Shape**: ``(N_grid,)``
+        * **Shape**: ``(num_grid_nodes,)``
     title : str or None, optional
         Optional figure title.
     vrange : tuple[float, float] or None, optional
diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py
index fcb02ded0..a6395cacd 100644
--- a/neural_lam/weather_dataset.py
+++ b/neural_lam/weather_dataset.py
@@ -19,27 +19,6 @@ class WeatherDataset(torch.utils.data.Dataset):
     """Dataset class for weather data.
 
     This class loads and processes weather data from a given datastore.
-
-    Parameters
-    ----------
-    datastore : BaseDatastore
-        The datastore to load the data from (e.g. mdp).
-    split : str, optional
-        The data split to use ("train", "val" or "test"). Default is "train".
-    ar_steps : int, optional
-        The number of autoregressive steps. Default is 3.
-    num_past_forcing_steps: int, optional
-        Number of past time steps to include in forcing input. If set to i,
-        forcing from times t-i, t-i+1, ..., t-1, t (and potentially beyond,
-        given num_future_forcing_steps) are included as forcing inputs at time t
-        Default is 1.
-    num_future_forcing_steps: int, optional
-        Number of future time steps to include in forcing input. If set to j,
-        forcing from times t, t+1, ..., t+j-1, t+j (and potentially times before
-        t, given num_past_forcing_steps) are included as forcing inputs at time
-        t. Default is 1.
-    standardize : bool, optional
-        Whether to standardize the data. Default is True.
     """
 
     def __init__(
@@ -652,9 +631,10 @@ def __init__(
         datastore : BaseDatastore
             Datastore used for all splits.
         ar_steps_train : int, optional
-            Number of AR steps for training batches. Default ``3``.
+            Number of autoregressive steps for training batches. Default ``3``.
         ar_steps_eval : int, optional
-            Number of AR steps for validation/test batches. Default ``25``.
+            Number of autoregressive steps for validation/test batches.
+            Default ``25``.
         standardize : bool, optional
             If ``True``, datasets are returned standardized. Default ``True``.
         num_past_forcing_steps : int, optional
@@ -695,7 +675,8 @@ def setup(self, stage=None):
         ----------
         stage : str or None, optional
             Trainer stage identifier (``"fit"``/``"test"``/``None``). When
-            ``None``, both train and evaluation datasets are created.
+            ``None``, both the training split and the validation/test
+            evaluation splits are prepared.
         """
         if stage == "fit" or stage is None:
             self.train_dataset = WeatherDataset(

From 721eee56a191b4f69d4aff05e98ccd9b71c3530b Mon Sep 17 00:00:00 2001
From: Mohit-Lakra <mohitlakra2007777@gmail.com>
Date: Tue, 17 Mar 2026 08:50:09 +0530
Subject: [PATCH 08/16] docs: add dimension glossary to README

---
 README.md | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/README.md b/README.md
index 1815b6d7c..b27ec1504 100644
--- a/README.md
+++ b/README.md
@@ -568,6 +568,23 @@ In addition, hierarchical mesh graphs (`L > 1`) feature a few additional files w
 These files have the same list format as the ones above, but each list has length `L-1` (as these edges describe connections between levels).
 Entries 0 in these lists describe edges between the lowest levels 1 and 2.
 
+## Dimension Glossary
+
+Canonical dimension names used in tensor shape annotations throughout the codebase:
+
+- `B` — batch size
+- `pred_steps` — number of autoregressive prediction steps
+- `num_grid_nodes` — number of nodes in the flattened spatial grid
+- `num_mesh_nodes` — number of mesh nodes; indexed as `num_mesh_nodes[l]` for hierarchical level `l`
+- `num_state_vars` — number of atmospheric state variables
+- `num_forcing_vars` — number of forcing input variables
+- `num_variables` — generic variable dimension used in metric functions
+- `hidden_dim` — internal hidden representation size in GNN layers and MLPs
+- `input_dim` — input feature dimensionality to a layer before transformation
+- `num_edges` — number of edges in a graph (g2m, m2g, same-level, up, down)
+- `num_send` — number of sender nodes in a message-passing step
+- `num_rec` — number of receiver nodes in a message-passing step
+
 # Development and Contributing
 Any push or Pull-Request to the main branch will trigger a selection of pre-commit hooks.
 These hooks will run a series of checks on the code, like formatting and linting.

From 1e20d51765937f7838d3483c5fed08ac95bf331b Mon Sep 17 00:00:00 2001
From: Mohit-Lakra <mohitlakra2007777@gmail.com>
Date: Thu, 19 Mar 2026 08:03:06 +0530
Subject: [PATCH 09/16] docs: address reviewer feedback on PR #252 docstrings

---
 neural_lam/utils.py | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/neural_lam/utils.py b/neural_lam/utils.py
index cb039d6fc..27a0c7d59 100644
--- a/neural_lam/utils.py
+++ b/neural_lam/utils.py
@@ -362,8 +362,9 @@ def make_mlp(blueprint, layer_norm=True):
     ----------
     blueprint : list[int]
         Sequence of layer dimensions where ``blueprint[0]`` is the input size,
-        ``blueprint[-1]`` is the output size, and the intermediate entries
-        specify the hidden widths.
+        ``blueprint[-1]`` is the output size, the intermediate entries specify
+        the hidden layer widths, and ``len(blueprint) - 2`` is the number of
+        hidden layers.
     layer_norm : bool, optional
         If ``True``, append a ``LayerNorm`` to the output as in GraphCast.
 
@@ -600,9 +601,12 @@ def inverse_softplus(x, beta=1, threshold=20):
     """
     Inverse of :func:`torch.nn.functional.softplus`.
 
-    Input is clamped to approximately positive values of ``x`` for numerical
-    stability; everything above ``threshold / beta`` is treated as linear and
-    exactly matches the softplus inverse within numerical precision.
+    For most inputs this function is exact up to numerical precision. The
+    input is clamped to ensure numerical stability: values above
+    ``threshold / beta`` are treated as linear (which is exact in that
+    regime), and values near zero are clamped to avoid ``log`` of
+    non-positive numbers. Only near the lower clamping bound does the
+    result deviate from the true inverse.
 
     Parameters
     ----------
@@ -611,7 +615,8 @@ def inverse_softplus(x, beta=1, threshold=20):
     beta : float, optional
         Softplus ``beta`` parameter that controls the sharpness. Default ``1``.
     threshold : float, optional
-        Threshold applied to the input for numerical stability. Default ``20``.
+        Threshold above which the function is treated as linear for numerical
+        stability. Default ``20``.
 
     Returns
     -------
@@ -643,6 +648,10 @@ def inverse_sigmoid(x):
     Sigmoid output takes values in ``[0, 1]``; we clamp the input slightly
     within that open interval before applying ``log(x / (1 - x))``.
 
+    Note that ``torch.clamp`` will make gradients 0 near the bounds, but
+    this is not a problem as values of x that are this close to 0 or 1
+    have gradients of 0 anyhow.
+
     Parameters
     ----------
     x : torch.Tensor

From 39cbe30d607048ba95d0ecf2693a21a57bb13c9e Mon Sep 17 00:00:00 2001
From: Mohit-Lakra <mohitlakra2007777@gmail.com>
Date: Sun, 5 Apr 2026 00:03:43 +0530
Subject: [PATCH 10/16] fix: correct typos in docstrings and comments (refs
 #359)

---
 neural_lam/config.py          |  4 ++--
 neural_lam/utils.py           |  2 +-
 neural_lam/weather_dataset.py | 14 +++++++-------
 3 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/neural_lam/config.py b/neural_lam/config.py
index 5d449f545..cc5cd0edc 100644
--- a/neural_lam/config.py
+++ b/neural_lam/config.py
@@ -99,9 +99,9 @@ class TrainingConfig:
     Attributes
     ----------
     state_feature_weighting : Union[ManualStateFeatureWeighting,
-                                    UnformFeatureWeighting]
+                                    UniformFeatureWeighting]
         The method to use for weighting the state features in the loss
-        function. Defaults to uniform weighting (`UnformFeatureWeighting`, i.e.
+        function. Defaults to uniform weighting (`UniformFeatureWeighting`, i.e.
         all features are weighted equally).
     """
 
diff --git a/neural_lam/utils.py b/neural_lam/utils.py
index 27a0c7d59..e232e2820 100644
--- a/neural_lam/utils.py
+++ b/neural_lam/utils.py
@@ -267,7 +267,7 @@ def loads_file(fn):
     assert g2m_edge_index.min() >= 0, "Negative node index in g2m"
 
     n_levels = len(m2m_edge_index)
-    hierarchical = n_levels > 1  # Nor just single level mesh graph
+    hierarchical = n_levels > 1  # Not just single level mesh graph
 
     # Load static edge features
     # List of (M_m2m[l], d_edge_f)
diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py
index 7c07383c3..839f02f8f 100644
--- a/neural_lam/weather_dataset.py
+++ b/neural_lam/weather_dataset.py
@@ -249,11 +249,11 @@ def __len__(self):
     def _slice_state_time(self, da_state, idx, n_steps: int):
         """
         Produce a time slice of the given dataarray `da_state` (state) starting
-        at `idx` and with `n_steps` steps. An `offset`is calculated based on the
-        `num_past_forcing_steps` class attribute. `Offset` is used to offset the
-        start of the sample, to assert that enough previous time steps are
-        available for the 2 initial states and any corresponding forcings
-        (calculated in `_slice_forcing_time`).
+        at `idx` and with `n_steps` steps. An `offset` is calculated based on
+        the `num_past_forcing_steps` class attribute. `Offset` is used to offset
+        the start of the sample, to assert that enough previous time steps
+        are available for the 2 initial states and any corresponding
+        forcings (calculated in `_slice_forcing_time`).
 
         Parameters
         ----------
@@ -338,8 +338,8 @@ def _slice_forcing_time(self, da_forcing, idx, n_steps: int):
         """
         # The current implementation requires at least 2 time steps for the
         # initial state (see GraphCast). The forcing data is windowed around the
-        # current autregressive time step. The two `init_steps` can also be used
-        # as past forcings.
+        # current autoregressive time step. The two `init_steps` can also be
+        # used as past forcings.
         init_steps = 2
         da_list = []
 

From 31df714f9bf78ccfc17e840ccb31d0344f0cb331 Mon Sep 17 00:00:00 2001
From: sadamov <kode@mailbox.org>
Date: Tue, 9 Jun 2026 12:11:49 +0200
Subject: [PATCH 11/16] docs: align #252 with README dimension glossary and
 tidy RST for autoapi

Takeover diff for mllam/neural-lam#252. Folds the verified findings
from a multi-agent review pass (covering NumPy-style consistency,
shape-name alignment vs the README glossary, Sphinx/autoapi
readiness, docstring-vs-implementation accuracy, information
preservation across the rewrite, and completeness beyond raw
interrogate coverage) into one mechanical diff Mohit-Lakra can merge
or cherry-pick.

Headline changes:
- Rename non-glossary tensor-shape names to canonical glossary forms
  across docstrings and inline comments (d_h -> hidden_dim,
  d_f -> num_state_vars, d_forcing -> num_forcing_vars,
  N_grid -> num_grid_nodes, N_edges -> num_edges, M_g2m / M_m2g /
  M_up / M_down / M_same / M_mesh -> num_edges, d_edge_f ->
  input_dim).
- Tidy RST: wrap bare True/False in double backticks, convert # to
  RST section underlines, replace single backticks around shape
  strings with double backticks, remove em-dashes from README
  glossary.
- Restore information the rewrite dropped: filename templates in
  npyfilesmeps/store, num_rec constraint in InteractionNet.aggregate,
  vrange "inferred from data" and projection mention in
  plot_spatial_error, full list of units in utils, hidden-dimension
  context in utils, equation context in base_graph_model.
- Add missing Raises sections where the code raises (setup_training_logger,
  _get_heatmap_color_values, prepare_clamping_params, get_metric,
  WeatherDataset.create_dataarray_from_tensor, MDPDatastore
  __init__ / coords_projection, stack_grid_coords).
- Cross-cutting: add CHANGELOG entry referencing #252, add
  [tool.interrogate] block to pyproject.toml documenting the gate,
  switch the pre-commit hook to read that config (`-c pyproject.toml`)
  so local interrogate runs match CI.

Format choices locked in for the project:
- prose-Shape (`Shape ``(B, ...)``.`) over bullet-Shape; matches the
  dominant pattern already in the codebase.
- numpydoc literal-set (`aggr : {"sum", "mean"}`) over prose
  enumeration.
- Glossary stays at the current 12 entries; renames pick the closest
  existing canonical (num_edges rather than num_g2m_edges et al).

Not in this PR (open follow-ups, discussed in the PR description):
- utils.py docstring restructure / dropped-context restores (the
  load_graph Returns block, the Joel-flagged context drops). Apply
  agent died before completing this file due to session budget; the
  v2 review draft at /tmp/pr252_review_draft_v2.md has the exact
  text to restore.
- plot_graph.py and train_model.py docstring polish (low-priority,
  no behaviour change needed).
- whether to extend the interrogate gate to `tests/` (currently
  excluded via [tool.interrogate]).

refs mllam/neural-lam#252

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 .pre-commit-config.yaml                       |  2 +-
 CHANGELOG.md                                  |  2 +
 README.md                                     | 24 +++---
 neural_lam/config.py                          |  6 +-
 neural_lam/custom_loggers.py                  |  6 ++
 neural_lam/datastore/base.py                  | 19 +++--
 neural_lam/datastore/mdp.py                   | 12 +++
 .../compute_standardization_stats.py          | 62 +++++++++------
 neural_lam/datastore/npyfilesmeps/config.py   |  5 ++
 neural_lam/datastore/npyfilesmeps/store.py    | 50 +++++++++---
 neural_lam/interaction_net.py                 | 25 +++---
 neural_lam/loss_weighting.py                  | 12 +++
 neural_lam/metrics.py                         | 76 ++++++++++---------
 .../models/forecasters/autoregressive.py      | 34 ++++-----
 neural_lam/models/forecasters/base.py         | 27 +++----
 neural_lam/models/module.py                   | 21 +++--
 neural_lam/models/step_predictors/base.py     | 28 ++++---
 .../models/step_predictors/graph/base.py      | 47 ++++++------
 .../models/step_predictors/graph/graph_lam.py | 22 +++---
 .../models/step_predictors/graph/hi_lam.py    | 40 +++++-----
 .../step_predictors/graph/hi_lam_parallel.py  | 14 ++--
 .../step_predictors/graph/hierarchical.py     | 24 +++---
 neural_lam/plot_graph.py                      |  4 +-
 neural_lam/utils.py                           | 24 +++---
 neural_lam/vis.py                             | 55 +++++++++-----
 neural_lam/weather_dataset.py                 | 49 +++++-------
 pyproject.toml                                | 12 +++
 27 files changed, 398 insertions(+), 304 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 45bb18444..b75a383ff 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -42,7 +42,7 @@ repos:
       - id: interrogate
         description: Ensure documentation coverage stays perfect
         pass_filenames: false
-        args: ["--fail-under=100", "neural_lam"]
+        args: ["-c", "pyproject.toml", "neural_lam"]
   - repo: https://github.com/pre-commit/mirrors-mypy
     rev: v1.19.0
     hooks:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index c30dedf36..2c1b8470a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -35,6 +35,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Maintenance
 
+- Establish 100% docstring coverage across `neural_lam/` via an `interrogate` pre-commit hook, add a Dimension Glossary to the README for canonical tensor-shape names, and rewrite public docstrings in NumPy style to serve as the entry point for the autoapi pipeline (#196 / #272). [\#252](https://github.com/mllam/neural-lam/pull/252) @Mohit-Lakra
+
 - Fix GPU CI torch version resolution to query the target wheel index instead of PyPI [\#639](https://github.com/mllam/neural-lam/pull/639) @Sir-Sloth-The-Lazy
 
 - Add comprehensive type hints to all functions and class methods in `utils.py` [\#620](https://github.com/mllam/neural-lam/pull/620) @GiGiKoneti
diff --git a/README.md b/README.md
index ff25915fa..2f7f0adb9 100644
--- a/README.md
+++ b/README.md
@@ -571,18 +571,18 @@ Entries 0 in these lists describe edges between the lowest levels 1 and 2.
 
 Canonical dimension names used in tensor shape annotations throughout the codebase:
 
-- `B` — batch size
-- `pred_steps` — number of autoregressive prediction steps
-- `num_grid_nodes` — number of nodes in the flattened spatial grid
-- `num_mesh_nodes` — number of mesh nodes; indexed as `num_mesh_nodes[l]` for hierarchical level `l`
-- `num_state_vars` — number of atmospheric state variables
-- `num_forcing_vars` — number of forcing input variables
-- `num_variables` — generic variable dimension used in metric functions
-- `hidden_dim` — internal hidden representation size in GNN layers and MLPs
-- `input_dim` — input feature dimensionality to a layer before transformation
-- `num_edges` — number of edges in a graph (g2m, m2g, same-level, up, down)
-- `num_send` — number of sender nodes in a message-passing step
-- `num_rec` — number of receiver nodes in a message-passing step
+- `B` - batch size
+- `pred_steps` - number of autoregressive prediction steps
+- `num_grid_nodes` - number of nodes in the flattened spatial grid
+- `num_mesh_nodes` - number of mesh nodes; indexed as `num_mesh_nodes[l]` for hierarchical level `l`
+- `num_state_vars` - number of atmospheric state variables
+- `num_forcing_vars` - number of forcing input variables
+- `num_variables` - generic variable dimension used in metric functions
+- `hidden_dim` - internal hidden representation size in GNN layers and MLPs
+- `input_dim` - input feature dimensionality to a layer before transformation
+- `num_edges` - number of edges in a graph (g2m, m2g, same-level, up, down)
+- `num_send` - number of sender nodes in a message-passing step
+- `num_rec` - number of receiver nodes in a message-passing step
 
 # Development and Contributing
 Any push or Pull-Request to the main branch will trigger a selection of pre-commit hooks.
diff --git a/neural_lam/config.py b/neural_lam/config.py
index cc5cd0edc..1da43fff9 100644
--- a/neural_lam/config.py
+++ b/neural_lam/config.py
@@ -33,6 +33,7 @@ class DatastoreSelection:
     """
 
     kind: str
+    config_path: str
 
     def __post_init__(self):
         """
@@ -46,8 +47,6 @@ def __post_init__(self):
         if self.kind not in DATASTORES:
             raise ValueError(f"Datastore kind {self.kind} is not implemented")
 
-    config_path: str
-
 
 @dataclasses.dataclass
 class ManualStateFeatureWeighting:
@@ -103,6 +102,9 @@ class TrainingConfig:
         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).
     """
 
     state_feature_weighting: Union[
diff --git a/neural_lam/custom_loggers.py b/neural_lam/custom_loggers.py
index 3891e00eb..8af34df0a 100644
--- a/neural_lam/custom_loggers.py
+++ b/neural_lam/custom_loggers.py
@@ -29,6 +29,11 @@ def __init__(self, experiment_name, tracking_uri, run_name):
             MLflow tracking server URI.
         run_name : str
             Human-readable run name stored as ``mlflow.runName``.
+
+        Notes
+        -----
+        Starts the MLflow run with ``log_system_metrics=True`` and also
+        records ``run_id`` as an MLflow param.
         """
         super().__init__(
             experiment_name=experiment_name, tracking_uri=tracking_uri
@@ -65,6 +70,7 @@ def log_image(self, key, images, step=None):
         step : int or None, optional
             Step to associate with the log entry. ``None`` logs without
             a step suffix.
+
         Raises
         ------
         SystemExit
diff --git a/neural_lam/datastore/base.py b/neural_lam/datastore/base.py
index edbf25059..8e353d099 100644
--- a/neural_lam/datastore/base.py
+++ b/neural_lam/datastore/base.py
@@ -30,25 +30,28 @@ class BaseDatastore(abc.ABC):
     `weather_dataset.WeatherDataset` class (which inherits from
     `torch.utils.data.Dataset` and uses the datastore to access the data).
 
-    # Forecast vs analysis data
+    Forecast vs analysis data
+    -------------------------
     If the datastore is used to represent forecast rather than analysis data,
-    then the `is_forecast` attribute should be set to True, and returned data
-    from `get_dataarray` is assumed to have `analysis_time` and `forecast_time`
+    then the ``is_forecast`` attribute should be set to True, and returned data
+    from ``get_dataarray`` is assumed to have `analysis_time` and `forecast_time`
     dimensions (rather than just `time`).
 
-    # Ensemble vs deterministic data
+    Ensemble vs deterministic data
+    ------------------------------
     If the datastore is used to present an ensemble of state realisations, for
     example for forecast ensembles, then the `is_ensemble` attribute should be
-    set to `True` and returned state data from `get_dataarray` is expected to
+    set to `True` and returned state data from ``get_dataarray`` is expected to
     have an `ensemble_member` dimension. If each ensemble member has its own
     forcing values, then `has_ensemble_forcing` should be set to `True`, and
-    returned forcing data from `get_dataarray` is expected to have an
+    returned forcing data from ``get_dataarray`` is expected to have an
     `ensemble_member` dimension; otherwise forcing data is expected not to have
     one.
 
-    # Grid index
+    Grid index
+    ----------
     All methods that return data specific to a grid point (like
-    `get_dataarray`) should have a single dimension named `grid_index` that
+    ``get_dataarray``) should have a single dimension named `grid_index` that
     represents the spatial grid index of the data. The actual x, y coordinates
     of the grid points should be stored in the `x` and `y` coordinates of the
     dataarray or dataset with the `grid_index` dimension as the coordinate for
diff --git a/neural_lam/datastore/mdp.py b/neural_lam/datastore/mdp.py
index 7fba8305c..2906af331 100644
--- a/neural_lam/datastore/mdp.py
+++ b/neural_lam/datastore/mdp.py
@@ -54,6 +54,12 @@ def __init__(self, config_path, n_boundary_points=30, reuse_existing=True):
             Whether to reuse an existing dataset zarr file if it exists and its
             creation date is newer than the configuration file.
 
+        Raises
+        ------
+        ValueError
+            If the dataset does not contain all of the required
+            train/val/test splits.
+
         """
         self._config_path = Path(config_path)
         self._root_path = self._config_path.parent
@@ -409,6 +415,12 @@ def coords_projection(self) -> ccrs.Projection:
         ccrs.Projection
             The projection of the coordinates.
 
+        Raises
+        ------
+        ValueError
+            If the `projection` entry is missing from the `extra` section of
+            the config, or if its `class_name` or `kwargs` keys are missing.
+
         """
         if "projection" not in self._config.extra:
             raise ValueError(
diff --git a/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py b/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py
index 06dd0a97a..74f85aa23 100644
--- a/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py
+++ b/neural_lam/datastore/npyfilesmeps/compute_standardization_stats.py
@@ -150,24 +150,36 @@ def save_stats(
     static_dir_path : str or pathlib.Path
         Directory where ``*.pt`` files should be written.
     means : Sequence[torch.Tensor]
-        Batch-wise means with shape ``(N_batch, d_features)``.
+        Shape ``(B, num_state_vars)``. Batch-wise means. Sequence of one or
+        more tensors each of shape ``(B_i, num_state_vars)``; concatenated /
+        stacked along the batch dim before reduction. Typical callers pass a
+        single already-gathered tensor. The list can have length 0 to skip
+        saving.
     squares : Sequence[torch.Tensor]
-        Batch-wise second moments with shape ``(N_batch, d_features)``.
+        Shape ``(B, num_state_vars)``. Batch-wise second moments. Sequence of
+        one or more tensors each of shape ``(B_i, num_state_vars)``;
+        concatenated / stacked along the batch dim before reduction. Typical
+        callers pass a single already-gathered tensor. The list can have
+        length 0 to skip saving.
     flux_means : Sequence[torch.Tensor]
-        Flux means of shape ``(N_batch,)``; pass an empty sequence to skip
-        saving.
+        Shape ``(B,)``. Flux means. Sequence of one or more tensors each of
+        shape ``(B_i,)``; concatenated / stacked along the batch dim before
+        reduction. Typical callers pass a single already-gathered tensor.
+        The list can have length 0 to skip saving.
     flux_squares : Sequence[torch.Tensor]
-        Flux second moments of shape ``(N_batch,)``; pass an empty sequence to
-        skip saving.
+        Shape ``(B,)``. Flux second moments. Sequence of one or more tensors
+        each of shape ``(B_i,)``; concatenated / stacked along the batch dim
+        before reduction. Typical callers pass a single already-gathered
+        tensor. The list can have length 0 to skip saving.
     filename_prefix : str
         Prefix (e.g., ``"parameter"`` or ``"diff"``) for saved tensors.
     """
     means = (
         torch.stack(means) if len(means) > 1 else means[0]
-    )  # (N_batch, d_features,)
+    )  # (B, d_features,)
     squares = (
         torch.stack(squares) if len(squares) > 1 else squares[0]
-    )  # (N_batch, d_features,)
+    )  # (B, d_features,)
     mean = torch.mean(means, dim=0)  # (d_features,)
     second_moment = torch.mean(squares, dim=0)  # (d_features,)
     std = torch.sqrt(second_moment - mean**2)  # (d_features,)
@@ -186,10 +198,10 @@ def save_stats(
         return
     flux_means = (
         torch.stack(flux_means) if len(flux_means) > 1 else flux_means[0]
-    )  # (N_batch,)
+    )  # (B,)
     flux_squares = (
         torch.stack(flux_squares) if len(flux_squares) > 1 else flux_squares[0]
-    )  # (N_batch,)
+    )  # (B,)
     flux_mean = torch.mean(flux_means)  # (,)
     flux_second_moment = torch.mean(flux_squares)  # (,)
     flux_std = torch.sqrt(flux_second_moment - flux_mean**2)  # (,)
@@ -277,15 +289,15 @@ def main(
                 target_batch.to(device),
                 forcing_batch.to(device),
             )
-        # (N_batch, N_t, N_grid, d_features)
+        # (B, N_t, num_grid_nodes, d_features)
         batch = torch.cat((init_batch, target_batch), dim=1)
         # Flux at 1st windowed position is index 0 in forcing
         flux_batch = forcing_batch[:, :, :, 0]
-        # (N_batch, d_features,)
+        # (B, d_features,)
         means.append(torch.mean(batch, dim=(1, 2)).cpu())
         squares.append(
             torch.mean(batch**2, dim=(1, 2)).cpu()
-        )  # (N_batch, d_features,)
+        )  # (B, d_features,)
         flux_means.append(torch.mean(flux_batch).cpu())  # (,)
         flux_squares.append(torch.mean(flux_batch**2).cpu())  # (,)
 
@@ -330,10 +342,10 @@ def main(
                 )
             ]
     else:
-        means = [torch.cat(means, dim=0)]  # (N_batch, d_features,)
-        squares = [torch.cat(squares, dim=0)]  # (N_batch, d_features,)
-        flux_means = [torch.tensor(flux_means)]  # (N_batch,)
-        flux_squares = [torch.tensor(flux_squares)]  # (N_batch,)
+        means = [torch.cat(means, dim=0)]  # (B, d_features,)
+        squares = [torch.cat(squares, dim=0)]  # (B, d_features,)
+        flux_means = [torch.tensor(flux_means)]  # (B,)
+        flux_squares = [torch.tensor(flux_squares)]  # (B,)
 
     if rank == 0:
         save_stats(
@@ -391,7 +403,7 @@ def main(
             init_batch, target_batch = init_batch.to(device), target_batch.to(
                 device
             )
-        # (N_batch, N_t', N_grid, d_features)
+        # (B, N_t', num_grid_nodes, d_features)
         batch = torch.cat((init_batch, target_batch), dim=1)
         # Note: batch contains only 1h-steps
         stepped_batch = torch.cat(
@@ -401,14 +413,14 @@ def main(
             ],
             dim=0,
         )
-        # (N_batch', N_t, N_grid, d_features),
-        # N_batch' = step_length*N_batch
+        # (B', N_t, num_grid_nodes, d_features),
+        # B' = step_length*B
         batch_diffs = stepped_batch[:, 1:] - stepped_batch[:, :-1]
-        # (N_batch', N_t-1, N_grid, d_features)
+        # (B', N_t-1, num_grid_nodes, d_features)
         diff_means.append(torch.mean(batch_diffs, dim=(1, 2)).cpu())
-        # (N_batch', d_features,)
+        # (B', d_features,)
         diff_squares.append(torch.mean(batch_diffs**2, dim=(1, 2)).cpu())
-        # (N_batch', d_features,)
+        # (B', d_features,)
 
     if distributed and world_size > 1:
         dist.barrier()
@@ -432,8 +444,8 @@ def main(
             diff_means = [diff_means_gathered[:n_original_windows]]
             diff_squares = [diff_squares_gathered[:n_original_windows]]
 
-    diff_means = [torch.cat(diff_means, dim=0)]  # (N_batch', d_features,)
-    diff_squares = [torch.cat(diff_squares, dim=0)]  # (N_batch', d_features,)
+    diff_means = [torch.cat(diff_means, dim=0)]  # (B', d_features,)
+    diff_squares = [torch.cat(diff_squares, dim=0)]  # (B', d_features,)
 
     if rank == 0:
         save_stats(static_dir_path, diff_means, diff_squares, [], [], "diff")
diff --git a/neural_lam/datastore/npyfilesmeps/config.py b/neural_lam/datastore/npyfilesmeps/config.py
index b5f015b6b..b423b4a15 100644
--- a/neural_lam/datastore/npyfilesmeps/config.py
+++ b/neural_lam/datastore/npyfilesmeps/config.py
@@ -38,6 +38,11 @@ class Dataset:
         var_units: A list of units for each variable.
         var_longnames: A list of long, descriptive names for each variable.
         num_forcing_features: The number of forcing features in the dataset.
+        num_timesteps: The number of timesteps per analysis sample.
+        step_length: The time delta between consecutive timesteps.
+        num_ensemble_members: The number of ensemble members in the dataset.
+        remove_state_features_with_index: Indices of state features to drop
+        when loading the dataset.
 
     """
 
diff --git a/neural_lam/datastore/npyfilesmeps/store.py b/neural_lam/datastore/npyfilesmeps/store.py
index 120bdde0c..51d3ddf29 100644
--- a/neural_lam/datastore/npyfilesmeps/store.py
+++ b/neural_lam/datastore/npyfilesmeps/store.py
@@ -64,22 +64,22 @@ class NpyFilesDatastoreMEPS(BaseRegularGridDatastore):
     """
     Represents a dataset stored as numpy files on disk. The dataset is assumed
     to be stored in a directory structure where each sample is stored in a
-    separate file. The file-name format is assumed to be described by
-    ``STATE_FILENAME_FORMAT``.
+    separate file. The file-name format is assumed to be
+    `nwp_{analysis_time:%Y%m%d%H}_mbr{member_id:03d}.npy`.
 
     The MEPS dataset is organised into three splits: train, val, and test. Each
     split has a set of files which are:
 
-    - ``STATE_FILENAME_FORMAT``:
+    - `nwp_{analysis_time:%Y%m%d%H}_mbr{member_id:03d}.npy`:
         The state variables for a forecast started at `analysis_time` with
         member id `member_id`. The dimensions of the array are
         `[forecast_timestep, y, x, feature]`.
 
-    - ``TOA_SW_DOWN_FLUX_FILENAME_FORMAT``:
+    - `nwp_toa_downwelling_shortwave_flux_{analysis_time:%Y%m%d%H}.npy`:
         The top-of-atmosphere downwelling shortwave flux at `time`. The
         dimensions of the array are `[forecast_timestep, y, x]`.
 
-    - ``OPEN_WATER_FILENAME_FORMAT``:
+    - `wtr_{analysis_time:%Y%m%d%H}.npy`:
         The open water fraction at `time`. The dimensions of the array are
         `[y, x]`.
 
@@ -139,23 +139,51 @@ class NpyFilesDatastoreMEPS(BaseRegularGridDatastore):
         ├── parameter_weights.npy
         └── surface_geopotential.npy
 
+    Notes
+    -----
+    Folder structure::
+
+        meps_example_reduced
+        ├── data_config.yaml
+        ├── samples
+        │   ├── test
+        │   │   ├── nwp_2022090100_mbr000.npy
+        │   │   ├── ...
+        │   ├── train
+        │   │   ├── nwp_2022040100_mbr000.npy
+        │   │   ├── ...
+        │   └── val
+        │       ├── nwp_2022060500_mbr000.npy
+        │       └── ...
+        └── static
+            ├── border_mask.npy
+            ├── diff_mean.pt
+            ├── diff_std.pt
+            ├── flux_stats.pt
+            ├── grid_features.pt
+            ├── nwp_xy.npy
+            ├── parameter_mean.pt
+            ├── parameter_std.pt
+            ├── parameter_weights.npy
+            └── surface_geopotential.npy
+
     For the MEPS dataset:
     N_t' = 65
     N_t = 65//subsample_step (= 21 for 3h steps)
     dim_y = 268
     dim_x = 238
-    N_grid = 268x238 = 63784
-    d_features = 17 (d_features' = 18)
-    d_forcing = 5
+    num_grid_nodes = 268x238 = 63784
+    num_state_vars = 17 (num_state_vars' = 18)
+    num_forcing_vars = 5
 
     For the MEPS reduced dataset:
     N_t' = 65
     N_t = 65//subsample_step (= 21 for 3h steps)
     dim_y = 134
     dim_x = 119
-    N_grid = 134x119 = 15946
-    d_features = 8
-    d_forcing = 1
+    num_grid_nodes = 134x119 = 15946
+    num_state_vars = 8
+    num_forcing_vars = 1
     """
 
     SHORT_NAME = "npyfilesmeps"
diff --git a/neural_lam/interaction_net.py b/neural_lam/interaction_net.py
index 31bf86936..1e97b2136 100644
--- a/neural_lam/interaction_net.py
+++ b/neural_lam/interaction_net.py
@@ -36,8 +36,7 @@ def __init__(
         ----------
         edge_index : torch.Tensor
             Edge connectivity tensor in PyG format.
-
-            * **Shape**: ``(2, num_edges)``
+            Shape ``(2, num_edges)``.
         input_dim : int
             Dimensionality of both node and edge input representations.
         update_edges : bool, optional
@@ -58,7 +57,7 @@ def __init__(
 
         Raises
         ------
-        AssertionError
+        ValueError
             If ``aggr`` is not one of ``"sum"`` or ``"mean"``.
         """
         if aggr not in ("sum", "mean"):
@@ -114,23 +113,23 @@ def forward(self, send_rep, rec_rep, edge_rep):
         Parameters
         ----------
         send_rep : torch.Tensor
-            Shape ``(B, N_send, d_h)``. Sender node representations.
-            Dims: ``B`` is batch size, ``N_send`` is the number of
-            sender nodes, and ``d_h`` is the hidden dimension.
+            Sender node representations.
+            Shape ``(num_send, input_dim)``.
         rec_rep : torch.Tensor
-            Shape ``(B, N_rec, d_h)``. Receiver node representations.
-            Dims: ``N_rec`` is the number of receiver nodes.
+            Receiver node representations.
+            Shape ``(num_rec, input_dim)``.
         edge_rep : torch.Tensor
-            Shape ``(B, M, d_h)``. Edge representations. Dims: ``M``
-            is the number of edges.
+            Edge representations.
+            Shape ``(num_edges, input_dim)``.
 
         Returns
         -------
         rec_rep : torch.Tensor
-            Shape ``(B, N_rec, d_h)``. Updated receiver node
-            representations.
+            Updated receiver node representations.
+            Shape ``(num_rec, input_dim)``.
         edge_rep : torch.Tensor
-            Shape ``(B, M, d_h)``. Updated edge representations.
+            Updated edge representations.
+            Shape ``(num_edges, input_dim)``.
             Only returned when ``update_edges=True``.
         """
         # Always concatenate to [rec_nodes, send_nodes] for propagation,
diff --git a/neural_lam/loss_weighting.py b/neural_lam/loss_weighting.py
index 3335fb6d7..12564ca2f 100644
--- a/neural_lam/loss_weighting.py
+++ b/neural_lam/loss_weighting.py
@@ -27,6 +27,12 @@ def get_manual_state_feature_weights(
     -------
     list[float]
         List of floats containing the state feature weights.
+
+    Raises
+    ------
+    ValueError
+        If the set of feature names in ``weighting_config.weights`` does
+        not match the state feature names in the datastore.
     """
     state_feature_names = datastore.get_vars_names(category="state")
     feature_weight_names = weighting_config.weights.keys()
@@ -92,6 +98,12 @@ def get_state_feature_weighting(
     -------
     list[float]
         List of floats containing the state feature weights.
+
+    Raises
+    ------
+    NotImplementedError
+        If ``config.training.state_feature_weighting`` is not a
+        recognised weighting configuration type.
     """
     weighting_config = config.training.state_feature_weighting
 
diff --git a/neural_lam/metrics.py b/neural_lam/metrics.py
index 76c86b8c3..47bf3d58d 100644
--- a/neural_lam/metrics.py
+++ b/neural_lam/metrics.py
@@ -17,6 +17,12 @@ def get_metric(metric_name):
     -------
     callable
         Function implementing the requested metric.
+
+    Raises
+    ------
+    AssertionError
+        If ``metric_name`` (case-insensitive) is not a key in
+        ``DEFINED_METRICS``.
     """
     metric_name_lower = metric_name.lower()
     assert (
@@ -32,23 +38,23 @@ def mask_and_reduce_metric(metric_entry_vals, mask, average_grid, sum_vars):
     Parameters
     ----------
     metric_entry_vals : torch.Tensor
-        Shape ``(..., N, d_state)``. Per-entry metric values. ``(...)``
+        Shape ``(..., N, num_variables)``. Per-entry metric values. ``(...)``
         denotes any number of broadcastable batch dimensions, ``N`` is
-        the number of grid nodes, and ``d_state`` is the number of state
-        variables.
+        the number of grid nodes, and ``num_variables`` is the number of
+        variables in the gridded representation (e.g. state features).
     mask : torch.Tensor or None
         Shape ``(N,)``. Boolean mask selecting which grid nodes to
         include. ``None`` means all nodes are used.
     average_grid : bool
         If True, average over the grid dimension ``N``.
     sum_vars : bool
-        If True, sum over the variable dimension ``d_state``.
+        If True, sum over the variable dimension ``num_variables``.
 
     Returns
     -------
     torch.Tensor
         Reduced metric values. Shape is one of ``(...,)``,
-        ``(..., d_state)``, ``(..., N)``, or ``(..., N, d_state)``
+        ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)``
         depending on ``average_grid`` and ``sum_vars``.
     """
     # Only keep grid nodes in mask
@@ -77,14 +83,14 @@ def wmse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     Parameters
     ----------
     pred : torch.Tensor
-        Shape ``(..., N, d_state)``. Model prediction. ``(...)`` denotes
+        Shape ``(..., N, num_variables)``. Model prediction. ``(...)`` denotes
         any number of broadcastable batch dimensions, ``N`` is the number
-        of grid nodes, and ``d_state`` is the number of state variables.
+        of grid nodes, and ``num_variables`` is the number of state variables.
     target : torch.Tensor
-        Shape ``(..., N, d_state)``. Ground-truth target. Dims: same as
+        Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as
         ``pred``.
     pred_std : torch.Tensor
-        Shape ``(..., N, d_state)`` or ``(d_state,)``. Predicted
+        Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted
         standard deviation used as per-entry weight.
     mask : torch.Tensor or None, optional
         Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all
@@ -98,7 +104,7 @@ def wmse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     -------
     torch.Tensor
         Reduced metric values. Shape is one of ``(...,)``,
-        ``(..., d_state)``, ``(..., N)``, or ``(..., N, d_state)``
+        ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)``
         depending on ``average_grid`` and ``sum_vars``.
     """
     entry_mse = torch.nn.functional.mse_loss(
@@ -123,14 +129,14 @@ def mse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     Parameters
     ----------
     pred : torch.Tensor
-        Shape ``(..., N, d_state)``. Model prediction. ``(...)`` denotes
+        Shape ``(..., N, num_variables)``. Model prediction. ``(...)`` denotes
         any number of broadcastable batch dimensions, ``N`` is the number
-        of grid nodes, and ``d_state`` is the number of state variables.
+        of grid nodes, and ``num_variables`` is the number of state variables.
     target : torch.Tensor
-        Shape ``(..., N, d_state)``. Ground-truth target. Dims: same as
+        Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as
         ``pred``.
     pred_std : torch.Tensor
-        Shape ``(..., N, d_state)`` or ``(d_state,)``. Predicted
+        Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted
         standard deviation (unused; ``pred_std`` is replaced by ones
         internally).
     mask : torch.Tensor or None, optional
@@ -145,7 +151,7 @@ def mse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     -------
     torch.Tensor
         Reduced metric values. Shape is one of ``(...,)``,
-        ``(..., d_state)``, ``(..., N)``, or ``(..., N, d_state)``
+        ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)``
         depending on ``average_grid`` and ``sum_vars``.
     """
     # Replace pred_std with constant ones
@@ -161,14 +167,14 @@ def wmae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     Parameters
     ----------
     pred : torch.Tensor
-        Shape ``(..., N, d_state)``. Model prediction. ``(...)`` denotes
+        Shape ``(..., N, num_variables)``. Model prediction. ``(...)`` denotes
         any number of broadcastable batch dimensions, ``N`` is the number
-        of grid nodes, and ``d_state`` is the number of state variables.
+        of grid nodes, and ``num_variables`` is the number of state variables.
     target : torch.Tensor
-        Shape ``(..., N, d_state)``. Ground-truth target. Dims: same as
+        Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as
         ``pred``.
     pred_std : torch.Tensor
-        Shape ``(..., N, d_state)`` or ``(d_state,)``. Predicted
+        Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted
         standard deviation used as per-entry weight.
     mask : torch.Tensor or None, optional
         Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all
@@ -182,7 +188,7 @@ def wmae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     -------
     torch.Tensor
         Reduced metric values. Shape is one of ``(...,)``,
-        ``(..., d_state)``, ``(..., N)``, or ``(..., N, d_state)``
+        ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)``
         depending on ``average_grid`` and ``sum_vars``.
     """
     entry_mae = torch.nn.functional.l1_loss(
@@ -207,14 +213,14 @@ def mae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     Parameters
     ----------
     pred : torch.Tensor
-        Shape ``(..., N, d_state)``. Model prediction. ``(...)`` denotes
+        Shape ``(..., N, num_variables)``. Model prediction. ``(...)`` denotes
         any number of broadcastable batch dimensions, ``N`` is the number
-        of grid nodes, and ``d_state`` is the number of state variables.
+        of grid nodes, and ``num_variables`` is the number of state variables.
     target : torch.Tensor
-        Shape ``(..., N, d_state)``. Ground-truth target. Dims: same as
+        Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as
         ``pred``.
     pred_std : torch.Tensor
-        Shape ``(..., N, d_state)`` or ``(d_state,)``. Predicted
+        Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted
         standard deviation (unused; ``pred_std`` is replaced by ones
         internally).
     mask : torch.Tensor or None, optional
@@ -229,7 +235,7 @@ def mae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     -------
     torch.Tensor
         Reduced metric values. Shape is one of ``(...,)``,
-        ``(..., d_state)``, ``(..., N)``, or ``(..., N, d_state)``
+        ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)``
         depending on ``average_grid`` and ``sum_vars``.
     """
     # Replace pred_std with constant ones
@@ -245,14 +251,14 @@ def nll(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     Parameters
     ----------
     pred : torch.Tensor
-        Shape ``(..., N, d_state)``. Predicted mean. ``(...)`` denotes
+        Shape ``(..., N, num_variables)``. Predicted mean. ``(...)`` denotes
         any number of broadcastable batch dimensions, ``N`` is the number
-        of grid nodes, and ``d_state`` is the number of state variables.
+        of grid nodes, and ``num_variables`` is the number of state variables.
     target : torch.Tensor
-        Shape ``(..., N, d_state)``. Ground-truth target. Dims: same as
+        Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as
         ``pred``.
     pred_std : torch.Tensor
-        Shape ``(..., N, d_state)`` or ``(d_state,)``. Predicted
+        Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted
         standard deviation of the Gaussian.
     mask : torch.Tensor or None, optional
         Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all
@@ -266,7 +272,7 @@ def nll(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True):
     -------
     torch.Tensor
         Reduced metric values. Shape is one of ``(...,)``,
-        ``(..., d_state)``, ``(..., N)``, or ``(..., N, d_state)``
+        ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)``
         depending on ``average_grid`` and ``sum_vars``.
     """
     # Broadcast pred_std if shaped (num_variables,) via distribution internals
@@ -290,14 +296,14 @@ def crps_gauss(
     Parameters
     ----------
     pred : torch.Tensor
-        Shape ``(..., N, d_state)``. Predicted mean. ``(...)`` denotes
+        Shape ``(..., N, num_variables)``. Predicted mean. ``(...)`` denotes
         any number of broadcastable batch dimensions, ``N`` is the number
-        of grid nodes, and ``d_state`` is the number of state variables.
+        of grid nodes, and ``num_variables`` is the number of state variables.
     target : torch.Tensor
-        Shape ``(..., N, d_state)``. Ground-truth target. Dims: same as
+        Shape ``(..., N, num_variables)``. Ground-truth target. Dims: same as
         ``pred``.
     pred_std : torch.Tensor
-        Shape ``(..., N, d_state)`` or ``(d_state,)``. Predicted
+        Shape ``(..., N, num_variables)`` or ``(num_variables,)``. Predicted
         standard deviation of the Gaussian.
     mask : torch.Tensor or None, optional
         Shape ``(N,)``. Boolean mask over grid nodes. ``None`` uses all
@@ -311,7 +317,7 @@ def crps_gauss(
     -------
     torch.Tensor
         Reduced metric values. Shape is one of ``(...,)``,
-        ``(..., d_state)``, ``(..., N)``, or ``(..., N, d_state)``
+        ``(..., num_variables)``, ``(..., N)``, or ``(..., N, num_variables)``
         depending on ``average_grid`` and ``sum_vars``.
     """
     std_normal = torch.distributions.Normal(
diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py
index c342dc115..502fb0ae6 100644
--- a/neural_lam/models/forecasters/autoregressive.py
+++ b/neural_lam/models/forecasters/autoregressive.py
@@ -1,6 +1,4 @@
-"""
-Forecaster that uses an auto-regressive strategy to unroll a forecast.
-"""
+"""Forecaster that uses an auto-regressive strategy to unroll a forecast."""
 
 # Standard library
 from typing import Optional
@@ -53,7 +51,7 @@ def predicts_std(self) -> bool:
         Returns
         -------
         bool
-            True if the forecaster predicts standard deviation, False otherwise.
+            ``True`` if the forecaster predicts standard deviation, ``False`` otherwise.
         """
         return self.predictor.predicts_std
 
@@ -71,36 +69,36 @@ def forward(
         Parameters
         ----------
         init_states : torch.Tensor
-            Shape ``(B, 2, num_grid_nodes, d_f)``. The two initial states
-            ``[X_{t-1}, X_t]`` used to start the rollout from. Dims: ``B`` is
-            batch size, ``2`` is the time index, ``num_grid_nodes`` is the
-            number of spatial nodes, and ``d_f`` is the state feature
-            dimension.
+            Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial
+            states ``[X_{t-1}, X_t]`` used to start the rollout from. Dims:
+            ``B`` is batch size, ``2`` initial time steps (``X_{t-1}, X_t``),
+            ``num_grid_nodes`` is the number of spatial nodes, and
+            ``num_state_vars`` is the number of state variables.
         forcing_features : torch.Tensor
-            Shape ``(B, pred_steps, num_grid_nodes, d_forcing)``. Forcing
-            features for each predicted step; ``pred_steps`` defines the
-            rollout length. Dims: ``B`` is batch size, ``pred_steps`` is
+            Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``.
+            Forcing features for each predicted step; ``pred_steps`` defines
+            the rollout length. Dims: ``B`` is batch size, ``pred_steps`` is
             the number of predicted steps, ``num_grid_nodes`` is the
-            number of spatial nodes, and ``d_forcing`` is the forcing
-            feature dimension (already concatenated past/current/future
+            number of spatial nodes, and ``num_forcing_vars`` is the number
+            of forcing variables (already concatenated past/current/future
             windows).
         boundary_states : torch.Tensor
-            Shape ``(B, pred_steps, num_grid_nodes, d_f)``. True state
+            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True state
             values used ONLY to overwrite boundary nodes at each AR step.
             The interior prediction at step ``i`` must not depend on
             ``boundary_states[:, i]`` in any other way. Dims: ``B`` is
             batch size, ``pred_steps`` is the number of predicted steps,
             ``num_grid_nodes`` is the number of spatial nodes, and
-            ``d_f`` is the state feature dimension.
+            ``num_state_vars`` is the state feature dimension.
 
         Returns
         -------
         prediction : torch.Tensor
-            Shape ``(B, pred_steps, num_grid_nodes, d_f)``. Stacked
+            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. Stacked
             per-step forecasts (with boundary overwritten by the true
             value). Dims: same as ``boundary_states``.
         pred_std : torch.Tensor or None
-            Shape ``(B, pred_steps, num_grid_nodes, d_f)`` when the
+            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)`` when the
             wrapped predictor outputs an std, otherwise ``None`` (in which
             case ``ForecasterModule`` substitutes the constant
             per-variable std). Dims: same as ``prediction``.
diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py
index 66b0d4c8c..8fdd80e80 100644
--- a/neural_lam/models/forecasters/base.py
+++ b/neural_lam/models/forecasters/base.py
@@ -1,6 +1,4 @@
-"""
-Base class for forecasters.
-"""
+"""Base class for forecasters."""
 
 # Standard library
 from abc import ABC, abstractmethod
@@ -27,7 +25,7 @@ def predicts_std(self) -> bool:
         Returns
         -------
         bool
-            True if the forecaster predicts standard deviation, False otherwise.
+            ``True`` if the forecaster predicts standard deviation, ``False`` otherwise.
         """
 
     @abstractmethod
@@ -44,25 +42,22 @@ def forward(
         Parameters
         ----------
         init_states : torch.Tensor
-            Shape ``(B, 2, num_grid_nodes, d_f)``. The two initial states
+            Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial states
             ``[X_{t-1}, X_t]`` used to start the forecast from. Dims: ``B`` is
             batch size, ``2`` is the time index (``[X_{t-1}, X_t]``),
-            ``num_grid_nodes`` is the number of spatial nodes, and ``d_f``
+            ``num_grid_nodes`` is the number of spatial nodes, and ``num_state_vars``
             is the state feature dimension.
         forcing_features : torch.Tensor
-            Shape ``(B, pred_steps, num_grid_nodes, d_forcing)``. External
+            Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. External
             forcings provided at each predicted step. Dims: ``B`` is batch
             size, ``pred_steps`` is the autoregressive rollout length,
             ``num_grid_nodes`` is the number of spatial nodes, and
-            ``d_forcing`` is the forcing feature dimension (already
+            ``num_forcing_vars`` is the forcing feature dimension (already
             concatenated past/current/future windows).
         boundary_states : torch.Tensor
-            Shape ``(B, pred_steps, num_grid_nodes, d_f)``. True state
-            values used ONLY to overwrite boundary nodes at each AR step
-            — interior predictions must not depend on ``boundary_states``
-            in any other way. Dims: ``B`` is batch size, ``pred_steps``
-            is the rollout length, ``num_grid_nodes`` is the number of
-            spatial nodes, and ``d_f`` is the state feature dimension.
+            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True state values used ONLY to overwrite boundary nodes at each AR step; interior predictions must not depend on ``boundary_states`` in any other way. Dims: ``B`` is batch size, ``pred_steps`` is the
+            rollout length, ``num_grid_nodes`` is the number of
+            spatial nodes, and ``num_state_vars`` is the state feature dimension.
             This is a temporary mechanism that mirrors the pre-refactor
             ARModel behavior; it will be replaced by a dedicated
             boundary-forcing input in #138 (training on interior +
@@ -72,11 +67,11 @@ def forward(
         Returns
         -------
         prediction : torch.Tensor
-            Shape ``(B, pred_steps, num_grid_nodes, d_f)``. Forecast of
+            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. Forecast of
             state at each predicted step. Dims: same as
             ``boundary_states``.
         pred_std : torch.Tensor or None
-            Shape ``(B, pred_steps, num_grid_nodes, d_f)`` when
+            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)`` when
             ``predicts_std`` is True, otherwise ``None``. Per-feature
             predicted standard deviation; when ``None``, the constant
             per-variable std is substituted upstream by
diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py
index 08babb612..958169c49 100644
--- a/neural_lam/models/module.py
+++ b/neural_lam/models/module.py
@@ -1,6 +1,4 @@
-"""
-Lightning module handling training, validation and testing loops.
-"""
+"""Lightning module handling training, validation and testing loops."""
 
 # Standard library
 import os
@@ -75,13 +73,20 @@ def __init__(
             Specific rollout steps to log during validation/testing.
         metrics_watch : list of str, optional
             List of metrics to watch and log specifically.
-        var_leads_metrics_watch : dict, optional
-            Dictionary mapping variable indices to lead times for watching.
+        var_leads_metrics_watch : dict of {int: list of int}, optional
+            Mapping from variable index to a list of rollout steps to log
+            individually for the configured metrics.
         args : argparse.Namespace, optional
-            Legacy arguments for backward compatibility.
+            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.
         """
         super().__init__()
-        # Pre-refactor ARModel checkpoints saved every hyperparameter nested
+        # Pre-refactor ``ARModel`` checkpoints saved every hyperparameter nested
         # inside an argparse Namespace under the single key 'args'. When
         # Lightning calls __init__ during load_from_checkpoint it would
         # otherwise drop 'args' (not in the new signature) and silently fall
@@ -862,7 +867,7 @@ def on_load_checkpoint(self, checkpoint):
         loaded_state_dict = checkpoint["state_dict"]
 
         # 1. Broad namespace remap: for pre-refactor checkpoints
-        # The old ARModel was a flat LightningModule. Everything that belonged
+        # The old ``ARModel`` was a flat LightningModule. Everything that belonged
         # to the predictor needs to be moved to 'forecaster.predictor.'
         old_keys = list(loaded_state_dict.keys())
         for key in old_keys:
diff --git a/neural_lam/models/step_predictors/base.py b/neural_lam/models/step_predictors/base.py
index e1f14700f..2d79ed72f 100644
--- a/neural_lam/models/step_predictors/base.py
+++ b/neural_lam/models/step_predictors/base.py
@@ -1,6 +1,4 @@
-"""
-Base class for step predictors.
-"""
+"""Base class for step predictors."""
 
 # Standard library
 from abc import ABC, abstractmethod
@@ -104,7 +102,7 @@ def predicts_std(self) -> bool:
         Returns
         -------
         bool
-            True if the predictor predicts standard deviation, False otherwise.
+            ``True`` if the predictor predicts standard deviation, ``False`` otherwise.
         """
         return self.output_std
 
@@ -141,28 +139,28 @@ def forward(
         Parameters
         ----------
         prev_state : torch.Tensor
-            Shape ``(B, num_grid_nodes, d_f)``. The current state
+            Shape ``(B, num_grid_nodes, num_state_vars)``. The current state
             ``X_t``. Dims: ``B`` is batch size, ``num_grid_nodes`` is the
-            number of spatial nodes, and ``d_f`` is the number of state
+            number of spatial nodes, and ``num_state_vars`` is the number of state
             variables.
         prev_prev_state : torch.Tensor
-            Shape ``(B, num_grid_nodes, d_f)``. The previous state
+            Shape ``(B, num_grid_nodes, num_state_vars)``. The previous state
             ``X_{t-1}``, used as additional conditioning. Dims: same as
             ``prev_state``.
         forcing : torch.Tensor
-            Shape ``(B, num_grid_nodes, d_forcing)``. External forcings
+            Shape ``(B, num_grid_nodes, num_forcing_vars)``. External forcings
             for this step (already concatenated past/current/future
             windows). Dims: ``B`` is batch size, ``num_grid_nodes`` is
-            the number of spatial nodes, and ``d_forcing`` is the
+            the number of spatial nodes, and ``num_forcing_vars`` is the
             forcing feature dimension.
 
         Returns
         -------
         pred_state : torch.Tensor
-            Shape ``(B, num_grid_nodes, d_f)``. The predicted next
+            Shape ``(B, num_grid_nodes, num_state_vars)``. The predicted next
             state ``X_{t+1}``. Dims: same as ``prev_state``.
         pred_std : torch.Tensor or None
-            Shape ``(B, num_grid_nodes, d_f)`` when ``output_std``
+            Shape ``(B, num_grid_nodes, num_state_vars)`` when ``output_std``
             is True, otherwise ``None``. Per-feature predicted standard
             deviation. Dims: same as ``prev_state``.
         """
@@ -324,19 +322,19 @@ def get_clamped_new_state(self, state_delta, prev_state):
         Parameters
         ----------
         state_delta : torch.Tensor
-            Shape ``(B, num_grid_nodes, d_f)``. Raw predicted state
+            Shape ``(B, num_grid_nodes, num_state_vars)``. Raw predicted state
             increment (network output, already rescaled). Dims: ``B`` is
             batch size, ``num_grid_nodes`` is the number of spatial nodes,
-            and ``d_f`` is the number of state variables.
+            and ``num_state_vars`` is the number of state variables.
         prev_state : torch.Tensor
-            Shape ``(B, num_grid_nodes, d_f)``. Current state ``X_t``
+            Shape ``(B, num_grid_nodes, num_state_vars)``. Current state ``X_t``
             used as the base for the clamped update. Dims: same as
             ``state_delta``.
 
         Returns
         -------
         torch.Tensor
-            Shape ``(B, num_grid_nodes, d_f)``. Clamped next state.
+            Shape ``(B, num_grid_nodes, num_state_vars)``. Clamped next state.
             Dims: same as ``state_delta``.
         """
 
diff --git a/neural_lam/models/step_predictors/graph/base.py b/neural_lam/models/step_predictors/graph/base.py
index c87d00fc5..d97500e0c 100644
--- a/neural_lam/models/step_predictors/graph/base.py
+++ b/neural_lam/models/step_predictors/graph/base.py
@@ -1,6 +1,4 @@
-"""
-Base class for graph-based step predictors.
-"""
+"""Base class for graph-based step predictors."""
 
 # Standard library
 from typing import Dict, Optional
@@ -41,7 +39,8 @@ def __init__(
         Parameters
         ----------
         datastore : BaseDatastore
-            The datastore providing grid metadata and data access.
+            Datastore supplying data and information about dataset and
+            forecast region.
         graph_name : str, default "multiscale"
             The name of the graph to load.
         hidden_dim : int, default 64
@@ -189,9 +188,9 @@ def embedd_mesh_nodes(self):
         Returns
         -------
         torch.Tensor
-            Shape ``(num_mesh_nodes, d_h)``. Embedded mesh node
+            Shape ``(num_mesh_nodes, hidden_dim)``. Embedded mesh node
             representations. Dims: ``num_mesh_nodes`` is the number of
-            mesh nodes and ``d_h`` is the hidden dimension.
+            mesh nodes and ``hidden_dim`` is the hidden dimension.
         """
         raise NotImplementedError("embedd_mesh_nodes not implemented")
 
@@ -203,15 +202,15 @@ def process_step(self, mesh_rep):
         Parameters
         ----------
         mesh_rep : torch.Tensor
-            Shape ``(B, num_mesh_nodes, d_h)``. Current mesh node
+            Shape ``(B, num_mesh_nodes, hidden_dim)``. Current mesh node
             representations. Dims: ``B`` is batch size,
-            ``num_mesh_nodes`` is the number of mesh nodes, and ``d_h``
+            ``num_mesh_nodes`` is the number of mesh nodes, and ``hidden_dim``
             is the hidden dimension.
 
         Returns
         -------
         torch.Tensor
-            Shape ``(B, num_mesh_nodes, d_h)``. Updated mesh node
+            Shape ``(B, num_mesh_nodes, hidden_dim)``. Updated mesh node
             representations. Dims: same as ``mesh_rep``.
         """
         raise NotImplementedError("process_step not implemented")
@@ -228,29 +227,29 @@ def forward(self, prev_state, prev_prev_state, forcing):
         Parameters
         ----------
         prev_state : torch.Tensor
-            Shape ``(B, num_grid_nodes, d_f)``. The current state
+            Shape ``(B, num_grid_nodes, num_state_vars)``. The current state
             ``X_t``. Dims: ``B`` is batch size, ``num_grid_nodes`` is the
-            number of spatial grid nodes, and ``d_f`` is the number of
+            number of spatial grid nodes, and ``num_state_vars`` is the number of
             state variables.
         prev_prev_state : torch.Tensor
-            Shape ``(B, num_grid_nodes, d_f)``. The previous state
+            Shape ``(B, num_grid_nodes, num_state_vars)``. The previous state
             ``X_{t-1}``, used as additional conditioning. Dims: same as
             ``prev_state``.
         forcing : torch.Tensor
-            Shape ``(B, num_grid_nodes, d_forcing)``. External forcings
+            Shape ``(B, num_grid_nodes, num_forcing_vars)``. External forcings
             for this step (already concatenated past/current/future
             windows). Dims: ``B`` is batch size, ``num_grid_nodes`` is
-            the number of spatial grid nodes, and ``d_forcing`` is the
+            the number of spatial grid nodes, and ``num_forcing_vars`` is the
             forcing feature dimension.
 
         Returns
         -------
         new_state : torch.Tensor
-            Shape ``(B, num_grid_nodes, d_f)``. The predicted next state
+            Shape ``(B, num_grid_nodes, num_state_vars)``. The predicted next state
             ``X_{t+1}`` after delta-add and clamping. Dims: same as
             ``prev_state``.
         pred_std : torch.Tensor or None
-            Shape ``(B, num_grid_nodes, d_f)`` when ``output_std`` is
+            Shape ``(B, num_grid_nodes, num_state_vars)`` when ``output_std`` is
             True, otherwise ``None``. Per-feature predicted standard
             deviation (raw softplus output, not rescaled by diff
             statistics). Dims: same as ``prev_state``.
@@ -269,25 +268,25 @@ def forward(self, prev_state, prev_prev_state, forcing):
         )
 
         # Embed all features
-        grid_emb = self.grid_embedder(grid_features)  # (B, num_grid_nodes, d_h)
-        g2m_emb = self.g2m_embedder(self.g2m_features)  # (M_g2m, d_h)
-        m2g_emb = self.m2g_embedder(self.m2g_features)  # (M_m2g, d_h)
+        grid_emb = self.grid_embedder(grid_features)  # (B, num_grid_nodes, hidden_dim)
+        g2m_emb = self.g2m_embedder(self.g2m_features)  # (num_edges, hidden_dim)
+        m2g_emb = self.m2g_embedder(self.m2g_features)  # (num_edges, hidden_dim)
         mesh_emb = self.embedd_mesh_nodes()
 
         # Map from grid to mesh
         mesh_emb_expanded = self.expand_to_batch(
             mesh_emb, batch_size
-        )  # (B, num_mesh_nodes, d_h)
+        )  # (B, num_mesh_nodes, hidden_dim)
         g2m_emb_expanded = self.expand_to_batch(g2m_emb, batch_size)
 
         # This also splits representation into grid and mesh
         mesh_rep = self.g2m_gnn(
             grid_emb, mesh_emb_expanded, g2m_emb_expanded
-        )  # (B, num_mesh_nodes, d_h)
+        )  # (B, num_mesh_nodes, hidden_dim)
         # Also MLP with residual for grid representation
         grid_rep = grid_emb + self.encoding_grid_mlp(
             grid_emb
-        )  # (B, num_grid_nodes, d_h)
+        )  # (B, num_grid_nodes, hidden_dim)
 
         # Run processor step
         mesh_rep = self.process_step(mesh_rep)
@@ -296,7 +295,7 @@ def forward(self, prev_state, prev_prev_state, forcing):
         m2g_emb_expanded = self.expand_to_batch(m2g_emb, batch_size)
         grid_rep = self.m2g_gnn(
             mesh_rep, grid_rep, m2g_emb_expanded
-        )  # (B, num_grid_nodes, d_h)
+        )  # (B, num_grid_nodes, hidden_dim)
 
         # Map to output dimension, only for grid
         net_output = self.output_map(
@@ -306,7 +305,7 @@ def forward(self, prev_state, prev_prev_state, forcing):
         if self.output_std:
             pred_delta_mean, pred_std_raw = net_output.chunk(
                 2, dim=-1
-            )  # both (B, num_grid_nodes, d_f)
+            )  # both (B, num_grid_nodes, num_state_vars)
             # NOTE: The predicted std. is not scaled in any way here
             # linter for some reason does not think softplus is callable
             # pylint: disable-next=not-callable
diff --git a/neural_lam/models/step_predictors/graph/graph_lam.py b/neural_lam/models/step_predictors/graph/graph_lam.py
index 11716cf3b..0159e1fd7 100644
--- a/neural_lam/models/step_predictors/graph/graph_lam.py
+++ b/neural_lam/models/step_predictors/graph/graph_lam.py
@@ -1,6 +1,4 @@
-"""
-Graph-based LAM model with a flat mesh.
-"""
+"""Graph-based LAM model with a flat mesh."""
 
 # Standard library
 from typing import Dict, Optional
@@ -136,13 +134,13 @@ def embedd_mesh_nodes(self):
         Returns
         -------
         torch.Tensor
-            Shape ``(num_mesh_nodes, d_h)``. Embedded mesh node representations.
-            Dims: ``num_mesh_nodes`` is the number of mesh nodes and ``d_h`` is
+            Shape ``(num_mesh_nodes, hidden_dim)``. Embedded mesh node representations.
+            Dims: ``num_mesh_nodes`` is the number of mesh nodes and ``hidden_dim`` is
             the hidden dimension.
         """
         return self.mesh_embedder(
             self.mesh_static_features
-        )  # (num_mesh_nodes, d_h)
+        )  # (num_mesh_nodes, hidden_dim)
 
     def process_step(self, mesh_rep):
         """
@@ -152,25 +150,25 @@ def process_step(self, mesh_rep):
         Parameters
         ----------
         mesh_rep : torch.Tensor
-            Shape ``(B, num_mesh_nodes, d_h)``. Current mesh node
+            Shape ``(B, num_mesh_nodes, hidden_dim)``. Current mesh node
             representations. Dims: ``B`` is batch size, ``num_mesh_nodes`` is
-            the number of mesh nodes, and ``d_h`` is the hidden
+            the number of mesh nodes, and ``hidden_dim`` is the hidden
             dimension.
 
         Returns
         -------
         torch.Tensor
-            Shape ``(B, num_mesh_nodes, d_h)``. Updated mesh node
+            Shape ``(B, num_mesh_nodes, hidden_dim)``. Updated mesh node
             representations. Dims: same as ``mesh_rep``.
         """
         # Embed m2m here first
         batch_size = mesh_rep.shape[0]
-        m2m_emb = self.m2m_embedder(self.m2m_features)  # (M_mesh, d_h)
+        m2m_emb = self.m2m_embedder(self.m2m_features)  # (num_edges, hidden_dim)
         m2m_emb_expanded = self.expand_to_batch(
             m2m_emb, batch_size
-        )  # (B, M_mesh, d_h)
+        )  # (B, num_edges, hidden_dim)
 
         mesh_rep, _ = self.processor(
             mesh_rep, m2m_emb_expanded
-        )  # (B, num_mesh_nodes, d_h)
+        )  # (B, num_mesh_nodes, hidden_dim)
         return mesh_rep
diff --git a/neural_lam/models/step_predictors/graph/hi_lam.py b/neural_lam/models/step_predictors/graph/hi_lam.py
index 89531dfbb..79ec95cd2 100644
--- a/neural_lam/models/step_predictors/graph/hi_lam.py
+++ b/neural_lam/models/step_predictors/graph/hi_lam.py
@@ -168,16 +168,16 @@ def mesh_down_step(
         Parameters
         ----------
         mesh_rep_levels : list of torch.Tensor
-            One tensor per level, each of shape ``(B, num_mesh_nodes[l], d_h)``.
+            One tensor per level, each of shape ``(B, num_mesh_nodes[l], hidden_dim)``.
             Node representations at each hierarchy level. Dims: ``B`` is
             batch size, ``num_mesh_nodes[l]`` is the node count at level ``l``,
-            and ``d_h`` is the hidden dimension.
+            and ``hidden_dim`` is the hidden dimension.
         mesh_same_rep : list of torch.Tensor
-            One tensor per level, each of shape ``(B, M_same[l], d_h)``.
+            One tensor per level, each of shape ``(B, num_edges[l], hidden_dim)``.
             Same-level edge representations.
         mesh_down_rep : list of torch.Tensor
             One tensor per inter-level gap, each of shape
-            ``(B, M_down[l], d_h)``. Downward edge representations from
+            ``(B, num_edges[l], hidden_dim)``. Downward edge representations from
             level ``l+1`` to ``l``.
         down_gnns : nn.ModuleList
             GNNs for downward edges, one per inter-level gap.
@@ -203,10 +203,10 @@ def mesh_down_step(
             # Extract representations
             send_node_rep = mesh_rep_levels[
                 level_l + 1
-            ]  # (B, num_mesh_nodes[l+1], d_h)
+            ]  # (B, num_mesh_nodes[l+1], hidden_dim)
             rec_node_rep = mesh_rep_levels[
                 level_l
-            ]  # (B, num_mesh_nodes[l], d_h)
+            ]  # (B, num_mesh_nodes[l], hidden_dim)
             down_edge_rep = mesh_down_rep[level_l]
             same_edge_rep = mesh_same_rep[level_l]
 
@@ -219,7 +219,7 @@ def mesh_down_step(
             mesh_rep_levels[level_l], mesh_same_rep[level_l] = same_gnn(
                 new_node_rep, new_node_rep, same_edge_rep
             )
-            # (B, num_mesh_nodes[l], d_h) and (B, M_same[l], d_h)
+            # (B, num_mesh_nodes[l], hidden_dim) and (B, num_edges[l], hidden_dim)
 
         return mesh_rep_levels, mesh_same_rep, mesh_down_rep
 
@@ -233,16 +233,16 @@ def mesh_up_step(
         Parameters
         ----------
         mesh_rep_levels : list of torch.Tensor
-            One tensor per level, each of shape ``(B, num_mesh_nodes[l], d_h)``.
+            One tensor per level, each of shape ``(B, num_mesh_nodes[l], hidden_dim)``.
             Node representations at each hierarchy level. Dims: ``B`` is
             batch size, ``num_mesh_nodes[l]`` is the node count at level ``l``,
-            and ``d_h`` is the hidden dimension.
+            and ``hidden_dim`` is the hidden dimension.
         mesh_same_rep : list of torch.Tensor
-            One tensor per level, each of shape ``(B, M_same[l], d_h)``.
+            One tensor per level, each of shape ``(B, num_edges[l], hidden_dim)``.
             Same-level edge representations.
         mesh_up_rep : list of torch.Tensor
             One tensor per inter-level gap, each of shape
-            ``(B, M_up[l], d_h)``. Upward edge representations from
+            ``(B, num_edges[l], hidden_dim)``. Upward edge representations from
             level ``l`` to ``l+1``.
         up_gnns : nn.ModuleList
             GNNs for upward edges, one per inter-level gap.
@@ -267,10 +267,10 @@ def mesh_up_step(
             # Extract representations
             send_node_rep = mesh_rep_levels[
                 level_l - 1
-            ]  # (B, num_mesh_nodes[l-1], d_h)
+            ]  # (B, num_mesh_nodes[l-1], hidden_dim)
             rec_node_rep = mesh_rep_levels[
                 level_l
-            ]  # (B, num_mesh_nodes[l], d_h)
+            ]  # (B, num_mesh_nodes[l], hidden_dim)
             up_edge_rep = mesh_up_rep[level_l - 1]
             same_edge_rep = mesh_same_rep[level_l]
 
@@ -278,13 +278,13 @@ def mesh_up_step(
             new_node_rep, mesh_up_rep[level_l - 1] = up_gnn(
                 send_node_rep, rec_node_rep, up_edge_rep
             )
-            # (B, num_mesh_nodes[l], d_h) and (B, M_up[l-1], d_h)
+            # (B, num_mesh_nodes[l], hidden_dim) and (B, num_edges[l-1], hidden_dim)
 
             # Run same level processing on level l
             mesh_rep_levels[level_l], mesh_same_rep[level_l] = same_gnn(
                 new_node_rep, new_node_rep, same_edge_rep
             )
-            # (B, num_mesh_nodes[l], d_h) and (B, M_same[l], d_h)
+            # (B, num_mesh_nodes[l], hidden_dim) and (B, num_edges[l], hidden_dim)
 
         return mesh_rep_levels, mesh_same_rep, mesh_up_rep
 
@@ -298,19 +298,19 @@ def hi_processor_step(
         Parameters
         ----------
         mesh_rep_levels : list of torch.Tensor
-            One tensor per level, each of shape ``(B, num_mesh_nodes[l], d_h)``.
+            One tensor per level, each of shape ``(B, num_mesh_nodes[l], hidden_dim)``.
             Node representations at each hierarchy level. Dims: ``B`` is
             batch size, ``num_mesh_nodes[l]`` is the node count at level ``l``,
-            and ``d_h`` is the hidden dimension.
+            and ``hidden_dim`` is the hidden dimension.
         mesh_same_rep : list of torch.Tensor
-            One tensor per level, each of shape ``(B, M_same[l], d_h)``.
+            One tensor per level, each of shape ``(B, num_edges[l], hidden_dim)``.
             Same-level edge representations.
         mesh_up_rep : list of torch.Tensor
             One tensor per inter-level gap, each of shape
-            ``(B, M_up[l], d_h)``. Upward edge representations.
+            ``(B, num_edges[l], hidden_dim)``. Upward edge representations.
         mesh_down_rep : list of torch.Tensor
             One tensor per inter-level gap, each of shape
-            ``(B, M_down[l], d_h)``. Downward edge representations.
+            ``(B, num_edges[l], hidden_dim)``. Downward edge representations.
 
         Returns
         -------
diff --git a/neural_lam/models/step_predictors/graph/hi_lam_parallel.py b/neural_lam/models/step_predictors/graph/hi_lam_parallel.py
index 340367ef8..b0ab43ad4 100644
--- a/neural_lam/models/step_predictors/graph/hi_lam_parallel.py
+++ b/neural_lam/models/step_predictors/graph/hi_lam_parallel.py
@@ -121,19 +121,19 @@ def hi_processor_step(
         Parameters
         ----------
         mesh_rep_levels : list of torch.Tensor
-            One tensor per level, each of shape ``(B, num_mesh_nodes[l], d_h)``.
+            One tensor per level, each of shape ``(B, num_mesh_nodes[l], hidden_dim)``.
             Node representations at each hierarchy level. Dims: ``B`` is
             batch size, ``num_mesh_nodes[l]`` is the node count at level ``l``,
-            and ``d_h`` is the hidden dimension.
+            and ``hidden_dim`` is the hidden dimension.
         mesh_same_rep : list of torch.Tensor
-            One tensor per level, each of shape ``(B, M_same[l], d_h)``.
+            One tensor per level, each of shape ``(B, num_edges[l], hidden_dim)``.
             Same-level edge representations.
         mesh_up_rep : list of torch.Tensor
             One tensor per inter-level gap, each of shape
-            ``(B, M_up[l], d_h)``. Upward edge representations.
+            ``(B, num_edges[l], hidden_dim)``. Upward edge representations.
         mesh_down_rep : list of torch.Tensor
             One tensor per inter-level gap, each of shape
-            ``(B, M_down[l], d_h)``. Downward edge representations.
+            ``(B, num_edges[l], hidden_dim)``. Downward edge representations.
 
         Returns
         -------
@@ -143,10 +143,10 @@ def hi_processor_step(
         """
 
         # First join all node and edge representations to single tensors
-        mesh_rep = torch.cat(mesh_rep_levels, dim=1)  # (B, num_mesh_nodes, d_h)
+        mesh_rep = torch.cat(mesh_rep_levels, dim=1)  # (B, num_mesh_nodes, hidden_dim)
         mesh_edge_rep = torch.cat(
             mesh_same_rep + mesh_up_rep + mesh_down_rep, axis=1
-        )  # (B, M_mesh, d_h)
+        )  # (B, num_edges, hidden_dim)
 
         # Here, update mesh_*_rep and mesh_rep
         mesh_rep, mesh_edge_rep = self.processor(mesh_rep, mesh_edge_rep)
diff --git a/neural_lam/models/step_predictors/graph/hierarchical.py b/neural_lam/models/step_predictors/graph/hierarchical.py
index 8408e794f..266baa516 100644
--- a/neural_lam/models/step_predictors/graph/hierarchical.py
+++ b/neural_lam/models/step_predictors/graph/hierarchical.py
@@ -157,9 +157,9 @@ def embedd_mesh_nodes(self):
         Returns
         -------
         torch.Tensor
-            Shape ``(num_mesh_nodes[0], d_h)``. Embedded bottom-level
+            Shape ``(num_mesh_nodes[0], hidden_dim)``. Embedded bottom-level
             mesh node representations. Dims: ``num_mesh_nodes[0]`` is
-            the number of nodes at level 0 and ``d_h`` is the hidden
+            the number of nodes at level 0 and ``hidden_dim`` is the hidden
             dimension.
         """
         return self.mesh_embedders[0](self.mesh_static_features[0])
@@ -172,15 +172,15 @@ def process_step(self, mesh_rep):
         Parameters
         ----------
         mesh_rep : torch.Tensor
-            Shape ``(B, num_mesh_nodes[0], d_h)``. Bottom-level mesh
+            Shape ``(B, num_mesh_nodes[0], hidden_dim)``. Bottom-level mesh
             node representations from the encoder. Dims: ``B`` is batch
             size, ``num_mesh_nodes[0]`` is the number of nodes at
-            level 0, and ``d_h`` is the hidden dimension.
+            level 0, and ``hidden_dim`` is the hidden dimension.
 
         Returns
         -------
         torch.Tensor
-            Shape ``(B, num_mesh_nodes[0], d_h)``. Updated bottom-level
+            Shape ``(B, num_mesh_nodes[0], hidden_dim)``. Updated bottom-level
             mesh node representations. Dims: same as ``mesh_rep``.
         """
         batch_size = mesh_rep.shape[0]
@@ -239,7 +239,7 @@ def process_step(self, mesh_rep):
                 new_node_rep  # (B, num_mesh_nodes[l], hidden_dim)
             )
             mesh_up_rep[level_l - 1] = (
-                new_edge_rep  # (B, M_up[l-1], hidden_dim)
+                new_edge_rep  # (B, num_edges[l-1], hidden_dim)
             )
 
         # - PROCESSOR -
@@ -282,21 +282,21 @@ def hi_processor_step(
         ----------
         mesh_rep_levels : list of torch.Tensor
             One tensor per level, each of shape
-            ``(B, num_mesh_nodes[l], d_h)``. Node representations at
+            ``(B, num_mesh_nodes[l], hidden_dim)``. Node representations at
             each hierarchy level. Dims: ``B`` is batch size,
             ``num_mesh_nodes[l]`` is the node count at level ``l``, and
-            ``d_h`` is the hidden dimension.
+            ``hidden_dim`` is the hidden dimension.
         mesh_same_rep : list of torch.Tensor
-            One tensor per level, each of shape ``(B, M_same[l], d_h)``.
-            Same-level edge representations. ``M_same[l]`` is the edge
+            One tensor per level, each of shape ``(B, num_edges[l], hidden_dim)``.
+            Same-level edge representations. ``num_edges[l]`` is the edge
             count at level ``l``.
         mesh_up_rep : list of torch.Tensor
             One tensor per inter-level gap, each of shape
-            ``(B, M_up[l], d_h)``. Upward edge representations from
+            ``(B, num_edges[l], hidden_dim)``. Upward edge representations from
             level ``l`` to ``l+1``.
         mesh_down_rep : list of torch.Tensor
             One tensor per inter-level gap, each of shape
-            ``(B, M_down[l], d_h)``. Downward edge representations from
+            ``(B, num_edges[l], hidden_dim)``. Downward edge representations from
             level ``l+1`` to ``l``.
 
         Returns
diff --git a/neural_lam/plot_graph.py b/neural_lam/plot_graph.py
index c8ae0c9a8..35788343c 100644
--- a/neural_lam/plot_graph.py
+++ b/neural_lam/plot_graph.py
@@ -30,7 +30,7 @@ def plot_graph(
     Parameters
     ----------
     grid_pos : np.ndarray
-        Grid node positions, shape (N_grid, 2).
+        Grid node positions, shape (num_grid_nodes, 2).
     hierarchical : bool
         Whether the loaded graph is hierarchical.
     graph_ldict : dict
@@ -265,7 +265,7 @@ def main():
         config_path=args.datastore_config_path
     )
 
-    xy = datastore.get_xy("state", stacked=True)  # (N_grid, 2)
+    xy = datastore.get_xy("state", stacked=True)  # (num_grid_nodes, 2)
     pos_max = np.max(np.abs(xy))
     grid_pos = xy / pos_max  # Divide by maximum coordinate
 
diff --git a/neural_lam/utils.py b/neural_lam/utils.py
index 2f6933431..4708b13d6 100644
--- a/neural_lam/utils.py
+++ b/neural_lam/utils.py
@@ -106,7 +106,7 @@ def zero_index_edge_index(edge_index: torch.Tensor) -> torch.Tensor:
     Parameters
     ----------
     edge_index : torch.Tensor
-        Edge index tensor of shape (2, N_edges).
+        Edge index tensor of shape (2, num_edges).
 
     Returns
     -------
@@ -130,7 +130,7 @@ def zero_index_m2g(
     Parameters
     ----------
     m2g_edge_index : torch.Tensor
-        Edge index tensor of shape (2, N_edges).
+        Edge index tensor of shape (2, num_edges).
     mesh_static_features : list of torch.Tensor
         Mesh node feature tensors.
     mesh_first : bool
@@ -182,7 +182,7 @@ def zero_index_g2m(
     Parameters
     ----------
     g2m_edge_index : torch.Tensor
-        Edge index tensor of shape (2, N_edges).
+        Edge index tensor of shape (2, num_edges).
     mesh_static_features : list of torch.Tensor
         Mesh node feature tensors.
     mesh_first : bool
@@ -299,8 +299,8 @@ def loads_file(fn: str) -> Any:
         [zero_index_edge_index(ei) for ei in loads_file("m2m_edge_index.pt")],
         persistent=False,
     )  # List of (2, M_m2m[l])
-    g2m_edge_index = loads_file("g2m_edge_index.pt")  # (2, M_g2m)
-    m2g_edge_index = loads_file("m2g_edge_index.pt")  # (2, M_m2g)
+    g2m_edge_index = loads_file("g2m_edge_index.pt")  # (2, num_edges)
+    m2g_edge_index = loads_file("m2g_edge_index.pt")  # (2, num_edges)
 
     # Change first indices to 0
     # m2g and g2m has to be handled specially as not all mesh nodes
@@ -321,10 +321,10 @@ def loads_file(fn: str) -> Any:
     hierarchical = n_levels > 1  # Not just single level mesh graph
 
     # Load static edge features
-    # List of (M_m2m[l], d_edge_f)
+    # List of (M_m2m[l], input_dim)
     m2m_features = loads_file("m2m_features.pt")
-    g2m_features = loads_file("g2m_features.pt")  # (M_g2m, d_edge_f)
-    m2g_features = loads_file("m2g_features.pt")  # (M_m2g, d_edge_f)
+    g2m_features = loads_file("g2m_features.pt")  # (num_edges, input_dim)
+    m2g_features = loads_file("m2g_features.pt")  # (num_edges, input_dim)
 
     # Normalize by dividing with longest edge (found in m2m)
     longest_edge = max(
@@ -352,21 +352,21 @@ def loads_file(fn: str) -> Any:
                 for ei in loads_file("mesh_up_edge_index.pt")
             ],
             persistent=False,
-        )  # List of (2, M_up[l])
+        )  # List of (2, num_edges[l])
         mesh_down_edge_index = BufferList(
             [
                 zero_index_edge_index(ei)
                 for ei in loads_file("mesh_down_edge_index.pt")
             ],
             persistent=False,
-        )  # List of (2, M_down[l])
+        )  # List of (2, num_edges[l])
 
         mesh_up_features = loads_file(
             "mesh_up_features.pt"
-        )  # List of (M_up[l], d_edge_f)
+        )  # List of (num_edges[l], input_dim)
         mesh_down_features = loads_file(
             "mesh_down_features.pt"
-        )  # List of (M_down[l], d_edge_f)
+        )  # List of (num_edges[l], input_dim)
 
         # Rescale
         mesh_up_features = BufferList(mesh_up_features, persistent=False)
diff --git a/neural_lam/vis.py b/neural_lam/vis.py
index 919ec6a88..2d6088e63 100644
--- a/neural_lam/vis.py
+++ b/neural_lam/vis.py
@@ -2,12 +2,16 @@
 
 # Standard library
 import warnings
+from typing import Optional, Union
 
 # Third-party
 import cartopy.crs as ccrs
 import cartopy.feature as cfeature
 import matplotlib
+import matplotlib.axes
+import matplotlib.collections
 import matplotlib.colors
+import matplotlib.figure
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
@@ -120,21 +124,24 @@ def _get_heatmap_var_labels(datastore: BaseRegularGridDatastore) -> list[str]:
 
 def _to_heatmap_matrix(values) -> np.ndarray:
     """
-    Convert heatmap inputs to a `(d_f, pred_steps)` matrix.
+    Convert heatmap inputs to a ``(num_state_vars, pred_steps)`` matrix.
 
-    A single-step tensor may arrive as one-dimensional `(d_f,)`, especially in
-    single-GPU or focused metric logging paths. In that case we first treat it
-    as one row of `(pred_steps=1, d_f)` before transposing.
+    A single-step tensor may arrive as one-dimensional ``(num_state_vars,)``,
+    especially in single-GPU or focused metric logging paths. In that case we
+    first treat it as one row of ``(pred_steps=1, num_state_vars)`` before
+    transposing.
 
     Parameters
     ----------
     values : array-like
         The input values to convert.
+        Shape ``(num_state_vars,)`` or ``(pred_steps, num_state_vars)``.
 
     Returns
     -------
     np.ndarray
-        The converted heatmap matrix with shape `(d_f, pred_steps)`.
+        The converted heatmap matrix with shape
+        ``(num_state_vars, pred_steps)``.
     """
     if hasattr(values, "detach"):
         values = values.detach().cpu().numpy()
@@ -214,6 +221,11 @@ def _get_heatmap_color_values(
         - color_values: The normalized values for the colormap.
         - colorbar_label: The label for the colorbar.
         - cmap: The colormap to use.
+
+    Raises
+    ------
+    ValueError
+        If ``normalization`` is not one of ``'state_std'`` or ``'diff_std'``.
     """
 
     def _per_var_fallback():
@@ -328,24 +340,25 @@ def _get_annotation_text_color(
 
 
 def plot_on_axis(
-    ax,
-    da,
-    datastore,
-    vmin=None,
-    vmax=None,
-    ax_title=None,
-    cmap="plasma",
-    boundary_alpha=None,
-    crop_to_interior=False,
-):
-    """Plot weather state on a projection-aware axis using datastore metadata.
+    ax: matplotlib.axes.Axes,
+    da: xr.DataArray,
+    datastore: BaseRegularGridDatastore,
+    vmin: Optional[float] = None,
+    vmax: Optional[float] = None,
+    ax_title: Optional[str] = None,
+    cmap: Union[str, matplotlib.colors.Colormap] = "plasma",
+    boundary_alpha: Optional[float] = None,
+    crop_to_interior: bool = False,
+) -> matplotlib.collections.QuadMesh:
+    """
+    Plot weather state on a projection-aware axis using datastore metadata.
 
     Parameters
     ----------
     ax : matplotlib.axes.Axes
         The axis to plot on. Should have a cartopy projection.
     da : xarray.DataArray
-        The data to plot. Should have shape (N_grid,).
+        The data to plot. Should have shape (num_grid_nodes,).
     datastore : BaseRegularGridDatastore
         The datastore containing metadata about the grid.
     vmin : float, optional
@@ -465,7 +478,7 @@ def plot_error_heatmap(
     Parameters
     ----------
     errors : torch.Tensor
-        Shape ``(pred_steps, d_f)``. Per-step, per-variable errors. These
+        Shape ``(pred_steps, num_state_vars)``. Per-step, per-variable errors. These
         values are used for the numeric annotations in each cell.
     datastore : BaseRegularGridDatastore
         Datastore providing variable names, units, and step length.
@@ -614,9 +627,9 @@ def plot_prediction(
     datastore : BaseRegularGridDatastore
        Datastore providing grid metadata and projection.
     da_prediction : xarray.DataArray
-        Shape ``(N_grid,)``. Predicted field values.
+        Shape ``(num_grid_nodes,)``. Predicted field values.
     da_target : xarray.DataArray
-        Shape ``(N_grid,)``. Ground-truth field values.
+        Shape ``(num_grid_nodes,)``. Ground-truth field values.
     title : str, optional
         Overall figure title.
     vrange : tuple of (float, float), optional
@@ -697,7 +710,7 @@ def plot_spatial_error(
     ----------
     error : torch.Tensor
         Error magnitudes on the flattened grid.
-        * **Shape**: ``(N_grid,)``
+        * **Shape**: ``(num_grid_nodes,)``
     datastore : BaseRegularGridDatastore
         Datastore providing coordinate metadata and boundary masks.
     title : str or None, optional
diff --git a/neural_lam/weather_dataset.py b/neural_lam/weather_dataset.py
index 3ad4bf03b..23c23d6b4 100644
--- a/neural_lam/weather_dataset.py
+++ b/neural_lam/weather_dataset.py
@@ -19,33 +19,8 @@
 class WeatherDataset(torch.utils.data.Dataset):
     """Dataset class for weather data.
 
-    This class loads and processes weather data from a given datastore.
-
-    Parameters
-    ----------
-    datastore : BaseDatastore
-        The datastore to load the data from (e.g. mdp).
-    split : str, optional
-        The data split to use ("train", "val" or "test"). Default is "train".
-    ar_steps : int, optional
-        The number of autoregressive steps. Default is 3.
-    num_past_forcing_steps: int, optional
-        Number of past time steps to include in forcing input. If set to i,
-        forcing from times t-i, t-i+1, ..., t-1, t (and potentially beyond,
-        given num_future_forcing_steps) are included as forcing inputs at time t
-        Default is 1.
-    num_future_forcing_steps: int, optional
-        Number of future time steps to include in forcing input. If set to j,
-        forcing from times t, t+1, ..., t+j-1, t+j (and potentially times before
-        t, given num_past_forcing_steps) are included as forcing inputs at time
-        t. Default is 1.
-    load_single_member : bool, optional
-        If `False` and the datastore returns an ensemble of state
-        realisations, treat each state ensemble member as an independent
-        sample. If `True`, only ensemble member 0 is used. Default is False,
-        so all members are used when available.
-    standardize : bool, optional
-        Whether to standardize the data. Default is True.
+    Loads and processes weather data from a given datastore. See
+    :meth:`__init__` for the full parameter list.
     """
 
     def __init__(
@@ -74,8 +49,22 @@ def __init__(
         num_future_forcing_steps : int, optional
             Future forcing window length ``j`` so that ``[t, ..., t+j]``
             forcings are available. Default ``1``.
+        load_single_member : bool, optional
+            If ``False`` and the datastore returns an ensemble of state
+            realisations, treat each state ensemble member as an independent
+            sample. If ``True``, only ensemble member 0 is used. Default
+            ``False``.
         standardize : bool, optional
             If ``True``, normalize state/forcing arrays via datastore stats.
+            Default ``True``.
+
+        Raises
+        ------
+        ValueError
+            If the datastore does not provide state data, if the configured
+            ``ar_steps`` and forcing windows leave zero samples in ``split``,
+            or if the state/forcing dimension order does not match the
+            datastore's expected dimension order.
         """
 
         super().__init__()
@@ -568,9 +557,9 @@ def __getitem__(self, idx):
 
         forcing = torch.tensor(da_forcing_windowed.values, dtype=tensor_dtype)
 
-        # init_states: (2, N_grid, d_features)
-        # target_states: (ar_steps, N_grid, d_features)
-        # forcing: (ar_steps, N_grid, d_windowed_forcing)
+        # init_states: (2, num_grid_nodes, d_features)
+        # target_states: (ar_steps, num_grid_nodes, d_features)
+        # forcing: (ar_steps, num_grid_nodes, d_windowed_forcing)
         # target_times: (ar_steps,)
 
         return init_states, target_states, forcing, target_times
diff --git a/pyproject.toml b/pyproject.toml
index bccec0828..2c1be2c79 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -108,6 +108,18 @@ allow-any-import-level = "neural_lam"
 [tool.pylint.SIMILARITIES]
 min-similarity-lines = 10
 
+[tool.interrogate]
+fail-under = 100
+verbose = 1
+quiet = false
+ignore-init-method = false
+ignore-init-module = false
+ignore-magic = false
+ignore-private = false
+ignore-nested-functions = false
+ignore-property-decorators = false
+exclude = ["tests", "docs", "build"]
+
 [build-system]
 requires = ["hatchling>=1.27.0", "hatch-vcs"]
 build-backend = "hatchling.build"

From 47632c2cf3c9ba0c9ec46c6bb04f113ca1f42867 Mon Sep 17 00:00:00 2001
From: sadamov <kode@mailbox.org>
Date: Tue, 9 Jun 2026 12:27:36 +0200
Subject: [PATCH 12/16] precommits

---
 neural_lam/datastore/base.py                  |  8 +-
 .../models/forecasters/autoregressive.py      | 20 +++--
 neural_lam/models/forecasters/base.py         | 39 +++++----
 neural_lam/models/module.py                   | 15 ++--
 neural_lam/models/step_predictors/base.py     | 11 +--
 .../models/step_predictors/graph/base.py      | 26 +++---
 .../models/step_predictors/graph/graph_lam.py | 15 ++--
 .../models/step_predictors/graph/hi_lam.py    | 84 ++++++++++---------
 .../step_predictors/graph/hi_lam_parallel.py  | 18 ++--
 .../step_predictors/graph/hierarchical.py     | 11 +--
 neural_lam/vis.py                             |  5 +-
 11 files changed, 140 insertions(+), 112 deletions(-)

diff --git a/neural_lam/datastore/base.py b/neural_lam/datastore/base.py
index 8e353d099..69799fd36 100644
--- a/neural_lam/datastore/base.py
+++ b/neural_lam/datastore/base.py
@@ -32,10 +32,10 @@ class BaseDatastore(abc.ABC):
 
     Forecast vs analysis data
     -------------------------
-    If the datastore is used to represent forecast rather than analysis data,
-    then the ``is_forecast`` attribute should be set to True, and returned data
-    from ``get_dataarray`` is assumed to have `analysis_time` and `forecast_time`
-    dimensions (rather than just `time`).
+    If the datastore is used to represent forecast rather than analysis
+    data, then the ``is_forecast`` attribute should be set to True, and
+    returned data from ``get_dataarray`` is assumed to have `analysis_time`
+    and `forecast_time` dimensions (rather than just `time`).
 
     Ensemble vs deterministic data
     ------------------------------
diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py
index 502fb0ae6..7ac0902d4 100644
--- a/neural_lam/models/forecasters/autoregressive.py
+++ b/neural_lam/models/forecasters/autoregressive.py
@@ -51,7 +51,8 @@ def predicts_std(self) -> bool:
         Returns
         -------
         bool
-            ``True`` if the forecaster predicts standard deviation, ``False`` otherwise.
+            ``True`` if the forecaster predicts standard deviation,
+            ``False`` otherwise.
         """
         return self.predictor.predicts_std
 
@@ -63,8 +64,8 @@ def forward(
     ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
         """
         Unroll the autoregressive model: at each step ``i`` call
-        ``self.predictor`` to produce the next state, then overwrite
-        boundary nodes with the true value from ``boundary_states[:, i]``.
+        ``self.predictor`` to produce the next state, then overwrite boundary
+        nodes with the true value from ``boundary_states[:, i]``.
 
         Parameters
         ----------
@@ -83,12 +84,13 @@ def forward(
             of forcing variables (already concatenated past/current/future
             windows).
         boundary_states : torch.Tensor
-            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True state
-            values used ONLY to overwrite boundary nodes at each AR step.
-            The interior prediction at step ``i`` must not depend on
-            ``boundary_states[:, i]`` in any other way. Dims: ``B`` is
-            batch size, ``pred_steps`` is the number of predicted steps,
-            ``num_grid_nodes`` is the number of spatial nodes, and
+            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``.
+            True state values used ONLY to overwrite boundary nodes at
+            each AR step. The interior prediction at step ``i`` must not
+            depend on ``boundary_states[:, i]`` in any other way. Dims:
+            ``B`` is batch size, ``pred_steps`` is the number of
+            predicted steps, ``num_grid_nodes`` is the number of spatial
+            nodes, and
             ``num_state_vars`` is the state feature dimension.
 
         Returns
diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py
index 8fdd80e80..e314424f5 100644
--- a/neural_lam/models/forecasters/base.py
+++ b/neural_lam/models/forecasters/base.py
@@ -25,7 +25,8 @@ def predicts_std(self) -> bool:
         Returns
         -------
         bool
-            ``True`` if the forecaster predicts standard deviation, ``False`` otherwise.
+            ``True`` if the forecaster predicts standard deviation,
+            ``False`` otherwise.
         """
 
     @abstractmethod
@@ -42,33 +43,35 @@ def forward(
         Parameters
         ----------
         init_states : torch.Tensor
-            Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial states
-            ``[X_{t-1}, X_t]`` used to start the forecast from. Dims: ``B`` is
-            batch size, ``2`` is the time index (``[X_{t-1}, X_t]``),
-            ``num_grid_nodes`` is the number of spatial nodes, and ``num_state_vars``
-            is the state feature dimension.
-        forcing_features : torch.Tensor
-            Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``. External
-            forcings provided at each predicted step. Dims: ``B`` is batch
-            size, ``pred_steps`` is the autoregressive rollout length,
+            Shape ``(B, 2, num_grid_nodes, num_state_vars)``. The two initial
+            states ``[X_{t-1}, X_t]`` used to start the forecast from. Dims:
+            ``B`` is batch size, ``2`` is the time index (``[X_{t-1}, X_t]``),
             ``num_grid_nodes`` is the number of spatial nodes, and
+            ``num_state_vars`` is the state feature dimension.
+        forcing_features : torch.Tensor
+            Shape ``(B, pred_steps, num_grid_nodes, num_forcing_vars)``.
+            External forcings provided at each predicted step. Dims: ``B``
+            is batch size, ``pred_steps`` is the autoregressive rollout
+            length, ``num_grid_nodes`` is the number of spatial nodes, and
             ``num_forcing_vars`` is the forcing feature dimension (already
             concatenated past/current/future windows).
         boundary_states : torch.Tensor
-            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True state values used ONLY to overwrite boundary nodes at each AR step; interior predictions must not depend on ``boundary_states`` in any other way. Dims: ``B`` is batch size, ``pred_steps`` is the
-            rollout length, ``num_grid_nodes`` is the number of
-            spatial nodes, and ``num_state_vars`` is the state feature dimension.
+            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. True
+            state values used ONLY to overwrite boundary nodes at each AR
+            step; interior predictions must not depend on ``boundary_states``
+            in any other way. Dims: ``B`` is batch size, ``pred_steps`` is
+            the rollout length, ``num_grid_nodes`` is the number of spatial
+            nodes, and ``num_state_vars`` is the state feature dimension.
             This is a temporary mechanism that mirrors the pre-refactor
             ARModel behavior; it will be replaced by a dedicated
-            boundary-forcing input in #138 (training on interior +
-            boundary datastore), at which point this parameter will be
-            removed.
+            boundary-forcing input in #138 (training on interior + boundary
+            datastore), at which point this parameter will be removed.
 
         Returns
         -------
         prediction : torch.Tensor
-            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``. Forecast of
-            state at each predicted step. Dims: same as
+            Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)``.
+            Forecast of state at each predicted step. Dims: same as
             ``boundary_states``.
         pred_std : torch.Tensor or None
             Shape ``(B, pred_steps, num_grid_nodes, num_state_vars)`` when
diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py
index 958169c49..5e5f784bb 100644
--- a/neural_lam/models/module.py
+++ b/neural_lam/models/module.py
@@ -116,11 +116,11 @@ def __init__(
         if var_leads_metrics_watch is None:
             var_leads_metrics_watch = {}
 
-        # datastore and forecaster are excluded from saved hparams and must be
-        # provided explicitly when calling load_from_checkpoint. Saving args
-        # makes the checkpoint self-describing: it carries model, graph_name,
-        # hidden_dim, etc. so the caller can reconstruct the exact forecaster
-        # architecture from the checkpoint alone.
+        # datastore and forecaster are excluded from saved hparams and must
+        # be provided explicitly when calling load_from_checkpoint. Saving
+        # args makes the checkpoint self-describing: it carries model,
+        # graph_name, hidden_dim, etc. so the caller can reconstruct the
+        # exact forecaster architecture from the checkpoint alone.
         self.save_hyperparameters(ignore=["datastore", "forecaster"])
         self.datastore = datastore
         self.forecaster = forecaster
@@ -867,8 +867,9 @@ def on_load_checkpoint(self, checkpoint):
         loaded_state_dict = checkpoint["state_dict"]
 
         # 1. Broad namespace remap: for pre-refactor checkpoints
-        # The old ``ARModel`` was a flat LightningModule. Everything that belonged
-        # to the predictor needs to be moved to 'forecaster.predictor.'
+        # The old ``ARModel`` was a flat LightningModule. Everything that
+        # belonged to the predictor needs to be moved to
+        # 'forecaster.predictor.'
         old_keys = list(loaded_state_dict.keys())
         for key in old_keys:
             if not key.startswith("forecaster.") and key not in (
diff --git a/neural_lam/models/step_predictors/base.py b/neural_lam/models/step_predictors/base.py
index 2d79ed72f..632b3d84d 100644
--- a/neural_lam/models/step_predictors/base.py
+++ b/neural_lam/models/step_predictors/base.py
@@ -102,7 +102,8 @@ def predicts_std(self) -> bool:
         Returns
         -------
         bool
-            ``True`` if the predictor predicts standard deviation, ``False`` otherwise.
+            ``True`` if the predictor predicts standard deviation,
+            ``False`` otherwise.
         """
         return self.output_std
 
@@ -113,8 +114,8 @@ def expand_to_batch(self, x: torch.Tensor, batch_size: int) -> torch.Tensor:
         Parameters
         ----------
         x : torch.Tensor
-            Shape ``(N, d)``. Tensor to expand. Dims: ``N`` is the number
-            of nodes and ``d`` is the feature dimension.
+            Shape ``(N, d)``. Tensor to expand. Dims: ``N`` is the number of
+            nodes and ``d`` is the feature dimension.
         batch_size : int
             Target batch size ``B``.
 
@@ -141,8 +142,8 @@ def forward(
         prev_state : torch.Tensor
             Shape ``(B, num_grid_nodes, num_state_vars)``. The current state
             ``X_t``. Dims: ``B`` is batch size, ``num_grid_nodes`` is the
-            number of spatial nodes, and ``num_state_vars`` is the number of state
-            variables.
+            number of spatial nodes, and ``num_state_vars`` is the number of
+            state variables.
         prev_prev_state : torch.Tensor
             Shape ``(B, num_grid_nodes, num_state_vars)``. The previous state
             ``X_{t-1}``, used as additional conditioning. Dims: same as
diff --git a/neural_lam/models/step_predictors/graph/base.py b/neural_lam/models/step_predictors/graph/base.py
index d97500e0c..f541b6105 100644
--- a/neural_lam/models/step_predictors/graph/base.py
+++ b/neural_lam/models/step_predictors/graph/base.py
@@ -227,10 +227,10 @@ def forward(self, prev_state, prev_prev_state, forcing):
         Parameters
         ----------
         prev_state : torch.Tensor
-            Shape ``(B, num_grid_nodes, num_state_vars)``. The current state
-            ``X_t``. Dims: ``B`` is batch size, ``num_grid_nodes`` is the
-            number of spatial grid nodes, and ``num_state_vars`` is the number of
-            state variables.
+            Shape ``(B, num_grid_nodes, num_state_vars)``. The current
+            state ``X_t``. Dims: ``B`` is batch size,
+            ``num_grid_nodes`` is the number of spatial grid nodes, and
+            ``num_state_vars`` is the number of state variables.
         prev_prev_state : torch.Tensor
             Shape ``(B, num_grid_nodes, num_state_vars)``. The previous state
             ``X_{t-1}``, used as additional conditioning. Dims: same as
@@ -245,9 +245,9 @@ def forward(self, prev_state, prev_prev_state, forcing):
         Returns
         -------
         new_state : torch.Tensor
-            Shape ``(B, num_grid_nodes, num_state_vars)``. The predicted next state
-            ``X_{t+1}`` after delta-add and clamping. Dims: same as
-            ``prev_state``.
+            Shape ``(B, num_grid_nodes, num_state_vars)``. The predicted
+            next state ``X_{t+1}`` after delta-add and clamping. Dims:
+            same as ``prev_state``.
         pred_std : torch.Tensor or None
             Shape ``(B, num_grid_nodes, num_state_vars)`` when ``output_std`` is
             True, otherwise ``None``. Per-feature predicted standard
@@ -268,9 +268,15 @@ def forward(self, prev_state, prev_prev_state, forcing):
         )
 
         # Embed all features
-        grid_emb = self.grid_embedder(grid_features)  # (B, num_grid_nodes, hidden_dim)
-        g2m_emb = self.g2m_embedder(self.g2m_features)  # (num_edges, hidden_dim)
-        m2g_emb = self.m2g_embedder(self.m2g_features)  # (num_edges, hidden_dim)
+        grid_emb = self.grid_embedder(
+            grid_features
+        )  # (B, num_grid_nodes, hidden_dim)
+        g2m_emb = self.g2m_embedder(
+            self.g2m_features
+        )  # (num_edges, hidden_dim)
+        m2g_emb = self.m2g_embedder(
+            self.m2g_features
+        )  # (num_edges, hidden_dim)
         mesh_emb = self.embedd_mesh_nodes()
 
         # Map from grid to mesh
diff --git a/neural_lam/models/step_predictors/graph/graph_lam.py b/neural_lam/models/step_predictors/graph/graph_lam.py
index 0159e1fd7..8a9171c07 100644
--- a/neural_lam/models/step_predictors/graph/graph_lam.py
+++ b/neural_lam/models/step_predictors/graph/graph_lam.py
@@ -108,7 +108,10 @@ def __init__(
         self.processor = pyg.nn.Sequential(
             "mesh_rep, edge_rep",
             [
-                (net, "mesh_rep, mesh_rep, edge_rep -> mesh_rep, edge_rep")
+                (
+                    net,
+                    "mesh_rep, mesh_rep, edge_rep -> mesh_rep, edge_rep",
+                )
                 for net in processor_nets
             ],
         )
@@ -134,9 +137,9 @@ def embedd_mesh_nodes(self):
         Returns
         -------
         torch.Tensor
-            Shape ``(num_mesh_nodes, hidden_dim)``. Embedded mesh node representations.
-            Dims: ``num_mesh_nodes`` is the number of mesh nodes and ``hidden_dim`` is
-            the hidden dimension.
+            Shape ``(num_mesh_nodes, hidden_dim)``. Embedded mesh node
+            representations. Dims: ``num_mesh_nodes`` is the number of
+            mesh nodes and ``hidden_dim`` is the hidden dimension.
         """
         return self.mesh_embedder(
             self.mesh_static_features
@@ -163,7 +166,9 @@ def process_step(self, mesh_rep):
         """
         # Embed m2m here first
         batch_size = mesh_rep.shape[0]
-        m2m_emb = self.m2m_embedder(self.m2m_features)  # (num_edges, hidden_dim)
+        m2m_emb = self.m2m_embedder(
+            self.m2m_features
+        )  # (num_edges, hidden_dim)
         m2m_emb_expanded = self.expand_to_batch(
             m2m_emb, batch_size
         )  # (B, num_edges, hidden_dim)
diff --git a/neural_lam/models/step_predictors/graph/hi_lam.py b/neural_lam/models/step_predictors/graph/hi_lam.py
index 79ec95cd2..7976ffff1 100644
--- a/neural_lam/models/step_predictors/graph/hi_lam.py
+++ b/neural_lam/models/step_predictors/graph/hi_lam.py
@@ -168,17 +168,19 @@ def mesh_down_step(
         Parameters
         ----------
         mesh_rep_levels : list of torch.Tensor
-            One tensor per level, each of shape ``(B, num_mesh_nodes[l], hidden_dim)``.
-            Node representations at each hierarchy level. Dims: ``B`` is
-            batch size, ``num_mesh_nodes[l]`` is the node count at level ``l``,
-            and ``hidden_dim`` is the hidden dimension.
+            One tensor per level, each of shape
+            ``(B, num_mesh_nodes[l], hidden_dim)``. Node representations
+            at each hierarchy level. Dims: ``B`` is batch size,
+            ``num_mesh_nodes[l]`` is the node count at level ``l``, and
+            ``hidden_dim`` is the hidden dimension.
         mesh_same_rep : list of torch.Tensor
-            One tensor per level, each of shape ``(B, num_edges[l], hidden_dim)``.
-            Same-level edge representations.
+            One tensor per level, each of shape
+            ``(B, num_edges[l], hidden_dim)``. Same-level edge
+            representations.
         mesh_down_rep : list of torch.Tensor
             One tensor per inter-level gap, each of shape
-            ``(B, num_edges[l], hidden_dim)``. Downward edge representations from
-            level ``l+1`` to ``l``.
+            ``(B, num_edges[l], hidden_dim)``. Downward edge
+            representations from level ``l+1`` to ``l``.
         down_gnns : nn.ModuleList
             GNNs for downward edges, one per inter-level gap.
         same_gnns : nn.ModuleList
@@ -191,7 +193,9 @@ def mesh_down_step(
         """
         # Run same level processing on level L
         mesh_rep_levels[-1], mesh_same_rep[-1] = same_gnns[-1](
-            mesh_rep_levels[-1], mesh_rep_levels[-1], mesh_same_rep[-1]
+            mesh_rep_levels[-1],
+            mesh_rep_levels[-1],
+            mesh_same_rep[-1],
         )
 
         # Let level_l go from L-1 to 0
@@ -201,25 +205,22 @@ def mesh_down_step(
             reversed(same_gnns[:-1]),
         ):
             # Extract representations
-            send_node_rep = mesh_rep_levels[
-                level_l + 1
-            ]  # (B, num_mesh_nodes[l+1], hidden_dim)
-            rec_node_rep = mesh_rep_levels[
-                level_l
-            ]  # (B, num_mesh_nodes[l], hidden_dim)
+            send_node_rep = mesh_rep_levels[level_l + 1]
+            rec_node_rep = mesh_rep_levels[level_l]
             down_edge_rep = mesh_down_rep[level_l]
             same_edge_rep = mesh_same_rep[level_l]
 
             # Apply down GNN
             new_node_rep, mesh_down_rep[level_l] = down_gnn(
-                send_node_rep, rec_node_rep, down_edge_rep
+                send_node_rep,
+                rec_node_rep,
+                down_edge_rep,
             )
 
             # Run same level processing on level l
             mesh_rep_levels[level_l], mesh_same_rep[level_l] = same_gnn(
                 new_node_rep, new_node_rep, same_edge_rep
             )
-            # (B, num_mesh_nodes[l], hidden_dim) and (B, num_edges[l], hidden_dim)
 
         return mesh_rep_levels, mesh_same_rep, mesh_down_rep
 
@@ -233,13 +234,15 @@ def mesh_up_step(
         Parameters
         ----------
         mesh_rep_levels : list of torch.Tensor
-            One tensor per level, each of shape ``(B, num_mesh_nodes[l], hidden_dim)``.
-            Node representations at each hierarchy level. Dims: ``B`` is
-            batch size, ``num_mesh_nodes[l]`` is the node count at level ``l``,
-            and ``hidden_dim`` is the hidden dimension.
+            One tensor per level, each of shape
+            ``(B, num_mesh_nodes[l], hidden_dim)``. Node representations
+            at each hierarchy level. Dims: ``B`` is batch size,
+            ``num_mesh_nodes[l]`` is the node count at level ``l``, and
+            ``hidden_dim`` is the hidden dimension.
         mesh_same_rep : list of torch.Tensor
-            One tensor per level, each of shape ``(B, num_edges[l], hidden_dim)``.
-            Same-level edge representations.
+            One tensor per level, each of shape
+            ``(B, num_edges[l], hidden_dim)``. Same-level edge
+            representations.
         mesh_up_rep : list of torch.Tensor
             One tensor per inter-level gap, each of shape
             ``(B, num_edges[l], hidden_dim)``. Upward edge representations from
@@ -257,34 +260,33 @@ def mesh_up_step(
 
         # Run same level processing on level 0
         mesh_rep_levels[0], mesh_same_rep[0] = same_gnns[0](
-            mesh_rep_levels[0], mesh_rep_levels[0], mesh_same_rep[0]
+            mesh_rep_levels[0],
+            mesh_rep_levels[0],
+            mesh_same_rep[0],
         )
 
         # Let level_l go from 1 to L
         for level_l, (up_gnn, same_gnn) in enumerate(
-            zip(up_gnns, same_gnns[1:]), start=1
+            zip(up_gnns, same_gnns[1:]),
+            start=1,
         ):
             # Extract representations
-            send_node_rep = mesh_rep_levels[
-                level_l - 1
-            ]  # (B, num_mesh_nodes[l-1], hidden_dim)
-            rec_node_rep = mesh_rep_levels[
-                level_l
-            ]  # (B, num_mesh_nodes[l], hidden_dim)
+            send_node_rep = mesh_rep_levels[level_l - 1]
+            rec_node_rep = mesh_rep_levels[level_l]
             up_edge_rep = mesh_up_rep[level_l - 1]
             same_edge_rep = mesh_same_rep[level_l]
 
             # Apply up GNN
             new_node_rep, mesh_up_rep[level_l - 1] = up_gnn(
-                send_node_rep, rec_node_rep, up_edge_rep
+                send_node_rep,
+                rec_node_rep,
+                up_edge_rep,
             )
-            # (B, num_mesh_nodes[l], hidden_dim) and (B, num_edges[l-1], hidden_dim)
 
             # Run same level processing on level l
             mesh_rep_levels[level_l], mesh_same_rep[level_l] = same_gnn(
                 new_node_rep, new_node_rep, same_edge_rep
             )
-            # (B, num_mesh_nodes[l], hidden_dim) and (B, num_edges[l], hidden_dim)
 
         return mesh_rep_levels, mesh_same_rep, mesh_up_rep
 
@@ -298,13 +300,15 @@ def hi_processor_step(
         Parameters
         ----------
         mesh_rep_levels : list of torch.Tensor
-            One tensor per level, each of shape ``(B, num_mesh_nodes[l], hidden_dim)``.
-            Node representations at each hierarchy level. Dims: ``B`` is
-            batch size, ``num_mesh_nodes[l]`` is the node count at level ``l``,
-            and ``hidden_dim`` is the hidden dimension.
+            One tensor per level, each of shape
+            ``(B, num_mesh_nodes[l], hidden_dim)``. Node representations
+            at each hierarchy level. Dims: ``B`` is batch size,
+            ``num_mesh_nodes[l]`` is the node count at level ``l``, and
+            ``hidden_dim`` is the hidden dimension.
         mesh_same_rep : list of torch.Tensor
-            One tensor per level, each of shape ``(B, num_edges[l], hidden_dim)``.
-            Same-level edge representations.
+            One tensor per level, each of shape
+            ``(B, num_edges[l], hidden_dim)``. Same-level edge
+            representations.
         mesh_up_rep : list of torch.Tensor
             One tensor per inter-level gap, each of shape
             ``(B, num_edges[l], hidden_dim)``. Upward edge representations.
diff --git a/neural_lam/models/step_predictors/graph/hi_lam_parallel.py b/neural_lam/models/step_predictors/graph/hi_lam_parallel.py
index b0ab43ad4..c83dbb5bf 100644
--- a/neural_lam/models/step_predictors/graph/hi_lam_parallel.py
+++ b/neural_lam/models/step_predictors/graph/hi_lam_parallel.py
@@ -121,13 +121,15 @@ def hi_processor_step(
         Parameters
         ----------
         mesh_rep_levels : list of torch.Tensor
-            One tensor per level, each of shape ``(B, num_mesh_nodes[l], hidden_dim)``.
-            Node representations at each hierarchy level. Dims: ``B`` is
-            batch size, ``num_mesh_nodes[l]`` is the node count at level ``l``,
-            and ``hidden_dim`` is the hidden dimension.
+            One tensor per level, each of shape
+            ``(B, num_mesh_nodes[l], hidden_dim)``. Node representations at
+            each hierarchy level. Dims: ``B`` is batch size,
+            ``num_mesh_nodes[l]`` is the node count at level ``l``, and
+            ``hidden_dim`` is the hidden dimension.
         mesh_same_rep : list of torch.Tensor
-            One tensor per level, each of shape ``(B, num_edges[l], hidden_dim)``.
-            Same-level edge representations.
+            One tensor per level, each of shape
+            ``(B, num_edges[l], hidden_dim)``. Same-level edge
+            representations.
         mesh_up_rep : list of torch.Tensor
             One tensor per inter-level gap, each of shape
             ``(B, num_edges[l], hidden_dim)``. Upward edge representations.
@@ -143,7 +145,9 @@ def hi_processor_step(
         """
 
         # First join all node and edge representations to single tensors
-        mesh_rep = torch.cat(mesh_rep_levels, dim=1)  # (B, num_mesh_nodes, hidden_dim)
+        mesh_rep = torch.cat(
+            mesh_rep_levels, dim=1
+        )  # (B, num_mesh_nodes, hidden_dim)
         mesh_edge_rep = torch.cat(
             mesh_same_rep + mesh_up_rep + mesh_down_rep, axis=1
         )  # (B, num_edges, hidden_dim)
diff --git a/neural_lam/models/step_predictors/graph/hierarchical.py b/neural_lam/models/step_predictors/graph/hierarchical.py
index 266baa516..dba528447 100644
--- a/neural_lam/models/step_predictors/graph/hierarchical.py
+++ b/neural_lam/models/step_predictors/graph/hierarchical.py
@@ -287,17 +287,18 @@ def hi_processor_step(
             ``num_mesh_nodes[l]`` is the node count at level ``l``, and
             ``hidden_dim`` is the hidden dimension.
         mesh_same_rep : list of torch.Tensor
-            One tensor per level, each of shape ``(B, num_edges[l], hidden_dim)``.
-            Same-level edge representations. ``num_edges[l]`` is the edge
-            count at level ``l``.
+            One tensor per level, each of shape
+            ``(B, num_edges[l], hidden_dim)``. Same-level edge
+            representations. ``num_edges[l]`` is the edge count at
+            level ``l``.
         mesh_up_rep : list of torch.Tensor
             One tensor per inter-level gap, each of shape
             ``(B, num_edges[l], hidden_dim)``. Upward edge representations from
             level ``l`` to ``l+1``.
         mesh_down_rep : list of torch.Tensor
             One tensor per inter-level gap, each of shape
-            ``(B, num_edges[l], hidden_dim)``. Downward edge representations from
-            level ``l+1`` to ``l``.
+            ``(B, num_edges[l], hidden_dim)``. Downward edge
+            representations from level ``l+1`` to ``l``.
 
         Returns
         -------
diff --git a/neural_lam/vis.py b/neural_lam/vis.py
index 2d6088e63..9346000ce 100644
--- a/neural_lam/vis.py
+++ b/neural_lam/vis.py
@@ -478,8 +478,9 @@ def plot_error_heatmap(
     Parameters
     ----------
     errors : torch.Tensor
-        Shape ``(pred_steps, num_state_vars)``. Per-step, per-variable errors. These
-        values are used for the numeric annotations in each cell.
+        Shape ``(pred_steps, num_state_vars)``. Per-step, per-variable
+        errors. These values are used for the numeric annotations in each
+        cell.
     datastore : BaseRegularGridDatastore
         Datastore providing variable names, units, and step length.
     title : str, optional

From 7ffb5d816567996018180191699dcd456263bef5 Mon Sep 17 00:00:00 2001
From: sadamov <kode@mailbox.org>
Date: Tue, 9 Jun 2026 12:48:12 +0200
Subject: [PATCH 13/16] flake8

---
 neural_lam/datastore/base.py                  | 56 +++++--------
 .../models/forecasters/autoregressive.py      |  6 +-
 neural_lam/models/forecasters/base.py         |  4 +-
 neural_lam/models/module.py                   | 25 +++---
 neural_lam/models/step_predictors/base.py     | 23 +++--
 .../models/step_predictors/graph/base.py      |  5 +-
 .../models/step_predictors/graph/graph_lam.py | 11 ++-
 .../models/step_predictors/graph/hi_lam.py    | 83 +++++++++----------
 .../step_predictors/graph/hi_lam_parallel.py  |  5 +-
 .../step_predictors/graph/hierarchical.py     | 83 ++++++++-----------
 neural_lam/vis.py                             | 12 ++-
 11 files changed, 137 insertions(+), 176 deletions(-)

diff --git a/neural_lam/datastore/base.py b/neural_lam/datastore/base.py
index 69799fd36..dabd15d6e 100644
--- a/neural_lam/datastore/base.py
+++ b/neural_lam/datastore/base.py
@@ -8,7 +8,6 @@
 from datetime import timedelta
 from functools import cached_property
 from pathlib import Path
-from typing import List, Optional, Union
 
 # Third-party
 import cartopy.crs as ccrs
@@ -75,7 +74,6 @@ def root_path(self) -> Path:
             The root path to the datastore.
 
         """
-        pass
 
     @property
     @abc.abstractmethod
@@ -89,7 +87,6 @@ def config(self) -> collections.abc.Mapping:
             returned.
 
         """
-        pass
 
     @property
     @abc.abstractmethod
@@ -102,10 +99,9 @@ def step_length(self) -> timedelta:
             The step length of the dataset.
 
         """
-        pass
 
     @abc.abstractmethod
-    def get_vars_units(self, category: str) -> List[str]:
+    def get_vars_units(self, category: str) -> list[str]:
         """Get the units of the variables in the given category.
 
         Parameters
@@ -119,10 +115,9 @@ def get_vars_units(self, category: str) -> List[str]:
             The units of the variables.
 
         """
-        pass
 
     @abc.abstractmethod
-    def get_vars_names(self, category: str) -> List[str]:
+    def get_vars_names(self, category: str) -> list[str]:
         """Get the names of the variables in the given category.
 
         Parameters
@@ -136,10 +131,9 @@ def get_vars_names(self, category: str) -> List[str]:
             The names of the variables.
 
         """
-        pass
 
     @abc.abstractmethod
-    def get_vars_long_names(self, category: str) -> List[str]:
+    def get_vars_long_names(self, category: str) -> list[str]:
         """Get the long names of the variables in the given category.
 
         Parameters
@@ -153,7 +147,6 @@ def get_vars_long_names(self, category: str) -> List[str]:
             The long names of the variables.
 
         """
-        pass
 
     @abc.abstractmethod
     def get_num_data_vars(self, category: str) -> int:
@@ -170,7 +163,6 @@ def get_num_data_vars(self, category: str) -> int:
             The number of data variables.
 
         """
-        pass
 
     @abc.abstractmethod
     def get_standardization_dataarray(self, category: str) -> xr.Dataset:
@@ -199,7 +191,6 @@ def get_standardization_dataarray(self, category: str) -> xr.Dataset:
             differences for state variables).
 
         """
-        pass
 
     def _standardize_datarray(
         self, da: xr.DataArray, category: str
@@ -233,9 +224,9 @@ def _standardize_datarray(
     def get_dataarray(
         self,
         category: str,
-        split: Optional[str],
+        split: str | None,
         standardize: bool = False,
-    ) -> Union[xr.DataArray, None]:
+    ) -> xr.DataArray | None:
         """
         Return the processed data (as a single `xr.DataArray`) for the given
         category of data and test/train/val-split that covers all the data (in
@@ -277,7 +268,6 @@ def get_dataarray(
             The xarray DataArray object with processed dataset.
 
         """
-        pass
 
     @cached_property
     @abc.abstractmethod
@@ -294,7 +284,6 @@ def boundary_mask(self) -> xr.DataArray:
             `('grid_index',)`.
 
         """
-        pass
 
     @abc.abstractmethod
     def get_xy(self, category: str, stacked: bool) -> np.ndarray:
@@ -329,10 +318,9 @@ def coords_projection(self) -> ccrs.Projection:
             The projection object.
 
         """
-        pass
 
     @functools.lru_cache
-    def get_xy_extent(self, category: str) -> List[float]:
+    def get_xy_extent(self, category: str) -> list[float]:
         """
         Return the extent of the x, y coordinates for a given category of data.
         The extent should be returned as a list of 4 floats with `[xmin, xmax,
@@ -386,11 +374,10 @@ def num_grid_points(self) -> int:
             The number of grid points in the dataset.
 
         """
-        pass
 
     @cached_property
     @abc.abstractmethod
-    def state_feature_weights_values(self) -> List[float]:
+    def state_feature_weights_values(self) -> list[float]:
         """
         Return the weights for each state feature as a list of floats.
 
@@ -404,12 +391,11 @@ def state_feature_weights_values(self) -> List[float]:
         list of float
             The weight for each state feature.
         """
-        pass
 
     @functools.lru_cache
     def expected_dim_order(
         self,
-        category: Optional[str] = None,
+        category: str | None = None,
     ) -> tuple[str, ...]:
         """
         Return the expected dimension order for the dataarray or dataset
@@ -447,15 +433,19 @@ def expected_dim_order(
             if category != "static":
                 # static data does not vary in time
                 if self.is_forecast:
-                    dim_order.extend(
-                        ["analysis_time", "elapsed_forecast_duration"]
-                    )
+                    dim_order.extend([
+                        "analysis_time",
+                        "elapsed_forecast_duration",
+                    ])
                 elif not self.is_forecast:
                     dim_order.append("time")
 
-            if category == "state" and self.is_ensemble:
-                dim_order.append("ensemble_member")
-            elif category == "forcing" and self.has_ensemble_forcing:
+            if (
+                category == "state"
+                and self.is_ensemble
+                or category == "forcing"
+                and self.has_ensemble_forcing
+            ):
                 dim_order.append("ensemble_member")
 
         dim_order.append("grid_index")
@@ -518,7 +508,6 @@ def grid_shape_state(self) -> CartesianGridShape:
             `y` attributes.
 
         """
-        pass
 
     @abc.abstractmethod
     def get_xy(self, category: str, stacked: bool) -> np.ndarray:
@@ -540,11 +529,10 @@ def get_xy(self, category: str, stacked: bool) -> np.ndarray:
                                n_grid_points=N_x*N_y.
             - `stacked==False`: shape `(N_x, N_y, 2)`
         """
-        pass
 
     def unstack_grid_coords(
-        self, da_or_ds: Union[xr.DataArray, xr.Dataset]
-    ) -> Union[xr.DataArray, xr.Dataset]:
+        self, da_or_ds: xr.DataArray | xr.Dataset
+    ) -> xr.DataArray | xr.Dataset:
         """
         Unstack the spatial grid coordinates from `grid_index` into separate `x`
         and `y` dimensions to create a 2D grid (if the spatial coordinates have
@@ -598,8 +586,8 @@ def unstack_grid_coords(
         return da_or_ds_unstacked
 
     def stack_grid_coords(
-        self, da_or_ds: Union[xr.DataArray, xr.Dataset]
-    ) -> Union[xr.DataArray, xr.Dataset]:
+        self, da_or_ds: xr.DataArray | xr.Dataset
+    ) -> xr.DataArray | xr.Dataset:
         """
         Stack the spatial grid coordinates (x and y) into a single `grid_index`
         dimension. Only performs stacking if the data is currently unstacked
diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py
index 7ac0902d4..7b0dc9138 100644
--- a/neural_lam/models/forecasters/autoregressive.py
+++ b/neural_lam/models/forecasters/autoregressive.py
@@ -1,7 +1,6 @@
 """Forecaster that uses an auto-regressive strategy to unroll a forecast."""
 
 # Standard library
-from typing import Optional
 
 # Third-party
 import torch
@@ -34,7 +33,8 @@ def __init__(self, predictor: StepPredictor, datastore: BaseDatastore):
 
         # Register boundary/interior masks on the forecaster, not the predictor
         boundary_mask = (
-            torch.tensor(datastore.boundary_mask.values, dtype=torch.float32)
+            torch
+            .tensor(datastore.boundary_mask.values, dtype=torch.float32)
             .unsqueeze(0)
             .unsqueeze(-1)
         )
@@ -61,7 +61,7 @@ def forward(
         init_states: torch.Tensor,
         forcing_features: torch.Tensor,
         boundary_states: torch.Tensor,
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
         """
         Unroll the autoregressive model: at each step ``i`` call
         ``self.predictor`` to produce the next state, then overwrite boundary
diff --git a/neural_lam/models/forecasters/base.py b/neural_lam/models/forecasters/base.py
index e314424f5..4d9579168 100644
--- a/neural_lam/models/forecasters/base.py
+++ b/neural_lam/models/forecasters/base.py
@@ -2,7 +2,6 @@
 
 # Standard library
 from abc import ABC, abstractmethod
-from typing import Optional
 
 # Third-party
 import torch
@@ -35,7 +34,7 @@ def forward(
         init_states: torch.Tensor,
         forcing_features: torch.Tensor,
         boundary_states: torch.Tensor,
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
         """
         Produce a forecast of length ``pred_steps`` from two initial states,
         the per-step forcing features, and the per-step true boundary states.
@@ -80,4 +79,3 @@ def forward(
             per-variable std is substituted upstream by
             ``ForecasterModule``. Dims: same as ``prediction``.
         """
-        pass
diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py
index 5e5f784bb..7ea4c81c9 100644
--- a/neural_lam/models/module.py
+++ b/neural_lam/models/module.py
@@ -3,7 +3,7 @@
 # Standard library
 import os
 import warnings
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 # Third-party
 import matplotlib.pyplot as plt
@@ -43,9 +43,9 @@ def __init__(
         restore_opt: bool = False,
         n_example_pred: int = 1,
         create_gif: bool = False,
-        val_steps_to_log: Optional[List[int]] = None,
-        metrics_watch: Optional[List[str]] = None,
-        var_leads_metrics_watch: Optional[Dict[int, List[int]]] = None,
+        val_steps_to_log: list[int] | None = None,
+        metrics_watch: list[str] | None = None,
+        var_leads_metrics_watch: dict[int, list[int]] | None = None,
         args=None,
     ):
         """
@@ -128,7 +128,8 @@ def __init__(
 
         # Compute interior_mask_bool directly from datastore
         boundary_mask = (
-            torch.tensor(datastore.boundary_mask.values, dtype=torch.float32)
+            torch
+            .tensor(datastore.boundary_mask.values, dtype=torch.float32)
             .unsqueeze(0)
             .unsqueeze(-1)
         )  # (1, num_grid_nodes, 1)
@@ -165,10 +166,10 @@ def __init__(
         # Instantiate loss function
         self.loss = metrics.get_metric(loss)
 
-        self.val_metrics: Dict[str, List] = {
+        self.val_metrics: dict[str, list] = {
             "mse": [],
         }
-        self.test_metrics: Dict[str, List] = {
+        self.test_metrics: dict[str, list] = {
             "mse": [],
             "mae": [],
         }
@@ -184,7 +185,7 @@ def __init__(
         self.plotted_examples = 0
 
         # For storing spatial loss maps during evaluation
-        self.spatial_loss_maps: List[Any] = []
+        self.spatial_loss_maps: list[Any] = []
 
         # Warn once per phase if val_steps_to_log exceeds the actual rollout
         self._val_steps_warn_issued = False
@@ -557,7 +558,8 @@ def plot_examples(self, batch, n_examples, split, prediction):
             ).unstack("grid_index")
 
             var_vmin = (
-                torch.minimum(
+                torch
+                .minimum(
                     pred_slice.flatten(0, 1).min(dim=0)[0],
                     target_slice.flatten(0, 1).min(dim=0)[0],
                 )
@@ -565,7 +567,8 @@ def plot_examples(self, batch, n_examples, split, prediction):
                 .numpy()
             )
             var_vmax = (
-                torch.maximum(
+                torch
+                .maximum(
                     pred_slice.flatten(0, 1).max(dim=0)[0],
                     target_slice.flatten(0, 1).max(dim=0)[0],
                 )
@@ -582,7 +585,7 @@ def plot_examples(self, batch, n_examples, split, prediction):
                     f"example_plots_{example_i}",
                 )
                 os.makedirs(plot_dir_path, exist_ok=True)
-                png_frames: Dict[str, List[str]] = {
+                png_frames: dict[str, list[str]] = {
                     var_name: []
                     for var_name in self.datastore.get_vars_names("state")
                 }
diff --git a/neural_lam/models/step_predictors/base.py b/neural_lam/models/step_predictors/base.py
index 632b3d84d..4d9b070ea 100644
--- a/neural_lam/models/step_predictors/base.py
+++ b/neural_lam/models/step_predictors/base.py
@@ -2,7 +2,6 @@
 
 # Standard library
 from abc import ABC, abstractmethod
-from typing import Dict, Optional
 
 # Third-party
 import torch
@@ -23,8 +22,8 @@ def __init__(
         self,
         datastore: BaseDatastore,
         output_std: bool = False,
-        output_clamping_lower: Optional[Dict[str, float]] = None,
-        output_clamping_upper: Optional[Dict[str, float]] = None,
+        output_clamping_lower: dict[str, float] | None = None,
+        output_clamping_upper: dict[str, float] | None = None,
     ):
         """
         Initialize the StepPredictor.
@@ -41,10 +40,10 @@ def __init__(
             Upper clamping limits for state variables.
         """
         super().__init__()
-        self._output_clamping_lower: Dict[str, float] = (
+        self._output_clamping_lower: dict[str, float] = (
             dict(output_clamping_lower) if output_clamping_lower else {}
         )
-        self._output_clamping_upper: Dict[str, float] = (
+        self._output_clamping_upper: dict[str, float] = (
             dict(output_clamping_upper) if output_clamping_upper else {}
         )
 
@@ -132,7 +131,7 @@ def forward(
         prev_state: torch.Tensor,
         prev_prev_state: torch.Tensor,
         forcing: torch.Tensor,
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
         """
         Advance the state by one step:
         ``(X_{t-1}, X_t, forcing_t) -> X_{t+1}``.
@@ -165,7 +164,6 @@ def forward(
             is True, otherwise ``None``. Per-feature predicted standard
             deviation. Dims: same as ``prev_state``.
         """
-        pass
 
     def prepare_clamping_params(self, datastore: BaseDatastore):
         """
@@ -206,9 +204,8 @@ def prepare_clamping_params(self, datastore: BaseDatastore):
         sigmoid_center = 0
         softplus_center = 0
 
-        normalize_clamping_lim = (
-            lambda x, feature_idx: (x - self.state_mean[feature_idx])
-            / self.state_std[feature_idx]
+        normalize_clamping_lim = lambda x, feature_idx: (
+            (x - self.state_mean[feature_idx]) / self.state_std[feature_idx]
         )
 
         # Check which clamping functions to use for each feature
@@ -224,11 +221,11 @@ def prepare_clamping_params(self, datastore: BaseDatastore):
 
         for feature_idx, feature in enumerate(state_feature_names):
             if feature in lower_lims and feature in upper_lims:
-                assert (
-                    lower_lims[feature] < upper_lims[feature]
-                ), f'Invalid clamping limits for feature "{feature}",\
+                assert lower_lims[feature] < upper_lims[feature], (
+                    f'Invalid clamping limits for feature "{feature}",\
                      lower: {lower_lims[feature]}, larger than\
                      upper: {upper_lims[feature]}'
+                )
                 sigmoid_lower_upper_idx.append(feature_idx)
                 sigmoid_lower_lims.append(
                     normalize_clamping_lim(lower_lims[feature], feature_idx)
diff --git a/neural_lam/models/step_predictors/graph/base.py b/neural_lam/models/step_predictors/graph/base.py
index f541b6105..87679221c 100644
--- a/neural_lam/models/step_predictors/graph/base.py
+++ b/neural_lam/models/step_predictors/graph/base.py
@@ -1,7 +1,6 @@
 """Base class for graph-based step predictors."""
 
 # Standard library
-from typing import Dict, Optional
 
 # Third-party
 import torch
@@ -30,8 +29,8 @@ def __init__(
         num_past_forcing_steps: int = 1,
         num_future_forcing_steps: int = 1,
         output_std: bool = False,
-        output_clamping_lower: Optional[Dict[str, float]] = None,
-        output_clamping_upper: Optional[Dict[str, float]] = None,
+        output_clamping_lower: dict[str, float] | None = None,
+        output_clamping_upper: dict[str, float] | None = None,
     ):
         """
         Initialize the BaseGraphModel.
diff --git a/neural_lam/models/step_predictors/graph/graph_lam.py b/neural_lam/models/step_predictors/graph/graph_lam.py
index 8a9171c07..60ed434f2 100644
--- a/neural_lam/models/step_predictors/graph/graph_lam.py
+++ b/neural_lam/models/step_predictors/graph/graph_lam.py
@@ -1,7 +1,6 @@
 """Graph-based LAM model with a flat mesh."""
 
 # Standard library
-from typing import Dict, Optional
 
 # Third-party
 import torch_geometric as pyg
@@ -32,8 +31,8 @@ def __init__(
         num_past_forcing_steps: int = 1,
         num_future_forcing_steps: int = 1,
         output_std: bool = False,
-        output_clamping_lower: Optional[Dict[str, float]] = None,
-        output_clamping_upper: Optional[Dict[str, float]] = None,
+        output_clamping_lower: dict[str, float] | None = None,
+        output_clamping_upper: dict[str, float] | None = None,
     ):
         """
         Initialize the GraphLAM model.
@@ -77,9 +76,9 @@ def __init__(
             output_clamping_upper=output_clamping_upper,
         )
 
-        assert (
-            not self.hierarchical
-        ), "GraphLAM does not use a hierarchical mesh graph"
+        assert not self.hierarchical, (
+            "GraphLAM does not use a hierarchical mesh graph"
+        )
 
         # grid_dim from data + static + batch_static
         mesh_dim = self.mesh_static_features.shape[1]
diff --git a/neural_lam/models/step_predictors/graph/hi_lam.py b/neural_lam/models/step_predictors/graph/hi_lam.py
index 7976ffff1..6db85a706 100644
--- a/neural_lam/models/step_predictors/graph/hi_lam.py
+++ b/neural_lam/models/step_predictors/graph/hi_lam.py
@@ -3,7 +3,6 @@
 """
 
 # Standard library
-from typing import Dict, Optional
 
 # Third-party
 from torch import nn
@@ -32,8 +31,8 @@ def __init__(
         num_past_forcing_steps: int = 1,
         num_future_forcing_steps: int = 1,
         output_std: bool = False,
-        output_clamping_lower: Optional[Dict[str, float]] = None,
-        output_clamping_upper: Optional[Dict[str, float]] = None,
+        output_clamping_lower: dict[str, float] | None = None,
+        output_clamping_upper: dict[str, float] | None = None,
     ):
         """
         Initialize the HiLAM model.
@@ -78,20 +77,20 @@ def __init__(
         )
 
         # Make down GNNs, both for down edges and same level
-        self.mesh_down_gnns = nn.ModuleList(
-            [self.make_down_gnns() for _ in range(processor_layers)]
-        )  # Nested lists (proc_steps, num_levels-1)
-        self.mesh_down_same_gnns = nn.ModuleList(
-            [self.make_same_gnns() for _ in range(processor_layers)]
-        )  # Nested lists (proc_steps, num_levels)
+        self.mesh_down_gnns = nn.ModuleList([
+            self.make_down_gnns() for _ in range(processor_layers)
+        ])  # Nested lists (proc_steps, num_levels-1)
+        self.mesh_down_same_gnns = nn.ModuleList([
+            self.make_same_gnns() for _ in range(processor_layers)
+        ])  # Nested lists (proc_steps, num_levels)
 
         # Make up GNNs, both for up edges and same level
-        self.mesh_up_gnns = nn.ModuleList(
-            [self.make_up_gnns() for _ in range(processor_layers)]
-        )  # Nested lists (proc_steps, num_levels-1)
-        self.mesh_up_same_gnns = nn.ModuleList(
-            [self.make_same_gnns() for _ in range(processor_layers)]
-        )  # Nested lists (proc_steps, num_levels)
+        self.mesh_up_gnns = nn.ModuleList([
+            self.make_up_gnns() for _ in range(processor_layers)
+        ])  # Nested lists (proc_steps, num_levels-1)
+        self.mesh_up_same_gnns = nn.ModuleList([
+            self.make_same_gnns() for _ in range(processor_layers)
+        ])  # Nested lists (proc_steps, num_levels)
 
     def make_same_gnns(self):
         """
@@ -102,16 +101,14 @@ def make_same_gnns(self):
         nn.ModuleList
             List of GNNs for each level.
         """
-        return nn.ModuleList(
-            [
-                InteractionNet(
-                    edge_index,
-                    self.hidden_dim,
-                    hidden_layers=self.hidden_layers,
-                )
-                for edge_index in self.m2m_edge_index
-            ]
-        )
+        return nn.ModuleList([
+            InteractionNet(
+                edge_index,
+                self.hidden_dim,
+                hidden_layers=self.hidden_layers,
+            )
+            for edge_index in self.m2m_edge_index
+        ])
 
     def make_up_gnns(self):
         """
@@ -122,16 +119,14 @@ def make_up_gnns(self):
         nn.ModuleList
             List of GNNs for each inter-level gap (upwards).
         """
-        return nn.ModuleList(
-            [
-                InteractionNet(
-                    edge_index,
-                    self.hidden_dim,
-                    hidden_layers=self.hidden_layers,
-                )
-                for edge_index in self.mesh_up_edge_index
-            ]
-        )
+        return nn.ModuleList([
+            InteractionNet(
+                edge_index,
+                self.hidden_dim,
+                hidden_layers=self.hidden_layers,
+            )
+            for edge_index in self.mesh_up_edge_index
+        ])
 
     def make_down_gnns(self):
         """
@@ -142,16 +137,14 @@ def make_down_gnns(self):
         nn.ModuleList
             List of GNNs for each inter-level gap (downwards).
         """
-        return nn.ModuleList(
-            [
-                InteractionNet(
-                    edge_index,
-                    self.hidden_dim,
-                    hidden_layers=self.hidden_layers,
-                )
-                for edge_index in self.mesh_down_edge_index
-            ]
-        )
+        return nn.ModuleList([
+            InteractionNet(
+                edge_index,
+                self.hidden_dim,
+                hidden_layers=self.hidden_layers,
+            )
+            for edge_index in self.mesh_down_edge_index
+        ])
 
     def mesh_down_step(
         self,
diff --git a/neural_lam/models/step_predictors/graph/hi_lam_parallel.py b/neural_lam/models/step_predictors/graph/hi_lam_parallel.py
index c83dbb5bf..8f5609df8 100644
--- a/neural_lam/models/step_predictors/graph/hi_lam_parallel.py
+++ b/neural_lam/models/step_predictors/graph/hi_lam_parallel.py
@@ -3,7 +3,6 @@
 """
 
 # Standard library
-from typing import Dict, Optional
 
 # Third-party
 import torch
@@ -35,8 +34,8 @@ def __init__(
         num_past_forcing_steps: int = 1,
         num_future_forcing_steps: int = 1,
         output_std: bool = False,
-        output_clamping_lower: Optional[Dict[str, float]] = None,
-        output_clamping_upper: Optional[Dict[str, float]] = None,
+        output_clamping_lower: dict[str, float] | None = None,
+        output_clamping_upper: dict[str, float] | None = None,
     ):
         """
         Initialize the HiLAMParallel model.
diff --git a/neural_lam/models/step_predictors/graph/hierarchical.py b/neural_lam/models/step_predictors/graph/hierarchical.py
index dba528447..8ee30f217 100644
--- a/neural_lam/models/step_predictors/graph/hierarchical.py
+++ b/neural_lam/models/step_predictors/graph/hierarchical.py
@@ -1,7 +1,6 @@
 """Base implementations for hierarchical (multi-level) graph models."""
 
 # Standard library
-from typing import Dict, Optional
 
 # Third-party
 from torch import nn
@@ -29,8 +28,8 @@ def __init__(
         num_past_forcing_steps: int = 1,
         num_future_forcing_steps: int = 1,
         output_std: bool = False,
-        output_clamping_lower: Optional[Dict[str, float]] = None,
-        output_clamping_upper: Optional[Dict[str, float]] = None,
+        output_clamping_lower: dict[str, float] | None = None,
+        output_clamping_upper: dict[str, float] | None = None,
     ):
         """Extend :class:`BaseGraphModel` with hierarchical mesh structures."""
         super().__init__(
@@ -80,56 +79,44 @@ def __init__(
         mesh_down_dim = self.mesh_down_features[0].shape[1]
 
         # Separate mesh node embedders for each level
-        self.mesh_embedders = nn.ModuleList(
-            [
-                utils.make_mlp([mesh_dim] + self.mlp_blueprint_end)
-                for _ in range(self.num_levels)
-            ]
-        )
-        self.mesh_same_embedders = nn.ModuleList(
-            [
-                utils.make_mlp([mesh_same_dim] + self.mlp_blueprint_end)
-                for _ in range(self.num_levels)
-            ]
-        )
-        self.mesh_up_embedders = nn.ModuleList(
-            [
-                utils.make_mlp([mesh_up_dim] + self.mlp_blueprint_end)
-                for _ in range(self.num_levels - 1)
-            ]
-        )
-        self.mesh_down_embedders = nn.ModuleList(
-            [
-                utils.make_mlp([mesh_down_dim] + self.mlp_blueprint_end)
-                for _ in range(self.num_levels - 1)
-            ]
-        )
+        self.mesh_embedders = nn.ModuleList([
+            utils.make_mlp([mesh_dim] + self.mlp_blueprint_end)
+            for _ in range(self.num_levels)
+        ])
+        self.mesh_same_embedders = nn.ModuleList([
+            utils.make_mlp([mesh_same_dim] + self.mlp_blueprint_end)
+            for _ in range(self.num_levels)
+        ])
+        self.mesh_up_embedders = nn.ModuleList([
+            utils.make_mlp([mesh_up_dim] + self.mlp_blueprint_end)
+            for _ in range(self.num_levels - 1)
+        ])
+        self.mesh_down_embedders = nn.ModuleList([
+            utils.make_mlp([mesh_down_dim] + self.mlp_blueprint_end)
+            for _ in range(self.num_levels - 1)
+        ])
 
         # Instantiate GNNs
         # Init GNNs
-        self.mesh_init_gnns = nn.ModuleList(
-            [
-                InteractionNet(
-                    edge_index,
-                    hidden_dim,
-                    hidden_layers=hidden_layers,
-                )
-                for edge_index in self.mesh_up_edge_index
-            ]
-        )
+        self.mesh_init_gnns = nn.ModuleList([
+            InteractionNet(
+                edge_index,
+                hidden_dim,
+                hidden_layers=hidden_layers,
+            )
+            for edge_index in self.mesh_up_edge_index
+        ])
 
         # Read out GNNs
-        self.mesh_read_gnns = nn.ModuleList(
-            [
-                InteractionNet(
-                    edge_index,
-                    hidden_dim,
-                    hidden_layers=hidden_layers,
-                    update_edges=False,
-                )
-                for edge_index in self.mesh_down_edge_index
-            ]
-        )
+        self.mesh_read_gnns = nn.ModuleList([
+            InteractionNet(
+                edge_index,
+                hidden_dim,
+                hidden_layers=hidden_layers,
+                update_edges=False,
+            )
+            for edge_index in self.mesh_down_edge_index
+        ])
 
     def get_num_mesh(self):
         """
diff --git a/neural_lam/vis.py b/neural_lam/vis.py
index 9346000ce..4b1888893 100644
--- a/neural_lam/vis.py
+++ b/neural_lam/vis.py
@@ -2,7 +2,6 @@
 
 # Standard library
 import warnings
-from typing import Optional, Union
 
 # Third-party
 import cartopy.crs as ccrs
@@ -11,7 +10,6 @@
 import matplotlib.axes
 import matplotlib.collections
 import matplotlib.colors
-import matplotlib.figure
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
@@ -343,11 +341,11 @@ def plot_on_axis(
     ax: matplotlib.axes.Axes,
     da: xr.DataArray,
     datastore: BaseRegularGridDatastore,
-    vmin: Optional[float] = None,
-    vmax: Optional[float] = None,
-    ax_title: Optional[str] = None,
-    cmap: Union[str, matplotlib.colors.Colormap] = "plasma",
-    boundary_alpha: Optional[float] = None,
+    vmin: float | None = None,
+    vmax: float | None = None,
+    ax_title: str | None = None,
+    cmap: str | matplotlib.colors.Colormap = "plasma",
+    boundary_alpha: float | None = None,
     crop_to_interior: bool = False,
 ) -> matplotlib.collections.QuadMesh:
     """

From 949332f6ec273e461013e9e021ae54477bf9c013 Mon Sep 17 00:00:00 2001
From: sadamov <kode@mailbox.org>
Date: Tue, 9 Jun 2026 12:54:49 +0200
Subject: [PATCH 14/16] format

---
 neural_lam/datastore/base.py                  | 10 ++-
 .../models/forecasters/autoregressive.py      |  3 +-
 neural_lam/models/module.py                   |  9 +--
 neural_lam/models/step_predictors/base.py     |  7 +-
 .../models/step_predictors/graph/graph_lam.py |  6 +-
 .../models/step_predictors/graph/hi_lam.py    | 78 ++++++++++---------
 .../step_predictors/graph/hierarchical.py     | 78 +++++++++++--------
 7 files changed, 104 insertions(+), 87 deletions(-)

diff --git a/neural_lam/datastore/base.py b/neural_lam/datastore/base.py
index dabd15d6e..8376d38de 100644
--- a/neural_lam/datastore/base.py
+++ b/neural_lam/datastore/base.py
@@ -433,10 +433,12 @@ def expected_dim_order(
             if category != "static":
                 # static data does not vary in time
                 if self.is_forecast:
-                    dim_order.extend([
-                        "analysis_time",
-                        "elapsed_forecast_duration",
-                    ])
+                    dim_order.extend(
+                        [
+                            "analysis_time",
+                            "elapsed_forecast_duration",
+                        ]
+                    )
                 elif not self.is_forecast:
                     dim_order.append("time")
 
diff --git a/neural_lam/models/forecasters/autoregressive.py b/neural_lam/models/forecasters/autoregressive.py
index 7b0dc9138..daa2081de 100644
--- a/neural_lam/models/forecasters/autoregressive.py
+++ b/neural_lam/models/forecasters/autoregressive.py
@@ -33,8 +33,7 @@ def __init__(self, predictor: StepPredictor, datastore: BaseDatastore):
 
         # Register boundary/interior masks on the forecaster, not the predictor
         boundary_mask = (
-            torch
-            .tensor(datastore.boundary_mask.values, dtype=torch.float32)
+            torch.tensor(datastore.boundary_mask.values, dtype=torch.float32)
             .unsqueeze(0)
             .unsqueeze(-1)
         )
diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py
index 7ea4c81c9..d67975bdc 100644
--- a/neural_lam/models/module.py
+++ b/neural_lam/models/module.py
@@ -128,8 +128,7 @@ def __init__(
 
         # Compute interior_mask_bool directly from datastore
         boundary_mask = (
-            torch
-            .tensor(datastore.boundary_mask.values, dtype=torch.float32)
+            torch.tensor(datastore.boundary_mask.values, dtype=torch.float32)
             .unsqueeze(0)
             .unsqueeze(-1)
         )  # (1, num_grid_nodes, 1)
@@ -558,8 +557,7 @@ def plot_examples(self, batch, n_examples, split, prediction):
             ).unstack("grid_index")
 
             var_vmin = (
-                torch
-                .minimum(
+                torch.minimum(
                     pred_slice.flatten(0, 1).min(dim=0)[0],
                     target_slice.flatten(0, 1).min(dim=0)[0],
                 )
@@ -567,8 +565,7 @@ def plot_examples(self, batch, n_examples, split, prediction):
                 .numpy()
             )
             var_vmax = (
-                torch
-                .maximum(
+                torch.maximum(
                     pred_slice.flatten(0, 1).max(dim=0)[0],
                     target_slice.flatten(0, 1).max(dim=0)[0],
                 )
diff --git a/neural_lam/models/step_predictors/base.py b/neural_lam/models/step_predictors/base.py
index 4d9b070ea..61ce9b660 100644
--- a/neural_lam/models/step_predictors/base.py
+++ b/neural_lam/models/step_predictors/base.py
@@ -204,9 +204,10 @@ def prepare_clamping_params(self, datastore: BaseDatastore):
         sigmoid_center = 0
         softplus_center = 0
 
-        normalize_clamping_lim = lambda x, feature_idx: (
-            (x - self.state_mean[feature_idx]) / self.state_std[feature_idx]
-        )
+        def normalize_clamping_lim(x, feature_idx):
+            return (x - self.state_mean[feature_idx]) / self.state_std[
+                feature_idx
+            ]
 
         # Check which clamping functions to use for each feature
         sigmoid_lower_upper_idx = []
diff --git a/neural_lam/models/step_predictors/graph/graph_lam.py b/neural_lam/models/step_predictors/graph/graph_lam.py
index 60ed434f2..15d90565b 100644
--- a/neural_lam/models/step_predictors/graph/graph_lam.py
+++ b/neural_lam/models/step_predictors/graph/graph_lam.py
@@ -76,9 +76,9 @@ def __init__(
             output_clamping_upper=output_clamping_upper,
         )
 
-        assert not self.hierarchical, (
-            "GraphLAM does not use a hierarchical mesh graph"
-        )
+        assert (
+            not self.hierarchical
+        ), "GraphLAM does not use a hierarchical mesh graph"
 
         # grid_dim from data + static + batch_static
         mesh_dim = self.mesh_static_features.shape[1]
diff --git a/neural_lam/models/step_predictors/graph/hi_lam.py b/neural_lam/models/step_predictors/graph/hi_lam.py
index 6db85a706..2e86ba465 100644
--- a/neural_lam/models/step_predictors/graph/hi_lam.py
+++ b/neural_lam/models/step_predictors/graph/hi_lam.py
@@ -77,20 +77,20 @@ def __init__(
         )
 
         # Make down GNNs, both for down edges and same level
-        self.mesh_down_gnns = nn.ModuleList([
-            self.make_down_gnns() for _ in range(processor_layers)
-        ])  # Nested lists (proc_steps, num_levels-1)
-        self.mesh_down_same_gnns = nn.ModuleList([
-            self.make_same_gnns() for _ in range(processor_layers)
-        ])  # Nested lists (proc_steps, num_levels)
+        self.mesh_down_gnns = nn.ModuleList(
+            [self.make_down_gnns() for _ in range(processor_layers)]
+        )  # Nested lists (proc_steps, num_levels-1)
+        self.mesh_down_same_gnns = nn.ModuleList(
+            [self.make_same_gnns() for _ in range(processor_layers)]
+        )  # Nested lists (proc_steps, num_levels)
 
         # Make up GNNs, both for up edges and same level
-        self.mesh_up_gnns = nn.ModuleList([
-            self.make_up_gnns() for _ in range(processor_layers)
-        ])  # Nested lists (proc_steps, num_levels-1)
-        self.mesh_up_same_gnns = nn.ModuleList([
-            self.make_same_gnns() for _ in range(processor_layers)
-        ])  # Nested lists (proc_steps, num_levels)
+        self.mesh_up_gnns = nn.ModuleList(
+            [self.make_up_gnns() for _ in range(processor_layers)]
+        )  # Nested lists (proc_steps, num_levels-1)
+        self.mesh_up_same_gnns = nn.ModuleList(
+            [self.make_same_gnns() for _ in range(processor_layers)]
+        )  # Nested lists (proc_steps, num_levels)
 
     def make_same_gnns(self):
         """
@@ -101,14 +101,16 @@ def make_same_gnns(self):
         nn.ModuleList
             List of GNNs for each level.
         """
-        return nn.ModuleList([
-            InteractionNet(
-                edge_index,
-                self.hidden_dim,
-                hidden_layers=self.hidden_layers,
-            )
-            for edge_index in self.m2m_edge_index
-        ])
+        return nn.ModuleList(
+            [
+                InteractionNet(
+                    edge_index,
+                    self.hidden_dim,
+                    hidden_layers=self.hidden_layers,
+                )
+                for edge_index in self.m2m_edge_index
+            ]
+        )
 
     def make_up_gnns(self):
         """
@@ -119,14 +121,16 @@ def make_up_gnns(self):
         nn.ModuleList
             List of GNNs for each inter-level gap (upwards).
         """
-        return nn.ModuleList([
-            InteractionNet(
-                edge_index,
-                self.hidden_dim,
-                hidden_layers=self.hidden_layers,
-            )
-            for edge_index in self.mesh_up_edge_index
-        ])
+        return nn.ModuleList(
+            [
+                InteractionNet(
+                    edge_index,
+                    self.hidden_dim,
+                    hidden_layers=self.hidden_layers,
+                )
+                for edge_index in self.mesh_up_edge_index
+            ]
+        )
 
     def make_down_gnns(self):
         """
@@ -137,14 +141,16 @@ def make_down_gnns(self):
         nn.ModuleList
             List of GNNs for each inter-level gap (downwards).
         """
-        return nn.ModuleList([
-            InteractionNet(
-                edge_index,
-                self.hidden_dim,
-                hidden_layers=self.hidden_layers,
-            )
-            for edge_index in self.mesh_down_edge_index
-        ])
+        return nn.ModuleList(
+            [
+                InteractionNet(
+                    edge_index,
+                    self.hidden_dim,
+                    hidden_layers=self.hidden_layers,
+                )
+                for edge_index in self.mesh_down_edge_index
+            ]
+        )
 
     def mesh_down_step(
         self,
diff --git a/neural_lam/models/step_predictors/graph/hierarchical.py b/neural_lam/models/step_predictors/graph/hierarchical.py
index 8ee30f217..72bad419b 100644
--- a/neural_lam/models/step_predictors/graph/hierarchical.py
+++ b/neural_lam/models/step_predictors/graph/hierarchical.py
@@ -79,44 +79,56 @@ def __init__(
         mesh_down_dim = self.mesh_down_features[0].shape[1]
 
         # Separate mesh node embedders for each level
-        self.mesh_embedders = nn.ModuleList([
-            utils.make_mlp([mesh_dim] + self.mlp_blueprint_end)
-            for _ in range(self.num_levels)
-        ])
-        self.mesh_same_embedders = nn.ModuleList([
-            utils.make_mlp([mesh_same_dim] + self.mlp_blueprint_end)
-            for _ in range(self.num_levels)
-        ])
-        self.mesh_up_embedders = nn.ModuleList([
-            utils.make_mlp([mesh_up_dim] + self.mlp_blueprint_end)
-            for _ in range(self.num_levels - 1)
-        ])
-        self.mesh_down_embedders = nn.ModuleList([
-            utils.make_mlp([mesh_down_dim] + self.mlp_blueprint_end)
-            for _ in range(self.num_levels - 1)
-        ])
+        self.mesh_embedders = nn.ModuleList(
+            [
+                utils.make_mlp([mesh_dim] + self.mlp_blueprint_end)
+                for _ in range(self.num_levels)
+            ]
+        )
+        self.mesh_same_embedders = nn.ModuleList(
+            [
+                utils.make_mlp([mesh_same_dim] + self.mlp_blueprint_end)
+                for _ in range(self.num_levels)
+            ]
+        )
+        self.mesh_up_embedders = nn.ModuleList(
+            [
+                utils.make_mlp([mesh_up_dim] + self.mlp_blueprint_end)
+                for _ in range(self.num_levels - 1)
+            ]
+        )
+        self.mesh_down_embedders = nn.ModuleList(
+            [
+                utils.make_mlp([mesh_down_dim] + self.mlp_blueprint_end)
+                for _ in range(self.num_levels - 1)
+            ]
+        )
 
         # Instantiate GNNs
         # Init GNNs
-        self.mesh_init_gnns = nn.ModuleList([
-            InteractionNet(
-                edge_index,
-                hidden_dim,
-                hidden_layers=hidden_layers,
-            )
-            for edge_index in self.mesh_up_edge_index
-        ])
+        self.mesh_init_gnns = nn.ModuleList(
+            [
+                InteractionNet(
+                    edge_index,
+                    hidden_dim,
+                    hidden_layers=hidden_layers,
+                )
+                for edge_index in self.mesh_up_edge_index
+            ]
+        )
 
         # Read out GNNs
-        self.mesh_read_gnns = nn.ModuleList([
-            InteractionNet(
-                edge_index,
-                hidden_dim,
-                hidden_layers=hidden_layers,
-                update_edges=False,
-            )
-            for edge_index in self.mesh_down_edge_index
-        ])
+        self.mesh_read_gnns = nn.ModuleList(
+            [
+                InteractionNet(
+                    edge_index,
+                    hidden_dim,
+                    hidden_layers=hidden_layers,
+                    update_edges=False,
+                )
+                for edge_index in self.mesh_down_edge_index
+            ]
+        )
 
     def get_num_mesh(self):
         """

From 58fa6efecd3c73cb4505c3f2eaa8c24aedd1e58d Mon Sep 17 00:00:00 2001
From: sadamov <kode@mailbox.org>
Date: Tue, 9 Jun 2026 12:57:20 +0200
Subject: [PATCH 15/16] format

---
 neural_lam/models/step_predictors/base.py | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/neural_lam/models/step_predictors/base.py b/neural_lam/models/step_predictors/base.py
index 61ce9b660..e56869586 100644
--- a/neural_lam/models/step_predictors/base.py
+++ b/neural_lam/models/step_predictors/base.py
@@ -205,6 +205,17 @@ def prepare_clamping_params(self, datastore: BaseDatastore):
         softplus_center = 0
 
         def normalize_clamping_lim(x, feature_idx):
+            """Normalize a clamping limit from the original feature space to the
+            standardized space of the model's output.
+
+            Parameters
+            ----------
+            x : float
+                The clamping limit in the original feature space.
+            feature_idx : int
+                The index of the feature this limit applies to, used to look up
+                the mean and std for normalization.
+            """
             return (x - self.state_mean[feature_idx]) / self.state_std[
                 feature_idx
             ]
@@ -222,11 +233,11 @@ def normalize_clamping_lim(x, feature_idx):
 
         for feature_idx, feature in enumerate(state_feature_names):
             if feature in lower_lims and feature in upper_lims:
-                assert lower_lims[feature] < upper_lims[feature], (
-                    f'Invalid clamping limits for feature "{feature}",\
+                assert (
+                    lower_lims[feature] < upper_lims[feature]
+                ), f'Invalid clamping limits for feature "{feature}",\
                      lower: {lower_lims[feature]}, larger than\
                      upper: {upper_lims[feature]}'
-                )
                 sigmoid_lower_upper_idx.append(feature_idx)
                 sigmoid_lower_lims.append(
                     normalize_clamping_lim(lower_lims[feature], feature_idx)

From 760c9053717c04cd589f70bded7f7c5ad623e5db Mon Sep 17 00:00:00 2001
From: sadamov <kode@mailbox.org>
Date: Thu, 11 Jun 2026 06:24:28 +0200
Subject: [PATCH 16/16] docs: drop duplicate folder-tree block in
 NpyFilesDatastoreMEPS docstring

@observingClouds flagged that the second (shorter) ASCII folder tree
is an incomplete repetition of the detailed one above. The detailed
version is the source of truth; remove the duplicate.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 neural_lam/datastore/npyfilesmeps/store.py | 28 ----------------------
 1 file changed, 28 deletions(-)

diff --git a/neural_lam/datastore/npyfilesmeps/store.py b/neural_lam/datastore/npyfilesmeps/store.py
index df0ed2659..e14bbc740 100644
--- a/neural_lam/datastore/npyfilesmeps/store.py
+++ b/neural_lam/datastore/npyfilesmeps/store.py
@@ -139,34 +139,6 @@ class NpyFilesDatastoreMEPS(BaseRegularGridDatastore):
         ├── parameter_weights.npy
         └── surface_geopotential.npy
 
-    Notes
-    -----
-    Folder structure::
-
-        meps_example_reduced
-        ├── data_config.yaml
-        ├── samples
-        │   ├── test
-        │   │   ├── nwp_2022090100_mbr000.npy
-        │   │   ├── ...
-        │   ├── train
-        │   │   ├── nwp_2022040100_mbr000.npy
-        │   │   ├── ...
-        │   └── val
-        │       ├── nwp_2022060500_mbr000.npy
-        │       └── ...
-        └── static
-            ├── border_mask.npy
-            ├── diff_mean.pt
-            ├── diff_std.pt
-            ├── flux_stats.pt
-            ├── grid_features.pt
-            ├── nwp_xy.npy
-            ├── parameter_mean.pt
-            ├── parameter_std.pt
-            ├── parameter_weights.npy
-            └── surface_geopotential.npy
-
     For the MEPS dataset:
     N_t' = 65
     N_t = 65//subsample_step (= 21 for 3h steps)