Can not create flat multiscale graphs with >= 3 levels

[joeloskarsson]

There seems to be a bug in https://github.com/mllam/weather-model-graphs/blob/2456064543200af1adfd17db0d1ccbb5f7441ca8/src/weather_model_graphs/create/mesh/kinds/flat.py#L8

that causes a crash whenever you try to create a multiscale graph by collapsing >= 3 levels of flat graphs.

Minimum example to reproduce

import weather_model_graphs as wmg
import numpy as np

x_coords = np.linspace(0, 15, 50)
y_coords = np.linspace(0, 5, 50)
meshgridded = np.meshgrid(x_coords, y_coords)
xy_grid = np.stack(meshgridded, axis=0)

graph = wmg.create.archetype.create_graphcast_graph(
    xy_grid,
    grid_refinement_factor=1,
    level_refinement_factor=3,
)

This gives output

DEBUG    | weather_model_graphs.create.mesh.mesh:create_multirange_2d_mesh_graphs:134 - mesh_levels: 3, nleaf: [27 27]

but then crashes with

File "/home/joel/repos/weather-model-graphs/src/weather_model_graphs/create/mesh/kinds/flat.py", line 65, in create_flat_multiscale_mesh_graph
    .reshape((num_nodes_x, num_nodes_y, 2))[
ValueError: cannot reshape array of size 162 into shape (26,26,2)

Problem

The issue seems to relate to the loop https://github.com/mllam/weather-model-graphs/blob/2456064543200af1adfd17db0d1ccbb5f7441ca8/src/weather_model_graphs/create/mesh/kinds/flat.py#L58-L76 where the variables num_nodes_x and num_nodes_y are not set correctly on the second iteration. I am a bit suspicious of the lines https://github.com/mllam/weather-model-graphs/blob/2456064543200af1adfd17db0d1ccbb5f7441ca8/src/weather_model_graphs/create/mesh/kinds/flat.py#L73-L75 which modify the graph G_all_levels[lev], which is what is used in the next iteration of the loop. It could be that this graph is being overwritten before it is used.

I encountered this when working on #32, but since this is an orthogonal issue I will not fix it in there.