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

CHANGELOG.md

@@ -5,10 +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
 
+- 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 `__version__` attribute to the package init
   [\#56](https://github.com/mllam/weather-model-graphs/pull/56) @AdMub
 

src/weather_model_graphs/create/archetype.py

@@ -57,18 +57,16 @@ 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),
+        m2m_connectivity_kwargs=dict(pattern="8-star"),
         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,
     )
@@ -132,22 +130,20 @@ 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",
+        m2m_connectivity_kwargs=dict(pattern="8-star"),
         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,
     )
@@ -216,22 +212,23 @@ def create_oskarsson_hierarchical_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,
+            refinement_factor=level_refinement_factor,
+            max_num_refinement_levels=max_num_levels,
+        ),
         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,
+            intra_level=dict(pattern="8-star"),
+            inter_level=dict(pattern="nearest", k=1),
         ),
         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,
     )