diff --git a/src/weather_model_graphs/create/base.py b/src/weather_model_graphs/create/base.py index fbfa93b..792184b 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/weather_model_graphs/create/base.py @@ -157,9 +157,21 @@ def create_all_graph_components( G_target=grid_connect_graph, method=g2m_connectivity, **g2m_connectivity_kwargs, - ) +) graph_components["g2m"] = G_g2m + grid_nodes_in_g2m = {u for u, v in G_g2m.edges()} + all_grid_nodes = set(G_grid.nodes) + missing_nodes = all_grid_nodes - grid_nodes_in_g2m + + if missing_nodes: + raise ValueError( + f"Found {len(missing_nodes)} grid node(s) with no g2m connection! " + f"Missing nodes: {missing_nodes}. " + f"This means these grid nodes won't participate in the encode phase. " + f"Try increasing g2m rel_max_dist or max_dist." + ) + if decode_mask is None: # decode to all grid nodes decode_grid = G_grid @@ -407,6 +419,7 @@ def _find_neighbour_node_idxs_in_source_mesh(xy_target): for target_node in target_nodes_list: xy_target = G_target.nodes[target_node]["pos"] neigh_idxs = _find_neighbour_node_idxs_in_source_mesh(xy_target) + neigh_idxs = np.atleast_1d(neigh_idxs) for i in neigh_idxs: source_node = source_nodes_list[i] # add edge from source to target diff --git a/src/weather_model_graphs/create/mesh/kinds/flat.py b/src/weather_model_graphs/create/mesh/kinds/flat.py index 92f47b0..a371a4c 100644 --- a/src/weather_model_graphs/create/mesh/kinds/flat.py +++ b/src/weather_model_graphs/create/mesh/kinds/flat.py @@ -6,7 +6,7 @@ def create_flat_multiscale_mesh_graph( - xy, mesh_node_distance: float, level_refinement_factor: int, max_num_levels: int + xy, mesh_node_distance: float, level_refinement_factor: int, max_num_levels: int, ): """ Create flat mesh graph by merging the single-level mesh diff --git a/tests/test_spherical_coords.py b/tests/test_spherical_coords.py new file mode 100644 index 0000000..baaf273 --- /dev/null +++ b/tests/test_spherical_coords.py @@ -0,0 +1,47 @@ +import numpy as np +import pytest +import weather_model_graphs as wmg +from weather_model_graphs.create.base import create_all_graph_components + +def test_flat_coords(): + coords=np.array([[0,0],[0,1],[1,0],[1,1]]) + graph=wmg.create.archetype.create_keisler_graph( + coords=coords, + mesh_node_distance=0.5 + ) + assert len(graph.nodes) > 0 + assert len(graph.edges) > 0 + print(f"Success flat graph has {len(graph.nodes)}nodes and {len(graph.edges)}edges") + + +def test_spherical_coords(): + coords=np.array([ + [-10.0,0.0], + [10.0,0.0], + [0.0,10.0], + [0.0,-10.0], + ]) + graph=wmg.create.archetype.create_keisler_graph( + coords=coords, + coords_crs="EPSG:4326", + graph_crs="EPSG:3857", + mesh_node_distance=500_000 + ) + assert len(graph.nodes) > 0 + assert len(graph.edges) > 0 + for _, _, data in graph.edges(data=True): + if 'len' in data: + print(f"Edge length in metres:{data['len']:.1f}m") + print(f"Success! Spherical graph has {len(graph.nodes)} nodes and {len(graph.edges)} edges") + + + +if __name__ == "__main__": + print("Running Test 1 --Flat Cartesian...") + test_flat_coords() + + print("\nRunning Test 2 --Spherical....") + test_spherical_coords() + + print("\nAll Tests Passed!") +