Add graph creation functionality using weather-model-graphs

[sadamov]

Overview

Suggestion to add functionality to create and save different graph architectures using the weather-model-graphs package. The implementation includes support for three graph types (currently):

• Keisler graphs (single-scale mesh)
• GraphCast (multi-scale)
• Oskarsson hierarchical graphs

Reference https://github.com/mllam/neural-lam/pull/66#issuecomment-2434482120

New Features

• Script to generate graphs from datastore xy coordinates
• Integration with weather-model-graphs package
• Support for multiple datastore types (MDP, NPYFilesMEPS)
• Graph serialization to PyTorch format

Implementation Details

from pathlib import Path
import weather_model_graphs as wmg
from neural_lam.datastore import init_datastore

def create_graphs(datastore):
    xy = datastore.get_xy("state", stacked=False)
    graphs = {
        'keisler': wmg.create.archetype.create_keisler_graph(xy_grid=xy),
        'graphcast': wmg.create.archetype.create_graphcast_graph(xy_grid=xy),
        'hierarchical': wmg.create.archetype.create_oscarsson_hierarchical_graph(xy_grid=xy)
    }
    return graphs

Known Issues

1. Graph serialization with wmg.save.to_pyg() not working correctly
2. Need to determine correct edge/node feature naming convention

TODO

• [ ] Expand/simplify PyG serialization
• [ ] Adapt tests for graph creation with wgm
• [ ] Document graph properties and usage

Questions

1. What's the preferred format for storing generated graphs?
2. Should we include visualization tools?
3. How closely do we want to interate wgm into nl (slang for neural-lam, obviously)

[joeloskarsson]

I will take this on, as there is a need to more tightly integrate neural-lam and wmg for #84 . Or, in particular, neural-lam will have to call wmg with specific arguments, so we can add a dependency from neural-lam to wmg and then graph-creation scripts in neural-lam that call wmg.

For #84 this will be needed to impose a specific ordering between interior grid nodes and boundary grid nodes. By calling wmg directly from neural-lam we can assure that these calls always place all interior nodes before all boundary nodes.