Can not create flat multiscale graphs with >= 3 levels

[joeloskarsson]

There seems to be a bug in

weather-model-graphs/src/weather_model_graphs/create/mesh/kinds/flat.py

Line 8 in 2456064

def create_flat_multiscale_mesh_graph(

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

weather-model-graphs/src/weather_model_graphs/create/mesh/kinds/flat.py

Lines 58 to 76 in 2456064

	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))[
	level_offset::level_refinement_factor,
	level_offset::level_refinement_factor,
	:,
	]
	.reshape(int(num_nodes_x * num_nodes_y / (level_refinement_factor**2)), 2)
	)
	ij = [tuple(x) for x in ij]
	G_all_levels[lev] = networkx.relabel_nodes(
	G_all_levels[lev], dict(zip(G_all_levels[lev].nodes, ij))
	)
	G_tot = networkx.compose(G_tot, G_all_levels[lev])

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

weather-model-graphs/src/weather_model_graphs/create/mesh/kinds/flat.py

Lines 73 to 75 in 2456064

	G_all_levels[lev] = networkx.relabel_nodes(
	G_all_levels[lev], dict(zip(G_all_levels[lev].nodes, ij))
	)

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.