From 30bfb82b34619640698f55be43a807c15da94ab7 Mon Sep 17 00:00:00 2001
From: joeloskarsson <joel.oskarsson@liu.se>
Date: Thu, 23 Jan 2025 21:34:34 +0100
Subject: [PATCH 1/3] Fix computation of number of nodex in x- and y-directions

---
 src/weather_model_graphs/create/mesh/kinds/flat.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/weather_model_graphs/create/mesh/kinds/flat.py b/src/weather_model_graphs/create/mesh/kinds/flat.py
index ab27a08..cf14fc5 100644
--- a/src/weather_model_graphs/create/mesh/kinds/flat.py
+++ b/src/weather_model_graphs/create/mesh/kinds/flat.py
@@ -52,11 +52,14 @@ def create_flat_multiscale_mesh_graph(
     G_tot = G_all_levels[0]
     # First node at level l+1 share position with node (offset, offset) at level l
     level_offset = level_refinement_factor // 2
+
+    first_level_nodes = list(G_all_levels[0].nodes)
+    # Last nodes in first layer has pos (nx-1, ny-1)
+    num_nodes_x = first_level_nodes[-1][0] + 1
+    num_nodes_y = first_level_nodes[-1][1] + 1
+
     for lev in range(1, len(G_all_levels)):
         nodes = list(G_all_levels[lev - 1].nodes)
-        # Last nodes always has pos (nx-1, ny-1)
-        num_nodes_x = nodes[-1][0] + 1
-        num_nodes_y = nodes[-1][1] + 1
         ij = (
             np.array(nodes)
             .reshape((num_nodes_x, num_nodes_y, 2))[
@@ -72,6 +75,10 @@ def create_flat_multiscale_mesh_graph(
         )
         G_tot = networkx.compose(G_tot, G_all_levels[lev])
 
+        # Update number of nodes in x- and y-direction for next iteraion
+        num_nodes_x //= level_refinement_factor
+        num_nodes_y //= level_refinement_factor
+
     # Relabel mesh nodes to start with 0
     G_tot = prepend_node_index(G_tot, 0)
 

From 3dc21ce95599cd15b1bea7e769f2b9a86e779896 Mon Sep 17 00:00:00 2001
From: joeloskarsson <joel.oskarsson@liu.se>
Date: Fri, 24 Jan 2025 09:10:18 +0100
Subject: [PATCH 2/3] Add test testing graph creation with more than 2 levels

---
 tests/test_graph_creation.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/tests/test_graph_creation.py b/tests/test_graph_creation.py
index ea63d1c..8d2b057 100644
--- a/tests/test_graph_creation.py
+++ b/tests/test_graph_creation.py
@@ -192,3 +192,24 @@ def test_create_lat_lon(kind):
         coords_crs=coords_crs,
         graph_crs=graph_crs,
     )
+
+
+@pytest.mark.parametrize("kind", ["graphcast", "oskarsson_hierarchical"])
+def test_create_many_levels(kind):
+    """Test that mesh graph creation methods that work with many levels
+    can handle more than 2 levels
+    """
+    # Test 4 levels at lrf=3
+    grid_coord_range = 10
+    level_refinement_factor = 3
+    mesh_node_distance = grid_coord_range / 3**4
+
+    xy = test_utils.create_fake_xy(N=grid_coord_range)
+    fn_name = f"create_{kind}_graph"
+    fn = getattr(wmg.create.archetype, fn_name)
+
+    fn(
+        coords=xy,
+        mesh_node_distance=mesh_node_distance,
+        level_refinement_factor=level_refinement_factor,
+    )

From ace4f8a6907052d58f07b6b3cba84f34a544913b Mon Sep 17 00:00:00 2001
From: joeloskarsson <joel.oskarsson@liu.se>
Date: Fri, 24 Jan 2025 09:40:24 +0100
Subject: [PATCH 3/3] Add changelog entry

---
 CHANGELOG.md | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a2fef38..2ce1af6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -20,6 +20,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   [2, Ny, Nx]), to allow for non-regularly gridded coordinates
   [\#32](https://github.com/mllam/weather-model-graphs/pull/32), @joeloskarsson
 
+### Fixed
+
+- Fix crash when trying to create flat multiscale graphs with >= 3 levels
+  [\#41](https://github.com/mllam/weather-model-graphs/pull/41), @joeloskarsson
+
 ## [v0.2.0](https://github.com/mllam/weather-model-graphs/releases/tag/v0.2.0)
 
 ### Added