From d7b6b745ae3f8ed373a83a9da63a4387effcefc7 Mon Sep 17 00:00:00 2001 From: Sidhant185 <74181527+Sidhant185@users.noreply.github.com> Date: Fri, 20 Mar 2026 12:37:12 +0530 Subject: [PATCH 1/4] docs: set up MkDocs documentation structure with mkdocstrings --- docs/api/datastore.md | 3 ++ docs/api/metrics.md | 3 ++ docs/api/models.md | 5 ++++ docs/getting_started/installation.md | 21 ++++++++++++++ docs/getting_started/quickstart.md | 13 +++++++++ docs/index.md | 23 ++++++++++++++++ mkdocs.yml | 41 ++++++++++++++++++++++++++++ pyproject.toml | 1 + 8 files changed, 110 insertions(+) create mode 100644 docs/api/datastore.md create mode 100644 docs/api/metrics.md create mode 100644 docs/api/models.md create mode 100644 docs/getting_started/installation.md create mode 100644 docs/getting_started/quickstart.md create mode 100644 docs/index.md create mode 100644 mkdocs.yml diff --git a/docs/api/datastore.md b/docs/api/datastore.md new file mode 100644 index 000000000..377414e73 --- /dev/null +++ b/docs/api/datastore.md @@ -0,0 +1,3 @@ +# Datastore + +::: neural_lam.datastore.base.BaseDatastore diff --git a/docs/api/metrics.md b/docs/api/metrics.md new file mode 100644 index 000000000..247ccd08a --- /dev/null +++ b/docs/api/metrics.md @@ -0,0 +1,3 @@ +# Metrics + +::: neural_lam.metrics diff --git a/docs/api/models.md b/docs/api/models.md new file mode 100644 index 000000000..e8a0dfe8a --- /dev/null +++ b/docs/api/models.md @@ -0,0 +1,5 @@ +# Models + +::: neural_lam.models.graph_lam.GraphLAM + +::: neural_lam.models.hi_lam.HiLAM diff --git a/docs/getting_started/installation.md b/docs/getting_started/installation.md new file mode 100644 index 000000000..002594887 --- /dev/null +++ b/docs/getting_started/installation.md @@ -0,0 +1,21 @@ +# Installation + +## Requirements + +- Python >= 3.10 +- PyTorch >= 2.3.0 + +## Install from source + +```bash +git clone https://github.com/mllam/neural-lam.git +cd neural-lam +uv sync +``` + +## Development install + +```bash +uv sync --group dev +pre-commit install +``` diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md new file mode 100644 index 000000000..a52301b01 --- /dev/null +++ b/docs/getting_started/quickstart.md @@ -0,0 +1,13 @@ +# Quick Start + +## Training a model + +```bash +python -m neural_lam.train_model --config config.yaml +``` + +## Running tests + +```bash +pytest tests/ +``` diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 000000000..fbc1d49e0 --- /dev/null +++ b/docs/index.md @@ -0,0 +1,23 @@ +# neural-lam + +Graph-based neural weather prediction for Limited Area Modeling (LAM). + +## Overview + +neural-lam is a Python package for training and evaluating graph-based +neural network models for weather forecasting. It is developed by the +[MLLAM community](https://github.com/mllam). + +## Features + +- Graph Neural Network-based weather prediction +- Support for Limited Area Modeling (LAM) +- Flexible datastore interface +- Multiple model architectures (GraphLAM, HiLAM, HiLAMParallel) + +## Quick Links + +- [Installation](getting_started/installation.md) +- [Quick Start](getting_started/quickstart.md) +- [API Reference](api/metrics.md) +- [GitHub Repository](https://github.com/mllam/neural-lam) diff --git a/mkdocs.yml b/mkdocs.yml new file mode 100644 index 000000000..6bf3bb2c8 --- /dev/null +++ b/mkdocs.yml @@ -0,0 +1,41 @@ +site_name: neural-lam +site_description: Graph-based neural weather prediction for Limited Area Modeling +site_url: https://mllam.github.io/neural-lam +repo_url: https://github.com/mllam/neural-lam +repo_name: mllam/neural-lam + +theme: + name: material + palette: + primary: blue + accent: cyan + features: + - navigation.tabs + - navigation.sections + - toc.integrate + - search.suggest + +plugins: + - search + - mkdocstrings: + handlers: + python: + options: + docstring_style: numpy + show_source: true + show_root_heading: true + +nav: + - Home: index.md + - Getting Started: + - Installation: getting_started/installation.md + - Quick Start: getting_started/quickstart.md + - API Reference: + - Metrics: api/metrics.md + - Models: api/models.md + - Datastore: api/datastore.md + +markdown_extensions: + - admonition + - pymdownx.highlight + - pymdownx.superfences diff --git a/pyproject.toml b/pyproject.toml index 0a7e8ec1a..860729c04 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,6 +42,7 @@ requires-python = ">=3.10" [dependency-groups] dev = ["pre-commit>=3.8.0", "pytest>=8.3.2", "pooch>=1.8.2"] +docs = ["mkdocs>=1.5.0", "mkdocs-material>=9.0.0", "mkdocstrings[python]>=0.24.0"] [tool.black] line-length = 80 From fa1da9241278cbf126fb57cbb68dbeb19e393fcf Mon Sep 17 00:00:00 2001 From: Sidhant185 <74181527+Sidhant185@users.noreply.github.com> Date: Fri, 20 Mar 2026 13:02:46 +0530 Subject: [PATCH 2/4] docs: convert metrics.py docstrings to NumPy style for mkdocstrings --- docs/api/metrics.md | 7 +- neural_lam/metrics.py | 269 +++++++++++++++++++++++++----------------- 2 files changed, 166 insertions(+), 110 deletions(-) diff --git a/docs/api/metrics.md b/docs/api/metrics.md index 247ccd08a..cd9cf84b3 100644 --- a/docs/api/metrics.md +++ b/docs/api/metrics.md @@ -1,3 +1,6 @@ -# Metrics - ::: neural_lam.metrics + options: + show_root_heading: true + show_source: true + docstring_style: numpy + members_order: source diff --git a/neural_lam/metrics.py b/neural_lam/metrics.py index 7db2cca6d..d8dbeae2e 100644 --- a/neural_lam/metrics.py +++ b/neural_lam/metrics.py @@ -4,12 +4,17 @@ def get_metric(metric_name): """ - Get a defined metric with given name + Get a defined metric with given name. - metric_name: str, name of the metric + Parameters + ---------- + metric_name : str + Name of the metric. - Returns: - metric: function implementing the metric + Returns + ------- + callable + Function implementing the metric. """ metric_name_lower = metric_name.lower() assert ( @@ -20,19 +25,25 @@ 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. + Masks and (optionally) reduces entry-wise metric values. + + Parameters + ---------- + metric_entry_vals : torch.Tensor + Shape (..., N, d_state), array of unaggregated metric values. (...,) + is any number of batch dimensions, potentially different but broadcastable. + mask : torch.Tensor, optional + Shape (N,), boolean mask for which grid nodes to use. + average_grid : bool, optional + If True, reduce grid dimension -2 by mean over N. + sum_vars : bool, optional + If True, reduce variable dimension -1 by sum over d_state. + + Returns + ------- + torch.Tensor + One of (...,), (..., d_state), (..., N), (..., N, d_state), + depending on reduction arguments. """ # Only keep grid nodes in mask if mask is not None: @@ -55,21 +66,28 @@ 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. + Weighted Mean Squared Error. + + Parameters + ---------- + pred : torch.Tensor + Shape (..., N, d_state), prediction tensor. + target : torch.Tensor + Shape (..., N, d_state), target tensor. + pred_std : torch.Tensor + Shape (..., N, d_state) or (d_state,), predicted std.-dev. + mask : torch.Tensor, optional + Shape (N,), boolean mask for which grid nodes to use. + average_grid : bool, optional + If True, reduce grid dimension -2 by mean over N. + sum_vars : bool, optional + If True, reduce variable dimension -1 by sum over d_state. + + Returns + ------- + torch.Tensor + One of (...,), (..., d_state), (..., N), (..., N, d_state), + depending on reduction arguments. """ entry_mse = torch.nn.functional.mse_loss( pred, target, reduction="none" @@ -86,21 +104,28 @@ 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. + (Unweighted) Mean Squared Error. + + Parameters + ---------- + pred : torch.Tensor + Shape (..., N, d_state), prediction tensor. + target : torch.Tensor + Shape (..., N, d_state), target tensor. + pred_std : torch.Tensor + Shape (..., N, d_state) or (d_state,), predicted std.-dev. + mask : torch.Tensor, optional + Shape (N,), boolean mask for which grid nodes to use. + average_grid : bool, optional + If True, reduce grid dimension -2 by mean over N. + sum_vars : bool, optional + If True, reduce variable dimension -1 by sum over d_state. + + Returns + ------- + torch.Tensor + One of (...,), (..., d_state), (..., N), (..., N, d_state), + depending on reduction arguments. """ # Replace pred_std with constant ones return wmse( @@ -110,21 +135,28 @@ 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. + Weighted Mean Absolute Error. + + Parameters + ---------- + pred : torch.Tensor + Shape (..., N, d_state), prediction tensor. + target : torch.Tensor + Shape (..., N, d_state), target tensor. + pred_std : torch.Tensor + Shape (..., N, d_state) or (d_state,), predicted std.-dev. + mask : torch.Tensor, optional + Shape (N,), boolean mask for which grid nodes to use. + average_grid : bool, optional + If True, reduce grid dimension -2 by mean over N. + sum_vars : bool, optional + If True, reduce variable dimension -1 by sum over d_state. + + Returns + ------- + torch.Tensor + One of (...,), (..., d_state), (..., N), (..., N, d_state), + depending on reduction arguments. """ entry_mae = torch.nn.functional.l1_loss( pred, target, reduction="none" @@ -141,21 +173,28 @@ 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. + (Unweighted) Mean Absolute Error. + + Parameters + ---------- + pred : torch.Tensor + Shape (..., N, d_state), prediction tensor. + target : torch.Tensor + Shape (..., N, d_state), target tensor. + pred_std : torch.Tensor + Shape (..., N, d_state) or (d_state,), predicted std.-dev. + mask : torch.Tensor, optional + Shape (N,), boolean mask for which grid nodes to use. + average_grid : bool, optional + If True, reduce grid dimension -2 by mean over N. + sum_vars : bool, optional + If True, reduce variable dimension -1 by sum over d_state. + + Returns + ------- + torch.Tensor + One of (...,), (..., d_state), (..., N), (..., N, d_state), + depending on reduction arguments. """ # Replace pred_std with constant ones return wmae( @@ -165,21 +204,28 @@ 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. + Negative Log Likelihood loss, for isotropic Gaussian likelihood. + + Parameters + ---------- + pred : torch.Tensor + Shape (..., N, d_state), prediction tensor. + target : torch.Tensor + Shape (..., N, d_state), target tensor. + pred_std : torch.Tensor + Shape (..., N, d_state) or (d_state,), predicted std.-dev. + mask : torch.Tensor, optional + Shape (N,), boolean mask for which grid nodes to use. + average_grid : bool, optional + If True, reduce grid dimension -2 by mean over N. + sum_vars : bool, optional + If True, reduce variable dimension -1 by sum over d_state. + + Returns + ------- + torch.Tensor + One of (...,), (..., d_state), (..., N), (..., N, d_state), + depending on reduction arguments. """ # Broadcast pred_std if shaped (d_state,), done internally in Normal class dist = torch.distributions.Normal(pred, pred_std) # (..., N, d_state) @@ -194,22 +240,29 @@ 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. + (Negative) Continuous Ranked Probability Score (CRPS). + Closed-form expression based on Gaussian predictive distribution. + + Parameters + ---------- + pred : torch.Tensor + Shape (..., N, d_state), prediction tensor. + target : torch.Tensor + Shape (..., N, d_state), target tensor. + pred_std : torch.Tensor + Shape (..., N, d_state) or (d_state,), predicted std.-dev. + mask : torch.Tensor, optional + Shape (N,), boolean mask for which grid nodes to use. + average_grid : bool, optional + If True, reduce grid dimension -2 by mean over N. + sum_vars : bool, optional + If True, reduce variable dimension -1 by sum over d_state. + + Returns + ------- + torch.Tensor + One of (...,), (..., d_state), (..., N), (..., N, d_state), + depending on reduction arguments. """ std_normal = torch.distributions.Normal( torch.zeros((), device=pred.device), torch.ones((), device=pred.device) From 6d74519f5ea30daa7ff65eda22130a8d47356cd7 Mon Sep 17 00:00:00 2001 From: Sidhant185 <74181527+Sidhant185@users.noreply.github.com> Date: Fri, 20 Mar 2026 13:05:53 +0530 Subject: [PATCH 3/4] docs: add config.py API page and improve all API reference pages --- docs/api/config.md | 9 +++++++++ docs/api/datastore.md | 5 +++++ docs/api/models.md | 17 +++++++++++++++++ mkdocs.yml | 1 + 4 files changed, 32 insertions(+) create mode 100644 docs/api/config.md diff --git a/docs/api/config.md b/docs/api/config.md new file mode 100644 index 000000000..5c0a2e197 --- /dev/null +++ b/docs/api/config.md @@ -0,0 +1,9 @@ +# Configuration + +::: neural_lam.config + options: + show_root_heading: true + show_source: true + docstring_style: numpy + members_order: source + show_submodules: true diff --git a/docs/api/datastore.md b/docs/api/datastore.md index 377414e73..d3b9f8084 100644 --- a/docs/api/datastore.md +++ b/docs/api/datastore.md @@ -1,3 +1,8 @@ # Datastore ::: neural_lam.datastore.base.BaseDatastore + options: + show_root_heading: true + show_source: true + docstring_style: numpy + members_order: source diff --git a/docs/api/models.md b/docs/api/models.md index e8a0dfe8a..bb2f05548 100644 --- a/docs/api/models.md +++ b/docs/api/models.md @@ -1,5 +1,22 @@ # Models ::: neural_lam.models.graph_lam.GraphLAM + options: + show_root_heading: true + show_source: true + docstring_style: numpy + members_order: source ::: neural_lam.models.hi_lam.HiLAM + options: + show_root_heading: true + show_source: true + docstring_style: numpy + members_order: source + +::: neural_lam.models.hi_lam_parallel.HiLAMParallel + options: + show_root_heading: true + show_source: true + docstring_style: numpy + members_order: source diff --git a/mkdocs.yml b/mkdocs.yml index 6bf3bb2c8..22bfb9510 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -34,6 +34,7 @@ nav: - Metrics: api/metrics.md - Models: api/models.md - Datastore: api/datastore.md + - Configuration: api/config.md markdown_extensions: - admonition From c98509020273088113205d09655de40cb9f281f7 Mon Sep 17 00:00:00 2001 From: Sidhant185 <74181527+Sidhant185@users.noreply.github.com> Date: Fri, 20 Mar 2026 13:07:14 +0530 Subject: [PATCH 4/4] docs: add Changelog page linking to CHANGELOG.md --- docs/changelog.md | 5 +++++ mkdocs.yml | 2 ++ 2 files changed, 7 insertions(+) create mode 100644 docs/changelog.md diff --git a/docs/changelog.md b/docs/changelog.md new file mode 100644 index 000000000..a3736a26f --- /dev/null +++ b/docs/changelog.md @@ -0,0 +1,5 @@ +# Changelog + +This page documents all notable changes to neural-lam. + +--8<-- "CHANGELOG.md" diff --git a/mkdocs.yml b/mkdocs.yml index 22bfb9510..2f416e0bd 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -35,8 +35,10 @@ nav: - Models: api/models.md - Datastore: api/datastore.md - Configuration: api/config.md + - Changelog: changelog.md markdown_extensions: - admonition - pymdownx.highlight - pymdownx.superfences + - pymdownx.snippets