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 new file mode 100644 index 000000000..d3b9f8084 --- /dev/null +++ b/docs/api/datastore.md @@ -0,0 +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/metrics.md b/docs/api/metrics.md new file mode 100644 index 000000000..cd9cf84b3 --- /dev/null +++ b/docs/api/metrics.md @@ -0,0 +1,6 @@ +::: neural_lam.metrics + 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 new file mode 100644 index 000000000..bb2f05548 --- /dev/null +++ b/docs/api/models.md @@ -0,0 +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/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/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..2f416e0bd --- /dev/null +++ b/mkdocs.yml @@ -0,0 +1,44 @@ +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 + - Configuration: api/config.md + - Changelog: changelog.md + +markdown_extensions: + - admonition + - pymdownx.highlight + - pymdownx.superfences + - pymdownx.snippets 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) 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