Ctrl+K

Design principles

Design principles#

The key design principle of weather-model-graphs is to work with networkx.DiGraph objects as the primary data structure for the graph representation right until the graph is to be stored on disk into a specific format. Using only networkx.DiGraph objects as the intermediate representations makes it possible to

  1. easily modularise the whole generation process, with every step outputting a networkx.DiGraph object,

  2. easily visualise the graph resulting from any step and

  3. easily connect graph nodes across graph components, combine graphs and split graphs based on node and edge attributes.

The graph generation in weather-model-graphs is split into to the following steps:

  1. Create the three graph components of the message-passing graph that constitute the auto-regressive atmospheric flow model, all represented by networkx.DiGraph objects:

    • grid-to-mesh (g2m): the encoding compenent, where edges represent the encoding of physical variables into the latent space of the model

    • mesh-to-mesh (m2m): the processing component, where edges represent information flow between nodes through the time evolution of the atmospheric state

    • mesh-to-grid (m2g): the decoding component, where edges represent the decoding of the latent space back into physical variables

  2. Combine all three graph components into a single networkx.DiGraph object and create a unique node identifier for each node in the combined graph.

  3. Split the combined graph into the three output graph components again (or more if the specific graph architecture requires it).

  4. Store each of the output graph components in the desired format, for example:

    • networkx .pickle file: save networkx.DiGraph objects using pickle to disk (weather_model_graphs.save.to_pickle(...))

    • pytorch-geometric for neural-lam: edges indexes and features are stored in separate torch.Tensor objects serialised to disk that can then be loaded into torch_geometric.data.Data objects (weather_model_graphs.save.to_pyg(...))

Diagram of the graph generation process:#

Below, the graph generation process is visualised in weather-model-graphs for the example given above:

graph TB xy_grid["(x,y) grid coordinates [numpy.ndarray]"] subgraph weather_model_graphs["weather-model-graphs"] subgraph graph_components["Graph components"] G_g2m["G_g2m[networkx.DiGraph]"] G_m2m["G_m2m[networkx.DiGraph]"] G_m2g["G_m2g[networkx.DiGraph]"] end G_g2m --merged into --> G_full G_m2m --merged into --> G_full G_m2g --merged into --> G_full G_full["complete graph\nG_full[networkx.Digraph]"] G_full --replace node labels with unique integer id --> G_full_int["G_int[networkx.Digraph]"] G_full_int --split and converted into --> pyg_g2m["pyg_g2m[pyg.data.Data]"] G_full_int --split and converted into --> pyg_m2m["pyg_m2m[pyg.data.Data]"] G_full_int --split and converted into --> pyg_m2g["pyg_m2g[pyg.data.Data]"] end xy_grid --used to create node coordinates in --> G_g2m xy_grid --used to create node coordinates in --> G_m2g xy_grid --used to create node coordinates in --> G_m2m subgraph stored_on_disk["stored on disk"] subgraph hidden_graph1[ ] pyg_g2m_file["g2m_edge_index.pt\ng2m_features.pt"] pyg_m2m_file["m2m_edge_index.pt\nm2m_features.pt"] pyg_m2g_file["m2g_edge_index.pt\nm2g_features.pt"] end end pyg_g2m --stored in--> pyg_g2m_file pyg_m2m --stored in--> pyg_m2m_file pyg_m2g --stored in--> pyg_m2g_file subgraph pyg_loaded["Loaded into model"] subgraph hidden_graph2["(e.g. torch.nn.Module)"] pyg_g2m_loaded["pyg_g2m[pyg.data.Data]"] pyg_m2m_loaded["pyg_m2m[pyg.data.Data]"] pyg_m2g_loaded["pyg_m2g[pyg.data.Data]"] end end pyg_g2m_file --loaded into --> pyg_g2m_loaded pyg_m2m_file --loaded into --> pyg_m2m_loaded pyg_m2g_file --loaded into --> pyg_m2g_loaded classDef title_left margin-right:250px,display:block classDef hidden stroke-width:0px,margin-right:250px,display:block class hidden_graph1 hidden class hidden_graph2 hidden class weather_model_graphs title_left class graph_components title_left class stored_on_disk title_left class pyg_loaded title_left

Node and edge attributes#

There are a number of node and edge attributes with special meanings in weather-model-graphs which enable the splitting and visualisation of the graph components.

Node attributes#

  • pos: the (x,y) coordinates of the node in the grid

  • type: the type of node, either grid or mesh

Edge attributes#

  • component: the component of the graph the edge belongs to, either g2m, m2m or m2g

  • level: for multi-range mesh graphs this denotes the refinement level of mesh connection. For hierarchical graphs the different ranges of connections are split into different levels and so here level also denotes the level in the hierarchy that the edge belongs to.

  • len: the length of the edge in the (x,y) coordinate space of the grid nodes, i.e. the distance between the two nodes in the grid

  • vdiff: the vector spanning between the (x,y) coordinates of the two nodes

  • direction: for hierarchical graphs this denotes the direction of the edge, either up, down and same

Splitting graphs#

The splitting of the graph is done with by utilising the edge attributes, and thus it is easy to split the complete graph by either which component the edge belongs to, or by the level of the edge in the graph. This is done using the weather_model_graphs.split_graph_by_edge_attribute(...) function.

Code layout#

The code layout of weather-model-graphs is organised into submodules by the functionality they provide. The main submodules are:

weather_model_graphs
    .create
        .archetype:
            for creating specific archetype graph
            architectures (e.g. Keisler 2021, Lam et al 2023,
            Oscarsson et al 2023)
        .base
            general interface for creating graph architectures
            (here you define the g2m, m2m and m2g connectivity directly)
        .mesh
            for creating the mesh nodes and edges
        .grid
            for creating the grid nodes
    .visualise
        for plotting graphs, allowing for easy visualisation using any
        edge or node attribute for colouring
    .save
        for saving the graph to specific formats (e.g. pytorch-geometric)
Contents