Support masking of unconnected nodes (LAM)

[JPXKQX]

This PR adds support for using masked graphs, resulting in a performance improvement in LAM by dropping unused nodes in the boundary forcing dataset.

Description

The current release of anemoi-graphs (0.4.1) supports dropping unconnected nodes from the limited area graph, with the following code:

post_processors:
  - _target_: anemoi.graphs.processors.RemoveUnconnectedNodes
    nodes_name: data
    ignore: cutout_mask # optional
    save_mask_indices_to_attr: indices_connected_nodes # optional

anemoi-graphs returns the masked graph with an additional node attribute indices_connected_nodes (in this example) representing the indices used to mask the input dataset to the new graph data nodes.

Solution

We need to specify these spatial indices in the dataloader config.

grid_indices:
    _target_: anemoi.training.data.grid_indices.MaskedGrid
    nodes_name: data
    node_attribute_name: indices_connected_nodes

The performance improvement is of 66 % (from 1,92 it/s to 3,19 it/s) for the 1 GPU default limited area model.

This PR has been tested successfully on 1 GPU (with plot: detailed without PlotSpectrum included):

• Limited Area Model with reader_group_size=1
• Stretched grid with reader_group_size=1
• Global model with reader_group_size=1

This PR has been tested successfully on 2 GPU (with plot: detailed without PlotSpectrum included):

• Limited Area Model with reader_group_size=2
• Stretched grid with reader_group_size=2
• Global model with reader_group_size=2

W.I.P.: This mask must be saved to the checkpoint so that it can be used in anemoi-inference.

[theissenhelen]

Just a typo to fix.

[theissenhelen]

Great work!