docs: 100 percent coverage of docstrings

.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: ["-c", "pyproject.toml", "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

CHANGELOG.md

@@ -74,6 +74,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
+
 - Register a `slow` pytest marker and apply it to `test_training` and `test_training_output_std` so contributors can skip long-running training tests during local iteration via `pytest -m "not slow"`. [\#651](https://github.com/mllam/neural-lam/pull/651) @sadamov
 
 - Add a short README pointer to [\#163](https://github.com/mllam/neural-lam/issues/163) for DGX Spark / PyTorch container compatibility notes, so users hitting `torch_scatter` errors know where to find the known-working / known-failing combos [\#266](https://github.com/mllam/neural-lam/pull/266) @Jayant-kernel

README.md

@@ -578,6 +578,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.

neural_lam/__init__.py

@@ -1,3 +1,5 @@
+"""Neural-LAM: graph-based neural weather prediction models."""
+
 # Standard library
 import importlib.metadata
 

neural_lam/config.py

@@ -1,3 +1,5 @@
+"""Configuration dataclasses and helpers for Neural-LAM experiments."""
+
 # Standard library
 import dataclasses
 from pathlib import Path
@@ -31,13 +33,20 @@ class DatastoreSelection:
     """
 
     kind: str
+    config_path: 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")
 
-    config_path: str
-
 
 @dataclasses.dataclass
 class ManualStateFeatureWeighting:
@@ -89,10 +98,13 @@ 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).
+    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[
@@ -107,15 +119,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 +167,8 @@ class _(dataclass_wizard.JSONWizard.Meta):
 
 
 class InvalidConfigError(Exception):
+    """Raised when the Neural-LAM configuration file is invalid or malformed."""
+
     pass
 
 

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
@@ -22,6 +24,21 @@
 def plot_graph(
     graph: pyg.data.Data, title: Optional[str] = None
 ) -> tuple[matplotlib.figure.Figure, matplotlib.axes.Axes]:
+    """
+    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
@@ -73,6 +90,19 @@ def plot_graph(
 
 
 def sort_nodes_internally(nx_graph: networkx.Graph) -> networkx.DiGraph:
+    """
+    Return a copy of ``nx_graph`` with deterministically ordered nodes.
+
+    Parameters
+    ----------
+    nx_graph : networkx.Graph
+        The input graph to sort nodes for.
+
+    Returns
+    -------
+    networkx.DiGraph
+        A directed graph with nodes sorted alphabetically by their labels.
+    """
     # 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.
@@ -83,6 +113,18 @@ def sort_nodes_internally(nx_graph: networkx.Graph) -> networkx.DiGraph:
 
 
 def save_edges(graph: pyg.data.Data, name: str, base_path: str) -> None:
+    """
+    Persist edge indices/features for a PyG graph under ``base_path``.
+
+    Parameters
+    ----------
+    graph : torch_geometric.data.Data
+        The graph containing edge data.
+    name : str
+        The name prefix for the saved files.
+    base_path : str
+        The directory path where files should be saved.
+    """
     torch.save(
         graph.edge_index, os.path.join(base_path, f"{name}_edge_index.pt")
     )
@@ -95,6 +137,18 @@ def save_edges(graph: pyg.data.Data, name: str, base_path: str) -> None:
 def save_edges_list(
     graphs: list[pyg.data.Data], name: str, base_path: str
 ) -> None:
+    """
+    Persist edge indices/features for a list of graphs.
+
+    Parameters
+    ----------
+    graphs : list of torch_geometric.data.Data
+        The list of graphs containing edge data.
+    name : str
+        The name prefix for the saved files.
+    base_path : str
+        The directory path where files should be saved.
+    """
     torch.save(
         [graph.edge_index for graph in graphs],
         os.path.join(base_path, f"{name}_edge_index.pt"),
@@ -111,12 +165,44 @@ def save_edges_list(
 def from_networkx_with_start_index(
     nx_graph: networkx.Graph, start_index: int
 ) -> pyg.data.Data:
+    """
+    Convert a NetworkX graph to PyG and offset node indices.
+
+    Parameters
+    ----------
+    nx_graph : networkx.Graph
+        The NetworkX graph to convert.
+    start_index : int
+        The value to add to each node index.
+
+    Returns
+    -------
+    pyg.data.Data
+        The converted PyG graph.
+    """
     pyg_graph = from_networkx(nx_graph)
     pyg_graph.edge_index += start_index
     return pyg_graph
 
 
 def mk_2d_graph(xy: np.ndarray, nx: int, ny: int) -> networkx.DiGraph:
+    """
+    Create a diagonal 2-D grid graph over the ``xy`` positions.
+
+    Parameters
+    ----------
+    xy : np.ndarray
+        The grid coordinates.
+    nx : int
+        Number of nodes in the x-dimension.
+    ny : int
+        Number of nodes in the y-dimension.
+
+    Returns
+    -------
+    networkx.DiGraph
+        The constructed directed 2-D grid graph.
+    """
     xm, xM = np.amin(xy[:, :, 0][:, 0]), np.amax(xy[:, :, 0][:, 0])
     ym, yM = np.amin(xy[:, :, 1][0, :]), np.amax(xy[:, :, 1][0, :])
 
@@ -156,6 +242,21 @@ def mk_2d_graph(xy: np.ndarray, nx: int, ny: int) -> networkx.DiGraph:
 
 
 def prepend_node_index(graph: networkx.Graph, new_index: int) -> networkx.Graph:
+    """
+    Relabel each node by prepending ``new_index`` to its tuple identifier.
+
+    Parameters
+    ----------
+    graph : networkx.Graph
+        The graph to relabel.
+    new_index : int
+        The value to prepend to each node identifier.
+
+    Returns
+    -------
+    networkx.Graph
+        The relabeled graph.
+    """
     # 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))
@@ -550,6 +651,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:
@@ -567,6 +684,15 @@ def create_graph_from_datastore(
 
 
 def cli(input_args: Optional[list[str]] = None) -> 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,

neural_lam/custom_loggers.py

@@ -1,3 +1,5 @@
+"""Custom logging utilities (e.g., MLFlow wrappers) used in Neural-LAM."""
+
 # Standard library
 import os
 from typing import Optional
@@ -24,10 +26,26 @@ def __init__(
         run_name: str,
         save_dir: str,
     ) -> None:
-        """Initialize the logger and ensure ``save_dir`` exists on disk.
+        """
+        Initialize the logger, ensure ``save_dir`` exists, and start the
+        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``.
+        save_dir : str
+            Directory where ``log_image`` writes temporary figure files.
+            Created eagerly with ``exist_ok=True``.
 
-        ``save_dir`` is created eagerly (with ``exist_ok=True``) so that
-        subsequent ``log_image`` calls can write temporary files there.
+        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
@@ -75,6 +93,11 @@ def log_image(
         step : int or None, optional
             Step to associate with the log entry. ``None`` logs without
             a step suffix.
+
+        Raises
+        ------
+        SystemExit
+            If AWS credentials for the MLflow artifact store are missing.
         """
         # Third-party
         from botocore.exceptions import NoCredentialsError