Graph Creator

This module is used to create custom graphs for data-driven weather models. The graphs are built using a recipe that defines the structure of the graph.

class anemoi.graphs.create.GraphCreator(config: str | Path | DotDict | DictConfig)

Bases: object

Graph creator.

update_graph(graph: HeteroData) → HeteroData

Update the graph.

It instantiates the node builders and edge builders defined in the configuration file and applies them to the graph.

Parameters:

graph (HeteroData) – The input graph to be updated.

Returns:

The updated graph with new nodes and edges added based on the configuration.

Return type:

HeteroData

clean(graph: HeteroData) → HeteroData

Remove private attributes used during creation from the graph.

Parameters:

graph (HeteroData) – Generated graph

Returns:

Cleaned graph

Return type:

HeteroData

post_process(graph: HeteroData) → HeteroData

Allow post-processing of the resulting graph.

This method applies any configured post-processors to the graph, which can modify or enhance the graph structure or attributes.

Parameters:

graph (HeteroData) – The graph to be post-processed.

Returns:

The post-processed graph.

Return type:

HeteroData

Notes

Post-processors are applied in the order they are specified in the configuration. Each post-processor should implement an update_graph method that takes and returns a HeteroData object.

save(graph: HeteroData, save_path: Path, overwrite: bool = False) → None

Save the generated graph to the output path.

Parameters:

  • graph (HeteroData) – generated graph

  • save_path (Path) – location to save the graph

  • overwrite (bool, optional) – whether to overwrite existing graph file, by default False

create(save_path: Path | None = None, overwrite: bool = False) → HeteroData

Create the graph and save it to the output path.

Parameters:

  • save_path (Path, optional) – location to save the graph, by default None

  • overwrite (bool, optional) – whether to overwrite existing graph file, by default False

Returns:

created graph object

Return type:

HeteroData