Docs: Initial documentation

docs/index.rst

@@ -19,17 +19,30 @@ framework it seeks to handle many of the complexities that
 meteorological organisations will share, allowing them to easily train
 models from existing recipes but with their own data.
 
-This package provides a series of utility functions for used by the rest
-of the *Anemoi* packages.
+*Anemoi models* provides the core model components for used by the rest
+of the *Anemoi* packages to train graph neural networks for data-driven
+weather forecasting.
 
+-  :doc:`overview`
 -  :doc:`installing`
 
 .. toctree::
    :maxdepth: 1
    :hidden:
 
+   overview
    installing
 
+*********
+ Modules
+*********
+
+.. toctree::
+   :maxdepth: 1
+   :glob:
+
+   modules/*
+
 *****************
  Anemoi packages
 *****************

docs/modules/data_indices.rst

@@ -0,0 +1,124 @@
+##############
+ Data Indices
+##############
+
+Data routing with a DataIndex system that knows where data is passed
+based on config entries.
+
+This documentation aims to first explain the user side (configs), and
+then goes into detail for developers.
+
+Good news first: You don't need to know the index of data in your Zarr
+
+*******************
+ Usage Information
+*******************
+
+The routing information can be changed in the config in data.
+
+There are currently three types of data with respect to the model:
+
+-  Prognostic: Data that is input into and output from the model
+
+-  Forcing: Data that informs the model as auxiliary information but
+   isn't output nor fed to the loss
+
+-  Diagnostic: Data that is not input into the model, but is produced in
+   the inference. (It's pseudo-diagnostic, as it is fed to the loss
+   during training, so the model is still conditioned on this data.)
+
+The default in Anemoi Models is that data is prognostic. But data can be
+switched to forcing or diagnostic by adding the value to the existing
+config entry:
+
+.. code:: yaml
+
+   data:
+     forcing:
+       - "lsm"
+       - "sdor"
+     diagnostic:
+       - "tp"
+       - "cp"
+
+.. figure:: ../schemas/data_indices.png
+   :alt: Schematic of IndexCollection with Data Indexing on Data and Model levels.
+   :align: center
+
+The are two Index-levels:
+
+-  Data: The data at "Zarr"-level provided by Anemoi-Datasets
+-  Model: The "squeezed" tensors with irrelevant parts missing.
+
+These are both split into two versions:
+
+-  Input: The data going into training / model
+-  Output: The data produced by training / model
+
+Each of these four indices has four integer Tensors available to slice
+into the data:
+
+-  full: The indices of all "relevant" data, i.e. prognostic and forcing
+   for the input data.
+-  diagnostic: The indices for data that is "output only" and not fed to
+   the model.
+-  prognostic: The indices for data that is in both the input and output
+   for the training and model.
+-  forcing: The indices for data that is only input into the model but
+   doesn't exist in the forecast state.
+
+The data can be accessed via dot notation in such a way that:
+
+.. code:: python
+
+   data_indices.<model/data>.<input/output>.<full/diagnostic/prognostic/forcing>
+
+Examples:
+
+.. code:: python
+
+   data_indices.data.input.full
+   data_indices.model.output.diagnostic
+
+The name_to_index dictionary is available for each of the four indices,
+which provides a mapping from variable name to index at that level.
+These are useful for providing users an interface in the config that
+does not rely on the knowledge of index-locations in the data.
+Generally, hard-coded locations are to be avoided.
+
+*******************
+ Index Collections
+*******************
+
+The Index Collections module provides a collection of indices, which are
+used to index and slice the data.
+
+.. automodule:: anemoi.models.data_indices.collection
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+*******
+ Index
+*******
+
+The Index module provides all data indices and methods for indexing and
+slicing the data.
+
+.. automodule:: anemoi.models.data_indices.index
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+********
+ Tensor
+********
+
+This provides the underlying data structure for the indices. It creates
+a tensor of the appropriate size and shape, and provides methods for
+indexing and slicing the data.
+
+.. automodule:: anemoi.models.data_indices.tensor
+   :members:
+   :no-undoc-members:
+   :show-inheritance:

docs/modules/distributed.rst

@@ -0,0 +1,69 @@
+#############
+ Distributed
+#############
+
+*******
+ graph
+*******
+
+.. automodule:: anemoi.models.distributed.graph
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+************
+ khop_edges
+************
+
+.. automodule:: anemoi.models.distributed.khop_edges
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+..
+   *************
+
+..
+   primitives
+
+..
+   *************
+
+..
+   .. automodule:: anemoi.models.distributed.primitives
+
+..
+   :members:
+
+..
+   :no-undoc-members:
+
+..
+   :show-inheritance:
+
+********
+ shapes
+********
+
+.. automodule:: anemoi.models.distributed.shapes
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+*************
+ transformer
+*************
+
+.. automodule:: anemoi.models.distributed.transformer
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+*******
+ utils
+*******
+
+.. automodule:: anemoi.models.distributed.utils
+   :members:
+   :no-undoc-members:
+   :show-inheritance:

docs/modules/interface.rst

@@ -0,0 +1,8 @@
+###########
+ Interface
+###########
+
+.. automodule:: anemoi.models.interface
+   :members:
+   :no-undoc-members:
+   :show-inheritance:

docs/modules/layers.rst

@@ -0,0 +1,84 @@
+########
+ Layers
+########
+
+*********
+ Mappers
+*********
+
+.. automodule:: anemoi.models.layers.mapper
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+************
+ Processors
+************
+
+.. automodule:: anemoi.models.layers.processor
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+********
+ Chunks
+********
+
+.. automodule:: anemoi.models.layers.chunk
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+********
+ Blocks
+********
+
+.. automodule:: anemoi.models.layers.block
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+*******
+ Graph
+*******
+
+.. automodule:: anemoi.models.layers.graph
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+******
+ Conv
+******
+
+.. automodule:: anemoi.models.layers.conv
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+***********
+ Attention
+***********
+
+.. automodule:: anemoi.models.layers.attention
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+************************
+ Multi-Layer Perceptron
+************************
+
+.. automodule:: anemoi.models.layers.mlp
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+*******
+ Utils
+*******
+
+.. automodule:: anemoi.models.layers.utils
+   :members:
+   :no-undoc-members:
+   :show-inheritance:

docs/modules/models.rst

@@ -0,0 +1,15 @@
+########
+ Models
+########
+
+*********************************
+ Encoder Processor Decoder Model
+*********************************
+
+The model defines a graph neural network architecture with configurable
+encoder, processor, and decoder.
+
+.. automodule:: anemoi.models.models.encoder_processor_decoder
+   :members:
+   :no-undoc-members:
+   :show-inheritance:

docs/modules/preprocessing.rst

@@ -0,0 +1,35 @@
+###############
+ Preprocessing
+###############
+
+The preprocessing module is used to pre- and post-process the data. The
+module contains the following classes:
+
+.. automodule:: anemoi.models.preprocessing
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+************
+ Normalizer
+************
+
+The normalizer module is used to normalize the data. The module contains
+the following classes:
+
+.. automodule:: anemoi.models.preprocessing.normalizer
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+*********
+ Imputer
+*********
+
+The imputer module is used to impute the data. The module contains the
+following classes:
+
+.. automodule:: anemoi.models.preprocessing.imputer
+   :members:
+   :no-undoc-members:
+   :show-inheritance: