Add Jupyter Book documentation infrastructure with auto-generated API reference via sphinx-autoapi

.github/workflows/docs.yml

@@ -0,0 +1,80 @@
+name: Documentation
+on:
+  push:
+    branches: ["main"]
+    paths:
+      - "neural_lam/**"
+      - "docs/**"
+      - "pyproject.toml"
+      - ".github/workflows/docs.yml"
+  pull_request:
+    branches: ["main"]
+    paths:
+      - "neural_lam/**"
+      - "docs/**"
+      - "pyproject.toml"
+      - ".github/workflows/docs.yml"
+  workflow_dispatch:
+
+concurrency:
+  group: "docs-${{ github.ref }}"
+  cancel-in-progress: true
+
+jobs:
+  build-docs:
+    name: Build docs
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Set up uv
+        uses: astral-sh/setup-uv@v5
+
+      - name: Set up Python 3.11
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install docs dependencies
+        run: uv sync --extra docs --group dev --no-cache
+
+
+
+      - name: Build docs with sphinx
+        run: uv run sphinx-build -W --keep-going -b html docs/ docs/_build/html/
+
+      - name: Upload docs artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: docs-html
+          path: docs/_build/html/
+
+  linkcheck:
+    name: Link check
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Set up uv
+        uses: astral-sh/setup-uv@v5
+
+      - name: Set up Python 3.11
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install docs dependencies
+        run: uv sync --extra docs --group dev --no-cache
+
+
+
+      - name: Run linkcheck
+        run: uv run sphinx-build -W --keep-going -b linkcheck docs/ docs/_build/linkcheck/

.gitignore

@@ -52,6 +52,11 @@ share/python-wheels/
 # mkdocs documentation
 /site
 
+# Jupyter Book / Sphinx documentation
+docs/_build/
+docs/api/
+
+
 ### Vim ###
 # Swap
 [._]*.s[a-v][a-z]

.pre-commit-config.yaml

@@ -17,6 +17,7 @@ repos:
     hooks:
       - id: codespell
         description: Check for spelling errors
+        args: ["--ignore-words-list", "astroid"]
 
   - repo: https://github.com/psf/black
     rev: 25.11.0

.readthedocs.yaml

@@ -0,0 +1,11 @@
+version: 2
+
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "3.11"
+  commands:
+    - pip install uv
+    - uv sync --extra docs --group dev
+    - uv run sphinx-build -W --keep-going -b html docs/ docs/_build/html/
+    - cp -r docs/_build/html/* $READTHEDOCS_OUTPUT/html/

CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Add a Sphinx + `sphinx-autoapi` documentation site (intro, install, quickstart, auto-generated API reference) on `myst-nb` and `sphinx-book-theme`, with a `docs` CI build/linkcheck job and a ReadTheDocs config. closes #61 [\#272](https://github.com/mllam/neural-lam/pull/272) @Mohit-Lakra
+
 - Add `PropagationNet` GNN layer that incentivises directional message
   propagation from sender to receiver nodes, and expose it alongside
   `InteractionNet` through four new CLI arguments (`--g2m_gnn_type`,

docs/conf.py

@@ -0,0 +1,114 @@
+# Standard library
+import os
+import sys
+
+# Add the project root to sys.path so sphinx can find neural_lam
+sys.path.insert(0, os.path.abspath(".."))
+
+project = "Neural-LAM"
+copyright = "2024–2026, MLLAM Community"
+author = "MLLAM Community"
+
+# General Sphinx configuration
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
+
+extensions = [
+    "autoapi.extension",
+    "myst_nb",
+    "sphinx.ext.napoleon",
+    "sphinx.ext.viewcode",
+    "sphinx.ext.intersphinx",
+    "sphinx_copybutton",
+    "sphinxext.opengraph",
+]
+
+# AutoAPI settings
+autoapi_dirs = ["../neural_lam"]
+autoapi_root = "api"
+autoapi_type = "python"
+autoapi_options = [
+    "members",
+    "undoc-members",
+    "show-inheritance",
+    "show-module-summary",
+]
+autoapi_python_class_content = "both"
+autoapi_member_order = "groupwise"
+autoapi_python_use_implicit_namespaces = False
+autoapi_keep_files = True
+autoapi_add_toctree_entry = True
+autoapi_ignore = ["**/tests/**", "**/conftest.py"]
+
+# Napoleon settings
+napoleon_numpy_docstring = True
+napoleon_google_docstring = False
+napoleon_include_init_with_doc = True
+napoleon_include_private_with_doc = False
+napoleon_include_special_with_doc = False
+napoleon_use_param = True
+napoleon_use_rtype = True
+napoleon_use_ivar = True
+
+# Intersphinx settings
+intersphinx_mapping = {
+    "python": ("https://docs.python.org/3", None),
+    "numpy": ("https://numpy.org/doc/stable", None),
+    "torch": ("https://pytorch.org/docs/stable", None),
+    "pytorch_lightning": ("https://lightning.ai/docs/pytorch/stable/", None),
+    "torch_geometric": (
+        "https://pytorch-geometric.readthedocs.io/en/latest/",
+        None,
+    ),
+}
+
+# MyST / Notebook settings
+myst_enable_extensions = [
+    "colon_fence",
+    "dollarmath",
+    "linkify",
+    "substitution",
+    "tasklist",
+    "deflist",
+    "fieldlist",
+    "html_admonition",
+    "html_image",
+    "smartquotes",
+    "attrs_inline",
+]
+nb_execution_mode = "off"
+
+# HTML Theme
+html_theme = "sphinx_book_theme"
+html_logo = "_static/logo.png"
+html_static_path = ["_static"]
+
+html_theme_options = {
+    "repository_url": "https://github.com/mllam/neural-lam",
+    "use_repository_button": True,
+    "use_issues_button": True,
+    "use_edit_page_button": True,
+    "repository_branch": "main",
+    "path_to_docs": "docs",
+    "show_navbar_depth": 2,
+    "show_toc_level": 2,
+    "extra_footer": (
+        '<p>Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> | '
+        '<a href="https://github.com/mllam/neural-lam">Source</a></p>'
+    ),
+}
+
+# OpenGraph settings
+ogp_site_url = "https://neural-lam.readthedocs.io/en/latest/"
+ogp_image = "_static/logo.png"
+ogp_use_first_image = True
+
+suppress_warnings = [
+    "autoapi.python_import_resolution",
+    "myst.xref_missing",
+]
+
+# Linkcheck settings
+linkcheck_ignore = [
+    r"https://kutt\.to/mllam",  # Returns 403 Forbidden for bots
+    r"https://docs\.pytorch\.org/.*",  # Flaky anchors in intersphinx
+]

docs/getting-started/installation.md

@@ -0,0 +1,29 @@
+# Installation
+
+## Prerequisites
+- Python >=3.10
+- Git
+- (Optional) CUDA-capable GPU for training
+
+```{include} ../../README.md
+:start-after: "# Installing Neural-LAM"
+:end-before: "# Using Neural-LAM"
+```
+
+## Building Documentation Locally
+
+```{code-block} bash
+# Install docs dependencies
+uv sync --extra docs --group dev
+
+# Build the documentation
+sphinx-build -b html docs/ docs/_build/html/
+
+# Open in browser
+open docs/_build/html/index.html  # macOS
+# xdg-open docs/_build/html/index.html  # Linux
+```
+
+```{seealso}
+See the {doc}`quickstart` guide to run your first training.
+```

docs/getting-started/quickstart.md

@@ -0,0 +1,52 @@
+# Quickstart
+
+## Overview
+
+This guide walks through the minimum steps to get started with Neural-LAM:
+1. Set up example data
+2. Create a graph
+3. Train a model
+4. Evaluate the model
+
+## Step 1: Download Example Data
+
+When running tests for the first time, example data is automatically downloaded from S3. You can trigger this download by running a minimal test.
+
+```{code-block} bash
+# Run a minimal test to trigger the data download
+pytest tests/test_training.py -vv -s -k "test_training"
+```
+
+## Step 2: Create a Graph
+
+Before training, you must construct a graph mesh for your data.
+
+```{code-block} bash
+python -m neural_lam.create_graph --config_path <path-to-config> --name <graph-name>
+```
+
+This script builds the mesh graph required by the GNN models.
+
+## Step 3: Train a Model
+
+Now you can train a model using the graph and configuration.
+
+```{code-block} bash
+python -m neural_lam.train_model --config_path <path-to-config> --model graph_lam --graph <graph-name>
+```
+
+Neural-LAM supports several models like `graph_lam`, `hi_lam`, and `hi_lam_parallel`.
+
+## Step 4: Evaluate
+
+After training, you can evaluate the model on the test set by loading the saved checkpoint.
+
+```{code-block} bash
+python -m neural_lam.train_model --eval test --config_path <path-to-config> --load <checkpoint-path>
+```
+
+## Configuration
+
+Neural-LAM uses a YAML configuration system powered by `dataclass-wizard`. The configuration defines the dataset paths, training parameters, and model hyperparameters. For complete details, see the API reference for {py:class}`neural_lam.config.NeuralLAMConfig`.
+
+- {doc}`../api/index` for complete API reference

docs/graph_storage_spec.md

@@ -1,7 +1,7 @@
 <!-- AUTO-GENERATED by docs/validate_graph.py; do not edit by hand. -->
 <!-- Regenerate using the validator script's write-spec-to-path option. -->
 
-# Neural-LAM Graph Storage Specification
+# Graph Storage Specification
 
 Version: 0.1.0
 

docs/index.md

@@ -0,0 +1,44 @@
+# Neural-LAM
+
+Neural-LAM is a PyTorch and PyTorch Lightning framework for high-resolution weather prediction using Graph Neural Networks. It provides a modular approach to Limited Area Modeling, supporting multiple graph-based architectures such as GraphLAM, HiLAM, and HiLAMParallel to process and predict meteorological data efficiently.
+
+```{toctree}
+:maxdepth: 2
+:hidden:
+
+getting-started/installation
+getting-started/quickstart
+```
+
+```{toctree}
+:maxdepth: 2
+:caption: Tutorials
+:hidden:
+
+notebooks/create_reduced_meps_dataset
+```
+
+```{toctree}
+:maxdepth: 2
+:caption: Reference
+:hidden:
+
+api/index
+graph_storage_spec
+```
+
+- **[🚀 Getting Started](getting-started/installation.md)**: Installation guide and quickstart tutorial to get you up and running.
+- **[📚 API Reference](api/index)**: Auto-generated reference for all modules, classes, and functions.
+
+## Key Features
+
+- **Modular design**: Swap datastores, models, and graph structures independently
+- **Multiple model architectures**: GraphLAM (flat), HiLAM (hierarchical), HiLAMParallel (parallel hierarchical)
+- **Flexible data handling**: Abstract datastore interface supporting zarr, numpy, and custom formats via mllam-data-prep
+- **Production-ready**: PyTorch Lightning for training, W&B/MLflow logging, checkpoint management
+
+## Quick Links
+
+- [GitHub Repository](https://github.com/mllam/neural-lam)
+- [Issue Tracker](https://github.com/mllam/neural-lam/issues)
+- [MLLAM Community Slack](https://kutt.to/mllam)

docs/validate_graph.py

@@ -1273,7 +1273,7 @@ def validate_graph_directory(
 
     spec_text = textwrap.dedent(
         f"""\
-    # Neural-LAM Graph Storage Specification
+    # Graph Storage Specification
 
     Version: {CURRENT_GRAPH_FORMAT_SPEC_VERSION}