Use anemoi-graphs HeteroData

[JPXKQX]

This code implements the use of the new HeteroData object generated with anemoi-graphs package.

Checklist before requesting a review

• I have performed a self-review of my code
• My code follows the style guidelines of this project
• I have commented my code, particularly in hard-to-understand areas
• I have updated the documentation and docstrings to reflect the changes
• I have added tests that prove my fix is effective or that my feature works
• I have ensured that the code is still pip-installable after the changes and runs
• I have ran this on single GPU
• I have ran this on multi-GPU or multi-node

Issue ticket number and link

This work would not have been possible without @theissenhelen and @JesperDramsch and the rest of the AIFS team!

[floriankrb]

From a style point of view, these changes definitely make the code easier to read.

[JesperDramsch]

Thanks for the PR.

This seems to modify parts of the edge registration on mapper components, but doesn't modify edge registration on processors at all. Were these tested? I doubt they would work if the mappers needed a change.

[JPXKQX]

Thanks for the PR.

This seems to modify parts of the edge registration on mapper components, but doesn't modify edge registration on processors at all. Were these tested? I doubt they would work if the mappers needed a change.

Yes, edge registration is modified in the GraphEdgeMixin class defined in mapper.py but used in both mappers and processors.

[JesperDramsch]

This re-introduces graph instantiation to non-graph components to define "which graph attributes should be defined" globally, when really this should not be the case.

[mchantry]

Can you tag anemoi-graphs as 0.2 and then specify anemoi-graphs>=0.2 in the toml please.

[JPXKQX]

Can you tag anemoi-graphs as 0.2 and then specify anemoi-graphs>=0.2 in the toml please.

anemoi-graphs is not a direct dependency of anemoi-model.s

[mchantry]

LGTM, great work, please merge.

Requested changes already addressed by Jesper