neural_lam.create_graph_with_wmg#

Create neural-lam graphs by delegating construction to weather-model-graphs.

Builds the g2m/m2m/m2g graph components with weather-model-graphs (wmg) and saves them to disk in neural-lam’s tensor-on-disk format, replacing the duplicated logic in create_graph.py.

Attributes#

Functions#

cli(→ None)

Command-line interface for graph creation using weather-model-graphs.

create_graph_from_datastore(→ None)

Create graph using weather-model-graphs and save in neural-lam format.

Module Contents#

neural_lam.create_graph_with_wmg.cli(input_args: list[str] | None = None) → None[source]#

Command-line interface for graph creation using weather-model-graphs.

Parameters:

input_args (list of str or None, optional) – Arguments to parse instead of sys.argv.

neural_lam.create_graph_with_wmg.create_graph_from_datastore(datastore: neural_lam.datastore.base.BaseDatastore, output_root_path: str, archetype: str = 'keisler', mesh_node_distance: float | None = None, mesh_grid_distance_ratio: float = 3.0, level_refinement_factor: int = 3, max_num_levels: int | None = None) → None[source]#

Create graph using weather-model-graphs and save in neural-lam format.

Parameters:
  • datastore (BaseDatastore) – Datastore providing grid coordinates.

  • output_root_path (str) – Directory where the .pt graph files will be saved.

  • archetype (str) – Graph archetype to create: "keisler", "graphcast", or "hierarchical".

  • mesh_node_distance (float or None) – Distance between created mesh nodes (in coordinate units). If None, the grid node distance is estimated automatically from the grid coordinates and multiplied by mesh_grid_distance_ratio.

  • mesh_grid_distance_ratio (float) – Ratio of mesh node distance to grid node distance. Only used when mesh_node_distance is None. Default is 3.0.

  • level_refinement_factor (int) – Refinement factor between mesh hierarchy levels. Only used for "graphcast" and "hierarchical" archetypes.

  • max_num_levels (int or None) – Maximum number of mesh hierarchy levels. Only used for "graphcast" and "hierarchical" archetypes.

neural_lam.create_graph_with_wmg.ARCHETYPE_FUNCTIONS#