docs: lam documentation WIP

docs/graphs/node_attributes.rst

@@ -1,4 +1,4 @@
-.. _node-attributes:
+.. _graphs-node_attributes:
 
 ####################
  Nodes - Attributes

docs/graphs/node_coordinates.rst

@@ -1,4 +1,4 @@
-.. _node-coordinates:
+.. _graphs-node_coordinates:
 
 #####################
  Nodes - Coordinates

docs/graphs/node_coordinates/hex_refined_icosahedron.rst

@@ -10,6 +10,13 @@ hexagons.
 To define the `node coordinates` based on the hexagonal refinements of
 an icosahedron, you can use the following YAML configuration:
 
+***************
+ Global graphs
+***************
+
+The class `HexNodes` allows us to define the nodes over the entire
+globe.
+
 .. code:: yaml
 
    nodes:
@@ -19,7 +26,29 @@ an icosahedron, you can use the following YAML configuration:
          resolution: 4
        attributes: ...
 
-where resolution is the number of refinements to be applied.
+where `resolution` is the number of refinements to be applied.
+
+*********************
+ Limited Area graphs
+*********************
+
+The class `LimitedAreaHexNodes` allows us to define the nodes only for a
+specific area of interest.
+
+.. code:: yaml
+
+   nodes:
+     data:
+       node_builder:
+         _target_: anemoi.graphs.nodes.LimitedAreaHexNodes
+         resolution: 4
+         reference_node_name: nodes_name
+         mask_attr_name: mask_name  # optional
+         margin_radius_km: 100  # optional
+       attributes: ...
+
+where `reference_node_name` is the name of the nodes to define the area
+of interest.
 
 .. csv-table:: Hexagonal Hierarchical refinements specifications
    :file: ./hex_refined.csv

docs/graphs/node_coordinates/tri_refined_icosahedron.rst

@@ -6,7 +6,13 @@ This class allows us to define nodes based on iterative refinements of
 an icoshaedron with triangles.
 
 To define the `node coordinates` based on icosahedral refinements of an
-icosahedron, you can use the following YAML configuration:
+icosahedron, you can use the following YAML configurations:
+
+***************
+ Global graphs
+***************
+
+The class `TriNodes` allows us to define the nodes over the entire globe
 
 .. code:: yaml
 
@@ -17,9 +23,55 @@ icosahedron, you can use the following YAML configuration:
          resolution: 4
        attributes: ...
 
-where resolution is the number of refinements to be applied to the
+where `resolution` is the number of refinements to be applied to the
 icosahedron.
 
+*********************
+ Limited Area graphs
+*********************
+
+The class `LimitedAreaTriNodes` allows us to define the nodes only for a
+specific area of interest.
+
+.. code:: yaml
+
+   nodes:
+     data:
+       node_builder:
+         _target_: anemoi.graphs.nodes.LimitedAreaHexNodes
+         resolution: 4
+         reference_node_name: nodes_name
+         mask_attr_name: mask_name  # optional
+         margin_radius_km: 100  # optional
+       attributes: ...
+
+where `reference_node_name` is the name of the nodes to define the area
+of interest.
+
+*****************
+ Stretched graph
+*****************
+
+The class `StretchedTriNodes` allows us to define the nodes with a
+different resolution for inside and outside the area of interest.
+
+.. code:: yaml
+
+   nodes:
+     data:
+       node_builder:
+         _target_: anemoi.graphs.nodes.StretchedTriNodes
+         global_resolution: 3
+         lam_resolution: 5
+         reference_node_name: nodes_name
+         mask_attr_name: mask_name  # optional
+         margin_radius_km: 100  # optional
+       attributes: ...
+
+where `resolution` argument is dropped divided into `global_resolution`
+and `lam_resolution`, which are the number of refinements to be applied
+to the icosahedron outside and inside the area of interest respectively.
+
 .. warning::
 
    This class will require the `trimesh <https://trimesh.org>`_ package

docs/index.rst

@@ -39,6 +39,7 @@ of the *Anemoi* packages.
 -  :doc:`graphs/node_attributes`
 -  :doc:`graphs/edges`
 -  :doc:`graphs/edge_attributes`
+-  :doc:`graphs/masks`
 
 .. toctree::
    :maxdepth: 1
@@ -50,6 +51,7 @@ of the *Anemoi* packages.
    graphs/node_attributes
    graphs/edges
    graphs/edge_attributes
+   graphs/masks
 
 *********
  Modules
@@ -111,13 +113,15 @@ of the *Anemoi* packages.
 ***********
 
 -  :doc:`usage/getting_started`
+-  :doc:`usage/limited_area`
 
 .. toctree::
    :maxdepth: 1
    :hidden:
    :caption: Usage
 
    usage/getting_started
+   usage/limited_area
 
 *****************
  Anemoi packages

docs/modules/masks.rst

@@ -0,0 +1,10 @@
+.. _modules-masks:
+
+#######
+ Masks
+#######
+
+.. automodule:: anemoi.graphs.generate.masks
+   :members:
+   :no-undoc-members:
+   :show-inheritance:

docs/usage/getting_started.rst

@@ -4,17 +4,19 @@
  Getting started
 #################
 
+The simplest use case is to build an encoder-processor-decoder graph for
+a global weather model.
+
 **************
  First recipe
 **************
 
-The simplest use case is to build an encoder-processor-decoder graph for
-a global weather model. In this case, the recipe must contain a
-``nodes`` section where the keys will be the names of the sets of
-`nodes`, that will later be used to build the connections. Each `nodes`
-configuration must include a ``node_builder`` section describing how to
-generate the `nodes`, and it may include an optional ``attributes``
-section to define additional attributes (weights, mask, ...).
+In this case, the recipe must contain a ``nodes`` section where the keys
+will be the names of the sets of `nodes`, that will later be used to
+build the connections. Each `nodes` configuration must include a
+``node_builder`` section describing how to generate the `nodes`, and it
+may include an optional ``attributes`` section to define additional
+attributes (weights, mask, ...).
 
 .. literalinclude:: yaml/nodes.yaml
    :language: yaml

docs/usage/limited_area.rst

@@ -0,0 +1,77 @@
+.. _usage-limited_area:
+
+#############################
+ Limited Area Modeling (LAM)
+#############################
+
+Anemoi Graphs brings another level of flexibility to the user by
+allowing the definition of limited area graphs.
+
+************************************************
+ Cut out regional dataset into a global dataset
+************************************************
+
+In this case, the user can use 2 datasets to build the data nodes
+combining nodes from the LAM dataset and the global dataset (as boundary
+forcings). The class `CutOutZarrDatasetNodes` allows this functionality:
+
+.. literalinclude:: yaml/cutout_zarr.yaml
+   :language: yaml
+
+The `CutOutZarrDatasetNodes` supports an optional `thinning` argument
+which can be used to sampling points from the regional dataset to reduce
+computation during development stage.
+
+In addition, this node builder class will create an additional node
+attribute with a mask showing which node correspond to each of the 2
+datasets.
+
+.. code:: console
+
+   >>> graph
+   HeteroData(
+      data={
+         x=[40320, 2],
+         node_type='ZarrDatasetNodes',
+         area_weight=[40320, 1],
+         cutout=[40320, 1],
+      }
+   )
+
+*********************************************
+ Define hidden nodes over region of interest
+*********************************************
+
+Once the ``data`` nodes are defined, the user can define the hidden
+nodes only over the region of interest. In this case, the area of
+interest is defined by the ``data`` nodes masked by the ``cutout``
+attribute.
+
+.. literalinclude:: yaml/limited_area_nodes.yaml
+   :language: yaml
+
+.. code:: console
+
+   >>> graph
+   HeteroData(
+      data={
+         x=[40320, 2],
+         node_type='ZarrDatasetNodes',
+         area_weight=[40320, 1],
+         cutout=[40320, 1],
+      },
+      hidden={
+         x=[10242, 2],
+         node_type='TriNodes',
+      }
+   )
+
+**************
+ Adding edges
+**************
+
+The user may define the edges using the same configuration as for the
+global graphs.
+
+.. literalinclude:: yaml/global.yaml
+   :language: yaml

docs/usage/yaml/cutout_zarr.yaml

@@ -0,0 +1,11 @@
+nodes:
+  data:
+    node_builder:
+      _target_: anemoi.graphs.nodes.CutOutZarrDatasetNodes
+      lam_dataset: /path/to/regional-dataset.zarr
+      forcing_dataset: /path/to/global-dataset.zarr
+      thinning: 25  # optional
+    attributes: ...
+  hidden: ...
+
+edges: ...

docs/usage/yaml/global.yaml

@@ -13,8 +13,8 @@ edges:
   - source_name: hidden
     target_name: hidden
     edge_builder:
-      _target_: anemoi.graphs.edges.KNNEdges
-      nearest_neighbours: 3
+      _target_: anemoi.graphs.edges.MultiScaleEdges
+      x_hops: 1
  # C) Processor connections
   - source_name: hidden
     target_name: data

docs/usage/yaml/limited_area_nodes.yaml

@@ -0,0 +1,10 @@
+nodes:
+  data: ...
+  hidden:
+    node_builder:
+      _target_: anemoi.graphs.nodes.LimitedAreaTriNodes
+      resolution: 5
+      reference_node_name: data
+      mask_attr_name: cutout
+
+edges: ...

docs/usage/yaml/nodes.yaml

@@ -5,6 +5,5 @@ nodes:
       dataset: /path/to/dataset.zarr
   hidden:
     node_builder:
-      _target_: anemoi.graphs.nodes.NPZFileNodes
-      grid_definition_path: /path/to/grids/
-      resolution: o48
+      _target_: anemoi.graphs.nodes.TriNodes
+      resolution: 5