diff --git a/docs/modules/data_indices.rst b/docs/modules/data_indices.rst index d776aab..59605ee 100644 --- a/docs/modules/data_indices.rst +++ b/docs/modules/data_indices.rst @@ -2,7 +2,102 @@ Data Indices ############## -.. automodule:: anemoi.models.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-block:: 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-block:: python + + data_indices... + +Examples: + + +.. code-block:: 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: diff --git a/docs/modules/distributed.rst b/docs/modules/distributed.rst index 6aa8811..de8a98a 100644 --- a/docs/modules/distributed.rst +++ b/docs/modules/distributed.rst @@ -2,6 +2,9 @@ Distributed ############# + + + .. automodule:: anemoi.models.distributed :members: :no-undoc-members: diff --git a/docs/modules/layers.rst b/docs/modules/layers.rst index d93d2c7..5bb25cb 100644 --- a/docs/modules/layers.rst +++ b/docs/modules/layers.rst @@ -2,7 +2,83 @@ Layers ######## -.. automodule:: anemoi.models.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: diff --git a/docs/modules/models.rst b/docs/modules/models.rst index 745046b..329791f 100644 --- a/docs/modules/models.rst +++ b/docs/modules/models.rst @@ -2,7 +2,13 @@ Models ######## -.. automodule:: anemoi.models.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: diff --git a/docs/modules/preprocessing.rst b/docs/modules/preprocessing.rst index bccb0d1..eb83d9a 100644 --- a/docs/modules/preprocessing.rst +++ b/docs/modules/preprocessing.rst @@ -2,7 +2,32 @@ 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: diff --git a/docs/modules/utils.rst b/docs/modules/utils.rst deleted file mode 100644 index daf5133..0000000 --- a/docs/modules/utils.rst +++ /dev/null @@ -1,8 +0,0 @@ -####### - Utils -####### - -.. automodule:: anemoi.models.utils - :members: - :no-undoc-members: - :show-inheritance: diff --git a/docs/schemas/data_indices.drawio b/docs/schemas/data_indices.drawio new file mode 100644 index 0000000..536c59e --- /dev/null +++ b/docs/schemas/data_indices.drawio @@ -0,0 +1,307 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docs/schemas/data_indices.png b/docs/schemas/data_indices.png new file mode 100644 index 0000000..0076e90 Binary files /dev/null and b/docs/schemas/data_indices.png differ