Edge builder

class anemoi.graphs.edges.builder.KNNEdges(source_name: str, target_name: str, num_nearest_neighbours: int, source_mask_attr_name: str | None = None, target_mask_attr_name: str | None = None)

Bases: BaseEdgeBuilder, NodeMaskingMixin

Computes KNN based edges and adds them to the graph.

source_name

The name of the source nodes.

Type:: str

target_name

The name of the target nodes.

Type:: str

num_nearest_neighbours

Number of nearest neighbours.

Type:: int

source_mask_attr_name

The name of the source mask attribute to filter edge connections.

Type:: str | None

target_mask_attr_name

The name of the target mask attribute to filter edge connections.

Type:: str | None

register_edges(graph): Register the edges in the graph.

register_attributes(graph, config): Register attributes in the edges of the graph.

update_graph(graph, attrs_config): Update the graph with the edges.

get_adjacency_matrix(source_nodes: NodeStorage, target_nodes: NodeStorage) → ndarray

Compute the adjacency matrix for the KNN method.

Parameters:

source_nodes (NodeStorage) – The source nodes.
target_nodes (NodeStorage) – The target nodes.

Returns:

The adjacency matrix.

Return type:

np.ndarray

class anemoi.graphs.edges.builder.CutOffEdges(source_name: str, target_name: str, cutoff_factor: float, source_mask_attr_name: str | None = None, target_mask_attr_name: str | None = None)

Bases: BaseEdgeBuilder, NodeMaskingMixin

Computes cut-off based edges and adds them to the graph.

source_name

The name of the source nodes.

Type:: str

target_name

The name of the target nodes.

Type:: str

cutoff_factor

Factor to multiply the grid reference distance to get the cut-off radius.

Type:: float

source_mask_attr_name

The name of the source mask attribute to filter edge connections.

Type:: str | None

target_mask_attr_name

The name of the target mask attribute to filter edge connections.

Type:: str | None

register_edges(graph): Register the edges in the graph.

register_attributes(graph, config): Register attributes in the edges of the graph.

update_graph(graph, attrs_config): Update the graph with the edges.

get_cutoff_radius(graph: HeteroData, mask_attr: Tensor | None = None) → float

Compute the cut-off radius.

The cut-off radius is computed as the product of the target nodes reference distance and the cut-off factor.

Parameters:

graph (HeteroData) – The graph.
mask_attr (torch.Tensor) – The mask attribute.

Returns:

The cut-off radius.

Return type:

float

prepare_node_data(graph: HeteroData) → tuple[NodeStorage, NodeStorage]: Prepare node information and get source and target nodes.

get_adjacency_matrix(source_nodes: NodeStorage, target_nodes: NodeStorage) → ndarray

Get the adjacency matrix for the cut-off method.

Parameters:

source_nodes (NodeStorage) – The source nodes.
target_nodes (NodeStorage) – The target nodes.

Returns:

The adjacency matrix.

Return type:

np.ndarray

class anemoi.graphs.edges.builder.MultiScaleEdges(source_name: str, target_name: str, x_hops: int, **kwargs)

Bases: BaseEdgeBuilder

Base class for multi-scale edges in the nodes of a graph.

source_name

The name of the source nodes.

Type:: str

target_name

The name of the target nodes.

Type:: str

x_hops

Number of hops (in the refined icosahedron) between two nodes to connect them with an edge.

Type:: int

register_edges(graph): Register the edges in the graph.

register_attributes(graph, config): Register attributes in the edges of the graph.

update_graph(graph, attrs_config): Update the graph with the edges.

update_graph(graph: HeteroData, attrs_config: DotDict | None = None) → HeteroData

Update the graph with the edges.

Parameters:

graph (HeteroData) – The graph.
attrs_config (DotDict) – The configuration of the edge attributes.

Returns:

The graph with the edges.

Return type:

HeteroData

class anemoi.graphs.edges.builder.ICONTopologicalBaseEdgeBuilder(source_name: str, target_name: str, icon_mesh: str, source_mask_attr_name: str | None = None, target_mask_attr_name: str | None = None)

Bases: BaseEdgeBuilder

Base class for computing edges based on ICON grid topology.

source_name

The name of the source nodes.

Type:: str

target_name

The name of the target nodes.

Type:: str

icon_mesh

The name of the ICON mesh (defines both the processor mesh and the data)

Type:: str

update_graph(graph: HeteroData, attrs_config: DotDict = None) → HeteroData: Update the graph with the edges.

get_adjacency_matrix(source_nodes: NodeStorage, target_nodes: NodeStorage)

Parameters:

source_nodes (NodeStorage) – The source nodes.
target_nodes (NodeStorage) – The target nodes.

class anemoi.graphs.edges.builder.ICONTopologicalProcessorEdges(source_name: str, target_name: str, icon_mesh: str, source_mask_attr_name: str | None = None, target_mask_attr_name: str | None = None)

Bases: ICONTopologicalBaseEdgeBuilder

Computes edges based on ICON grid topology: processor grid built from ICON grid vertices.

class anemoi.graphs.edges.builder.ICONTopologicalEncoderEdges(source_name: str, target_name: str, icon_mesh: str, source_mask_attr_name: str | None = None, target_mask_attr_name: str | None = None)

Bases: ICONTopologicalBaseEdgeBuilder

Computes encoder edges based on ICON grid topology: ICON cell circumcenters for mapped onto processor grid built from ICON grid vertices.

class anemoi.graphs.edges.builder.ICONTopologicalDecoderEdges(source_name: str, target_name: str, icon_mesh: str, source_mask_attr_name: str | None = None, target_mask_attr_name: str | None = None)

Bases: ICONTopologicalBaseEdgeBuilder

Computes encoder edges based on ICON grid topology: mapping from processor grid built from ICON grid vertices onto ICON cell circumcenters.