graph.yaml

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

graph_create.py

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

def create_verification_graph(save_path: Path):
    # Data Nodes (o48)
    data_nodes = ReducedGaussianGridNodes(
        grid="o48",
        name="data",
        attributes=[
            SphericalAreaWeights(name="area_weight", norm="unit-max")
        ]
    )

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

    # Edge Builders
    # 1. Data -> Hidden (CutOff)
    edge_data_hidden = CutOffEdges(
        source_name="data",
        target_name="hidden",
        cutoff_factor=0.6,
        attributes=[
            EdgeLength(name="edge_length", norm="unit-std"),
            EdgeDirection(name="edge_dirs", norm="unit-std"),
        ]
    )

    # 2. Hidden -> Hidden (KNN=8)
    edge_hidden_hidden = KNNEdges(
        source_name="hidden",
        target_name="hidden",
        num_nearest_neighbours=8,
        attributes=[
            EdgeLength(name="edge_length", norm="unit-std"),
            EdgeDirection(name="edge_dirs", norm="unit-std"),
        ]
    )

    # 3. Hidden -> Data (KNN=3)
    edge_hidden_data = KNNEdges(
        source_name="hidden",
        target_name="data",
        num_nearest_neighbours=3,
        attributes=[
            EdgeLength(name="edge_length", norm="unit-std"),
        ]
    )

    creator = GraphCreator(
        nodes=[data_nodes, hidden_nodes],
        edges=[edge_data_hidden, edge_hidden_hidden, edge_hidden_data]
    )

    graph = creator.create(save_path=save_path, overwrite=True)
    return graph