Edge attributes

class anemoi.graphs.edges.attributes.BaseEdgeAttributeBuilder(norm: str | None = None, dtype: str = 'float32')

Bases: MessagePassing, NormaliserMixin, ABC

Base class for edge attribute builders.

forward(x: tuple[NodeStorage, NodeStorage], edge_index: Tensor | SparseTensor, size: Tuple[int, int] | None = None) → Tensor

Runs the forward pass of the module.

message(x_i: Tensor, x_j: Tensor) → Tensor

Constructs messages from node \(j\) to node \(i\) in analogy to \(\phi_{\mathbf{\Theta}}\) for each edge in edge_index. This function can take any argument as input which was initially passed to propagate(). Furthermore, tensors passed to propagate() can be mapped to the respective nodes \(i\) and \(j\) by appending _i or _j to the variable name, .e.g. x_i and x_j.

aggregate(edge_features: Tensor) → Tensor

Aggregates messages from neighbors as \(\bigoplus_{j \in \mathcal{N}(i)}\).

Takes in the output of message computation as first argument and any argument which was initially passed to propagate().

By default, this function will delegate its call to the underlying Aggregation module to reduce messages as specified in __init__() by the aggr argument.

class anemoi.graphs.edges.attributes.BasePositionalBuilder(norm: str | None = None, dtype: str = 'float32')

Bases: BaseEdgeAttributeBuilder, ABC

class anemoi.graphs.edges.attributes.EdgeLength(norm: str | None = None, dtype: str = 'float32')

Bases: BasePositionalBuilder

Computes edge length for bipartite graphs.

class anemoi.graphs.edges.attributes.EdgeDirection(norm: str | None = None, dtype: str = 'float32')

Bases: BasePositionalBuilder

Computes edge direction for bipartite graphs.

class anemoi.graphs.edges.attributes.Azimuth(norm: str | None = None, dtype: str = 'float32')

Bases: BasePositionalBuilder

Compute the azimuth of the edge.

norm

Normalisation method. Options: None, “l1”, “l2”, “unit-max”, “unit-range”, “unit-std”.

Type:

Optional[str]

invert

Whether to invert the edge lengths, i.e. 1 - edge_length. Defaults to False.

Type:

bool

compute(x_i, x_j)

Compute edge lengths attributes.

References

• https://www.movable-type.co.uk/scripts/latlong.html

class anemoi.graphs.edges.attributes.BaseBooleanEdgeAttributeBuilder

Bases: BaseEdgeAttributeBuilder, ABC

Base class for boolean edge attributes.

class anemoi.graphs.edges.attributes.BaseEdgeAttributeFromNodeBuilder(node_attr_name: str)

Bases: BaseBooleanEdgeAttributeBuilder, ABC

Base class for propagating an attribute from the nodes to the edges.

class anemoi.graphs.edges.attributes.AttributeFromSourceNode(node_attr_name: str)

Bases: BaseEdgeAttributeFromNodeBuilder

Copy an attribute of the source node to the edge. Used for example to identify if an encoder edge originates from a LAM or global node.

node_attr_name

Name of the node attribute to propagate.

Type:

str

class anemoi.graphs.edges.attributes.AttributeFromTargetNode(node_attr_name: str)

Bases: BaseEdgeAttributeFromNodeBuilder

Copy an attribute of the target node to the edge.

Used for example to identify if an encoder edge ends at a LAM or global node.

node_attr_name

Name of the node attribute to propagate.

Type:

str