Create your first graph

The simplest use case is to build an encoder-processor-decoder graph for a global weather model.

Encoder-processor-decoder graph

The following is an expanded version of the examples given in the graphs introduction.

Nodes and edges

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, …).

nodes:
  data:
    node_builder:
      _target_: anemoi.graphs.nodes.ZarrDatasetNodes
      dataset: /path/to/dataset.zarr
  hidden:
    node_builder:
      _target_: anemoi.graphs.nodes.TriNodes
      resolution: 5

Once the nodes have been defined, you need to create the edges between them through which information will flow. To this aim, the recipe file must contain a edges section. These connections are defined between pairs of nodes (source and target, specified by source_name and target_name).

There are several methods to build these edges such as cutoff (CutOffEdges) or nearest neighbours (KNNEdges). For an encoder-processor-decoder graph you will need to build two sets of edges. The first set of edges will connect the data nodes with the hidden nodes to encode the input data into the latent space, normally referred to as the encoder edges and represented here by the first element of the edges section. The second set of edges will connect the hidden nodes with the data nodes to decode the latent space into the output data, normally referred to as decoder edges and represented here by the second element of the edges section.

edges:
  # A) Encoder connections
  - source_name: data
    target_name: hidden
    edge_builders:
    - _target_: anemoi.graphs.edges.CutOffEdges
      cutoff_factor: 0.7
  # B) Decoder connections
  - source_name: hidden
    target_name: data
    edge_builders:
    - _target_: anemoi.graphs.edges.KNNEdges
      num_nearest_neighbours: 3
Schema of global graph (without processor connections)

To create the graph, run the following command:

$ anemoi-graphs create recipe.yaml graph.pt

Once the build is complete, you can inspect the dataset using the following command:

$ anemoi-graphs inspect graph.pt

This will generate the following graph:

┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈
📦 Path          : graph.pt
🔢 Format version: 0.0.1

💽 Size       : 3.1 MiB (3,283,650)

   Nodes name  │ Num. nodes | Attribute dim | Min. latitude | Max. latitude | Min. longitude | Max. longitude
  ─────────────┼────────────┼───────────────┼───────────────┼───────────────┼────────────────┼────────────────
    data       |     10,840 |             4 |        -3.135 |         3.140 |           0.02 |          6.13
    hidden     |      6,200 |             4 |        -3.141 |         3.137 |           0.01 |          6.14
  ─────────────┴────────────┴───────────────┴───────────────┴───────────────┴────────────────┴────────────────


   Source      │ Destination  │ Num. edges  │ Attribute dim | Min. length │ Max. length │ Mean length │ Std dev
  ─────────────┼──────────────┼─────────────┼───────────────┼─────────────┼─────────────┼─────────────┼─────────
    data       │ hidden       │       13508 │             1 |      0.3116 │       25.79 │   11.059531 │   5.5856
    hidden     │ data         │       40910 │             1 |      0.2397 │      21.851 │   12.270924 │   4.2347
  ─────────────┴──────────────┴─────────────┴───────────────|─────────────┴─────────────┴─────────────┴─────────
🔋 Graph ready, last update 17 seconds ago.
📊 Statistics ready.

Note

Note that that the resulting graph will only work with a Transformer processor because there are no connections between the hidden nodes. This is the default behaviour for anemoi-training.

Adding processor connections

To add connections within the hidden nodes, to be used in the processor, you need to add a new set of edges to the recipe file. These connections are normally referred to as processor edges and are represented here by the third element of the edges section.

nodes:
  data: ...
  hidden: ...

edges:
  # A) Encoder connections
  - source_name: data
    target_name: hidden
    edge_builders:
    - _target_: anemoi.graphs.edges.CutOffEdges
      cutoff_factor: 0.7
  # B) Decoder connections
  - source_name: hidden
    target_name: data
    target_mask_attr_name: cutout
    edge_builders:
    - _target_: anemoi.graphs.edges.KNNEdges
      num_nearest_neighbours: 3
 # C) Processor connections
  - source_name: hidden
    target_name: hidden
    edge_builders:
    - _target_: anemoi.graphs.edges.MultiScaleEdges
      x_hops: 1
Schema of global graph

This will generate the following graph:

┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈
📦 Path          : graph.pt
🔢 Format version: 0.0.1

💽 Size       : 3.1 MiB (3,283,650)

   Nodes name  │ Num. nodes | Attribute dim | Min. latitude | Max. latitude | Min. longitude | Max. longitude
  ─────────────┼────────────┼───────────────┼───────────────┼───────────────┼────────────────┼────────────────
    data       |     10,840 |             0 |        -3.135 |         3.140 |           0.02 |          6.13
    hidden     |      6,200 |             0 |        -3.141 |         3.137 |           0.01 |          6.14
  ─────────────┴────────────┴───────────────┴───────────────┴───────────────┴────────────────┴────────────────


   Source      │ Destination  │ Num. edges  │ Attribute dim | Min. length │ Max. length │ Mean length │ Std dev
  ─────────────┼──────────────┼─────────────┼───────────────┼─────────────┼─────────────┼─────────────┼─────────
    data       │ hidden       │       13508 │             1 |      0.3116 │       25.79 │   11.059531 │   5.5856
    hidden     │ data         │       40910 │             1 |      0.2397 │      21.851 │   12.270924 │   4.2347
    hidden     │ hidden       │       32010 │             1 |      0.2397 │      21.851 │   10.270924 │   4.2347
  ─────────────┴──────────────┴─────────────┴───────────────|─────────────┴─────────────┴─────────────┴─────────
🔋 Graph ready, last update 7 seconds ago.
📊 Statistics ready.

Adding attributes

When training a data-driven weather model, it is common to add attributes to the nodes or edges. For example, you may want to add node attributes to weight the loss function, or add edge attributes to represent the direction of the edges.

To add attributes to the nodes, you must include the attributes section in the nodes configuration. The attributes can be defined as a list of dictionaries, where each dictionary contains the name of the attribute and the type of the attribute.

nodes:
  data:
    node_builder:
      _target_: anemoi.graphs.nodes.ZarrDatasetNodes
      dataset: /path/to/dataset.zarr
    attributes:
      weights:
        _target_: anemoi.graphs.nodes.attributes.AreaWeights
        norm: unit-max
  hidden: ...

To add the extra features to the edges of the graph, you need to set them in the attributes section.

nodes:
  data: ...
  hidden: ...

edges:
  # A) Encoder connections
  - source_name: data
    target_name: hidden
    edge_builders:
    - _target_: anemoi.graphs.edges.CutOffEdges
      cutoff_factor: 0.7
    attributes:
      edge_length:
        _target_: anemoi.graphs.edges.attributes.EdgeLength
  # B) Decoder connections
  - source_name: hidden
    target_name: data
    edge_builders:
    - _target_: anemoi.graphs.edges.KNNEdges
      num_nearest_neighbours: 3
    attributes:
      edge_length:
        _target_: anemoi.graphs.edges.attributes.EdgeLength
 # C) Processor connections
  - source_name: hidden
    target_name: hidden
    edge_builders:
    - _target_: anemoi.graphs.edges.MultiScaleEdges
      x_hops: 1
    attributes:
      edge_length:
        _target_: anemoi.graphs.edges.attributes.EdgeLength

This will generate the following graph:

┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈┈
📦 Path          : graph.pt
🔢 Format version: 0.0.1

💽 Size       : 3.1 MiB (3,283,650)

   Nodes name  │ Num. nodes | Attribute dim | Min. latitude | Max. latitude | Min. longitude | Max. longitude
  ─────────────┼────────────┼───────────────┼───────────────┼───────────────┼────────────────┼────────────────
    data       |     10,840 |             1 |        -3.135 |         3.140 |           0.02 |          6.13
    hidden     |      6,200 |             1 |        -3.141 |         3.137 |           0.01 |          6.14
  ─────────────┴────────────┴───────────────┴───────────────┴───────────────┴────────────────┴────────────────


   Source      │ Destination  │ Num. edges  │ Attribute dim | Min. length │ Max. length │ Mean length │ Std dev
  ─────────────┼──────────────┼─────────────┼───────────────┼─────────────┼─────────────┼─────────────┼─────────
    data       │ hidden       │       13508 │             3 |      0.3116 │       25.79 │   11.059531 │   5.5856
    hidden     │ data         │       40910 │             3 |      0.2397 │      21.851 │   12.270924 │   4.2347
    hidden     │ hidden       │       32010 │             3 |      0.2397 │      21.851 │   10.270924 │   4.2347
  ─────────────┴──────────────┴─────────────┴───────────────|─────────────┴─────────────┴─────────────┴─────────
🔋 Graph ready, last update 7 seconds ago.
📊 Statistics ready.