Add mesh_layout argument with rectilinear two-step coordinate/connectivity architecture

CHANGELOG.md

@@ -5,11 +5,15 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## [Unreleased]
+## [Unreleased](https://github.com/mllam/weather-model-graphs/compare/v0.3.0...HEAD)
 
 ### Added
 
-- Added a standalone graph consistency checking tool (`wmg.diagnostics.check_graph_consistency`) to ensure structural health, such as verifying all grid nodes successfully connect to the mesh (#42).
+- Add `mesh_layout` argument to mesh graph creation functions, with `rectilinear`
+  as the first supported layout. Uses a two-step architecture separating coordinate
+  creation from connectivity creation, enabling future alternative layouts (e.g. triangular).
+  [\#78](https://github.com/mllam/weather-model-graphs/issues/78), @prajwal-tech07
+- Add a standalone graph consistency checking tool (`wmg.diagnostics.check_graph_consistency`) to ensure structural health, such as verifying all grid nodes successfully connect to the mesh (#42).
 - Add Django-style graph filtering via `filter_graph`, for example to select
   nodes by type (`node__type="mesh"`), edges by component
   (`edge__component="g2m"`), long edges (`edge__len__gt=...`), and spatial

docs/creating_the_graph.ipynb

@@ -411,8 +411,11 @@
     "graph = wmg.create.create_all_graph_components(\n",
     "    m2m_connectivity=\"flat_multiscale\",\n",
     "    coords=xy,\n",
-    "    m2m_connectivity_kwargs=dict(\n",
-    "        mesh_node_distance=2, level_refinement_factor=3, max_num_levels=None\n",
+    "    mesh_layout=\"rectilinear\",\n",
+    "    mesh_layout_kwargs=dict(\n",
+    "        mesh_node_spacing=2,\n",
+    "        refinement_factor=3,\n",
+    "        max_num_refinement_levels=3,\n",
     "    ),\n",
     "    g2m_connectivity=\"nearest_neighbour\",\n",
     "    m2g_connectivity=\"nearest_neighbour\",\n",

src/weather_model_graphs/create/archetype.py

@@ -1,5 +1,6 @@
-from typing import Iterable
+from typing import Dict, Iterable, Union
 
+import networkx
 import pyproj
 
 from .base import create_all_graph_components
@@ -12,7 +13,7 @@ def create_keisler_graph(
     graph_crs: pyproj.crs.CRS | None = None,
     decode_mask: Iterable[bool] | None = None,
     return_components: bool = False,
-):
+) -> Union[networkx.DiGraph, Dict[str, networkx.DiGraph]]:
     """
     Create a flat LAM graph from Oskarsson et al (2023, https://arxiv.org/abs/2309.17370)
     This graph setup is inspired by the global graph used by Keisler (2022, https://arxiv.org/abs/2202.07575).
@@ -57,18 +58,15 @@ def create_keisler_graph(
     """
     return create_all_graph_components(
         coords=coords,
+        coords_crs=coords_crs,
+        graph_crs=graph_crs,
+        mesh_layout="rectilinear",
+        mesh_layout_kwargs=dict(mesh_node_spacing=mesh_node_distance),
         m2m_connectivity="flat",
-        m2m_connectivity_kwargs=dict(mesh_node_distance=mesh_node_distance),
         g2m_connectivity="within_radius",
+        g2m_connectivity_kwargs=dict(rel_max_dist=0.51),
         m2g_connectivity="nearest_neighbours",
-        g2m_connectivity_kwargs=dict(
-            rel_max_dist=0.51,
-        ),
-        m2g_connectivity_kwargs=dict(
-            max_num_neighbours=4,
-        ),
-        coords_crs=coords_crs,
-        graph_crs=graph_crs,
+        m2g_connectivity_kwargs=dict(max_num_neighbours=4),
         decode_mask=decode_mask,
         return_components=return_components,
     )
@@ -83,7 +81,7 @@ def create_graphcast_graph(
     graph_crs: pyproj.crs.CRS | None = None,
     decode_mask: Iterable[bool] | None = None,
     return_components: bool = False,
-):
+) -> Union[networkx.DiGraph, Dict[str, networkx.DiGraph]]:
     """
     Create a multiscale LAM graph from Oskarsson et al (2023, https://arxiv.org/abs/2309.17370)
     This graph setup is inspired by the global GraphCast graph used by Lam et al (2023, https://arxiv.org/abs/2212.12794)
@@ -132,22 +130,19 @@ def create_graphcast_graph(
     """
     return create_all_graph_components(
         coords=coords,
-        m2m_connectivity="flat_multiscale",
-        m2m_connectivity_kwargs=dict(
-            mesh_node_distance=mesh_node_distance,
-            level_refinement_factor=level_refinement_factor,
-            max_num_levels=max_num_levels,
+        coords_crs=coords_crs,
+        graph_crs=graph_crs,
+        mesh_layout="rectilinear",
+        mesh_layout_kwargs=dict(
+            mesh_node_spacing=mesh_node_distance,
+            refinement_factor=level_refinement_factor,
+            max_num_refinement_levels=max_num_levels,
         ),
+        m2m_connectivity="flat_multiscale",
         g2m_connectivity="within_radius",
+        g2m_connectivity_kwargs=dict(rel_max_dist=0.51),
         m2g_connectivity="nearest_neighbours",
-        g2m_connectivity_kwargs=dict(
-            rel_max_dist=0.51,
-        ),
-        m2g_connectivity_kwargs=dict(
-            max_num_neighbours=4,
-        ),
-        coords_crs=coords_crs,
-        graph_crs=graph_crs,
+        m2g_connectivity_kwargs=dict(max_num_neighbours=4),
         decode_mask=decode_mask,
         return_components=return_components,
     )
@@ -162,7 +157,7 @@ def create_oskarsson_hierarchical_graph(
     graph_crs: pyproj.crs.CRS | None = None,
     decode_mask: Iterable[bool] | None = None,
     return_components: bool = False,
-):
+) -> Union[networkx.DiGraph, Dict[str, networkx.DiGraph]]:
     """
     Create a LAM graph following Oskarsson et al (2023, https://arxiv.org/abs/2309.17370)
     hierarchical architecture.
@@ -216,22 +211,19 @@ def create_oskarsson_hierarchical_graph(
     """
     return create_all_graph_components(
         coords=coords,
-        m2m_connectivity="hierarchical",
-        m2m_connectivity_kwargs=dict(
-            mesh_node_distance=mesh_node_distance,
-            level_refinement_factor=level_refinement_factor,
-            max_num_levels=max_num_levels,
+        coords_crs=coords_crs,
+        graph_crs=graph_crs,
+        mesh_layout="rectilinear",
+        mesh_layout_kwargs=dict(
+            mesh_node_spacing=mesh_node_distance,
+            refinement_factor=level_refinement_factor,
+            max_num_refinement_levels=max_num_levels,
         ),
+        m2m_connectivity="hierarchical",
         g2m_connectivity="within_radius",
+        g2m_connectivity_kwargs=dict(rel_max_dist=0.51),
         m2g_connectivity="nearest_neighbours",
-        g2m_connectivity_kwargs=dict(
-            rel_max_dist=0.51,
-        ),
-        m2g_connectivity_kwargs=dict(
-            max_num_neighbours=4,
-        ),
-        coords_crs=coords_crs,
-        graph_crs=graph_crs,
+        m2g_connectivity_kwargs=dict(max_num_neighbours=4),
         decode_mask=decode_mask,
         return_components=return_components,
     )