Post processor

class anemoi.graphs.processors.post_process.BaseNodeMaskingProcessor(nodes_name: str, save_mask_indices_to_attr: str | None = None)

Bases: PostProcessor, ABC

Base class for mask based node processor.

removing_nodes(graph: HeteroData) → HeteroData

Remove nodes based on the mask passed.

update_edge_indices(graph: HeteroData) → HeteroData

Update the edge indices to the new position of the nodes.

add_attribute(graph: HeteroData) → HeteroData

Add an attribute of the mask indices as node attribute.

update_graph(graph: HeteroData, **kwargs: Any) → HeteroData

Post-process the graph.

Parameters:

• graph (HeteroData) – The graph to post-process.

• kwargs (Any) – Additional keyword arguments.

Returns:

The post-processed graph.

Return type:

HeteroData

class anemoi.graphs.processors.post_process.RemoveUnconnectedNodes(nodes_name: str, save_mask_indices_to_attr: str | None = None, ignore: str | None = None)

Bases: BaseNodeMaskingProcessor

Remove unconnected nodes in the graph.

nodes_name

Name of the unconnected nodes to remove.

Type:

str

ignore

Name of an attribute to ignore when removing nodes. Nodes with this attribute set to True will not be removed.

Type:

str, optional

save_mask_indices_to_attr

Name of the attribute to save the mask indices. If provided, the indices of the kept nodes will be saved in this attribute.

Type:

str, optional

compute_mask(graph)

Compute the mask of the connected nodes.

prune_graph(graph, mask)

Prune the nodes with the specified mask.

add_attribute(graph, mask)

Add an attribute of the mask indices as node attribute.

update_graph(graph)

Post-process the graph.

compute_mask(graph: HeteroData) → Tensor

Compute the mask of connected nodes.

class anemoi.graphs.processors.post_process.BaseEdgeMaskingProcessor(source_name: str, target_name: str)

Bases: PostProcessor, ABC

Base class for mask based edge processor.

source_name

Name of the source nodes of edges to remove.

Type:

str

target_name

Name of the target nodes of edges to remove.

Type:

str

removing_edges(graph: HeteroData) → HeteroData

Remove edges based on the mask passed.

recompute_attributes(graph: HeteroData, graph_config: dict) → HeteroData

Recompute attributes

update_graph(graph: HeteroData, **kwargs: Any) → HeteroData

Post-process the graph.

Parameters:

• graph (HeteroData) – The graph to post-process.

• kwargs (Any) – Additional keyword arguments.

Returns:

The post-processed graph.

Return type:

HeteroData

class anemoi.graphs.processors.post_process.RestrictEdgeLength(source_name: str, target_name: str, max_length_km: float, source_mask_attr_name: str | None = None, target_mask_attr_name: str | None = None)

Bases: BaseEdgeMaskingProcessor

Remove edges longer than a given treshold from the graph.

source_name

Name of the source nodes of edges to remove.

Type:

str

target_name

Name of the target nodes of edges to remove.

Type:

str

max_length_km

The maximal length of edges not to be removed.

Type:

float

source_mask_attr_name

the postprocessing will be restricted to edges with source node having True in this mask_attr

Type:

str , optional

target_mask_attr_name

the postprocessing will be restricted to edges with target node having True in this mask_attr

Type:

str, optional

compute_mask(graph)

Compute the mask of the relevant edges longer than max_length_km.