refactor(graphs): Expose Python API

[radiradev]

Description

This PR introduces an alternative way to construct graphs by instantiating objects directly, in addition to the existing YAML-based configuration approach.

Creating graphs via direct object instantiation improves transparency: docstrings are visible, IDE enables us to “go to definition”, and objects can be inspected directly during development. Furthermore this way we can utilise static type checking, reducing the need for Pydantic schema validation.

Sidenote: The GraphCreator could already be instantiated directly on main in Python, but this requires to pass in a DotDict which is basically a parsed config. This PR instead explores using the edge, node and attribute objects directly.

This is my first attempt as part of the broader plan in #768 and would serve as a test-bed for ideas and general feedback to a more pythonic anemoi.

Example graph with the new Python API:

from anemoi.graphs.create import GraphBuilder
from anemoi.graphs.nodes.builders.from_reduced_gaussian import ReducedGaussianGridNodes
from anemoi.graphs.nodes.attributes.area_weights import SphericalAreaWeights
from anemoi.graphs.edges.attributes import EdgeLength, EdgeDirection
from anemoi.graphs.edges import CutOffEdges, KNNEdges

def create_graph():
    data_nodes = ReducedGaussianGridNodes(
        grid="o48",
        name="data",
        attributes=[
            SphericalAreaWeights(norm="unit-max")
        ]
    )

    hidden_nodes = ReducedGaussianGridNodes(
        grid="o48",
        name="hidden",
        attributes=[
            SphericalAreaWeights(norm="unit-max")
        ]
    )

    edge_data_hidden = CutOffEdges(
        source_name="data",
        target_name="hidden",
        cutoff_factor=0.6,
        attributes=[
            EdgeLength(norm="unit-std"),
            EdgeDirection(norm="unit-std"),
        ]
    )

    edge_hidden_hidden = KNNEdges(
        source_name="hidden",
        target_name="hidden",
        num_nearest_neighbours=8,
        attributes=[
            EdgeLength(norm="unit-std"),
            EdgeDirection(norm="unit-std"),
        ]
    )

    edge_hidden_data = KNNEdges(
        source_name="hidden",
        target_name="data",
        num_nearest_neighbours=3,
        attributes=[
            EdgeLength(norm="unit-std"),
        ]
    )

    builder = GraphBuilder(
        nodes=[data_nodes, hidden_nodes],
        edges=[edge_data_hidden, edge_hidden_hidden, edge_hidden_data]
    )
    return builder.create()

YAML equivalent

---
nodes:
  data:
    node_builder:
      _target_: anemoi.graphs.nodes.ReducedGaussianGridNodes
      grid: o48
    attributes:
      area_weight:
        _target_: anemoi.graphs.nodes.attributes.SphericalAreaWeights
        norm: unit-max
  hidden:
    node_builder:
      _target_: anemoi.graphs.nodes.ReducedGaussianGridNodes
      grid: o48
    attributes:
      area_weight:
        _target_: anemoi.graphs.nodes.attributes.SphericalAreaWeights
        norm: unit-max

edges:
- source_name: data
  target_name: hidden
  edge_builders:
  - _target_: anemoi.graphs.edges.CutOffEdges
    cutoff_factor: 0.6
  attributes:
    edge_length:
      _target_: anemoi.graphs.edges.attributes.EdgeLength
      norm: unit-std
    edge_dirs:
      _target_: anemoi.graphs.edges.attributes.EdgeDirection
      norm: unit-std

- source_name: hidden
  target_name: hidden
  edge_builders:
  - _target_: anemoi.graphs.edges.KNNEdges
    num_nearest_neighbours: 8
  attributes:
    edge_length:
      _target_: anemoi.graphs.edges.attributes.EdgeLength
      norm: unit-std
    edge_dirs:
      _target_: anemoi.graphs.edges.attributes.EdgeDirection
      norm: unit-std

- 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
      norm: unit-std

What problem does this change solve?

This an experimental new feature that is fully backwards compatible with the YAML configuration

What issue or task does this change relate to?

#768

Additional notes

As a contributor to the Anemoi framework, please ensure that your changes include unit tests, updates to any affected dependencies and documentation, and have been tested in a parallel setting (i.e., with multiple GPUs). As a reviewer, you are also responsible for verifying these aspects and requesting changes if they are not adequately addressed. For guidelines about those please refer to https://anemoi.readthedocs.io/en/latest/

By opening this pull request, I affirm that all authors agree to the Contributor License Agreement.

[radiradev]

• Convert yaml configs and pass them directly to GraphCreator.
• Serialize Python graph creation into yaml

[anaprietonem]

If this is being removed, do we expect the yaml generated as a result of the API generation to work? I think a test checking that would be ideal but just wondering if a first step, you had tried to generate the graph with the config and that still works, and same for training?

[radiradev]

There is currently no YAML generated from the GraphBuilder (i. e the pure python API - ideas for a better name welcome). GraphCreator can still be instantiated with a yaml config as before. I have checked locally that the graphs created by the two are the same on a few cases but I agree I should add some tests to the parsers and to check that graphs are the same.

[anaprietonem]

@radiradev - current configs for graphs have interconnections with for example the dataloader section https://github.com/ecmwf/anemoi-core/blob/main/training/src/anemoi/training/config/graph/multi_scale.yaml

  data:
    node_builder:
      _target_: anemoi.graphs.nodes.AnemoiDatasetNodes # options: AnemoiDatasetNodes, NPZFileNodes
      dataset: ${dataloader.dataset}
    attributes: ${graph.attributes.nodes}

how would this be handled by the API? - I think this a useful feature but one that I think during the working group session in December we also touched on the fact that it brings some additional complexity, so not flagging this a must that we would need to keep, but rather curious on the interplay of the two features.

[radiradev]

@radiradev - current configs for graphs have interconnections with for example the dataloader section https://github.com/ecmwf/anemoi-core/blob/main/training/src/anemoi/training/config/graph/multi_scale.yaml

  data:
    node_builder:
      _target_: anemoi.graphs.nodes.AnemoiDatasetNodes # options: AnemoiDatasetNodes, NPZFileNodes
      dataset: ${dataloader.dataset}
    attributes: ${graph.attributes.nodes}

how would this be handled by the API? - I think this a useful feature but one that I think during the working group session in December we also touched on the fact that it brings some additional complexity, so not flagging this a must that we would need to keep, but rather curious on the interplay of the two features.

@anaprietonem currently the dataset: ... remains unchanged and expects a DictConfig. It can be changed to pass in directly the dataset or DictConfig to remain compatible with the rest of anemoi-training and still have a completely config-free option?

[radiradev]

fiddle has the option to serialize into YAML directly, here is what this looks like for the example above:

!fdl.Config
__fn_or_cls__:
  module: anemoi.graphs.create
  name: GraphBuilder
edges:
- !fdl.Config
  __fn_or_cls__:
    module: anemoi.graphs.edges.builders.cutoff
    name: CutOffEdges
  attributes:
  - !fdl.Config
    __fn_or_cls__:
      module: anemoi.graphs.edges.attributes
      name: EdgeLength
    norm: unit-std
  - !fdl.Config
    __fn_or_cls__:
      module: anemoi.graphs.edges.attributes
      name: EdgeDirection
    norm: unit-std
  cutoff_factor: 0.6
  source_name: data
  target_name: hidden
...

[leifdenby]

I've made some suggestions to your PR in radiradev#3, hopefully they are useful!