Post processor

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

Bases: BaseMaskingProcessor

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.