From 6f82d4341f9a1086cf1485087bc8f6cbbef97c12 Mon Sep 17 00:00:00 2001 From: prajwal Date: Mon, 2 Mar 2026 19:15:50 +0530 Subject: [PATCH 01/21] feat: introduce mesh_layout argument with two-step coordinate/connectivity architecture Implement the mesh_layout parameter and refactor mesh graph creation into a two-step process as discussed in #78: 1. Coordinate creation (mesh_layout + mesh_layout_kwargs): - create_single_level_2d_mesh_coordinates() returns nx.Graph with spatial adjacency edges annotated as 'cardinal' or 'diagonal' - create_multirange_2d_mesh_coordinates() returns list of nx.Graph for multi-level meshes with interlevel_refinement_factor stored as graph attr 2. Connectivity creation (m2m_connectivity + m2m_connectivity_kwargs): - create_directed_mesh_graph() converts nx.Graph to nx.DiGraph based on pattern ('4-star' or '8-star') - create_flat_singlescale_from_coordinates() for flat single-scale - create_flat_multiscale_from_coordinates() with intra_level/inter_level sub-dicts for explicit connectivity control - create_hierarchical_from_coordinates() with intra_level/inter_level sub-dicts supporting 'nearest' pattern with k parameter Parameter restructuring: - grid_spacing replaces mesh_node_distance in mesh_layout_kwargs - interlevel_refinement_factor replaces level_refinement_factor - max_num_levels moves to mesh_layout_kwargs - m2m_connectivity_kwargs restructured with intra_level/inter_level sub-dicts Backward compatibility: - Old-style kwargs (mesh_node_distance, level_refinement_factor, max_num_levels in m2m_connectivity_kwargs) are auto-migrated with deprecation warnings - All existing wrapper functions preserved - All existing tests pass unchanged Archetype functions updated to use new parameter scheme: - create_keisler_graph: pattern='8-star' - create_graphcast_graph: intra_level=8-star, inter_level=coincident - create_oskarsson_hierarchical_graph: intra_level=8-star, inter_level=nearest(k=1) Closes #78 --- src/weather_model_graphs/create/archetype.py | 26 +- src/weather_model_graphs/create/base.py | 211 ++++++++++++-- .../create/mesh/__init__.py | 7 +- .../create/mesh/kinds/flat.py | 176 +++++++++--- .../create/mesh/kinds/hierarchical.py | 168 ++++++++---- src/weather_model_graphs/create/mesh/mesh.py | 257 ++++++++++++++---- 6 files changed, 686 insertions(+), 159 deletions(-) diff --git a/src/weather_model_graphs/create/archetype.py b/src/weather_model_graphs/create/archetype.py index b4d717e..13a79ce 100644 --- a/src/weather_model_graphs/create/archetype.py +++ b/src/weather_model_graphs/create/archetype.py @@ -58,7 +58,9 @@ def create_keisler_graph( return create_all_graph_components( coords=coords, m2m_connectivity="flat", - m2m_connectivity_kwargs=dict(mesh_node_distance=mesh_node_distance), + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(grid_spacing=mesh_node_distance), + m2m_connectivity_kwargs=dict(pattern="8-star"), g2m_connectivity="within_radius", m2g_connectivity="nearest_neighbours", g2m_connectivity_kwargs=dict( @@ -133,11 +135,16 @@ 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, + mesh_layout="rectilinear", + mesh_layout_kwargs=dict( + grid_spacing=mesh_node_distance, + interlevel_refinement_factor=level_refinement_factor, max_num_levels=max_num_levels, ), + m2m_connectivity_kwargs=dict( + intra_level=dict(pattern="8-star"), + inter_level=dict(pattern="coincident"), + ), g2m_connectivity="within_radius", m2g_connectivity="nearest_neighbours", g2m_connectivity_kwargs=dict( @@ -217,11 +224,16 @@ def create_oskarsson_hierarchical_graph( return create_all_graph_components( coords=coords, m2m_connectivity="hierarchical", - m2m_connectivity_kwargs=dict( - mesh_node_distance=mesh_node_distance, - level_refinement_factor=level_refinement_factor, + mesh_layout="rectilinear", + mesh_layout_kwargs=dict( + grid_spacing=mesh_node_distance, + interlevel_refinement_factor=level_refinement_factor, max_num_levels=max_num_levels, ), + m2m_connectivity_kwargs=dict( + intra_level=dict(pattern="8-star"), + inter_level=dict(pattern="nearest", k=1), + ), g2m_connectivity="within_radius", m2g_connectivity="nearest_neighbours", g2m_connectivity_kwargs=dict( diff --git a/src/weather_model_graphs/create/base.py b/src/weather_model_graphs/create/base.py index f922b2d..3f2f231 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/weather_model_graphs/create/base.py @@ -8,7 +8,7 @@ function uses `connect_nodes_across_graphs` to connect nodes across the component graphs. """ - +import warnings from typing import Iterable import networkx @@ -25,10 +25,19 @@ ) from .grid import create_grid_graph_nodes from .mesh.kinds.flat import ( + create_flat_multiscale_from_coordinates, create_flat_multiscale_mesh_graph, + create_flat_singlescale_from_coordinates, create_flat_singlescale_mesh_graph, ) -from .mesh.kinds.hierarchical import create_hierarchical_multiscale_mesh_graph +from .mesh.kinds.hierarchical import ( + create_hierarchical_from_coordinates, + create_hierarchical_multiscale_mesh_graph, +) +from .mesh.mesh import ( + create_multirange_2d_mesh_coordinates, + create_single_level_2d_mesh_coordinates, +) def create_all_graph_components( @@ -36,7 +45,9 @@ def create_all_graph_components( m2m_connectivity: str, m2g_connectivity: str, g2m_connectivity: str, - m2m_connectivity_kwargs={}, + mesh_layout: str = "rectilinear", + mesh_layout_kwargs: dict = None, + m2m_connectivity_kwargs: dict = None, m2g_connectivity_kwargs={}, g2m_connectivity_kwargs={}, coords_crs: pyproj.crs.CRS | None = None, @@ -49,6 +60,14 @@ def create_all_graph_components( grid-to-mesh (g2m), mesh-to-mesh (m2m) and mesh-to-grid (m2g), representing the encode-process-decode respectively. + The mesh graph creation follows a two-step process: + 1. **Coordinate creation** (controlled by `mesh_layout` + `mesh_layout_kwargs`): + Creates an undirected graph (nx.Graph) with node positions and spatial + adjacency edges annotated with adjacency types. + 2. **Connectivity creation** (controlled by `m2m_connectivity` + `m2m_connectivity_kwargs`): + Converts the coordinate graph to directed connectivity (nx.DiGraph) + based on the specified pattern and connectivity method. + For each graph component, the method for connecting nodes across graphs should be specified (with the `*_connectivity` arguments, e.g. `m2g_connectivity`). And the method-specific arguments should be passed as keyword arguments using @@ -62,15 +81,23 @@ def create_all_graph_components( - "within_radius": Find all neighbours in grid within an absolute distance of `max_dist` or relative distance of `rel_max_dist` from each node in mesh + mesh_layout: + - "rectilinear": Regular rectilinear grid (default). Uses grid_spacing to + determine mesh node placement. Produces nodes with 4-star (cardinal) and + 8-star (cardinal + diagonal) spatial adjacency edges. + + mesh_layout_kwargs (for mesh_layout="rectilinear"): + - grid_spacing: float, distance between mesh nodes in coordinate units + - interlevel_refinement_factor: int, refinement factor between levels (for multi-level) + - max_num_levels: int, maximum number of mesh levels (for multi-level) + m2m_connectivity: - - "flat": Create a single-level 2D mesh graph with `mesh_node_distance`, - similar to Keisler et al. (2022) - - "flat_multiscale": Create a flat multiscale mesh graph with `max_num_levels`, - `mesh_node_distance` and `level_refinement_factor`, - similar to GraphCast, Lam et al. (2023) - - "hierarchical": Create a hierarchical mesh graph with `max_num_levels`, - `mesh_node_distance` and `level_refinement_factor`, - similar to Oskarsson et al. (2023) + - "flat": Create a single-level directed mesh graph. + m2m_connectivity_kwargs: pattern="4-star" or "8-star" (default: "8-star") + - "flat_multiscale": Create a flat multiscale mesh graph. + m2m_connectivity_kwargs: intra_level=dict(pattern=...), inter_level=dict(pattern=...) + - "hierarchical": Create a hierarchical mesh graph with up/down connections. + m2m_connectivity_kwargs: intra_level=dict(pattern=...), inter_level=dict(pattern=..., k=...) m2g_connectivity: - "nearest_neighbour": Find the nearest neighbour in mesh for each node in grid @@ -97,6 +124,50 @@ def create_all_graph_components( """ graph_components: dict[networkx.DiGraph] = {} + # Initialize mutable default arguments (and copy to avoid mutating caller's dicts) + if mesh_layout_kwargs is None: + mesh_layout_kwargs = {} + else: + mesh_layout_kwargs = dict(mesh_layout_kwargs) + if m2m_connectivity_kwargs is None: + m2m_connectivity_kwargs = {} + else: + m2m_connectivity_kwargs = dict(m2m_connectivity_kwargs) + + # Backward compatibility: migrate old-style kwargs where mesh_node_distance, + # level_refinement_factor, and max_num_levels were passed via + # m2m_connectivity_kwargs. In the new design these belong in mesh_layout_kwargs. + if "mesh_node_distance" in m2m_connectivity_kwargs and "grid_spacing" not in mesh_layout_kwargs: + warnings.warn( + "Passing 'mesh_node_distance' in m2m_connectivity_kwargs is deprecated. " + "Use mesh_layout_kwargs=dict(grid_spacing=...) instead.", + DeprecationWarning, + stacklevel=2, + ) + mesh_layout_kwargs["grid_spacing"] = m2m_connectivity_kwargs.pop( + "mesh_node_distance" + ) + if "level_refinement_factor" in m2m_connectivity_kwargs and "interlevel_refinement_factor" not in mesh_layout_kwargs: + warnings.warn( + "Passing 'level_refinement_factor' in m2m_connectivity_kwargs is deprecated. " + "Use mesh_layout_kwargs=dict(interlevel_refinement_factor=...) instead.", + DeprecationWarning, + stacklevel=2, + ) + mesh_layout_kwargs["interlevel_refinement_factor"] = ( + m2m_connectivity_kwargs.pop("level_refinement_factor") + ) + if "max_num_levels" in m2m_connectivity_kwargs and "max_num_levels" not in mesh_layout_kwargs: + warnings.warn( + "Passing 'max_num_levels' in m2m_connectivity_kwargs is deprecated. " + "Use mesh_layout_kwargs=dict(max_num_levels=...) instead.", + DeprecationWarning, + stacklevel=2, + ) + mesh_layout_kwargs["max_num_levels"] = m2m_connectivity_kwargs.pop( + "max_num_levels" + ) + assert ( len(coords.shape) == 2 and coords.shape[1] == 2 ), "Grid node coordinates should be given as an array of shape [num_grid_nodes, 2]." @@ -126,26 +197,124 @@ def create_all_graph_components( xy = np.stack(xy_tuple, axis=1) if m2m_connectivity == "flat": - graph_components["m2m"] = create_flat_singlescale_mesh_graph( - xy, - **m2m_connectivity_kwargs, + # --- Step 1: Coordinate creation based on mesh_layout --- + if mesh_layout == "rectilinear": + grid_spacing = mesh_layout_kwargs.get("grid_spacing") + if grid_spacing is None: + raise ValueError( + "mesh_layout='rectilinear' requires 'grid_spacing' in " + "mesh_layout_kwargs (or 'mesh_node_distance' in " + "m2m_connectivity_kwargs for backward compatibility)." + ) + # Compute number of mesh nodes from grid_spacing + range_x, range_y = np.ptp(xy, axis=0) + nx_mesh = int(range_x / grid_spacing) + ny_mesh = int(range_y / grid_spacing) + if nx_mesh == 0 or ny_mesh == 0: + raise ValueError( + "The given `grid_spacing` is too large for the provided " + f"coordinates. Got grid_spacing={grid_spacing}, but the " + f"x-range is {range_x} and y-range is {range_y}. Maybe you " + "want to decrease the `grid_spacing` so that the mesh nodes " + "are spaced closer together?" + ) + G_mesh_coords = create_single_level_2d_mesh_coordinates( + xy, nx_mesh, ny_mesh + ) + else: + raise NotImplementedError( + f"mesh_layout='{mesh_layout}' is not yet supported. " + "Currently only 'rectilinear' is implemented." + ) + + # --- Step 2: Connectivity creation --- + pattern = m2m_connectivity_kwargs.get("pattern", "8-star") + graph_components["m2m"] = create_flat_singlescale_from_coordinates( + G_mesh_coords, pattern=pattern ) grid_connect_graph = graph_components["m2m"] + elif m2m_connectivity == "hierarchical": + # --- Step 1: Coordinate creation based on mesh_layout --- + if mesh_layout == "rectilinear": + grid_spacing = mesh_layout_kwargs.get("grid_spacing") + interlevel_refinement_factor = mesh_layout_kwargs.get( + "interlevel_refinement_factor" + ) + max_num_levels = mesh_layout_kwargs.get("max_num_levels") + if grid_spacing is None: + raise ValueError( + "mesh_layout='rectilinear' with m2m_connectivity='hierarchical' " + "requires 'grid_spacing' in mesh_layout_kwargs." + ) + G_coords_list = create_multirange_2d_mesh_coordinates( + max_num_levels=max_num_levels, + xy=xy, + grid_spacing=grid_spacing, + interlevel_refinement_factor=interlevel_refinement_factor, + ) + else: + raise NotImplementedError( + f"mesh_layout='{mesh_layout}' is not yet supported. " + "Currently only 'rectilinear' is implemented." + ) + + # --- Step 2: Connectivity creation --- + intra_level = m2m_connectivity_kwargs.get( + "intra_level", {"pattern": "8-star"} + ) + inter_level = m2m_connectivity_kwargs.get( + "inter_level", {"pattern": "nearest", "k": 1} + ) # hierarchical mesh graph have three sub-graphs: - # `m2m` (mesh-to-mesh), `mesh_up` (up edge connections) and `mesh_down` (down edge connections) - graph_components["m2m"] = create_hierarchical_multiscale_mesh_graph( - xy=xy, - **m2m_connectivity_kwargs, + # `m2m` (mesh-to-mesh), `mesh_up` (up edge connections) and + # `mesh_down` (down edge connections) + graph_components["m2m"] = create_hierarchical_from_coordinates( + G_coords_list, + intra_level=intra_level, + inter_level=inter_level, ) # Only connect grid to bottom level of hierarchy grid_connect_graph = split_graph_by_edge_attribute( graph_components["m2m"], "level" )[0] + elif m2m_connectivity == "flat_multiscale": - graph_components["m2m"] = create_flat_multiscale_mesh_graph( - xy=xy, - **m2m_connectivity_kwargs, + # --- Step 1: Coordinate creation based on mesh_layout --- + if mesh_layout == "rectilinear": + grid_spacing = mesh_layout_kwargs.get("grid_spacing") + interlevel_refinement_factor = mesh_layout_kwargs.get( + "interlevel_refinement_factor" + ) + max_num_levels = mesh_layout_kwargs.get("max_num_levels") + if grid_spacing is None: + raise ValueError( + "mesh_layout='rectilinear' with m2m_connectivity='flat_multiscale' " + "requires 'grid_spacing' in mesh_layout_kwargs." + ) + G_coords_list = create_multirange_2d_mesh_coordinates( + max_num_levels=max_num_levels, + xy=xy, + grid_spacing=grid_spacing, + interlevel_refinement_factor=interlevel_refinement_factor, + ) + else: + raise NotImplementedError( + f"mesh_layout='{mesh_layout}' is not yet supported. " + "Currently only 'rectilinear' is implemented." + ) + + # --- Step 2: Connectivity creation --- + intra_level = m2m_connectivity_kwargs.get( + "intra_level", {"pattern": "8-star"} + ) + inter_level = m2m_connectivity_kwargs.get( + "inter_level", {"pattern": "coincident"} + ) + graph_components["m2m"] = create_flat_multiscale_from_coordinates( + G_coords_list, + intra_level=intra_level, + inter_level=inter_level, ) grid_connect_graph = graph_components["m2m"] else: diff --git a/src/weather_model_graphs/create/mesh/__init__.py b/src/weather_model_graphs/create/mesh/__init__.py index 573d9f9..09f3e7b 100644 --- a/src/weather_model_graphs/create/mesh/__init__.py +++ b/src/weather_model_graphs/create/mesh/__init__.py @@ -1 +1,6 @@ -from .mesh import create_single_level_2d_mesh_graph +from .mesh import ( + create_directed_mesh_graph, + create_multirange_2d_mesh_coordinates, + create_single_level_2d_mesh_coordinates, + create_single_level_2d_mesh_graph, +) diff --git a/src/weather_model_graphs/create/mesh/kinds/flat.py b/src/weather_model_graphs/create/mesh/kinds/flat.py index 92f47b0..b2cbca3 100644 --- a/src/weather_model_graphs/create/mesh/kinds/flat.py +++ b/src/weather_model_graphs/create/mesh/kinds/flat.py @@ -5,53 +5,77 @@ from .. import mesh as mesh_graph -def create_flat_multiscale_mesh_graph( - xy, mesh_node_distance: float, level_refinement_factor: int, max_num_levels: int +def create_flat_multiscale_from_coordinates( + G_coords_list, + intra_level=None, + inter_level=None, ): """ - Create flat mesh graph by merging the single-level mesh - graphs across all levels in `G_all_levels`. + Create flat multiscale mesh graph from a list of coordinate graphs. + + This is the connectivity creation step for flat multiscale meshes. + It takes undirected coordinate graphs (one per level) and produces a + single directed mesh graph with intra-level connectivity. Parameters ---------- - xy : np.ndarray [N_grid_points, 2] - Grid point coordinates, with first column representing - x coordinates and second column y coordinates. N_grid_points is the - total number of grid points. - mesh_node_distance: float - Distance (in x- and y-direction) between created mesh nodes, - in coordinate system of xy - level_refinement_factor: int - Refinement factor between grid points and bottom level of mesh hierarchy - NOTE: Must be an odd integer >1 to create proper multiscale graph - max_num_levels : int - Maximum number of levels in the multi-scale graph + G_coords_list : list of networkx.Graph + List of undirected coordinate graphs, one per level. Each should have + nodes with "pos" and "type" attributes, and edges with "adjacency_type" + attributes. Created by create_multirange_2d_mesh_coordinates. + intra_level : dict or None + Intra-level connectivity options. Supports: + - pattern: str, "4-star" or "8-star" (default: "8-star") + inter_level : dict or None + Inter-level connectivity options. Supports: + - pattern: str, "coincident" (default: "coincident") + Currently only "coincident" is supported (coarser level nodes are + exactly coincident with a subset of finer level nodes). + Returns ------- - G_tot : networkx.Graph - The merged mesh graph + G_tot : networkx.DiGraph + The merged flat multiscale mesh graph """ - # Check that level_refinement_factor is an odd integer + if intra_level is None: + intra_level = {"pattern": "8-star"} + if inter_level is None: + inter_level = {"pattern": "coincident"} + + intra_pattern = intra_level.get("pattern", "8-star") + inter_pattern = inter_level.get("pattern", "coincident") + + if inter_pattern != "coincident": + raise NotImplementedError( + f"Inter-level pattern '{inter_pattern}' is not yet supported. " + "Only 'coincident' is currently implemented." + ) + + # Retrieve interlevel_refinement_factor from graph attributes + interlevel_refinement_factor = G_coords_list[0].graph.get( + "interlevel_refinement_factor", 3 + ) + + # Check that interlevel_refinement_factor is an odd integer if ( - int(level_refinement_factor) != level_refinement_factor - or level_refinement_factor % 2 != 1 + int(interlevel_refinement_factor) != interlevel_refinement_factor + or interlevel_refinement_factor % 2 != 1 ): raise ValueError( - "The `level_refinement_factor` must be an odd integer. " - f"Given value: {level_refinement_factor}." + "The `interlevel_refinement_factor` must be an odd integer. " + f"Given value: {interlevel_refinement_factor}." ) - G_all_levels: list[networkx.DiGraph] = mesh_graph.create_multirange_2d_mesh_graphs( - max_num_levels=max_num_levels, - xy=xy, - mesh_node_distance=mesh_node_distance, - level_refinement_factor=level_refinement_factor, - ) + # Convert each level's coordinate graph to directed graph with chosen pattern + G_all_levels = [ + mesh_graph.create_directed_mesh_graph(g_coords, pattern=intra_pattern) + for g_coords in G_coords_list + ] # combine all levels to one 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 + level_offset = interlevel_refinement_factor // 2 first_level_nodes = list(G_all_levels[0].nodes) # Last nodes in first layer has pos (nx-1, ny-1) @@ -63,11 +87,18 @@ def create_flat_multiscale_mesh_graph( ij = ( np.array(nodes) .reshape((num_nodes_x, num_nodes_y, 2))[ - level_offset::level_refinement_factor, - level_offset::level_refinement_factor, + level_offset::interlevel_refinement_factor, + level_offset::interlevel_refinement_factor, :, ] - .reshape(int(num_nodes_x * num_nodes_y / (level_refinement_factor**2)), 2) + .reshape( + int( + num_nodes_x + * num_nodes_y + / (interlevel_refinement_factor**2) + ), + 2, + ) ) ij = [tuple(x) for x in ij] G_all_levels[lev] = networkx.relabel_nodes( @@ -75,9 +106,9 @@ 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 + # Update number of nodes in x- and y-direction for next iteration + num_nodes_x //= interlevel_refinement_factor + num_nodes_y //= interlevel_refinement_factor # Relabel mesh nodes to start with 0 G_tot = prepend_node_index(G_tot, 0) @@ -89,10 +120,83 @@ def create_flat_multiscale_mesh_graph( return G_tot +def create_flat_singlescale_from_coordinates(G_coords, pattern="8-star"): + """ + Create a flat single-scale directed mesh graph from a coordinate graph. + + This is the connectivity creation step for flat single-scale meshes. + It converts an undirected coordinate graph to a directed mesh graph + using the specified connectivity pattern. + + Parameters + ---------- + G_coords : networkx.Graph + Undirected coordinate graph with nodes having "pos" attributes and + edges having "adjacency_type" attributes. Created by + create_single_level_2d_mesh_coordinates. + pattern : str + Connectivity pattern: "4-star" or "8-star" (default: "8-star") + + Returns + ------- + networkx.DiGraph + The flat single-scale directed mesh graph + """ + return mesh_graph.create_directed_mesh_graph(G_coords, pattern=pattern) + + +def create_flat_multiscale_mesh_graph( + xy, mesh_node_distance: float, level_refinement_factor: int, max_num_levels: int +): + """ + Create flat mesh graph by merging the single-level mesh + graphs across all levels in `G_all_levels`. + + Internally uses the two-step process: + 1. create_multirange_2d_mesh_coordinates (coordinate creation) + 2. create_flat_multiscale_from_coordinates (connectivity creation) + + Parameters + ---------- + xy : np.ndarray [N_grid_points, 2] + Grid point coordinates, with first column representing + x coordinates and second column y coordinates. N_grid_points is the + total number of grid points. + mesh_node_distance: float + Distance (in x- and y-direction) between created mesh nodes, + in coordinate system of xy + level_refinement_factor: int + Refinement factor between grid points and bottom level of mesh hierarchy + NOTE: Must be an odd integer >1 to create proper multiscale graph + max_num_levels : int + Maximum number of levels in the multi-scale graph + Returns + ------- + G_tot : networkx.Graph + The merged mesh graph + """ + G_coords_list = mesh_graph.create_multirange_2d_mesh_coordinates( + max_num_levels=max_num_levels, + xy=xy, + grid_spacing=mesh_node_distance, + interlevel_refinement_factor=level_refinement_factor, + ) + + return create_flat_multiscale_from_coordinates( + G_coords_list, + intra_level={"pattern": "8-star"}, + inter_level={"pattern": "coincident"}, + ) + + def create_flat_singlescale_mesh_graph(xy, mesh_node_distance: float): """ Create flat mesh graph of single level + Internally uses the two-step process: + 1. create_single_level_2d_mesh_coordinates (coordinate creation) + 2. create_directed_mesh_graph (connectivity creation, pattern="8-star") + Parameters ---------- xy : np.ndarray [N_grid_points, 2] diff --git a/src/weather_model_graphs/create/mesh/kinds/hierarchical.py b/src/weather_model_graphs/create/mesh/kinds/hierarchical.py index b897693..d7d8c70 100644 --- a/src/weather_model_graphs/create/mesh/kinds/hierarchical.py +++ b/src/weather_model_graphs/create/mesh/kinds/hierarchical.py @@ -6,62 +6,78 @@ from .. import mesh as mesh_graph -def create_hierarchical_multiscale_mesh_graph( - xy, - mesh_node_distance: float, - level_refinement_factor: float, - max_num_levels: int, +def create_hierarchical_from_coordinates( + G_coords_list, + intra_level=None, + inter_level=None, ): """ - Create a hierarchical multiscale mesh graph with nearest neighbour - connections within each level (horizontally, vertically, and diagonally), and - connections between levels (coarse to fine and fine to coarse) using the - nearest neighbour connection. + Create a hierarchical multiscale mesh graph from a list of coordinate graphs. + + This is the connectivity creation step for hierarchical meshes. + It takes undirected coordinate graphs (one per level) and produces a + directed mesh graph with intra-level connectivity and inter-level + up/down connections. Parameters ---------- - xy: np.ndarray - 2D array of mesh point positions. - Distance (in x- and y-direction) between created mesh nodes in bottom level, - in coordinate system of xy - mesh_node_distance: float - Distance (in x- and y-direction) between created mesh nodes in bottom level, - in coordinate system of xy - level_refinement_factor: float - Refinement factor between grid points and bottom level of mesh hierarchy - max_num_levels: int - The number of levels in the hierarchical mesh graph. + G_coords_list : list of networkx.Graph + List of undirected coordinate graphs, one per level. Each should have + nodes with "pos" and "type" attributes, and edges with "adjacency_type" + attributes. Created by create_multirange_2d_mesh_coordinates. + intra_level : dict or None + Intra-level connectivity options. Supports: + - pattern: str, "4-star" or "8-star" (default: "8-star") + inter_level : dict or None + Inter-level connectivity options. Supports: + - pattern: str, "nearest" (default: "nearest") + - k: int, number of nearest neighbours (default: 1) Returns ------- - dict - A dictionary containing the hierarchical mesh graph, the mesh down graph, and - the mesh up graph, with keys "m2m", "mesh_down", and "mesh_up" respectively. + networkx.DiGraph + A directed graph containing the hierarchical mesh with intra-level + edges (direction="same"), inter-level down edges (direction="down"), + and inter-level up edges (direction="up"). """ - Gs_all_levels: list[networkx.DiGraph] = mesh_graph.create_multirange_2d_mesh_graphs( - max_num_levels=max_num_levels, - xy=xy, - mesh_node_distance=mesh_node_distance, - level_refinement_factor=level_refinement_factor, - ) + if intra_level is None: + intra_level = {"pattern": "8-star"} + if inter_level is None: + inter_level = {"pattern": "nearest", "k": 1} + + intra_pattern = intra_level.get("pattern", "8-star") + inter_pattern = inter_level.get("pattern", "nearest") + inter_k = inter_level.get("k", 1) + + if inter_pattern != "nearest": + raise NotImplementedError( + f"Inter-level pattern '{inter_pattern}' is not yet supported " + "for hierarchical graphs. Only 'nearest' is currently implemented." + ) + + # Convert each level's coordinate graph to directed graph with chosen pattern + Gs_all_levels = [ + mesh_graph.create_directed_mesh_graph(g_coords, pattern=intra_pattern) + for g_coords in G_coords_list + ] + n_mesh_levels = len(Gs_all_levels) if n_mesh_levels < 2: raise ValueError( "At least two mesh levels are required for hierarchical mesh graph. " "You may need to reduce the level refinement factor " - f"or increase the max number of levels {max_num_levels} " - f"or number of grid points {xy.shape[0]}." + f"or increase the max number of levels " + f"or number of grid points." ) # Relabel nodes of each level with level index first - Gs_all_levels = [ prepend_node_index(graph, level_i) for level_i, graph in enumerate(Gs_all_levels) ] - # add `direction` attribute to all edges with value `same`` + # add `direction` attribute to all edges with value `same` for i, G in enumerate(Gs_all_levels): for u, v in G.edges: G.edges[u, v]["direction"] = "same" @@ -92,21 +108,29 @@ def create_hierarchical_multiscale_mesh_graph( v_from_xy = np.array([xy for _, xy in G_from.nodes.data("pos")]) kdt_m = scipy.spatial.KDTree(v_from_xy) - # add edges from mesh to grid + # add edges from coarser to finer level for v in v_to_list: - # find 1(?) nearest neighbours (index to vm_xy) - neigh_idx = kdt_m.query(G_down.nodes[v]["pos"], 1)[1] - u = v_from_list[neigh_idx] - - # add edge from mesh to grid - G_down.add_edge(u, v) - d = np.sqrt(np.sum((G_down.nodes[u]["pos"] - G_down.nodes[v]["pos"]) ** 2)) - G_down.edges[u, v]["len"] = d - G_down.edges[u, v]["vdiff"] = ( - G_down.nodes[u]["pos"] - G_down.nodes[v]["pos"] - ) - G_down.edges[u, v]["levels"] = f"{from_level}>{to_level}" - G_down.edges[u, v]["direction"] = "down" + # find k nearest neighbours (index to vm_xy) + neigh_idx = kdt_m.query(G_down.nodes[v]["pos"], inter_k)[1] + if inter_k == 1: + neigh_idx = [neigh_idx] + + for idx in neigh_idx: + u = v_from_list[idx] + + # add edge from coarser to finer + G_down.add_edge(u, v) + d = np.sqrt( + np.sum( + (G_down.nodes[u]["pos"] - G_down.nodes[v]["pos"]) ** 2 + ) + ) + G_down.edges[u, v]["len"] = d + G_down.edges[u, v]["vdiff"] = ( + G_down.nodes[u]["pos"] - G_down.nodes[v]["pos"] + ) + G_down.edges[u, v]["levels"] = f"{from_level}>{to_level}" + G_down.edges[u, v]["direction"] = "down" G_up = networkx.DiGraph() G_up.add_nodes_from(G_down.nodes(data=True)) @@ -130,3 +154,53 @@ def create_hierarchical_multiscale_mesh_graph( G_m2m.graph[prop] = {i: g.graph[prop] for i, g in enumerate(Gs_all_levels)} return G_m2m + + +def create_hierarchical_multiscale_mesh_graph( + xy, + mesh_node_distance: float, + level_refinement_factor: float, + max_num_levels: int, +): + """ + Create a hierarchical multiscale mesh graph with nearest neighbour + connections within each level (horizontally, vertically, and diagonally), and + connections between levels (coarse to fine and fine to coarse) using the + nearest neighbour connection. + + Internally uses the two-step process: + 1. create_multirange_2d_mesh_coordinates (coordinate creation) + 2. create_hierarchical_from_coordinates (connectivity creation) + + Parameters + ---------- + xy: np.ndarray + 2D array of mesh point positions. + Distance (in x- and y-direction) between created mesh nodes in bottom level, + in coordinate system of xy + mesh_node_distance: float + Distance (in x- and y-direction) between created mesh nodes in bottom level, + in coordinate system of xy + level_refinement_factor: float + Refinement factor between grid points and bottom level of mesh hierarchy + max_num_levels: int + The number of levels in the hierarchical mesh graph. + + Returns + ------- + networkx.DiGraph + A directed graph containing the hierarchical mesh with intra-level, + up, and down edges. + """ + G_coords_list = mesh_graph.create_multirange_2d_mesh_coordinates( + max_num_levels=max_num_levels, + xy=xy, + grid_spacing=mesh_node_distance, + interlevel_refinement_factor=level_refinement_factor, + ) + + return create_hierarchical_from_coordinates( + G_coords_list, + intra_level={"pattern": "8-star"}, + inter_level={"pattern": "nearest", "k": 1}, + ) diff --git a/src/weather_model_graphs/create/mesh/mesh.py b/src/weather_model_graphs/create/mesh/mesh.py index 75c82d6..39620f6 100644 --- a/src/weather_model_graphs/create/mesh/mesh.py +++ b/src/weather_model_graphs/create/mesh/mesh.py @@ -3,22 +3,19 @@ from loguru import logger -def create_single_level_2d_mesh_graph(xy, nx, ny): +def create_single_level_2d_mesh_coordinates(xy, nx, ny): """ - Create directed graph with nx * ny nodes representing a 2D grid with - positions spanning the range of xy coordinate values (first dimension - is assumed to be x and y coordinate values respectively). Each nodes is - connected to its eight nearest neighbours, both horizontally, vertically - and diagonally as directed edges (which means that the graph contains two - edges between each pair of connected nodes). + Create an undirected mesh graph (nx.Graph) with node positions and spatial + adjacency edges, representing the coordinate creation step. - The nodes contain a "pos" attribute with the x and y - coordinates of the node, and an "type" attribute with the - type of the node (i.e. "mesh" for mesh nodes). + This produces a graph where: + - Nodes have "pos" (spatial coordinates) and "type" ("mesh") attributes + - Edges have "adjacency_type" attribute: "cardinal" for horizontal/vertical + neighbours (4-star) or "diagonal" for diagonal neighbours (8-star only) - The edges contain a "len" attribute with the length of the edge - and a "vdiff" attribute with the vector difference between the - nodes. + This is the first step in the two-step mesh creation process: + 1. Coordinate creation (this function) -> nx.Graph with spatial adjacency + 2. Connectivity creation (create_directed_mesh_graph) -> nx.DiGraph Parameters ---------- @@ -33,8 +30,8 @@ def create_single_level_2d_mesh_graph(xy, nx, ny): Returns ------- - networkx.DiGraph - Graph representing the 2D grid + networkx.Graph + Undirected graph with node positions and annotated spatial adjacency edges """ xm, xM = np.amin(xy[:, 0]), np.amax(xy[:, 0]) ym, yM = np.amin(xy[:, 1]), np.amax(xy[:, 1]) @@ -56,38 +53,149 @@ def create_single_level_2d_mesh_graph(xy, nx, ny): ) g.nodes[node]["type"] = "mesh" - # add diagonal edges - g.add_edges_from( - [((x, y), (x + 1, y + 1)) for x in range(nx - 1) for y in range(ny - 1)] - + [((x + 1, y), (x, y + 1)) for x in range(nx - 1) for y in range(ny - 1)] - ) - - # turn into directed graph - dg = networkx.DiGraph(g) + # Mark existing grid_2d_graph edges as cardinal (4-star adjacency) for u, v in g.edges(): - d = np.sqrt(np.sum((g.nodes[u]["pos"] - g.nodes[v]["pos"]) ** 2)) + g.edges[u, v]["adjacency_type"] = "cardinal" + + # Add diagonal edges (8-star adjacency) + diagonal_edges = [ + ((x, y), (x + 1, y + 1)) for x in range(nx - 1) for y in range(ny - 1) + ] + [((x + 1, y), (x, y + 1)) for x in range(nx - 1) for y in range(ny - 1)] + g.add_edges_from(diagonal_edges) + for u, v in diagonal_edges: + g.edges[u, v]["adjacency_type"] = "diagonal" + + g.graph["dx"] = dx + g.graph["dy"] = dy + + return g + + +def create_directed_mesh_graph(G_undirected, pattern="8-star"): + """ + Convert an undirected mesh graph with spatial adjacency edges to a + directed mesh graph (nx.DiGraph) based on the specified connectivity pattern. + + This is the second step in the two-step mesh creation process: + 1. Coordinate creation (create_single_level_2d_mesh_coordinates) -> nx.Graph + 2. Connectivity creation (this function) -> nx.DiGraph + + Parameters + ---------- + G_undirected : networkx.Graph + Undirected graph with nodes having "pos" attributes and edges having + "adjacency_type" attributes ("cardinal" or "diagonal"). + pattern : str + Connectivity pattern. Options: + - "4-star": only cardinal edges (horizontal/vertical neighbours) + - "8-star": all edges (cardinal + diagonal neighbours) + + Returns + ------- + networkx.DiGraph + Directed graph with bidirectional edges, each having "len" and "vdiff" + attributes. + """ + if pattern == "4-star": + # Filter to only cardinal edges (horizontal/vertical) + edges_to_use = [ + (u, v) + for u, v, d in G_undirected.edges(data=True) + if d.get("adjacency_type") == "cardinal" + ] + elif pattern == "8-star": + # Use all edges (cardinal + diagonal) + edges_to_use = list(G_undirected.edges()) + else: + raise ValueError( + f"Unknown connectivity pattern: '{pattern}'. " + "Choose '4-star' or '8-star'." + ) + + # Create filtered undirected graph with only selected edges + g_filtered = networkx.Graph() + g_filtered.add_nodes_from(G_undirected.nodes(data=True)) + g_filtered.add_edges_from(edges_to_use) + + # Convert to directed graph (creates edges in both directions) + dg = networkx.DiGraph(g_filtered) + for u, v in g_filtered.edges(): + d = np.sqrt( + np.sum( + (G_undirected.nodes[u]["pos"] - G_undirected.nodes[v]["pos"]) ** 2 + ) + ) dg.edges[u, v]["len"] = d - dg.edges[u, v]["vdiff"] = g.nodes[u]["pos"] - g.nodes[v]["pos"] - dg.add_edge(v, u) + dg.edges[u, v]["vdiff"] = ( + G_undirected.nodes[u]["pos"] - G_undirected.nodes[v]["pos"] + ) + # Ensure reverse edge exists and has attributes dg.edges[v, u]["len"] = d - dg.edges[v, u]["vdiff"] = g.nodes[v]["pos"] - g.nodes[u]["pos"] + dg.edges[v, u]["vdiff"] = ( + G_undirected.nodes[v]["pos"] - G_undirected.nodes[u]["pos"] + ) - dg.graph["dx"] = dx - dg.graph["dy"] = dy + # Preserve graph-level attributes (dx, dy, level, etc.) + dg.graph.update(G_undirected.graph) return dg -def create_multirange_2d_mesh_graphs( - max_num_levels, xy, mesh_node_distance=3, level_refinement_factor=3 +def create_single_level_2d_mesh_graph(xy, nx, ny): + """ + Create directed graph with nx * ny nodes representing a 2D grid with + positions spanning the range of xy coordinate values (first dimension + is assumed to be x and y coordinate values respectively). Each nodes is + connected to its eight nearest neighbours, both horizontally, vertically + and diagonally as directed edges (which means that the graph contains two + edges between each pair of connected nodes). + + The nodes contain a "pos" attribute with the x and y + coordinates of the node, and an "type" attribute with the + type of the node (i.e. "mesh" for mesh nodes). + + The edges contain a "len" attribute with the length of the edge + and a "vdiff" attribute with the vector difference between the + nodes. + + Internally, this uses the two-step process: + 1. create_single_level_2d_mesh_coordinates (coordinate creation) + 2. create_directed_mesh_graph (connectivity creation, pattern="8-star") + + Parameters + ---------- + xy : np.ndarray [N_grid_points, 2] + Grid point coordinates, with first column representing + x coordinates and second column y coordinates. N_grid_points is the + total number of grid points. + nx : int + Number of nodes in x direction + ny : int + Number of nodes in y direction + + Returns + ------- + networkx.DiGraph + Graph representing the 2D grid + """ + G_coords = create_single_level_2d_mesh_coordinates(xy, nx, ny) + return create_directed_mesh_graph(G_coords, pattern="8-star") + + +def create_multirange_2d_mesh_coordinates( + max_num_levels, xy, grid_spacing=3, interlevel_refinement_factor=3 ): """ - Create a list of 2D grid mesh graphs representing different levels of edge-length - scales spanning the spatial domain of the xy coordinates. - This list of graphs can then later be for example a) flattened into single graph - containing multiple ranges of connections or b) combined into a hierarchical graph. + Create a list of undirected coordinate graphs (nx.Graph) representing + different levels of mesh resolution spanning the spatial domain of the + xy coordinates. - Each graph in the list contains a "level" attribute with the level index of the graph. + This is the coordinate creation step for multi-level meshes. Each returned + graph contains nodes with spatial positions and edges annotated with + adjacency type ("cardinal" or "diagonal"). + + The graphs can be consumed by connectivity creation functions to produce + directed mesh graphs for flat_multiscale or hierarchical architectures. Parameters ---------- @@ -95,32 +203,35 @@ def create_multirange_2d_mesh_graphs( Number of edge-distance levels in mesh graph xy : np.ndarray Grid point coordinates, shaped [N_grid_points, 2] - mesh_node_distance: float + grid_spacing : float Distance (in x- and y-direction) between created mesh nodes, in coordinate system of xy - level_refinement_factor: float + interlevel_refinement_factor : float Refinement factor between grid points and bottom level of mesh hierarchy Returns ------- G_all_levels : list of networkx.Graph - List of networkx graphs for each level representing the connectivity - of the mesh within each level + List of undirected coordinate graphs for each level, each with + node positions and annotated spatial adjacency edges. + Each graph has "level" and "interlevel_refinement_factor" graph attributes. """ # Compute the size along x and y direction of area to cover with graph - # This is measured in the Cartesian coordiantes of xy + # This is measured in the Cartesian coordinates of xy coord_extent = np.ptp(xy, axis=0) # Number of nodes that would fit on bottom level of hierarchy, # in both directions - max_nodes_bottom = (coord_extent / mesh_node_distance).astype(int) + max_nodes_bottom = (coord_extent / grid_spacing).astype(int) # Find the number of mesh levels possible in x- and y-direction, # and the number of leaf nodes that would correspond to - # max_nodes_bottom/(level_refinement_factor^mesh_levels) = 1 - max_mesh_levels_float = np.log(max_nodes_bottom) / np.log(level_refinement_factor) + # max_nodes_bottom/(interlevel_refinement_factor^mesh_levels) = 1 + max_mesh_levels_float = np.log(max_nodes_bottom) / np.log( + interlevel_refinement_factor + ) max_mesh_levels = max_mesh_levels_float.astype(int) # (2,) - nleaf = level_refinement_factor**max_mesh_levels + nleaf = interlevel_refinement_factor**max_mesh_levels # leaves at the bottom in each direction, if using max_mesh_levels # As we can not instantiate different number of mesh levels in each @@ -137,14 +248,66 @@ def create_multirange_2d_mesh_graphs( G_all_levels = [] for lev in range(mesh_levels_to_create): # 0-index mesh levels # Compute number of nodes on level separate for each direction - nodes_x, nodes_y = (nleaf / (level_refinement_factor**lev)).astype(int) - g = create_single_level_2d_mesh_graph(xy, nodes_x, nodes_y) + nodes_x, nodes_y = ( + nleaf / (interlevel_refinement_factor**lev) + ).astype(int) + g = create_single_level_2d_mesh_coordinates(xy, nodes_x, nodes_y) # Add level information to nodes, edges and full graph for node in g.nodes: g.nodes[node]["level"] = lev for edge in g.edges: g.edges[edge]["level"] = lev g.graph["level"] = lev + # Store refinement factor so connectivity step can use it + g.graph["interlevel_refinement_factor"] = interlevel_refinement_factor G_all_levels.append(g) return G_all_levels + + +def create_multirange_2d_mesh_graphs( + max_num_levels, xy, mesh_node_distance=3, level_refinement_factor=3 +): + """ + Create a list of 2D grid mesh graphs representing different levels of edge-length + scales spanning the spatial domain of the xy coordinates. + This list of graphs can then later be for example a) flattened into single graph + containing multiple ranges of connections or b) combined into a hierarchical graph. + + Each graph in the list contains a "level" attribute with the level index of the graph. + + Internally uses the two-step process: + 1. create_multirange_2d_mesh_coordinates (coordinate creation) + 2. create_directed_mesh_graph (connectivity creation, pattern="8-star") + + Parameters + ---------- + max_num_levels : int + Number of edge-distance levels in mesh graph + xy : np.ndarray + Grid point coordinates, shaped [N_grid_points, 2] + mesh_node_distance: float + Distance (in x- and y-direction) between created mesh nodes, + in coordinate system of xy + level_refinement_factor: float + Refinement factor between grid points and bottom level of mesh hierarchy + + Returns + ------- + G_all_levels : list of networkx.Graph + List of networkx graphs for each level representing the connectivity + of the mesh within each level + """ + G_coords_list = create_multirange_2d_mesh_coordinates( + max_num_levels=max_num_levels, + xy=xy, + grid_spacing=mesh_node_distance, + interlevel_refinement_factor=level_refinement_factor, + ) + + G_all_levels = [] + for g_coords in G_coords_list: + g_directed = create_directed_mesh_graph(g_coords, pattern="8-star") + G_all_levels.append(g_directed) + + return G_all_levels From 63af99685af99c36ffbf66829476bdd8aa14fdab Mon Sep 17 00:00:00 2001 From: prajwal Date: Mon, 2 Mar 2026 19:34:43 +0530 Subject: [PATCH 02/21] test: add comprehensive tests for mesh_layout two-step architecture Add 46 new tests in test_mesh_layout.py covering: - Coordinate creation (nx.Graph with adjacency_type annotations) - Connectivity creation (4-star vs 8-star pattern filtering) - New API via create_all_graph_components (flat, flat_multiscale, hierarchical) - Backward compatibility with deprecation warnings - Caller dict non-mutation safety - Error handling (unsupported layouts, missing grid_spacing, etc.) - Equivalence between archetype functions and new API Also remove unused imports from base.py (old wrapper functions no longer called directly from the dispatch logic). --- src/weather_model_graphs/create/base.py | 3 - tests/test_mesh_layout.py | 747 ++++++++++++++++++++++++ 2 files changed, 747 insertions(+), 3 deletions(-) create mode 100644 tests/test_mesh_layout.py diff --git a/src/weather_model_graphs/create/base.py b/src/weather_model_graphs/create/base.py index 3f2f231..43e2ec0 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/weather_model_graphs/create/base.py @@ -26,13 +26,10 @@ from .grid import create_grid_graph_nodes from .mesh.kinds.flat import ( create_flat_multiscale_from_coordinates, - create_flat_multiscale_mesh_graph, create_flat_singlescale_from_coordinates, - create_flat_singlescale_mesh_graph, ) from .mesh.kinds.hierarchical import ( create_hierarchical_from_coordinates, - create_hierarchical_multiscale_mesh_graph, ) from .mesh.mesh import ( create_multirange_2d_mesh_coordinates, diff --git a/tests/test_mesh_layout.py b/tests/test_mesh_layout.py new file mode 100644 index 0000000..e78fd85 --- /dev/null +++ b/tests/test_mesh_layout.py @@ -0,0 +1,747 @@ +""" +Tests for the mesh_layout parameter and two-step coordinate/connectivity +architecture introduced in Issue #78. + +These tests verify: +1. The new API (mesh_layout, mesh_layout_kwargs, m2m_connectivity_kwargs with + intra_level/inter_level sub-dicts) +2. The two-step process (coordinate creation → connectivity creation) +3. The 4-star vs 8-star pattern functionality +4. Backward compatibility with old-style kwargs +5. Edge annotations on coordinate graphs +6. Error handling for invalid inputs +""" + +import warnings + +import networkx as nx +import numpy as np +import pytest + +import tests.utils as test_utils +import weather_model_graphs as wmg +from weather_model_graphs.create.mesh.mesh import ( + create_directed_mesh_graph, + create_multirange_2d_mesh_coordinates, + create_single_level_2d_mesh_coordinates, +) +from weather_model_graphs.create.mesh.kinds.flat import ( + create_flat_multiscale_from_coordinates, + create_flat_singlescale_from_coordinates, +) +from weather_model_graphs.create.mesh.kinds.hierarchical import ( + create_hierarchical_from_coordinates, +) + + +# ==================== +# Step 1: Coordinate creation tests +# ==================== + + +class TestSingleLevelCoordinateCreation: + """Tests for create_single_level_2d_mesh_coordinates.""" + + def test_returns_undirected_graph(self): + xy = test_utils.create_fake_xy(N=10) + G = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=5) + assert isinstance(G, nx.Graph) + assert not isinstance(G, nx.DiGraph) + + def test_nodes_have_pos_and_type(self): + xy = test_utils.create_fake_xy(N=10) + G = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + for node in G.nodes: + assert "pos" in G.nodes[node] + assert "type" in G.nodes[node] + assert G.nodes[node]["type"] == "mesh" + assert len(G.nodes[node]["pos"]) == 2 + + def test_correct_number_of_nodes(self): + xy = test_utils.create_fake_xy(N=10) + G = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=4) + assert len(G.nodes) == 5 * 4 + + def test_edges_have_adjacency_type(self): + xy = test_utils.create_fake_xy(N=10) + G = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=5) + cardinal_count = 0 + diagonal_count = 0 + for u, v, d in G.edges(data=True): + assert "adjacency_type" in d, f"Edge ({u}, {v}) missing adjacency_type" + assert d["adjacency_type"] in ("cardinal", "diagonal") + if d["adjacency_type"] == "cardinal": + cardinal_count += 1 + else: + diagonal_count += 1 + # For a 5x5 grid: cardinal = 2*(5*4) = 40, diagonal = 2*(4*4) = 32 + assert cardinal_count == 2 * (5 * 4) + assert diagonal_count == 2 * (4 * 4) + + def test_graph_has_dx_dy(self): + xy = test_utils.create_fake_xy(N=10) + G = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=5) + assert "dx" in G.graph + assert "dy" in G.graph + assert G.graph["dx"] > 0 + assert G.graph["dy"] > 0 + + +class TestMultirangeCoordinateCreation: + """Tests for create_multirange_2d_mesh_coordinates.""" + + def test_returns_list_of_undirected_graphs(self): + xy = test_utils.create_fake_xy(N=30) + G_list = create_multirange_2d_mesh_coordinates( + max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + assert isinstance(G_list, list) + assert len(G_list) > 0 + for G in G_list: + assert isinstance(G, nx.Graph) + assert not isinstance(G, nx.DiGraph) + + def test_each_level_has_level_attribute(self): + xy = test_utils.create_fake_xy(N=30) + G_list = create_multirange_2d_mesh_coordinates( + max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + for i, G in enumerate(G_list): + assert G.graph["level"] == i + for node in G.nodes: + assert G.nodes[node]["level"] == i + + def test_interlevel_refinement_factor_stored(self): + xy = test_utils.create_fake_xy(N=30) + G_list = create_multirange_2d_mesh_coordinates( + max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + for G in G_list: + assert G.graph["interlevel_refinement_factor"] == 3 + + def test_edges_have_adjacency_type(self): + xy = test_utils.create_fake_xy(N=30) + G_list = create_multirange_2d_mesh_coordinates( + max_num_levels=2, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + for G in G_list: + for u, v, d in G.edges(data=True): + assert "adjacency_type" in d + + def test_coarser_levels_have_fewer_nodes(self): + xy = test_utils.create_fake_xy(N=30) + G_list = create_multirange_2d_mesh_coordinates( + max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + if len(G_list) >= 2: + for i in range(len(G_list) - 1): + assert len(G_list[i].nodes) > len(G_list[i + 1].nodes) + + +# ==================== +# Step 2: Connectivity creation tests +# ==================== + + +class TestDirectedMeshGraph: + """Tests for create_directed_mesh_graph.""" + + def test_returns_directed_graph(self): + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G_directed = create_directed_mesh_graph(G_coords, pattern="8-star") + assert isinstance(G_directed, nx.DiGraph) + + def test_4star_has_fewer_edges_than_8star(self): + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=5) + G_4star = create_directed_mesh_graph(G_coords, pattern="4-star") + G_8star = create_directed_mesh_graph(G_coords, pattern="8-star") + assert len(G_4star.edges) < len(G_8star.edges) + + def test_4star_only_cardinal_edges(self): + """4-star should only include cardinal (horizontal/vertical) edges.""" + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G_4star = create_directed_mesh_graph(G_coords, pattern="4-star") + + # In a 4x4 grid, 4-star adjacency means each node connects only to + # horizontal/vertical neighbours + # For a 4x4 grid: 2 * (4*3 + 3*4) = 2 * 24 = 48 directed edges + expected_edges = 2 * (4 * 3 + 3 * 4) + assert len(G_4star.edges) == expected_edges + + def test_8star_includes_diagonal_edges(self): + """8-star should include both cardinal and diagonal edges.""" + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G_8star = create_directed_mesh_graph(G_coords, pattern="8-star") + + # Cardinal: 2 * (4*3 + 3*4) = 48 + # Diagonal: 2 * 2 * (3*3) = 36 + expected_edges = 48 + 36 + assert len(G_8star.edges) == expected_edges + + def test_edges_have_len_and_vdiff(self): + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G_directed = create_directed_mesh_graph(G_coords, pattern="8-star") + for u, v, d in G_directed.edges(data=True): + assert "len" in d + assert "vdiff" in d + assert d["len"] > 0 + + def test_invalid_pattern_raises_error(self): + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + with pytest.raises(ValueError, match="Unknown connectivity pattern"): + create_directed_mesh_graph(G_coords, pattern="6-star") + + def test_preserves_graph_attributes(self): + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G_directed = create_directed_mesh_graph(G_coords, pattern="8-star") + assert "dx" in G_directed.graph + assert "dy" in G_directed.graph + + def test_bidirectional_edges(self): + """Each undirected edge should produce two directed edges.""" + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_coordinates(xy, nx=3, ny=3) + G_directed = create_directed_mesh_graph(G_coords, pattern="4-star") + for u, v in G_directed.edges(): + assert G_directed.has_edge(v, u), f"Missing reverse edge ({v}, {u})" + + +class TestFlatSinglescaleFromCoordinates: + """Tests for create_flat_singlescale_from_coordinates.""" + + def test_basic_creation(self): + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=5) + G = create_flat_singlescale_from_coordinates(G_coords, pattern="8-star") + assert isinstance(G, nx.DiGraph) + + def test_4star_pattern(self): + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=5) + G = create_flat_singlescale_from_coordinates(G_coords, pattern="4-star") + assert isinstance(G, nx.DiGraph) + # Fewer edges than 8-star + G_8 = create_flat_singlescale_from_coordinates(G_coords, pattern="8-star") + assert len(G.edges) < len(G_8.edges) + + +class TestFlatMultiscaleFromCoordinates: + """Tests for create_flat_multiscale_from_coordinates.""" + + def test_basic_creation(self): + xy = test_utils.create_fake_xy(N=30) + G_coords_list = create_multirange_2d_mesh_coordinates( + max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + G = create_flat_multiscale_from_coordinates(G_coords_list) + assert isinstance(G, nx.DiGraph) + + def test_intra_level_pattern(self): + xy = test_utils.create_fake_xy(N=30) + G_coords_list = create_multirange_2d_mesh_coordinates( + max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + G_4star = create_flat_multiscale_from_coordinates( + G_coords_list, + intra_level={"pattern": "4-star"}, + inter_level={"pattern": "coincident"}, + ) + G_8star = create_flat_multiscale_from_coordinates( + G_coords_list, + intra_level={"pattern": "8-star"}, + inter_level={"pattern": "coincident"}, + ) + assert len(G_4star.edges) < len(G_8star.edges) + + def test_invalid_inter_level_pattern_raises(self): + xy = test_utils.create_fake_xy(N=30) + G_coords_list = create_multirange_2d_mesh_coordinates( + max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + with pytest.raises(NotImplementedError, match="Inter-level pattern"): + create_flat_multiscale_from_coordinates( + G_coords_list, + inter_level={"pattern": "some_unknown"}, + ) + + +class TestHierarchicalFromCoordinates: + """Tests for create_hierarchical_from_coordinates.""" + + def test_basic_creation(self): + xy = test_utils.create_fake_xy(N=30) + G_coords_list = create_multirange_2d_mesh_coordinates( + max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + G = create_hierarchical_from_coordinates(G_coords_list) + assert isinstance(G, nx.DiGraph) + + def test_has_up_down_same_edges(self): + xy = test_utils.create_fake_xy(N=30) + G_coords_list = create_multirange_2d_mesh_coordinates( + max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + G = create_hierarchical_from_coordinates(G_coords_list) + directions = set() + for u, v, d in G.edges(data=True): + if "direction" in d: + directions.add(d["direction"]) + assert "same" in directions + assert "up" in directions + assert "down" in directions + + def test_intra_level_pattern(self): + xy = test_utils.create_fake_xy(N=30) + G_coords_list = create_multirange_2d_mesh_coordinates( + max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + G_4star = create_hierarchical_from_coordinates( + G_coords_list, + intra_level={"pattern": "4-star"}, + ) + G_8star = create_hierarchical_from_coordinates( + G_coords_list, + intra_level={"pattern": "8-star"}, + ) + assert len(G_4star.edges) < len(G_8star.edges) + + def test_inter_level_k_parameter(self): + xy = test_utils.create_fake_xy(N=30) + G_coords_list = create_multirange_2d_mesh_coordinates( + max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + G_k1 = create_hierarchical_from_coordinates( + G_coords_list, + inter_level={"pattern": "nearest", "k": 1}, + ) + G_k3 = create_hierarchical_from_coordinates( + G_coords_list, + inter_level={"pattern": "nearest", "k": 3}, + ) + # More neighbours → more inter-level edges + assert len(G_k3.edges) > len(G_k1.edges) + + def test_invalid_inter_level_pattern_raises(self): + xy = test_utils.create_fake_xy(N=30) + G_coords_list = create_multirange_2d_mesh_coordinates( + max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + with pytest.raises(NotImplementedError, match="Inter-level pattern"): + create_hierarchical_from_coordinates( + G_coords_list, + inter_level={"pattern": "some_unknown"}, + ) + + +# ==================== +# New API via create_all_graph_components tests +# ==================== + + +class TestNewAPIFlat: + """Tests for create_all_graph_components with new mesh_layout API (flat).""" + + def test_flat_with_new_api(self): + xy = test_utils.create_fake_xy(N=32) + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(grid_spacing=3), + m2m_connectivity_kwargs=dict(pattern="8-star"), + g2m_connectivity="within_radius", + m2g_connectivity="nearest_neighbours", + g2m_connectivity_kwargs=dict(rel_max_dist=0.51), + m2g_connectivity_kwargs=dict(max_num_neighbours=4), + ) + assert isinstance(graph, nx.DiGraph) + + def test_flat_4star_pattern(self): + xy = test_utils.create_fake_xy(N=32) + graph_4 = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(grid_spacing=3), + m2m_connectivity_kwargs=dict(pattern="4-star"), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + graph_8 = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(grid_spacing=3), + m2m_connectivity_kwargs=dict(pattern="8-star"), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + assert isinstance(graph_4, nx.DiGraph) + assert isinstance(graph_8, nx.DiGraph) + assert len(graph_4.edges) < len(graph_8.edges) + + def test_missing_grid_spacing_raises(self): + xy = test_utils.create_fake_xy(N=32) + with pytest.raises(ValueError, match="grid_spacing"): + wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout="rectilinear", + mesh_layout_kwargs={}, + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + + +class TestNewAPIFlatMultiscale: + """Tests for create_all_graph_components with new API (flat_multiscale).""" + + def test_flat_multiscale_with_new_api(self): + xy = test_utils.create_fake_xy(N=32) + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat_multiscale", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict( + grid_spacing=3, + interlevel_refinement_factor=3, + max_num_levels=3, + ), + m2m_connectivity_kwargs=dict( + intra_level=dict(pattern="8-star"), + inter_level=dict(pattern="coincident"), + ), + g2m_connectivity="within_radius", + m2g_connectivity="nearest_neighbours", + g2m_connectivity_kwargs=dict(rel_max_dist=0.51), + m2g_connectivity_kwargs=dict(max_num_neighbours=4), + ) + assert isinstance(graph, nx.DiGraph) + + def test_flat_multiscale_4star_intra(self): + xy = test_utils.create_fake_xy(N=32) + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat_multiscale", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict( + grid_spacing=3, + interlevel_refinement_factor=3, + max_num_levels=3, + ), + m2m_connectivity_kwargs=dict( + intra_level=dict(pattern="4-star"), + inter_level=dict(pattern="coincident"), + ), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + assert isinstance(graph, nx.DiGraph) + + +class TestNewAPIHierarchical: + """Tests for create_all_graph_components with new API (hierarchical).""" + + def test_hierarchical_with_new_api(self): + xy = test_utils.create_fake_xy(N=32) + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="hierarchical", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict( + grid_spacing=3, + interlevel_refinement_factor=3, + max_num_levels=3, + ), + m2m_connectivity_kwargs=dict( + intra_level=dict(pattern="8-star"), + inter_level=dict(pattern="nearest", k=1), + ), + g2m_connectivity="within_radius", + m2g_connectivity="nearest_neighbours", + g2m_connectivity_kwargs=dict(rel_max_dist=0.51), + m2g_connectivity_kwargs=dict(max_num_neighbours=4), + ) + assert isinstance(graph, nx.DiGraph) + + def test_hierarchical_4star_intra(self): + xy = test_utils.create_fake_xy(N=32) + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="hierarchical", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict( + grid_spacing=3, + interlevel_refinement_factor=3, + max_num_levels=3, + ), + m2m_connectivity_kwargs=dict( + intra_level=dict(pattern="4-star"), + inter_level=dict(pattern="nearest", k=1), + ), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + assert isinstance(graph, nx.DiGraph) + + def test_hierarchical_k3_nearest(self): + xy = test_utils.create_fake_xy(N=32) + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="hierarchical", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict( + grid_spacing=3, + interlevel_refinement_factor=3, + max_num_levels=3, + ), + m2m_connectivity_kwargs=dict( + intra_level=dict(pattern="8-star"), + inter_level=dict(pattern="nearest", k=3), + ), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + assert isinstance(graph, nx.DiGraph) + + +# ==================== +# Backward compatibility tests +# ==================== + + +class TestBackwardCompatibility: + """Tests that old-style kwargs still work with deprecation warnings.""" + + def test_old_style_flat_with_mesh_node_distance(self): + xy = test_utils.create_fake_xy(N=32) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + m2m_connectivity_kwargs=dict(mesh_node_distance=3), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + # Should have deprecation warning + deprecation_warnings = [ + x for x in w if issubclass(x.category, DeprecationWarning) + ] + assert len(deprecation_warnings) >= 1 + assert "mesh_node_distance" in str(deprecation_warnings[0].message) + assert isinstance(graph, nx.DiGraph) + + def test_old_style_flat_multiscale(self): + xy = test_utils.create_fake_xy(N=32) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat_multiscale", + m2m_connectivity_kwargs=dict( + mesh_node_distance=3, + level_refinement_factor=3, + max_num_levels=3, + ), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + deprecation_warnings = [ + x for x in w if issubclass(x.category, DeprecationWarning) + ] + assert len(deprecation_warnings) >= 3 # 3 migrated kwargs + assert isinstance(graph, nx.DiGraph) + + def test_old_style_hierarchical(self): + xy = test_utils.create_fake_xy(N=32) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="hierarchical", + m2m_connectivity_kwargs=dict( + mesh_node_distance=3, + level_refinement_factor=3, + max_num_levels=3, + ), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + deprecation_warnings = [ + x for x in w if issubclass(x.category, DeprecationWarning) + ] + assert len(deprecation_warnings) >= 3 + assert isinstance(graph, nx.DiGraph) + + def test_kwargs_dict_not_mutated(self): + """Verify that passing dict kwargs doesn't mutate the caller's dict.""" + xy = test_utils.create_fake_xy(N=32) + original_kwargs = dict( + mesh_node_distance=3, + level_refinement_factor=3, + max_num_levels=3, + ) + kwargs_copy = original_kwargs.copy() + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="hierarchical", + m2m_connectivity_kwargs=original_kwargs, + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + + # The original dict should not be modified + assert original_kwargs == kwargs_copy + + +# ==================== +# Error handling tests +# ==================== + + +class TestErrorHandling: + """Tests for proper error handling.""" + + def test_unsupported_mesh_layout_raises(self): + xy = test_utils.create_fake_xy(N=32) + with pytest.raises(NotImplementedError, match="not yet supported"): + wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout="triangular", + mesh_layout_kwargs=dict(grid_spacing=3), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + + def test_unsupported_m2m_connectivity_raises(self): + xy = test_utils.create_fake_xy(N=32) + with pytest.raises(NotImplementedError, match="not implemented"): + wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="some_unknown", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(grid_spacing=3), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + + def test_grid_spacing_too_large_raises(self): + xy = test_utils.create_fake_xy(N=10) + with pytest.raises(ValueError, match="too large"): + wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(grid_spacing=100), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + + +# ==================== +# Equivalence tests: new API == old wrappers +# ==================== + + +class TestEquivalence: + """Verify that the new API produces equivalent results to the old wrappers.""" + + def test_keisler_archetype_matches_new_api(self): + """The keisler archetype function should produce the same result as + calling create_all_graph_components with the new API directly.""" + xy = test_utils.create_fake_xy(N=32) + + graph_archetype = wmg.create.archetype.create_keisler_graph( + coords=xy, mesh_node_distance=3 + ) + + graph_new_api = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(grid_spacing=3), + m2m_connectivity_kwargs=dict(pattern="8-star"), + g2m_connectivity="within_radius", + m2g_connectivity="nearest_neighbours", + g2m_connectivity_kwargs=dict(rel_max_dist=0.51), + m2g_connectivity_kwargs=dict(max_num_neighbours=4), + ) + + assert len(graph_archetype.nodes) == len(graph_new_api.nodes) + assert len(graph_archetype.edges) == len(graph_new_api.edges) + + def test_graphcast_archetype_matches_new_api(self): + xy = test_utils.create_fake_xy(N=32) + + graph_archetype = wmg.create.archetype.create_graphcast_graph( + coords=xy, + mesh_node_distance=3, + level_refinement_factor=3, + max_num_levels=3, + ) + + graph_new_api = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat_multiscale", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict( + grid_spacing=3, + interlevel_refinement_factor=3, + max_num_levels=3, + ), + m2m_connectivity_kwargs=dict( + intra_level=dict(pattern="8-star"), + inter_level=dict(pattern="coincident"), + ), + g2m_connectivity="within_radius", + m2g_connectivity="nearest_neighbours", + g2m_connectivity_kwargs=dict(rel_max_dist=0.51), + m2g_connectivity_kwargs=dict(max_num_neighbours=4), + ) + + assert len(graph_archetype.nodes) == len(graph_new_api.nodes) + assert len(graph_archetype.edges) == len(graph_new_api.edges) + + def test_oskarsson_archetype_matches_new_api(self): + xy = test_utils.create_fake_xy(N=32) + + graph_archetype = wmg.create.archetype.create_oskarsson_hierarchical_graph( + coords=xy, + mesh_node_distance=3, + level_refinement_factor=3, + max_num_levels=3, + ) + + graph_new_api = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="hierarchical", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict( + grid_spacing=3, + interlevel_refinement_factor=3, + max_num_levels=3, + ), + m2m_connectivity_kwargs=dict( + intra_level=dict(pattern="8-star"), + inter_level=dict(pattern="nearest", k=1), + ), + g2m_connectivity="within_radius", + m2g_connectivity="nearest_neighbours", + g2m_connectivity_kwargs=dict(rel_max_dist=0.51), + m2g_connectivity_kwargs=dict(max_num_neighbours=4), + ) + + assert len(graph_archetype.nodes) == len(graph_new_api.nodes) + assert len(graph_archetype.edges) == len(graph_new_api.edges) From f417962f45a228ba2fceb28ed35b16408bab05d4 Mon Sep 17 00:00:00 2001 From: prajwal Date: Mon, 2 Mar 2026 22:55:49 +0530 Subject: [PATCH 03/21] =?UTF-8?q?refactor:=20align=20API=20with=20Leif's?= =?UTF-8?q?=20final=20table=20=E2=80=94=20rename=20params,=20simplify=20fl?= =?UTF-8?q?at=5Fmultiscale=20kwargs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Rename interlevel_refinement_factor → refinement_factor, max_num_levels → max_num_refinement_levels - flat_multiscale uses simple pattern='8-star' (not intra_level/inter_level sub-dicts) - Only hierarchical uses intra_level/inter_level sub-dicts - Update archetypes and backward compat migration messages - Add 36 comprehensive edge case tests (81 total mesh layout tests pass) --- src/weather_model_graphs/create/archetype.py | 11 +- src/weather_model_graphs/create/base.py | 48 +- .../create/mesh/kinds/flat.py | 38 +- tests/test_mesh_layout.py | 651 ++++++++++++++++-- 4 files changed, 649 insertions(+), 99 deletions(-) diff --git a/src/weather_model_graphs/create/archetype.py b/src/weather_model_graphs/create/archetype.py index 13a79ce..ac76fcc 100644 --- a/src/weather_model_graphs/create/archetype.py +++ b/src/weather_model_graphs/create/archetype.py @@ -138,12 +138,11 @@ def create_graphcast_graph( mesh_layout="rectilinear", mesh_layout_kwargs=dict( grid_spacing=mesh_node_distance, - interlevel_refinement_factor=level_refinement_factor, - max_num_levels=max_num_levels, + refinement_factor=level_refinement_factor, + max_num_refinement_levels=max_num_levels, ), m2m_connectivity_kwargs=dict( - intra_level=dict(pattern="8-star"), - inter_level=dict(pattern="coincident"), + pattern="8-star", ), g2m_connectivity="within_radius", m2g_connectivity="nearest_neighbours", @@ -227,8 +226,8 @@ def create_oskarsson_hierarchical_graph( mesh_layout="rectilinear", mesh_layout_kwargs=dict( grid_spacing=mesh_node_distance, - interlevel_refinement_factor=level_refinement_factor, - max_num_levels=max_num_levels, + refinement_factor=level_refinement_factor, + max_num_refinement_levels=max_num_levels, ), m2m_connectivity_kwargs=dict( intra_level=dict(pattern="8-star"), diff --git a/src/weather_model_graphs/create/base.py b/src/weather_model_graphs/create/base.py index 43e2ec0..a6f81a6 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/weather_model_graphs/create/base.py @@ -85,14 +85,14 @@ def create_all_graph_components( mesh_layout_kwargs (for mesh_layout="rectilinear"): - grid_spacing: float, distance between mesh nodes in coordinate units - - interlevel_refinement_factor: int, refinement factor between levels (for multi-level) - - max_num_levels: int, maximum number of mesh levels (for multi-level) + - refinement_factor: int, refinement factor between levels (for multi-level) + - max_num_refinement_levels: int, maximum number of mesh levels (for multi-level) m2m_connectivity: - "flat": Create a single-level directed mesh graph. m2m_connectivity_kwargs: pattern="4-star" or "8-star" (default: "8-star") - "flat_multiscale": Create a flat multiscale mesh graph. - m2m_connectivity_kwargs: intra_level=dict(pattern=...), inter_level=dict(pattern=...) + m2m_connectivity_kwargs: pattern="4-star" or "8-star" (default: "8-star") - "hierarchical": Create a hierarchical mesh graph with up/down connections. m2m_connectivity_kwargs: intra_level=dict(pattern=...), inter_level=dict(pattern=..., k=...) @@ -144,24 +144,24 @@ def create_all_graph_components( mesh_layout_kwargs["grid_spacing"] = m2m_connectivity_kwargs.pop( "mesh_node_distance" ) - if "level_refinement_factor" in m2m_connectivity_kwargs and "interlevel_refinement_factor" not in mesh_layout_kwargs: + if "level_refinement_factor" in m2m_connectivity_kwargs and "refinement_factor" not in mesh_layout_kwargs: warnings.warn( "Passing 'level_refinement_factor' in m2m_connectivity_kwargs is deprecated. " - "Use mesh_layout_kwargs=dict(interlevel_refinement_factor=...) instead.", + "Use mesh_layout_kwargs=dict(refinement_factor=...) instead.", DeprecationWarning, stacklevel=2, ) - mesh_layout_kwargs["interlevel_refinement_factor"] = ( + mesh_layout_kwargs["refinement_factor"] = ( m2m_connectivity_kwargs.pop("level_refinement_factor") ) - if "max_num_levels" in m2m_connectivity_kwargs and "max_num_levels" not in mesh_layout_kwargs: + if "max_num_levels" in m2m_connectivity_kwargs and "max_num_refinement_levels" not in mesh_layout_kwargs: warnings.warn( "Passing 'max_num_levels' in m2m_connectivity_kwargs is deprecated. " - "Use mesh_layout_kwargs=dict(max_num_levels=...) instead.", + "Use mesh_layout_kwargs=dict(max_num_refinement_levels=...) instead.", DeprecationWarning, stacklevel=2, ) - mesh_layout_kwargs["max_num_levels"] = m2m_connectivity_kwargs.pop( + mesh_layout_kwargs["max_num_refinement_levels"] = m2m_connectivity_kwargs.pop( "max_num_levels" ) @@ -235,20 +235,20 @@ def create_all_graph_components( # --- Step 1: Coordinate creation based on mesh_layout --- if mesh_layout == "rectilinear": grid_spacing = mesh_layout_kwargs.get("grid_spacing") - interlevel_refinement_factor = mesh_layout_kwargs.get( - "interlevel_refinement_factor" + refinement_factor = mesh_layout_kwargs.get("refinement_factor") + max_num_refinement_levels = mesh_layout_kwargs.get( + "max_num_refinement_levels" ) - max_num_levels = mesh_layout_kwargs.get("max_num_levels") if grid_spacing is None: raise ValueError( "mesh_layout='rectilinear' with m2m_connectivity='hierarchical' " "requires 'grid_spacing' in mesh_layout_kwargs." ) G_coords_list = create_multirange_2d_mesh_coordinates( - max_num_levels=max_num_levels, + max_num_levels=max_num_refinement_levels, xy=xy, grid_spacing=grid_spacing, - interlevel_refinement_factor=interlevel_refinement_factor, + interlevel_refinement_factor=refinement_factor, ) else: raise NotImplementedError( @@ -280,20 +280,20 @@ def create_all_graph_components( # --- Step 1: Coordinate creation based on mesh_layout --- if mesh_layout == "rectilinear": grid_spacing = mesh_layout_kwargs.get("grid_spacing") - interlevel_refinement_factor = mesh_layout_kwargs.get( - "interlevel_refinement_factor" + refinement_factor = mesh_layout_kwargs.get("refinement_factor") + max_num_refinement_levels = mesh_layout_kwargs.get( + "max_num_refinement_levels" ) - max_num_levels = mesh_layout_kwargs.get("max_num_levels") if grid_spacing is None: raise ValueError( "mesh_layout='rectilinear' with m2m_connectivity='flat_multiscale' " "requires 'grid_spacing' in mesh_layout_kwargs." ) G_coords_list = create_multirange_2d_mesh_coordinates( - max_num_levels=max_num_levels, + max_num_levels=max_num_refinement_levels, xy=xy, grid_spacing=grid_spacing, - interlevel_refinement_factor=interlevel_refinement_factor, + interlevel_refinement_factor=refinement_factor, ) else: raise NotImplementedError( @@ -302,16 +302,10 @@ def create_all_graph_components( ) # --- Step 2: Connectivity creation --- - intra_level = m2m_connectivity_kwargs.get( - "intra_level", {"pattern": "8-star"} - ) - inter_level = m2m_connectivity_kwargs.get( - "inter_level", {"pattern": "coincident"} - ) + pattern = m2m_connectivity_kwargs.get("pattern", "8-star") graph_components["m2m"] = create_flat_multiscale_from_coordinates( G_coords_list, - intra_level=intra_level, - inter_level=inter_level, + pattern=pattern, ) grid_connect_graph = graph_components["m2m"] else: diff --git a/src/weather_model_graphs/create/mesh/kinds/flat.py b/src/weather_model_graphs/create/mesh/kinds/flat.py index b2cbca3..eca018a 100644 --- a/src/weather_model_graphs/create/mesh/kinds/flat.py +++ b/src/weather_model_graphs/create/mesh/kinds/flat.py @@ -7,15 +7,18 @@ def create_flat_multiscale_from_coordinates( G_coords_list, - intra_level=None, - inter_level=None, + pattern="8-star", ): """ Create flat multiscale mesh graph from a list of coordinate graphs. This is the connectivity creation step for flat multiscale meshes. It takes undirected coordinate graphs (one per level) and produces a - single directed mesh graph with intra-level connectivity. + single directed mesh graph where all levels are merged into one flat graph. + + In a flat multiscale graph, coarser levels are merged into the finer level + by coincident node positions (no separate inter-level connectivity needed). + The ``pattern`` controls the intra-level edge connectivity for each level. Parameters ---------- @@ -23,33 +26,15 @@ def create_flat_multiscale_from_coordinates( List of undirected coordinate graphs, one per level. Each should have nodes with "pos" and "type" attributes, and edges with "adjacency_type" attributes. Created by create_multirange_2d_mesh_coordinates. - intra_level : dict or None - Intra-level connectivity options. Supports: - - pattern: str, "4-star" or "8-star" (default: "8-star") - inter_level : dict or None - Inter-level connectivity options. Supports: - - pattern: str, "coincident" (default: "coincident") - Currently only "coincident" is supported (coarser level nodes are - exactly coincident with a subset of finer level nodes). + pattern : str + Connectivity pattern for intra-level edges: "4-star" or "8-star" + (default: "8-star") Returns ------- G_tot : networkx.DiGraph The merged flat multiscale mesh graph """ - if intra_level is None: - intra_level = {"pattern": "8-star"} - if inter_level is None: - inter_level = {"pattern": "coincident"} - - intra_pattern = intra_level.get("pattern", "8-star") - inter_pattern = inter_level.get("pattern", "coincident") - - if inter_pattern != "coincident": - raise NotImplementedError( - f"Inter-level pattern '{inter_pattern}' is not yet supported. " - "Only 'coincident' is currently implemented." - ) # Retrieve interlevel_refinement_factor from graph attributes interlevel_refinement_factor = G_coords_list[0].graph.get( @@ -68,7 +53,7 @@ def create_flat_multiscale_from_coordinates( # Convert each level's coordinate graph to directed graph with chosen pattern G_all_levels = [ - mesh_graph.create_directed_mesh_graph(g_coords, pattern=intra_pattern) + mesh_graph.create_directed_mesh_graph(g_coords, pattern=pattern) for g_coords in G_coords_list ] @@ -184,8 +169,7 @@ def create_flat_multiscale_mesh_graph( return create_flat_multiscale_from_coordinates( G_coords_list, - intra_level={"pattern": "8-star"}, - inter_level={"pattern": "coincident"}, + pattern="8-star", ) diff --git a/tests/test_mesh_layout.py b/tests/test_mesh_layout.py index e78fd85..b131925 100644 --- a/tests/test_mesh_layout.py +++ b/tests/test_mesh_layout.py @@ -3,8 +3,9 @@ architecture introduced in Issue #78. These tests verify: -1. The new API (mesh_layout, mesh_layout_kwargs, m2m_connectivity_kwargs with - intra_level/inter_level sub-dicts) +1. The new API (mesh_layout, mesh_layout_kwargs with refinement_factor and + max_num_refinement_levels, m2m_connectivity_kwargs with pattern for flat/ + flat_multiscale and intra_level/inter_level sub-dicts for hierarchical) 2. The two-step process (coordinate creation → connectivity creation) 3. The 4-star vs 8-star pattern functionality 4. Backward compatibility with old-style kwargs @@ -243,34 +244,21 @@ def test_basic_creation(self): G = create_flat_multiscale_from_coordinates(G_coords_list) assert isinstance(G, nx.DiGraph) - def test_intra_level_pattern(self): + def test_pattern_argument(self): xy = test_utils.create_fake_xy(N=30) G_coords_list = create_multirange_2d_mesh_coordinates( max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 ) G_4star = create_flat_multiscale_from_coordinates( G_coords_list, - intra_level={"pattern": "4-star"}, - inter_level={"pattern": "coincident"}, + pattern="4-star", ) G_8star = create_flat_multiscale_from_coordinates( G_coords_list, - intra_level={"pattern": "8-star"}, - inter_level={"pattern": "coincident"}, + pattern="8-star", ) assert len(G_4star.edges) < len(G_8star.edges) - def test_invalid_inter_level_pattern_raises(self): - xy = test_utils.create_fake_xy(N=30) - G_coords_list = create_multirange_2d_mesh_coordinates( - max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 - ) - with pytest.raises(NotImplementedError, match="Inter-level pattern"): - create_flat_multiscale_from_coordinates( - G_coords_list, - inter_level={"pattern": "some_unknown"}, - ) - class TestHierarchicalFromCoordinates: """Tests for create_hierarchical_from_coordinates.""" @@ -411,12 +399,11 @@ def test_flat_multiscale_with_new_api(self): mesh_layout="rectilinear", mesh_layout_kwargs=dict( grid_spacing=3, - interlevel_refinement_factor=3, - max_num_levels=3, + refinement_factor=3, + max_num_refinement_levels=3, ), m2m_connectivity_kwargs=dict( - intra_level=dict(pattern="8-star"), - inter_level=dict(pattern="coincident"), + pattern="8-star", ), g2m_connectivity="within_radius", m2g_connectivity="nearest_neighbours", @@ -425,7 +412,7 @@ def test_flat_multiscale_with_new_api(self): ) assert isinstance(graph, nx.DiGraph) - def test_flat_multiscale_4star_intra(self): + def test_flat_multiscale_4star_pattern(self): xy = test_utils.create_fake_xy(N=32) graph = wmg.create.create_all_graph_components( coords=xy, @@ -433,12 +420,11 @@ def test_flat_multiscale_4star_intra(self): mesh_layout="rectilinear", mesh_layout_kwargs=dict( grid_spacing=3, - interlevel_refinement_factor=3, - max_num_levels=3, + refinement_factor=3, + max_num_refinement_levels=3, ), m2m_connectivity_kwargs=dict( - intra_level=dict(pattern="4-star"), - inter_level=dict(pattern="coincident"), + pattern="4-star", ), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", @@ -457,8 +443,8 @@ def test_hierarchical_with_new_api(self): mesh_layout="rectilinear", mesh_layout_kwargs=dict( grid_spacing=3, - interlevel_refinement_factor=3, - max_num_levels=3, + refinement_factor=3, + max_num_refinement_levels=3, ), m2m_connectivity_kwargs=dict( intra_level=dict(pattern="8-star"), @@ -479,8 +465,8 @@ def test_hierarchical_4star_intra(self): mesh_layout="rectilinear", mesh_layout_kwargs=dict( grid_spacing=3, - interlevel_refinement_factor=3, - max_num_levels=3, + refinement_factor=3, + max_num_refinement_levels=3, ), m2m_connectivity_kwargs=dict( intra_level=dict(pattern="4-star"), @@ -499,8 +485,8 @@ def test_hierarchical_k3_nearest(self): mesh_layout="rectilinear", mesh_layout_kwargs=dict( grid_spacing=3, - interlevel_refinement_factor=3, - max_num_levels=3, + refinement_factor=3, + max_num_refinement_levels=3, ), m2m_connectivity_kwargs=dict( intra_level=dict(pattern="8-star"), @@ -698,12 +684,11 @@ def test_graphcast_archetype_matches_new_api(self): mesh_layout="rectilinear", mesh_layout_kwargs=dict( grid_spacing=3, - interlevel_refinement_factor=3, - max_num_levels=3, + refinement_factor=3, + max_num_refinement_levels=3, ), m2m_connectivity_kwargs=dict( - intra_level=dict(pattern="8-star"), - inter_level=dict(pattern="coincident"), + pattern="8-star", ), g2m_connectivity="within_radius", m2g_connectivity="nearest_neighbours", @@ -730,8 +715,8 @@ def test_oskarsson_archetype_matches_new_api(self): mesh_layout="rectilinear", mesh_layout_kwargs=dict( grid_spacing=3, - interlevel_refinement_factor=3, - max_num_levels=3, + refinement_factor=3, + max_num_refinement_levels=3, ), m2m_connectivity_kwargs=dict( intra_level=dict(pattern="8-star"), @@ -745,3 +730,591 @@ def test_oskarsson_archetype_matches_new_api(self): assert len(graph_archetype.nodes) == len(graph_new_api.nodes) assert len(graph_archetype.edges) == len(graph_new_api.edges) + + +# ==================== +# Edge case tests +# ==================== + + +class TestCoordinateCreationEdgeCases: + """Edge cases for coordinate creation step.""" + + def test_minimum_grid_2x2(self): + """Smallest possible grid: 2x2 nodes.""" + xy = test_utils.create_fake_xy(N=10) + G = create_single_level_2d_mesh_coordinates(xy, nx=2, ny=2) + assert len(G.nodes) == 4 + # 2x2 grid: cardinal edges = 2*(2*1) = 4, diagonal edges = 2*(1*1) = 2 + cardinal = sum( + 1 for _, _, d in G.edges(data=True) if d["adjacency_type"] == "cardinal" + ) + diagonal = sum( + 1 for _, _, d in G.edges(data=True) if d["adjacency_type"] == "diagonal" + ) + assert cardinal == 4 + assert diagonal == 2 + + def test_single_row_grid(self): + """Grid with only 1 row (nx=5, ny=1).""" + xy = test_utils.create_fake_xy(N=10) + G = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=1) + assert len(G.nodes) == 5 + # 5x1 grid: only horizontal cardinal edges, no diagonals + cardinal = sum( + 1 for _, _, d in G.edges(data=True) if d["adjacency_type"] == "cardinal" + ) + diagonal = sum( + 1 for _, _, d in G.edges(data=True) if d["adjacency_type"] == "diagonal" + ) + assert cardinal == 4 # 5-1 = 4 horizontal edges + assert diagonal == 0 + + def test_single_column_grid(self): + """Grid with only 1 column (nx=1, ny=5).""" + xy = test_utils.create_fake_xy(N=10) + G = create_single_level_2d_mesh_coordinates(xy, nx=1, ny=5) + assert len(G.nodes) == 5 + cardinal = sum( + 1 for _, _, d in G.edges(data=True) if d["adjacency_type"] == "cardinal" + ) + diagonal = sum( + 1 for _, _, d in G.edges(data=True) if d["adjacency_type"] == "diagonal" + ) + assert cardinal == 4 # 5-1 = 4 vertical edges + assert diagonal == 0 + + def test_1x1_grid_no_edges(self): + """Grid with a single node (1x1): should have no edges.""" + xy = test_utils.create_fake_xy(N=10) + G = create_single_level_2d_mesh_coordinates(xy, nx=1, ny=1) + assert len(G.nodes) == 1 + assert len(G.edges) == 0 + + def test_large_grid(self): + """Larger grid should still work correctly.""" + xy = test_utils.create_fake_xy(N=50) + G = create_single_level_2d_mesh_coordinates(xy, nx=10, ny=10) + assert len(G.nodes) == 100 + expected_cardinal = 2 * (10 * 9) # 180 + expected_diagonal = 2 * (9 * 9) # 162 + cardinal = sum( + 1 for _, _, d in G.edges(data=True) if d["adjacency_type"] == "cardinal" + ) + diagonal = sum( + 1 for _, _, d in G.edges(data=True) if d["adjacency_type"] == "diagonal" + ) + assert cardinal == expected_cardinal + assert diagonal == expected_diagonal + + def test_node_positions_within_bounds(self): + """Node positions should be within the xy bounds.""" + xy = test_utils.create_fake_xy(N=20) + G = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=5) + x_min, y_min = np.amin(xy, axis=0) + x_max, y_max = np.amax(xy, axis=0) + for node in G.nodes: + pos = G.nodes[node]["pos"] + assert pos[0] >= x_min and pos[0] <= x_max + assert pos[1] >= y_min and pos[1] <= y_max + + def test_multirange_with_max_levels_1(self): + """Multi-range with max_num_levels=1 should return single-level list.""" + xy = test_utils.create_fake_xy(N=30) + G_list = create_multirange_2d_mesh_coordinates( + max_num_levels=1, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + assert len(G_list) == 1 + assert G_list[0].graph["level"] == 0 + + def test_multirange_with_none_max_levels(self): + """max_num_levels=None should auto-compute levels.""" + xy = test_utils.create_fake_xy(N=30) + G_list = create_multirange_2d_mesh_coordinates( + max_num_levels=None, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + assert isinstance(G_list, list) + assert len(G_list) >= 1 + + def test_multirange_refinement_factor_5(self): + """Test with a different refinement factor.""" + xy = test_utils.create_fake_xy(N=50) + G_list = create_multirange_2d_mesh_coordinates( + max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=5 + ) + if len(G_list) >= 2: + for i in range(len(G_list) - 1): + assert len(G_list[i].nodes) > len(G_list[i + 1].nodes) + + +class TestConnectivityCreationEdgeCases: + """Edge cases for connectivity creation step.""" + + def test_directed_graph_from_1x1(self): + """Creating directed graph from a single-node coordinate graph.""" + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_coordinates(xy, nx=1, ny=1) + G_directed = create_directed_mesh_graph(G_coords, pattern="8-star") + assert isinstance(G_directed, nx.DiGraph) + assert len(G_directed.nodes) == 1 + assert len(G_directed.edges) == 0 + + def test_directed_graph_from_2x1(self): + """Creating directed graph from a 2x1 grid.""" + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_coordinates(xy, nx=2, ny=1) + G_4star = create_directed_mesh_graph(G_coords, pattern="4-star") + G_8star = create_directed_mesh_graph(G_coords, pattern="8-star") + # 2x1: 1 edge, both patterns should have same (no diagonals possible) + assert len(G_4star.edges) == 2 # bidirectional + assert len(G_8star.edges) == 2 + + def test_4star_is_subset_of_8star(self): + """All edges in 4-star should exist in 8-star.""" + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=5) + G_4star = create_directed_mesh_graph(G_coords, pattern="4-star") + G_8star = create_directed_mesh_graph(G_coords, pattern="8-star") + for u, v in G_4star.edges(): + assert G_8star.has_edge(u, v), f"4-star edge ({u},{v}) missing from 8-star" + + def test_edge_lengths_are_positive(self): + """All edge lengths should be positive.""" + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G = create_directed_mesh_graph(G_coords, pattern="8-star") + for u, v, d in G.edges(data=True): + assert d["len"] > 0, f"Edge ({u},{v}) has non-positive length" + + def test_vdiff_antisymmetric(self): + """vdiff(u,v) should be -vdiff(v,u).""" + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G = create_directed_mesh_graph(G_coords, pattern="8-star") + for u, v in G.edges(): + if G.has_edge(v, u): + np.testing.assert_allclose( + G.edges[u, v]["vdiff"], -G.edges[v, u]["vdiff"], + err_msg=f"vdiff not antisymmetric for ({u},{v})" + ) + + def test_edge_len_matches_vdiff_norm(self): + """Edge length should equal the norm of vdiff.""" + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G = create_directed_mesh_graph(G_coords, pattern="8-star") + for u, v, d in G.edges(data=True): + expected_len = np.sqrt(np.sum(d["vdiff"] ** 2)) + np.testing.assert_allclose( + d["len"], expected_len, + err_msg=f"Edge ({u},{v}) len doesn't match vdiff norm" + ) + + def test_flat_multiscale_single_level_input(self): + """Flat multiscale with a single-level list should work (degenerate case).""" + xy = test_utils.create_fake_xy(N=30) + G_coords_list = create_multirange_2d_mesh_coordinates( + max_num_levels=1, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + G = create_flat_multiscale_from_coordinates(G_coords_list, pattern="8-star") + assert isinstance(G, nx.DiGraph) + assert len(G.nodes) > 0 + + def test_flat_multiscale_4star_vs_8star_edge_count(self): + """4-star flat_multiscale should have fewer edges than 8-star.""" + xy = test_utils.create_fake_xy(N=30) + G_coords_list = create_multirange_2d_mesh_coordinates( + max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + G_4 = create_flat_multiscale_from_coordinates(G_coords_list, pattern="4-star") + G_8 = create_flat_multiscale_from_coordinates(G_coords_list, pattern="8-star") + assert len(G_4.edges) < len(G_8.edges) + + def test_hierarchical_single_level_raises(self): + """Hierarchical with only 1 level should raise ValueError.""" + xy = test_utils.create_fake_xy(N=30) + G_coords_list = create_multirange_2d_mesh_coordinates( + max_num_levels=1, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + with pytest.raises(ValueError, match="At least two mesh levels"): + create_hierarchical_from_coordinates(G_coords_list) + + def test_hierarchical_edge_direction_attributes(self): + """Every edge in hierarchical graph must have a 'direction' attribute.""" + xy = test_utils.create_fake_xy(N=30) + G_coords_list = create_multirange_2d_mesh_coordinates( + max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + G = create_hierarchical_from_coordinates(G_coords_list) + for u, v, d in G.edges(data=True): + assert "direction" in d, f"Edge ({u},{v}) missing 'direction'" + assert d["direction"] in ("same", "up", "down") + + def test_hierarchical_up_down_symmetry(self): + """For each 'down' edge (u,v), there should be an 'up' edge (v,u).""" + xy = test_utils.create_fake_xy(N=30) + G_coords_list = create_multirange_2d_mesh_coordinates( + max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + ) + G = create_hierarchical_from_coordinates(G_coords_list) + for u, v, d in G.edges(data=True): + if d.get("direction") == "down": + assert G.has_edge(v, u), f"Missing 'up' edge for 'down' ({u},{v})" + assert G.edges[v, u]["direction"] == "up" + + +class TestAPIEdgeCases: + """Edge cases for the public create_all_graph_components API.""" + + def test_flat_default_pattern_is_8star(self): + """When no m2m_connectivity_kwargs given, flat should default to 8-star.""" + xy = test_utils.create_fake_xy(N=32) + graph_default = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(grid_spacing=3), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + graph_8star = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(grid_spacing=3), + m2m_connectivity_kwargs=dict(pattern="8-star"), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + assert len(graph_default.edges) == len(graph_8star.edges) + + def test_flat_multiscale_default_pattern_is_8star(self): + """When no m2m_connectivity_kwargs given, flat_multiscale defaults to 8-star.""" + xy = test_utils.create_fake_xy(N=32) + graph_default = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat_multiscale", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict( + grid_spacing=3, + refinement_factor=3, + max_num_refinement_levels=3, + ), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + graph_8star = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat_multiscale", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict( + grid_spacing=3, + refinement_factor=3, + max_num_refinement_levels=3, + ), + m2m_connectivity_kwargs=dict(pattern="8-star"), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + assert len(graph_default.edges) == len(graph_8star.edges) + + def test_hierarchical_default_kwargs(self): + """When no m2m_connectivity_kwargs given, hierarchical has sensible defaults.""" + xy = test_utils.create_fake_xy(N=32) + graph_default = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="hierarchical", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict( + grid_spacing=3, + refinement_factor=3, + max_num_refinement_levels=3, + ), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + graph_explicit = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="hierarchical", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict( + grid_spacing=3, + refinement_factor=3, + max_num_refinement_levels=3, + ), + m2m_connectivity_kwargs=dict( + intra_level=dict(pattern="8-star"), + inter_level=dict(pattern="nearest", k=1), + ), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + assert len(graph_default.edges) == len(graph_explicit.edges) + + def test_mesh_layout_default_is_rectilinear(self): + """When mesh_layout not specified, it should default to rectilinear.""" + xy = test_utils.create_fake_xy(N=32) + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout_kwargs=dict(grid_spacing=3), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + assert isinstance(graph, nx.DiGraph) + + def test_return_components_flat(self): + """return_components=True should return dict with g2m, m2m, m2g.""" + xy = test_utils.create_fake_xy(N=32) + components = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(grid_spacing=3), + m2m_connectivity_kwargs=dict(pattern="8-star"), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + return_components=True, + ) + assert isinstance(components, dict) + assert "g2m" in components + assert "m2m" in components + assert "m2g" in components + for name, g in components.items(): + assert isinstance(g, nx.DiGraph) + + def test_return_components_hierarchical(self): + """return_components=True for hierarchical should contain 3 components.""" + xy = test_utils.create_fake_xy(N=32) + components = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="hierarchical", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict( + grid_spacing=3, + refinement_factor=3, + max_num_refinement_levels=3, + ), + m2m_connectivity_kwargs=dict( + intra_level=dict(pattern="8-star"), + inter_level=dict(pattern="nearest", k=1), + ), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + return_components=True, + ) + assert isinstance(components, dict) + assert "g2m" in components + assert "m2m" in components + assert "m2g" in components + + def test_return_components_flat_multiscale(self): + """return_components=True for flat_multiscale.""" + xy = test_utils.create_fake_xy(N=32) + components = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat_multiscale", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict( + grid_spacing=3, + refinement_factor=3, + max_num_refinement_levels=3, + ), + m2m_connectivity_kwargs=dict(pattern="8-star"), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + return_components=True, + ) + assert isinstance(components, dict) + assert set(components.keys()) == {"g2m", "m2m", "m2g"} + + def test_flat_multiscale_no_sub_dicts_interface(self): + """Ensure flat_multiscale accepts a simple pattern arg, not sub-dicts.""" + xy = test_utils.create_fake_xy(N=32) + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat_multiscale", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict( + grid_spacing=3, + refinement_factor=3, + max_num_refinement_levels=3, + ), + m2m_connectivity_kwargs=dict(pattern="8-star"), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + assert isinstance(graph, nx.DiGraph) + + def test_decode_mask_with_new_api(self): + """decode_mask should work correctly with the new API.""" + xy = test_utils.create_fake_xy(N=32) + n_points = len(xy) + mask = [True] * (n_points // 2) + [False] * (n_points - n_points // 2) + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(grid_spacing=3), + m2m_connectivity_kwargs=dict(pattern="8-star"), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + decode_mask=mask, + ) + assert isinstance(graph, nx.DiGraph) + + +class TestBackwardCompatEdgeCases: + """Advanced backward compatibility edge cases.""" + + def test_old_kwargs_with_flat_multiscale_compat(self): + """Old-style flat_multiscale kwargs should trigger deprecation warnings + and be migrated to the new names.""" + xy = test_utils.create_fake_xy(N=32) + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat_multiscale", + m2m_connectivity_kwargs=dict( + mesh_node_distance=3, + level_refinement_factor=3, + max_num_levels=3, + ), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + deprecation_warnings = [ + x for x in w if issubclass(x.category, DeprecationWarning) + ] + # Should have 3 deprecation warnings + assert len(deprecation_warnings) >= 3 + # Check the new names appear in the messages + msgs = " ".join(str(x.message) for x in deprecation_warnings) + assert "grid_spacing" in msgs + assert "refinement_factor" in msgs + assert "max_num_refinement_levels" in msgs + assert isinstance(graph, nx.DiGraph) + + +class TestGraphStructuralProperties: + """Tests verifying structural properties of generated graphs.""" + + def test_all_mesh_nodes_have_pos(self): + """Every node in the final graph should have 'pos' attribute.""" + xy = test_utils.create_fake_xy(N=32) + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(grid_spacing=3), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + for node in graph.nodes: + assert "pos" in graph.nodes[node], f"Node {node} missing 'pos'" + + def test_all_edges_have_component(self): + """Every edge should have a 'component' attribute ('g2m', 'm2m', 'm2g').""" + xy = test_utils.create_fake_xy(N=32) + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(grid_spacing=3), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + for u, v, d in graph.edges(data=True): + assert "component" in d, f"Edge ({u},{v}) missing 'component'" + assert d["component"] in ("g2m", "m2m", "m2g") + + def test_all_edges_have_len_and_vdiff(self): + """Every edge should have 'len' and 'vdiff' attributes.""" + xy = test_utils.create_fake_xy(N=32) + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(grid_spacing=3), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + for u, v, d in graph.edges(data=True): + assert "len" in d, f"Edge ({u},{v}) missing 'len'" + assert "vdiff" in d, f"Edge ({u},{v}) missing 'vdiff'" + + def test_graph_is_directed(self): + """Final graph should always be a DiGraph.""" + xy = test_utils.create_fake_xy(N=32) + for connectivity in ["flat"]: + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity=connectivity, + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(grid_spacing=3), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + assert isinstance(graph, nx.DiGraph) + + def test_flat_4star_strictly_fewer_m2m_edges(self): + """4-star flat should have strictly fewer m2m edges than 8-star.""" + xy = test_utils.create_fake_xy(N=32) + components_4 = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(grid_spacing=3), + m2m_connectivity_kwargs=dict(pattern="4-star"), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + return_components=True, + ) + components_8 = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(grid_spacing=3), + m2m_connectivity_kwargs=dict(pattern="8-star"), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + return_components=True, + ) + # m2m component specifically should differ + assert len(components_4["m2m"].edges) < len(components_8["m2m"].edges) + # g2m and m2g should be the same (same grid spacing, same connectivity) + assert len(components_4["g2m"].edges) == len(components_8["g2m"].edges) + assert len(components_4["m2g"].edges) == len(components_8["m2g"].edges) + + def test_hierarchical_has_same_up_down_edge_count(self): + """Hierarchical graph should have equal number of up and down edges.""" + xy = test_utils.create_fake_xy(N=32) + graph = wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="hierarchical", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict( + grid_spacing=3, + refinement_factor=3, + max_num_refinement_levels=3, + ), + m2m_connectivity_kwargs=dict( + intra_level=dict(pattern="8-star"), + inter_level=dict(pattern="nearest", k=1), + ), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + return_components=True, + ) + m2m = graph["m2m"] + up_count = sum( + 1 for _, _, d in m2m.edges(data=True) if d.get("direction") == "up" + ) + down_count = sum( + 1 for _, _, d in m2m.edges(data=True) if d.get("direction") == "down" + ) + assert up_count == down_count, ( + f"Up edges ({up_count}) != Down edges ({down_count})" + ) + assert up_count > 0, "Should have at least some up/down edges" \ No newline at end of file From 30ddb19ef6a01824e64b3cce03e8ccda49afe649 Mon Sep 17 00:00:00 2001 From: prajwal Date: Tue, 3 Mar 2026 00:05:54 +0530 Subject: [PATCH 04/21] docs: add CHANGELOG entry for mesh_layout feature --- CHANGELOG.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1b1d8e6..1d93fa2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +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](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 + ## [v0.3.0](https://github.com/mllam/weather-model-graphs/releases/tag/v0.3.0) ### Added From 802b66a9df6d2a1b143b8882f65b1d15468e88e0 Mon Sep 17 00:00:00 2001 From: prajwal Date: Wed, 4 Mar 2026 23:31:21 +0530 Subject: [PATCH 05/21] Address all review feedback from PR #81 - Rename mesh.py -> coords.py to reflect coordinate/primitive focus - Split mesh creation into two steps: primitive creation (nx.Graph) and directed connectivity creation (nx.DiGraph) - Add create_single_level_2d_mesh_primitive and create_directed_mesh_graph as new public API - Make mesh_layout a required parameter (no default) to force explicit choice - Extract _migrate_deprecated_kwargs() as a separate module-private helper function for planned removal - Use dict-based interface for intra_level/inter_level kwargs in hierarchical mesh creation - Remove redundant default values in base.py, let each mesh kind handle its own defaults - Add _check_required_graph_attributes() validation in flat.py - Group method+kwargs arguments together in archetype functions - Add comprehensive docstrings explaining 4-star/8-star patterns - Preserve backward compatibility with deprecation warnings - Update all tests to pass mesh_layout explicitly - All 81 mesh layout tests passing, 177/179 full suite passing (2 pre-existing Windows PermissionError on temp PNG files) --- src/weather_model_graphs/create/archetype.py | 44 ++-- src/weather_model_graphs/create/base.py | 196 ++++++++------ .../create/mesh/__init__.py | 6 +- .../create/mesh/{mesh.py => coords.py} | 101 +++++--- .../create/mesh/kinds/flat.py | 120 +++++++-- .../create/mesh/kinds/hierarchical.py | 97 ++++--- tests/test_graph_creation.py | 1 + tests/test_graph_plots.py | 1 + tests/test_mesh_layout.py | 243 +++++++++--------- 9 files changed, 483 insertions(+), 326 deletions(-) rename src/weather_model_graphs/create/mesh/{mesh.py => coords.py} (71%) diff --git a/src/weather_model_graphs/create/archetype.py b/src/weather_model_graphs/create/archetype.py index ac76fcc..d8ac806 100644 --- a/src/weather_model_graphs/create/archetype.py +++ b/src/weather_model_graphs/create/archetype.py @@ -57,20 +57,20 @@ def create_keisler_graph( """ return create_all_graph_components( coords=coords, - m2m_connectivity="flat", + coords_crs=coords_crs, + graph_crs=graph_crs, mesh_layout="rectilinear", - mesh_layout_kwargs=dict(grid_spacing=mesh_node_distance), - m2m_connectivity_kwargs=dict(pattern="8-star"), + mesh_layout_kwargs=dict(mesh_node_spacing=mesh_node_distance), g2m_connectivity="within_radius", - m2g_connectivity="nearest_neighbours", g2m_connectivity_kwargs=dict( rel_max_dist=0.51, ), + m2m_connectivity="flat", + m2m_connectivity_kwargs=dict(pattern="8-star"), + m2g_connectivity="nearest_neighbours", m2g_connectivity_kwargs=dict( max_num_neighbours=4, ), - coords_crs=coords_crs, - graph_crs=graph_crs, decode_mask=decode_mask, return_components=return_components, ) @@ -134,26 +134,26 @@ def create_graphcast_graph( """ return create_all_graph_components( coords=coords, - m2m_connectivity="flat_multiscale", + coords_crs=coords_crs, + graph_crs=graph_crs, mesh_layout="rectilinear", mesh_layout_kwargs=dict( - grid_spacing=mesh_node_distance, + mesh_node_spacing=mesh_node_distance, refinement_factor=level_refinement_factor, max_num_refinement_levels=max_num_levels, ), - m2m_connectivity_kwargs=dict( - pattern="8-star", - ), g2m_connectivity="within_radius", - m2g_connectivity="nearest_neighbours", g2m_connectivity_kwargs=dict( rel_max_dist=0.51, ), + m2m_connectivity="flat_multiscale", + m2m_connectivity_kwargs=dict( + pattern="8-star", + ), + m2g_connectivity="nearest_neighbours", m2g_connectivity_kwargs=dict( max_num_neighbours=4, ), - coords_crs=coords_crs, - graph_crs=graph_crs, decode_mask=decode_mask, return_components=return_components, ) @@ -222,27 +222,27 @@ def create_oskarsson_hierarchical_graph( """ return create_all_graph_components( coords=coords, - m2m_connectivity="hierarchical", + coords_crs=coords_crs, + graph_crs=graph_crs, mesh_layout="rectilinear", mesh_layout_kwargs=dict( - grid_spacing=mesh_node_distance, + mesh_node_spacing=mesh_node_distance, refinement_factor=level_refinement_factor, max_num_refinement_levels=max_num_levels, ), + g2m_connectivity="within_radius", + g2m_connectivity_kwargs=dict( + rel_max_dist=0.51, + ), + m2m_connectivity="hierarchical", m2m_connectivity_kwargs=dict( intra_level=dict(pattern="8-star"), inter_level=dict(pattern="nearest", k=1), ), - g2m_connectivity="within_radius", 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, decode_mask=decode_mask, return_components=return_components, ) diff --git a/src/weather_model_graphs/create/base.py b/src/weather_model_graphs/create/base.py index a6f81a6..ddee05b 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/weather_model_graphs/create/base.py @@ -31,22 +31,79 @@ from .mesh.kinds.hierarchical import ( create_hierarchical_from_coordinates, ) -from .mesh.mesh import ( - create_multirange_2d_mesh_coordinates, - create_single_level_2d_mesh_coordinates, +from .mesh.coords import ( + create_multirange_2d_mesh_primitives, + create_single_level_2d_mesh_primitive, ) +def _migrate_deprecated_kwargs(mesh_layout_kwargs, m2m_connectivity_kwargs): + """Migrate old-style kwargs to the new mesh_layout_kwargs structure. + + In the old API, ``mesh_node_distance``, ``level_refinement_factor``, and + ``max_num_levels`` were passed via ``m2m_connectivity_kwargs``. In the new + design these belong in ``mesh_layout_kwargs`` (as ``mesh_node_spacing``, + ``refinement_factor``, and ``max_num_refinement_levels`` respectively). + + This helper emits ``DeprecationWarning`` for each migrated key and moves + the value into *mesh_layout_kwargs*. It is intended to be removed once the + old API is no longer supported. + + Parameters + ---------- + mesh_layout_kwargs : dict + Mutable dict of mesh layout keyword arguments. + m2m_connectivity_kwargs : dict + Mutable dict of m2m connectivity keyword arguments. + + Returns + ------- + tuple[dict, dict] + Updated (mesh_layout_kwargs, m2m_connectivity_kwargs). + """ + if "mesh_node_distance" in m2m_connectivity_kwargs and "mesh_node_spacing" not in mesh_layout_kwargs: + warnings.warn( + "Passing 'mesh_node_distance' in m2m_connectivity_kwargs is deprecated. " + "Use mesh_layout_kwargs=dict(mesh_node_spacing=...) instead.", + DeprecationWarning, + stacklevel=3, + ) + mesh_layout_kwargs["mesh_node_spacing"] = m2m_connectivity_kwargs.pop( + "mesh_node_distance" + ) + if "level_refinement_factor" in m2m_connectivity_kwargs and "refinement_factor" not in mesh_layout_kwargs: + warnings.warn( + "Passing 'level_refinement_factor' in m2m_connectivity_kwargs is deprecated. " + "Use mesh_layout_kwargs=dict(refinement_factor=...) instead.", + DeprecationWarning, + stacklevel=3, + ) + mesh_layout_kwargs["refinement_factor"] = ( + m2m_connectivity_kwargs.pop("level_refinement_factor") + ) + if "max_num_levels" in m2m_connectivity_kwargs and "max_num_refinement_levels" not in mesh_layout_kwargs: + warnings.warn( + "Passing 'max_num_levels' in m2m_connectivity_kwargs is deprecated. " + "Use mesh_layout_kwargs=dict(max_num_refinement_levels=...) instead.", + DeprecationWarning, + stacklevel=3, + ) + mesh_layout_kwargs["max_num_refinement_levels"] = m2m_connectivity_kwargs.pop( + "max_num_levels" + ) + return mesh_layout_kwargs, m2m_connectivity_kwargs + + def create_all_graph_components( coords: np.ndarray, m2m_connectivity: str, m2g_connectivity: str, g2m_connectivity: str, - mesh_layout: str = "rectilinear", + mesh_layout: str, mesh_layout_kwargs: dict = None, m2m_connectivity_kwargs: dict = None, - m2g_connectivity_kwargs={}, - g2m_connectivity_kwargs={}, + m2g_connectivity_kwargs: dict = None, + g2m_connectivity_kwargs: dict = None, coords_crs: pyproj.crs.CRS | None = None, graph_crs: pyproj.crs.CRS | None = None, decode_mask: Iterable[bool] | None = None, @@ -79,20 +136,27 @@ def create_all_graph_components( of `max_dist` or relative distance of `rel_max_dist` from each node in mesh mesh_layout: - - "rectilinear": Regular rectilinear grid (default). Uses grid_spacing to - determine mesh node placement. Produces nodes with 4-star (cardinal) and + - "rectilinear": Uniform regular grid with ``mesh_node_spacing`` resolution. + Produces an undirected mesh primitive with 4-star (cardinal) and 8-star (cardinal + diagonal) spatial adjacency edges. mesh_layout_kwargs (for mesh_layout="rectilinear"): - - grid_spacing: float, distance between mesh nodes in coordinate units - - refinement_factor: int, refinement factor between levels (for multi-level) - - max_num_refinement_levels: int, maximum number of mesh levels (for multi-level) + - mesh_node_spacing: float, distance between mesh nodes in coordinate units. + - refinement_factor: int, refinement factor between levels + (for multi-level and hierarchical mesh graphs, default: 3) + - max_num_refinement_levels: int, maximum number of mesh levels + (for multi-level and hierarchical mesh graphs) + + Wherever the ``pattern`` argument appears below it defines the spatial + neighbourhood connectivity: + - ``"4-star"``: only cardinal directions (horizontal and vertical neighbours) + - ``"8-star"``: cardinal plus diagonal neighbours (all 8 surrounding nodes) m2m_connectivity: - "flat": Create a single-level directed mesh graph. - m2m_connectivity_kwargs: pattern="4-star" or "8-star" (default: "8-star") + m2m_connectivity_kwargs: pattern (default: "8-star") - "flat_multiscale": Create a flat multiscale mesh graph. - m2m_connectivity_kwargs: pattern="4-star" or "8-star" (default: "8-star") + m2m_connectivity_kwargs: pattern (default: "8-star") - "hierarchical": Create a hierarchical mesh graph with up/down connections. m2m_connectivity_kwargs: intra_level=dict(pattern=...), inter_level=dict(pattern=..., k=...) @@ -130,40 +194,19 @@ def create_all_graph_components( m2m_connectivity_kwargs = {} else: m2m_connectivity_kwargs = dict(m2m_connectivity_kwargs) + if m2g_connectivity_kwargs is None: + m2g_connectivity_kwargs = {} + else: + m2g_connectivity_kwargs = dict(m2g_connectivity_kwargs) + if g2m_connectivity_kwargs is None: + g2m_connectivity_kwargs = {} + else: + g2m_connectivity_kwargs = dict(g2m_connectivity_kwargs) - # Backward compatibility: migrate old-style kwargs where mesh_node_distance, - # level_refinement_factor, and max_num_levels were passed via - # m2m_connectivity_kwargs. In the new design these belong in mesh_layout_kwargs. - if "mesh_node_distance" in m2m_connectivity_kwargs and "grid_spacing" not in mesh_layout_kwargs: - warnings.warn( - "Passing 'mesh_node_distance' in m2m_connectivity_kwargs is deprecated. " - "Use mesh_layout_kwargs=dict(grid_spacing=...) instead.", - DeprecationWarning, - stacklevel=2, - ) - mesh_layout_kwargs["grid_spacing"] = m2m_connectivity_kwargs.pop( - "mesh_node_distance" - ) - if "level_refinement_factor" in m2m_connectivity_kwargs and "refinement_factor" not in mesh_layout_kwargs: - warnings.warn( - "Passing 'level_refinement_factor' in m2m_connectivity_kwargs is deprecated. " - "Use mesh_layout_kwargs=dict(refinement_factor=...) instead.", - DeprecationWarning, - stacklevel=2, - ) - mesh_layout_kwargs["refinement_factor"] = ( - m2m_connectivity_kwargs.pop("level_refinement_factor") - ) - if "max_num_levels" in m2m_connectivity_kwargs and "max_num_refinement_levels" not in mesh_layout_kwargs: - warnings.warn( - "Passing 'max_num_levels' in m2m_connectivity_kwargs is deprecated. " - "Use mesh_layout_kwargs=dict(max_num_refinement_levels=...) instead.", - DeprecationWarning, - stacklevel=2, - ) - mesh_layout_kwargs["max_num_refinement_levels"] = m2m_connectivity_kwargs.pop( - "max_num_levels" - ) + # Migrate deprecated kwargs (to be removed in a future version) + mesh_layout_kwargs, m2m_connectivity_kwargs = _migrate_deprecated_kwargs( + mesh_layout_kwargs, m2m_connectivity_kwargs + ) assert ( len(coords.shape) == 2 and coords.shape[1] == 2 @@ -196,26 +239,29 @@ def create_all_graph_components( if m2m_connectivity == "flat": # --- Step 1: Coordinate creation based on mesh_layout --- if mesh_layout == "rectilinear": - grid_spacing = mesh_layout_kwargs.get("grid_spacing") - if grid_spacing is None: + mesh_node_spacing = mesh_layout_kwargs.get("mesh_node_spacing") + # Backward compat: also check for old name "grid_spacing" + if mesh_node_spacing is None: + mesh_node_spacing = mesh_layout_kwargs.get("grid_spacing") + if mesh_node_spacing is None: raise ValueError( - "mesh_layout='rectilinear' requires 'grid_spacing' in " + "mesh_layout='rectilinear' requires 'mesh_node_spacing' in " "mesh_layout_kwargs (or 'mesh_node_distance' in " "m2m_connectivity_kwargs for backward compatibility)." ) - # Compute number of mesh nodes from grid_spacing + # Compute number of mesh nodes from mesh_node_spacing range_x, range_y = np.ptp(xy, axis=0) - nx_mesh = int(range_x / grid_spacing) - ny_mesh = int(range_y / grid_spacing) + nx_mesh = int(range_x / mesh_node_spacing) + ny_mesh = int(range_y / mesh_node_spacing) if nx_mesh == 0 or ny_mesh == 0: raise ValueError( - "The given `grid_spacing` is too large for the provided " - f"coordinates. Got grid_spacing={grid_spacing}, but the " + "The given `mesh_node_spacing` is too large for the provided " + f"coordinates. Got mesh_node_spacing={mesh_node_spacing}, but the " f"x-range is {range_x} and y-range is {range_y}. Maybe you " - "want to decrease the `grid_spacing` so that the mesh nodes " + "want to decrease the `mesh_node_spacing` so that the mesh nodes " "are spaced closer together?" ) - G_mesh_coords = create_single_level_2d_mesh_coordinates( + G_mesh_coords = create_single_level_2d_mesh_primitive( xy, nx_mesh, ny_mesh ) else: @@ -234,20 +280,22 @@ def create_all_graph_components( elif m2m_connectivity == "hierarchical": # --- Step 1: Coordinate creation based on mesh_layout --- if mesh_layout == "rectilinear": - grid_spacing = mesh_layout_kwargs.get("grid_spacing") - refinement_factor = mesh_layout_kwargs.get("refinement_factor") + mesh_node_spacing = mesh_layout_kwargs.get("mesh_node_spacing") + if mesh_node_spacing is None: + mesh_node_spacing = mesh_layout_kwargs.get("grid_spacing") + refinement_factor = mesh_layout_kwargs.get("refinement_factor", 3) max_num_refinement_levels = mesh_layout_kwargs.get( "max_num_refinement_levels" ) - if grid_spacing is None: + if mesh_node_spacing is None: raise ValueError( "mesh_layout='rectilinear' with m2m_connectivity='hierarchical' " - "requires 'grid_spacing' in mesh_layout_kwargs." + "requires 'mesh_node_spacing' in mesh_layout_kwargs." ) - G_coords_list = create_multirange_2d_mesh_coordinates( + G_coords_list = create_multirange_2d_mesh_primitives( max_num_levels=max_num_refinement_levels, xy=xy, - grid_spacing=grid_spacing, + mesh_node_spacing=mesh_node_spacing, interlevel_refinement_factor=refinement_factor, ) else: @@ -257,19 +305,13 @@ def create_all_graph_components( ) # --- Step 2: Connectivity creation --- - intra_level = m2m_connectivity_kwargs.get( - "intra_level", {"pattern": "8-star"} - ) - inter_level = m2m_connectivity_kwargs.get( - "inter_level", {"pattern": "nearest", "k": 1} - ) # hierarchical mesh graph have three sub-graphs: # `m2m` (mesh-to-mesh), `mesh_up` (up edge connections) and # `mesh_down` (down edge connections) graph_components["m2m"] = create_hierarchical_from_coordinates( G_coords_list, - intra_level=intra_level, - inter_level=inter_level, + intra_level=m2m_connectivity_kwargs.get("intra_level"), + inter_level=m2m_connectivity_kwargs.get("inter_level"), ) # Only connect grid to bottom level of hierarchy grid_connect_graph = split_graph_by_edge_attribute( @@ -279,20 +321,22 @@ def create_all_graph_components( elif m2m_connectivity == "flat_multiscale": # --- Step 1: Coordinate creation based on mesh_layout --- if mesh_layout == "rectilinear": - grid_spacing = mesh_layout_kwargs.get("grid_spacing") - refinement_factor = mesh_layout_kwargs.get("refinement_factor") + mesh_node_spacing = mesh_layout_kwargs.get("mesh_node_spacing") + if mesh_node_spacing is None: + mesh_node_spacing = mesh_layout_kwargs.get("grid_spacing") + refinement_factor = mesh_layout_kwargs.get("refinement_factor", 3) max_num_refinement_levels = mesh_layout_kwargs.get( "max_num_refinement_levels" ) - if grid_spacing is None: + if mesh_node_spacing is None: raise ValueError( "mesh_layout='rectilinear' with m2m_connectivity='flat_multiscale' " - "requires 'grid_spacing' in mesh_layout_kwargs." + "requires 'mesh_node_spacing' in mesh_layout_kwargs." ) - G_coords_list = create_multirange_2d_mesh_coordinates( + G_coords_list = create_multirange_2d_mesh_primitives( max_num_levels=max_num_refinement_levels, xy=xy, - grid_spacing=grid_spacing, + mesh_node_spacing=mesh_node_spacing, interlevel_refinement_factor=refinement_factor, ) else: diff --git a/src/weather_model_graphs/create/mesh/__init__.py b/src/weather_model_graphs/create/mesh/__init__.py index 09f3e7b..2d3d7cc 100644 --- a/src/weather_model_graphs/create/mesh/__init__.py +++ b/src/weather_model_graphs/create/mesh/__init__.py @@ -1,6 +1,6 @@ -from .mesh import ( +from .coords import ( create_directed_mesh_graph, - create_multirange_2d_mesh_coordinates, - create_single_level_2d_mesh_coordinates, + create_multirange_2d_mesh_primitives, create_single_level_2d_mesh_graph, + create_single_level_2d_mesh_primitive, ) diff --git a/src/weather_model_graphs/create/mesh/mesh.py b/src/weather_model_graphs/create/mesh/coords.py similarity index 71% rename from src/weather_model_graphs/create/mesh/mesh.py rename to src/weather_model_graphs/create/mesh/coords.py index 39620f6..47dd30c 100644 --- a/src/weather_model_graphs/create/mesh/mesh.py +++ b/src/weather_model_graphs/create/mesh/coords.py @@ -3,15 +3,22 @@ from loguru import logger -def create_single_level_2d_mesh_coordinates(xy, nx, ny): +def create_single_level_2d_mesh_primitive(xy: np.ndarray, nx: int, ny: int): """ - Create an undirected mesh graph (nx.Graph) with node positions and spatial - adjacency edges, representing the coordinate creation step. + Create an undirected mesh primitive graph (nx.Graph) with node positions + and spatial adjacency edges, representing the coordinate creation step. + + A mesh primitive is an undirected graph that encodes all potential + neighbourhood connectivity edges. It serves as a blueprint from which + directed connectivity graphs can later be built by selecting a subset + of edges (e.g. 4-star or 8-star pattern). This produces a graph where: - - Nodes have "pos" (spatial coordinates) and "type" ("mesh") attributes - - Edges have "adjacency_type" attribute: "cardinal" for horizontal/vertical - neighbours (4-star) or "diagonal" for diagonal neighbours (8-star only) + - Nodes have a ``"pos"`` attribute (np.ndarray of shape [2,] with x and y + coordinates) and a ``"type"`` attribute (str, always ``"mesh"``). + - Edges have an ``"adjacency_type"`` attribute (str): ``"cardinal"`` for + horizontal/vertical neighbours (4-star connectivity) or ``"diagonal"`` + for diagonal neighbours (additional edges in 8-star connectivity). This is the first step in the two-step mesh creation process: 1. Coordinate creation (this function) -> nx.Graph with spatial adjacency @@ -19,10 +26,9 @@ def create_single_level_2d_mesh_coordinates(xy, nx, ny): Parameters ---------- - xy : np.ndarray [N_grid_points, 2] - Grid point coordinates, with first column representing - x coordinates and second column y coordinates. N_grid_points is the - total number of grid points. + xy : np.ndarray + Grid point coordinates, shaped [N_grid_points, 2], with first column + representing x coordinates and second column y coordinates. nx : int Number of nodes in x direction ny : int @@ -31,7 +37,8 @@ def create_single_level_2d_mesh_coordinates(xy, nx, ny): Returns ------- networkx.Graph - Undirected graph with node positions and annotated spatial adjacency edges + Undirected mesh primitive graph with node positions and annotated + spatial adjacency edges. """ xm, xM = np.amin(xy[:, 0]), np.amax(xy[:, 0]) ym, yM = np.amin(xy[:, 1]), np.amax(xy[:, 1]) @@ -71,48 +78,59 @@ def create_single_level_2d_mesh_coordinates(xy, nx, ny): return g -def create_directed_mesh_graph(G_undirected, pattern="8-star"): +def create_directed_mesh_graph( + G_undirected: networkx.Graph, pattern: str = "8-star" +): """ - Convert an undirected mesh graph with spatial adjacency edges to a + Convert an undirected mesh primitive graph with spatial adjacency edges to a directed mesh graph (nx.DiGraph) based on the specified connectivity pattern. This is the second step in the two-step mesh creation process: - 1. Coordinate creation (create_single_level_2d_mesh_coordinates) -> nx.Graph + 1. Coordinate creation (create_single_level_2d_mesh_primitive) -> nx.Graph 2. Connectivity creation (this function) -> nx.DiGraph + The ``pattern`` argument defines the spatial neighbourhood connectivity: + - ``"4-star"``: only cardinal directions (horizontal and vertical neighbours) + - ``"8-star"``: cardinal directions plus diagonals (all 8 surrounding neighbours) + Parameters ---------- G_undirected : networkx.Graph - Undirected graph with nodes having "pos" attributes and edges having - "adjacency_type" attributes ("cardinal" or "diagonal"). + Undirected mesh primitive graph. Expected node attributes: + - ``"pos"``: np.ndarray of shape [2,], spatial coordinates. + Expected edge attributes: + - ``"adjacency_type"``: str, either ``"cardinal"`` or ``"diagonal"``. + Additional edge attributes (e.g. ``"level"``) are preserved in the + output directed graph. pattern : str Connectivity pattern. Options: - - "4-star": only cardinal edges (horizontal/vertical neighbours) - - "8-star": all edges (cardinal + diagonal neighbours) + - ``"4-star"``: only cardinal edges (horizontal/vertical neighbours) + - ``"8-star"``: all edges (cardinal + diagonal neighbours) Returns ------- networkx.DiGraph - Directed graph with bidirectional edges, each having "len" and "vdiff" - attributes. + Directed graph with bidirectional edges, each having ``"len"`` and + ``"vdiff"`` attributes. All original edge attributes from the + primitive graph are preserved. """ if pattern == "4-star": - # Filter to only cardinal edges (horizontal/vertical) + # Filter to only cardinal edges, preserving edge data edges_to_use = [ - (u, v) + (u, v, d) for u, v, d in G_undirected.edges(data=True) if d.get("adjacency_type") == "cardinal" ] elif pattern == "8-star": - # Use all edges (cardinal + diagonal) - edges_to_use = list(G_undirected.edges()) + # Use all edges with their data + edges_to_use = list(G_undirected.edges(data=True)) else: raise ValueError( f"Unknown connectivity pattern: '{pattern}'. " "Choose '4-star' or '8-star'." ) - # Create filtered undirected graph with only selected edges + # Create filtered undirected graph with only selected edges (preserving attrs) g_filtered = networkx.Graph() g_filtered.add_nodes_from(G_undirected.nodes(data=True)) g_filtered.add_edges_from(edges_to_use) @@ -159,7 +177,7 @@ def create_single_level_2d_mesh_graph(xy, nx, ny): nodes. Internally, this uses the two-step process: - 1. create_single_level_2d_mesh_coordinates (coordinate creation) + 1. create_single_level_2d_mesh_primitive (coordinate creation) 2. create_directed_mesh_graph (connectivity creation, pattern="8-star") Parameters @@ -178,21 +196,21 @@ def create_single_level_2d_mesh_graph(xy, nx, ny): networkx.DiGraph Graph representing the 2D grid """ - G_coords = create_single_level_2d_mesh_coordinates(xy, nx, ny) + G_coords = create_single_level_2d_mesh_primitive(xy, nx, ny) return create_directed_mesh_graph(G_coords, pattern="8-star") -def create_multirange_2d_mesh_coordinates( - max_num_levels, xy, grid_spacing=3, interlevel_refinement_factor=3 +def create_multirange_2d_mesh_primitives( + max_num_levels, xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ): """ - Create a list of undirected coordinate graphs (nx.Graph) representing + Create a list of undirected mesh primitive graphs (nx.Graph) representing different levels of mesh resolution spanning the spatial domain of the xy coordinates. - This is the coordinate creation step for multi-level meshes. Each returned - graph contains nodes with spatial positions and edges annotated with - adjacency type ("cardinal" or "diagonal"). + This is the coordinate creation step for multi-level and hierarchical mesh + graphs. Each returned graph contains nodes with spatial positions and edges + annotated with adjacency type (``"cardinal"`` or ``"diagonal"``). The graphs can be consumed by connectivity creation functions to produce directed mesh graphs for flat_multiscale or hierarchical architectures. @@ -203,7 +221,7 @@ def create_multirange_2d_mesh_coordinates( Number of edge-distance levels in mesh graph xy : np.ndarray Grid point coordinates, shaped [N_grid_points, 2] - grid_spacing : float + mesh_node_spacing : float Distance (in x- and y-direction) between created mesh nodes, in coordinate system of xy interlevel_refinement_factor : float @@ -212,16 +230,17 @@ def create_multirange_2d_mesh_coordinates( Returns ------- G_all_levels : list of networkx.Graph - List of undirected coordinate graphs for each level, each with + List of undirected mesh primitive graphs for each level, each with node positions and annotated spatial adjacency edges. - Each graph has "level" and "interlevel_refinement_factor" graph attributes. + Each graph has ``"level"`` and ``"interlevel_refinement_factor"`` + graph attributes. """ # Compute the size along x and y direction of area to cover with graph # This is measured in the Cartesian coordinates of xy coord_extent = np.ptp(xy, axis=0) # Number of nodes that would fit on bottom level of hierarchy, # in both directions - max_nodes_bottom = (coord_extent / grid_spacing).astype(int) + max_nodes_bottom = (coord_extent / mesh_node_spacing).astype(int) # Find the number of mesh levels possible in x- and y-direction, # and the number of leaf nodes that would correspond to @@ -251,7 +270,7 @@ def create_multirange_2d_mesh_coordinates( nodes_x, nodes_y = ( nleaf / (interlevel_refinement_factor**lev) ).astype(int) - g = create_single_level_2d_mesh_coordinates(xy, nodes_x, nodes_y) + g = create_single_level_2d_mesh_primitive(xy, nodes_x, nodes_y) # Add level information to nodes, edges and full graph for node in g.nodes: g.nodes[node]["level"] = lev @@ -277,7 +296,7 @@ def create_multirange_2d_mesh_graphs( Each graph in the list contains a "level" attribute with the level index of the graph. Internally uses the two-step process: - 1. create_multirange_2d_mesh_coordinates (coordinate creation) + 1. create_multirange_2d_mesh_primitives (coordinate creation) 2. create_directed_mesh_graph (connectivity creation, pattern="8-star") Parameters @@ -298,10 +317,10 @@ def create_multirange_2d_mesh_graphs( List of networkx graphs for each level representing the connectivity of the mesh within each level """ - G_coords_list = create_multirange_2d_mesh_coordinates( + G_coords_list = create_multirange_2d_mesh_primitives( max_num_levels=max_num_levels, xy=xy, - grid_spacing=mesh_node_distance, + mesh_node_spacing=mesh_node_distance, interlevel_refinement_factor=level_refinement_factor, ) diff --git a/src/weather_model_graphs/create/mesh/kinds/flat.py b/src/weather_model_graphs/create/mesh/kinds/flat.py index eca018a..4329389 100644 --- a/src/weather_model_graphs/create/mesh/kinds/flat.py +++ b/src/weather_model_graphs/create/mesh/kinds/flat.py @@ -1,46 +1,100 @@ +from typing import List + import networkx import numpy as np from ....networkx_utils import prepend_node_index -from .. import mesh as mesh_graph +from .. import coords as mesh_coords + + +def _check_required_graph_attributes(G: networkx.Graph, context: str): + """Check that a coordinate graph has the required node and edge attributes. + + Parameters + ---------- + G : networkx.Graph + The coordinate graph to validate. + context : str + Description of where this check is being called, for error messages. + + Raises + ------ + ValueError + If required attributes are missing. + """ + # Check at least one node has required attributes + if len(G.nodes) > 0: + sample_node = next(iter(G.nodes)) + if "pos" not in G.nodes[sample_node]: + raise ValueError( + f"{context}: coordinate graph nodes must have a 'pos' attribute " + "(np.ndarray of shape [2,])." + ) + if "type" not in G.nodes[sample_node]: + raise ValueError( + f"{context}: coordinate graph nodes must have a 'type' attribute." + ) + # Check at least one edge has required attributes + if len(G.edges) > 0: + sample_edge = next(iter(G.edges)) + if "adjacency_type" not in G.edges[sample_edge]: + raise ValueError( + f"{context}: coordinate graph edges must have an 'adjacency_type' " + "attribute ('cardinal' or 'diagonal')." + ) def create_flat_multiscale_from_coordinates( - G_coords_list, - pattern="8-star", + G_coords_list: List[networkx.Graph], + pattern: str = "8-star", ): """ Create flat multiscale mesh graph from a list of coordinate graphs. This is the connectivity creation step for flat multiscale meshes. - It takes undirected coordinate graphs (one per level) and produces a + It takes undirected mesh primitive graphs (one per level) and produces a single directed mesh graph where all levels are merged into one flat graph. In a flat multiscale graph, coarser levels are merged into the finer level by coincident node positions (no separate inter-level connectivity needed). - The ``pattern`` controls the intra-level edge connectivity for each level. + + The ``pattern`` argument defines the spatial neighbourhood connectivity: + - ``"4-star"``: only cardinal directions (horizontal and vertical neighbours) + - ``"8-star"``: cardinal directions plus diagonals (all 8 surrounding neighbours) Parameters ---------- G_coords_list : list of networkx.Graph - List of undirected coordinate graphs, one per level. Each should have - nodes with "pos" and "type" attributes, and edges with "adjacency_type" - attributes. Created by create_multirange_2d_mesh_coordinates. + List of undirected mesh primitive graphs, one per level. Each graph + must have: + - Node attributes: ``"pos"`` (np.ndarray of shape [2,]), ``"type"`` (str) + - Edge attributes: ``"adjacency_type"`` (str, ``"cardinal"`` or ``"diagonal"``) + - Graph attribute: ``"interlevel_refinement_factor"`` (int) + Created by ``create_multirange_2d_mesh_primitives``. pattern : str - Connectivity pattern for intra-level edges: "4-star" or "8-star" - (default: "8-star") + Connectivity pattern for intra-level edges: ``"4-star"`` or ``"8-star"`` + (default: ``"8-star"``) Returns ------- G_tot : networkx.DiGraph The merged flat multiscale mesh graph """ - - # Retrieve interlevel_refinement_factor from graph attributes - interlevel_refinement_factor = G_coords_list[0].graph.get( - "interlevel_refinement_factor", 3 + # Validate required attributes on first graph + _check_required_graph_attributes( + G_coords_list[0], "create_flat_multiscale_from_coordinates" ) + # Assert interlevel_refinement_factor is set (no silent default) + if "interlevel_refinement_factor" not in G_coords_list[0].graph: + raise ValueError( + "The coordinate graphs must have an 'interlevel_refinement_factor' " + "graph attribute. This is set by create_multirange_2d_mesh_primitives." + ) + interlevel_refinement_factor = G_coords_list[0].graph[ + "interlevel_refinement_factor" + ] + # Check that interlevel_refinement_factor is an odd integer if ( int(interlevel_refinement_factor) != interlevel_refinement_factor @@ -53,7 +107,7 @@ def create_flat_multiscale_from_coordinates( # Convert each level's coordinate graph to directed graph with chosen pattern G_all_levels = [ - mesh_graph.create_directed_mesh_graph(g_coords, pattern=pattern) + mesh_coords.create_directed_mesh_graph(g_coords, pattern=pattern) for g_coords in G_coords_list ] @@ -105,29 +159,39 @@ def create_flat_multiscale_from_coordinates( return G_tot -def create_flat_singlescale_from_coordinates(G_coords, pattern="8-star"): +def create_flat_singlescale_from_coordinates( + G_coords: networkx.Graph, pattern: str = "8-star" +): """ - Create a flat single-scale directed mesh graph from a coordinate graph. + Create a flat single-scale directed mesh graph from a mesh primitive graph. This is the connectivity creation step for flat single-scale meshes. - It converts an undirected coordinate graph to a directed mesh graph + It converts an undirected mesh primitive graph to a directed mesh graph using the specified connectivity pattern. + The ``pattern`` argument defines the spatial neighbourhood connectivity: + - ``"4-star"``: only cardinal directions (horizontal and vertical neighbours) + - ``"8-star"``: cardinal directions plus diagonals (all 8 surrounding neighbours) + Parameters ---------- G_coords : networkx.Graph - Undirected coordinate graph with nodes having "pos" attributes and - edges having "adjacency_type" attributes. Created by - create_single_level_2d_mesh_coordinates. + Undirected mesh primitive graph. Must have: + - Node attributes: ``"pos"`` (np.ndarray of shape [2,]), ``"type"`` (str) + - Edge attributes: ``"adjacency_type"`` (str, ``"cardinal"`` or ``"diagonal"``) + Created by ``create_single_level_2d_mesh_primitive``. pattern : str - Connectivity pattern: "4-star" or "8-star" (default: "8-star") + Connectivity pattern: ``"4-star"`` or ``"8-star"`` (default: ``"8-star"``) Returns ------- networkx.DiGraph The flat single-scale directed mesh graph """ - return mesh_graph.create_directed_mesh_graph(G_coords, pattern=pattern) + _check_required_graph_attributes( + G_coords, "create_flat_singlescale_from_coordinates" + ) + return mesh_coords.create_directed_mesh_graph(G_coords, pattern=pattern) def create_flat_multiscale_mesh_graph( @@ -138,7 +202,7 @@ def create_flat_multiscale_mesh_graph( graphs across all levels in `G_all_levels`. Internally uses the two-step process: - 1. create_multirange_2d_mesh_coordinates (coordinate creation) + 1. create_multirange_2d_mesh_primitives (coordinate creation) 2. create_flat_multiscale_from_coordinates (connectivity creation) Parameters @@ -160,10 +224,10 @@ def create_flat_multiscale_mesh_graph( G_tot : networkx.Graph The merged mesh graph """ - G_coords_list = mesh_graph.create_multirange_2d_mesh_coordinates( + G_coords_list = mesh_coords.create_multirange_2d_mesh_primitives( max_num_levels=max_num_levels, xy=xy, - grid_spacing=mesh_node_distance, + mesh_node_spacing=mesh_node_distance, interlevel_refinement_factor=level_refinement_factor, ) @@ -178,7 +242,7 @@ def create_flat_singlescale_mesh_graph(xy, mesh_node_distance: float): Create flat mesh graph of single level Internally uses the two-step process: - 1. create_single_level_2d_mesh_coordinates (coordinate creation) + 1. create_single_level_2d_mesh_primitive (coordinate creation) 2. create_directed_mesh_graph (connectivity creation, pattern="8-star") Parameters @@ -208,4 +272,4 @@ def create_flat_singlescale_mesh_graph(xy, mesh_node_distance: float): " so that the mesh nodes are spaced closer together?" ) - return mesh_graph.create_single_level_2d_mesh_graph(xy=xy, nx=nx, ny=ny) + return mesh_coords.create_single_level_2d_mesh_graph(xy, nx, ny) diff --git a/src/weather_model_graphs/create/mesh/kinds/hierarchical.py b/src/weather_model_graphs/create/mesh/kinds/hierarchical.py index d7d8c70..d1a9ca3 100644 --- a/src/weather_model_graphs/create/mesh/kinds/hierarchical.py +++ b/src/weather_model_graphs/create/mesh/kinds/hierarchical.py @@ -1,37 +1,49 @@ +from typing import Dict, List, Optional + import networkx import numpy as np import scipy from ....networkx_utils import prepend_node_index -from .. import mesh as mesh_graph +from .. import coords as mesh_coords def create_hierarchical_from_coordinates( - G_coords_list, - intra_level=None, - inter_level=None, + G_coords_list: List[networkx.Graph], + intra_level: Optional[Dict[str, object]] = None, + inter_level: Optional[Dict[str, object]] = None, ): """ - Create a hierarchical multiscale mesh graph from a list of coordinate graphs. + Create a hierarchical multiscale mesh graph from a list of mesh primitive + graphs. This is the connectivity creation step for hierarchical meshes. - It takes undirected coordinate graphs (one per level) and produces a + It takes undirected mesh primitive graphs (one per level) and produces a directed mesh graph with intra-level connectivity and inter-level up/down connections. + The ``intra_level["pattern"]`` defines the spatial neighbourhood connectivity + within each mesh level: + - ``"4-star"``: only cardinal directions (horizontal and vertical neighbours) + - ``"8-star"``: cardinal directions plus diagonals (all 8 surrounding neighbours) + Parameters ---------- G_coords_list : list of networkx.Graph - List of undirected coordinate graphs, one per level. Each should have - nodes with "pos" and "type" attributes, and edges with "adjacency_type" - attributes. Created by create_multirange_2d_mesh_coordinates. - intra_level : dict or None - Intra-level connectivity options. Supports: - - pattern: str, "4-star" or "8-star" (default: "8-star") - inter_level : dict or None - Inter-level connectivity options. Supports: - - pattern: str, "nearest" (default: "nearest") - - k: int, number of nearest neighbours (default: 1) + List of undirected mesh primitive graphs, one per level. Each graph + must have: + - Node attributes: ``"pos"`` (np.ndarray of shape [2,]), ``"type"`` (str) + - Edge attributes: ``"adjacency_type"`` (str, ``"cardinal"`` or ``"diagonal"``) + Created by ``create_multirange_2d_mesh_primitives``. + intra_level : dict, optional + Configuration for intra-level connectivity. Keys: + - ``"pattern"`` (str): ``"4-star"`` or ``"8-star"``. + Default: ``{"pattern": "8-star"}`` + inter_level : dict, optional + Configuration for inter-level connectivity. Keys: + - ``"pattern"`` (str): Currently only ``"nearest"`` is supported. + - ``"k"`` (int): Number of nearest neighbours for inter-level connections. + Default: ``{"pattern": "nearest", "k": 1}`` Returns ------- @@ -45,19 +57,21 @@ def create_hierarchical_from_coordinates( if inter_level is None: inter_level = {"pattern": "nearest", "k": 1} - intra_pattern = intra_level.get("pattern", "8-star") - inter_pattern = inter_level.get("pattern", "nearest") - inter_k = inter_level.get("k", 1) + intra_level_pattern = intra_level.get("pattern", "8-star") + inter_level_pattern = inter_level.get("pattern", "nearest") + inter_level_k = inter_level.get("k", 1) - if inter_pattern != "nearest": + if inter_level_pattern != "nearest": raise NotImplementedError( - f"Inter-level pattern '{inter_pattern}' is not yet supported " + f"Inter-level pattern '{inter_level_pattern}' is not yet supported " "for hierarchical graphs. Only 'nearest' is currently implemented." ) # Convert each level's coordinate graph to directed graph with chosen pattern Gs_all_levels = [ - mesh_graph.create_directed_mesh_graph(g_coords, pattern=intra_pattern) + mesh_coords.create_directed_mesh_graph( + g_coords, pattern=intra_level_pattern + ) for g_coords in G_coords_list ] @@ -111,8 +125,8 @@ def create_hierarchical_from_coordinates( # add edges from coarser to finer level for v in v_to_list: # find k nearest neighbours (index to vm_xy) - neigh_idx = kdt_m.query(G_down.nodes[v]["pos"], inter_k)[1] - if inter_k == 1: + neigh_idx = kdt_m.query(G_down.nodes[v]["pos"], inter_level_k)[1] + if inter_level_k == 1: neigh_idx = [neigh_idx] for idx in neigh_idx: @@ -157,10 +171,12 @@ def create_hierarchical_from_coordinates( def create_hierarchical_multiscale_mesh_graph( - xy, + xy: np.ndarray, mesh_node_distance: float, level_refinement_factor: float, max_num_levels: int, + intra_level: Optional[Dict[str, object]] = None, + inter_level: Optional[Dict[str, object]] = None, ): """ Create a hierarchical multiscale mesh graph with nearest neighbour @@ -169,22 +185,29 @@ def create_hierarchical_multiscale_mesh_graph( nearest neighbour connection. Internally uses the two-step process: - 1. create_multirange_2d_mesh_coordinates (coordinate creation) + 1. create_multirange_2d_mesh_primitives (coordinate creation) 2. create_hierarchical_from_coordinates (connectivity creation) Parameters ---------- - xy: np.ndarray - 2D array of mesh point positions. - Distance (in x- and y-direction) between created mesh nodes in bottom level, - in coordinate system of xy - mesh_node_distance: float + xy : np.ndarray + 2D array of mesh point positions, shaped [N_points, 2]. + mesh_node_distance : float Distance (in x- and y-direction) between created mesh nodes in bottom level, in coordinate system of xy - level_refinement_factor: float + level_refinement_factor : float Refinement factor between grid points and bottom level of mesh hierarchy - max_num_levels: int + max_num_levels : int The number of levels in the hierarchical mesh graph. + intra_level : dict, optional + Configuration for intra-level connectivity. Keys: + - ``"pattern"`` (str): ``"4-star"`` or ``"8-star"``. + Default: ``{"pattern": "8-star"}`` + inter_level : dict, optional + Configuration for inter-level connectivity. Keys: + - ``"pattern"`` (str): Currently only ``"nearest"`` is supported. + - ``"k"`` (int): Number of nearest neighbours. + Default: ``{"pattern": "nearest", "k": 1}`` Returns ------- @@ -192,15 +215,15 @@ def create_hierarchical_multiscale_mesh_graph( A directed graph containing the hierarchical mesh with intra-level, up, and down edges. """ - G_coords_list = mesh_graph.create_multirange_2d_mesh_coordinates( + G_coords_list = mesh_coords.create_multirange_2d_mesh_primitives( max_num_levels=max_num_levels, xy=xy, - grid_spacing=mesh_node_distance, + mesh_node_spacing=mesh_node_distance, interlevel_refinement_factor=level_refinement_factor, ) return create_hierarchical_from_coordinates( G_coords_list, - intra_level={"pattern": "8-star"}, - inter_level={"pattern": "nearest", "k": 1}, + intra_level=intra_level, + inter_level=inter_level, ) diff --git a/tests/test_graph_creation.py b/tests/test_graph_creation.py index b45d486..858299e 100644 --- a/tests/test_graph_creation.py +++ b/tests/test_graph_creation.py @@ -75,6 +75,7 @@ def test_create_graph_generic(m2g_connectivity, g2m_connectivity, m2m_connectivi graph = wmg.create.create_all_graph_components( coords=xy, m2m_connectivity=m2m_connectivity, + mesh_layout="rectilinear", m2m_connectivity_kwargs=m2m_kwargs, g2m_connectivity=g2m_connectivity, g2m_connectivity_kwargs=g2m_kwargs, diff --git a/tests/test_graph_plots.py b/tests/test_graph_plots.py index c5b1e10..eba23b5 100644 --- a/tests/test_graph_plots.py +++ b/tests/test_graph_plots.py @@ -14,6 +14,7 @@ def test_plot(): graph = wmg.create.create_all_graph_components( m2m_connectivity="flat_multiscale", coords=xy, + mesh_layout="rectilinear", m2m_connectivity_kwargs=dict( max_num_levels=3, mesh_node_distance=2, diff --git a/tests/test_mesh_layout.py b/tests/test_mesh_layout.py index b131925..a5c15ec 100644 --- a/tests/test_mesh_layout.py +++ b/tests/test_mesh_layout.py @@ -21,10 +21,10 @@ import tests.utils as test_utils import weather_model_graphs as wmg -from weather_model_graphs.create.mesh.mesh import ( +from weather_model_graphs.create.mesh.coords import ( create_directed_mesh_graph, - create_multirange_2d_mesh_coordinates, - create_single_level_2d_mesh_coordinates, + create_multirange_2d_mesh_primitives, + create_single_level_2d_mesh_primitive, ) from weather_model_graphs.create.mesh.kinds.flat import ( create_flat_multiscale_from_coordinates, @@ -41,17 +41,17 @@ class TestSingleLevelCoordinateCreation: - """Tests for create_single_level_2d_mesh_coordinates.""" + """Tests for create_single_level_2d_mesh_primitive.""" def test_returns_undirected_graph(self): xy = test_utils.create_fake_xy(N=10) - G = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=5) + G = create_single_level_2d_mesh_primitive(xy, nx=5, ny=5) assert isinstance(G, nx.Graph) assert not isinstance(G, nx.DiGraph) def test_nodes_have_pos_and_type(self): xy = test_utils.create_fake_xy(N=10) - G = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G = create_single_level_2d_mesh_primitive(xy, nx=4, ny=4) for node in G.nodes: assert "pos" in G.nodes[node] assert "type" in G.nodes[node] @@ -60,12 +60,12 @@ def test_nodes_have_pos_and_type(self): def test_correct_number_of_nodes(self): xy = test_utils.create_fake_xy(N=10) - G = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=4) + G = create_single_level_2d_mesh_primitive(xy, nx=5, ny=4) assert len(G.nodes) == 5 * 4 def test_edges_have_adjacency_type(self): xy = test_utils.create_fake_xy(N=10) - G = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=5) + G = create_single_level_2d_mesh_primitive(xy, nx=5, ny=5) cardinal_count = 0 diagonal_count = 0 for u, v, d in G.edges(data=True): @@ -81,7 +81,7 @@ def test_edges_have_adjacency_type(self): def test_graph_has_dx_dy(self): xy = test_utils.create_fake_xy(N=10) - G = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=5) + G = create_single_level_2d_mesh_primitive(xy, nx=5, ny=5) assert "dx" in G.graph assert "dy" in G.graph assert G.graph["dx"] > 0 @@ -89,12 +89,12 @@ def test_graph_has_dx_dy(self): class TestMultirangeCoordinateCreation: - """Tests for create_multirange_2d_mesh_coordinates.""" + """Tests for create_multirange_2d_mesh_primitives.""" def test_returns_list_of_undirected_graphs(self): xy = test_utils.create_fake_xy(N=30) - G_list = create_multirange_2d_mesh_coordinates( - max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_list = create_multirange_2d_mesh_primitives( + max_num_levels=3, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) assert isinstance(G_list, list) assert len(G_list) > 0 @@ -104,8 +104,8 @@ def test_returns_list_of_undirected_graphs(self): def test_each_level_has_level_attribute(self): xy = test_utils.create_fake_xy(N=30) - G_list = create_multirange_2d_mesh_coordinates( - max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_list = create_multirange_2d_mesh_primitives( + max_num_levels=3, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) for i, G in enumerate(G_list): assert G.graph["level"] == i @@ -114,16 +114,16 @@ def test_each_level_has_level_attribute(self): def test_interlevel_refinement_factor_stored(self): xy = test_utils.create_fake_xy(N=30) - G_list = create_multirange_2d_mesh_coordinates( - max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_list = create_multirange_2d_mesh_primitives( + max_num_levels=3, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) for G in G_list: assert G.graph["interlevel_refinement_factor"] == 3 def test_edges_have_adjacency_type(self): xy = test_utils.create_fake_xy(N=30) - G_list = create_multirange_2d_mesh_coordinates( - max_num_levels=2, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_list = create_multirange_2d_mesh_primitives( + max_num_levels=2, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) for G in G_list: for u, v, d in G.edges(data=True): @@ -131,8 +131,8 @@ def test_edges_have_adjacency_type(self): def test_coarser_levels_have_fewer_nodes(self): xy = test_utils.create_fake_xy(N=30) - G_list = create_multirange_2d_mesh_coordinates( - max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_list = create_multirange_2d_mesh_primitives( + max_num_levels=3, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) if len(G_list) >= 2: for i in range(len(G_list) - 1): @@ -149,13 +149,13 @@ class TestDirectedMeshGraph: def test_returns_directed_graph(self): xy = test_utils.create_fake_xy(N=10) - G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=4, ny=4) G_directed = create_directed_mesh_graph(G_coords, pattern="8-star") assert isinstance(G_directed, nx.DiGraph) def test_4star_has_fewer_edges_than_8star(self): xy = test_utils.create_fake_xy(N=10) - G_coords = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=5) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=5, ny=5) G_4star = create_directed_mesh_graph(G_coords, pattern="4-star") G_8star = create_directed_mesh_graph(G_coords, pattern="8-star") assert len(G_4star.edges) < len(G_8star.edges) @@ -163,7 +163,7 @@ def test_4star_has_fewer_edges_than_8star(self): def test_4star_only_cardinal_edges(self): """4-star should only include cardinal (horizontal/vertical) edges.""" xy = test_utils.create_fake_xy(N=10) - G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=4, ny=4) G_4star = create_directed_mesh_graph(G_coords, pattern="4-star") # In a 4x4 grid, 4-star adjacency means each node connects only to @@ -175,7 +175,7 @@ def test_4star_only_cardinal_edges(self): def test_8star_includes_diagonal_edges(self): """8-star should include both cardinal and diagonal edges.""" xy = test_utils.create_fake_xy(N=10) - G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=4, ny=4) G_8star = create_directed_mesh_graph(G_coords, pattern="8-star") # Cardinal: 2 * (4*3 + 3*4) = 48 @@ -185,7 +185,7 @@ def test_8star_includes_diagonal_edges(self): def test_edges_have_len_and_vdiff(self): xy = test_utils.create_fake_xy(N=10) - G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=4, ny=4) G_directed = create_directed_mesh_graph(G_coords, pattern="8-star") for u, v, d in G_directed.edges(data=True): assert "len" in d @@ -194,13 +194,13 @@ def test_edges_have_len_and_vdiff(self): def test_invalid_pattern_raises_error(self): xy = test_utils.create_fake_xy(N=10) - G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=4, ny=4) with pytest.raises(ValueError, match="Unknown connectivity pattern"): create_directed_mesh_graph(G_coords, pattern="6-star") def test_preserves_graph_attributes(self): xy = test_utils.create_fake_xy(N=10) - G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=4, ny=4) G_directed = create_directed_mesh_graph(G_coords, pattern="8-star") assert "dx" in G_directed.graph assert "dy" in G_directed.graph @@ -208,7 +208,7 @@ def test_preserves_graph_attributes(self): def test_bidirectional_edges(self): """Each undirected edge should produce two directed edges.""" xy = test_utils.create_fake_xy(N=10) - G_coords = create_single_level_2d_mesh_coordinates(xy, nx=3, ny=3) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=3, ny=3) G_directed = create_directed_mesh_graph(G_coords, pattern="4-star") for u, v in G_directed.edges(): assert G_directed.has_edge(v, u), f"Missing reverse edge ({v}, {u})" @@ -219,13 +219,13 @@ class TestFlatSinglescaleFromCoordinates: def test_basic_creation(self): xy = test_utils.create_fake_xy(N=10) - G_coords = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=5) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=5, ny=5) G = create_flat_singlescale_from_coordinates(G_coords, pattern="8-star") assert isinstance(G, nx.DiGraph) def test_4star_pattern(self): xy = test_utils.create_fake_xy(N=10) - G_coords = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=5) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=5, ny=5) G = create_flat_singlescale_from_coordinates(G_coords, pattern="4-star") assert isinstance(G, nx.DiGraph) # Fewer edges than 8-star @@ -238,16 +238,16 @@ class TestFlatMultiscaleFromCoordinates: def test_basic_creation(self): xy = test_utils.create_fake_xy(N=30) - G_coords_list = create_multirange_2d_mesh_coordinates( - max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_coords_list = create_multirange_2d_mesh_primitives( + max_num_levels=3, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) G = create_flat_multiscale_from_coordinates(G_coords_list) assert isinstance(G, nx.DiGraph) def test_pattern_argument(self): xy = test_utils.create_fake_xy(N=30) - G_coords_list = create_multirange_2d_mesh_coordinates( - max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_coords_list = create_multirange_2d_mesh_primitives( + max_num_levels=3, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) G_4star = create_flat_multiscale_from_coordinates( G_coords_list, @@ -265,16 +265,16 @@ class TestHierarchicalFromCoordinates: def test_basic_creation(self): xy = test_utils.create_fake_xy(N=30) - G_coords_list = create_multirange_2d_mesh_coordinates( - max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_coords_list = create_multirange_2d_mesh_primitives( + max_num_levels=3, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) G = create_hierarchical_from_coordinates(G_coords_list) assert isinstance(G, nx.DiGraph) def test_has_up_down_same_edges(self): xy = test_utils.create_fake_xy(N=30) - G_coords_list = create_multirange_2d_mesh_coordinates( - max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_coords_list = create_multirange_2d_mesh_primitives( + max_num_levels=3, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) G = create_hierarchical_from_coordinates(G_coords_list) directions = set() @@ -287,8 +287,8 @@ def test_has_up_down_same_edges(self): def test_intra_level_pattern(self): xy = test_utils.create_fake_xy(N=30) - G_coords_list = create_multirange_2d_mesh_coordinates( - max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_coords_list = create_multirange_2d_mesh_primitives( + max_num_levels=3, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) G_4star = create_hierarchical_from_coordinates( G_coords_list, @@ -302,8 +302,8 @@ def test_intra_level_pattern(self): def test_inter_level_k_parameter(self): xy = test_utils.create_fake_xy(N=30) - G_coords_list = create_multirange_2d_mesh_coordinates( - max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_coords_list = create_multirange_2d_mesh_primitives( + max_num_levels=3, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) G_k1 = create_hierarchical_from_coordinates( G_coords_list, @@ -318,8 +318,8 @@ def test_inter_level_k_parameter(self): def test_invalid_inter_level_pattern_raises(self): xy = test_utils.create_fake_xy(N=30) - G_coords_list = create_multirange_2d_mesh_coordinates( - max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_coords_list = create_multirange_2d_mesh_primitives( + max_num_levels=3, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) with pytest.raises(NotImplementedError, match="Inter-level pattern"): create_hierarchical_from_coordinates( @@ -342,7 +342,7 @@ def test_flat_with_new_api(self): coords=xy, m2m_connectivity="flat", mesh_layout="rectilinear", - mesh_layout_kwargs=dict(grid_spacing=3), + mesh_layout_kwargs=dict(mesh_node_spacing=3), m2m_connectivity_kwargs=dict(pattern="8-star"), g2m_connectivity="within_radius", m2g_connectivity="nearest_neighbours", @@ -357,7 +357,7 @@ def test_flat_4star_pattern(self): coords=xy, m2m_connectivity="flat", mesh_layout="rectilinear", - mesh_layout_kwargs=dict(grid_spacing=3), + mesh_layout_kwargs=dict(mesh_node_spacing=3), m2m_connectivity_kwargs=dict(pattern="4-star"), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", @@ -366,7 +366,7 @@ def test_flat_4star_pattern(self): coords=xy, m2m_connectivity="flat", mesh_layout="rectilinear", - mesh_layout_kwargs=dict(grid_spacing=3), + mesh_layout_kwargs=dict(mesh_node_spacing=3), m2m_connectivity_kwargs=dict(pattern="8-star"), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", @@ -375,9 +375,9 @@ def test_flat_4star_pattern(self): assert isinstance(graph_8, nx.DiGraph) assert len(graph_4.edges) < len(graph_8.edges) - def test_missing_grid_spacing_raises(self): + def test_missing_mesh_node_spacing_raises(self): xy = test_utils.create_fake_xy(N=32) - with pytest.raises(ValueError, match="grid_spacing"): + with pytest.raises(ValueError, match="mesh_node_spacing"): wmg.create.create_all_graph_components( coords=xy, m2m_connectivity="flat", @@ -398,7 +398,7 @@ def test_flat_multiscale_with_new_api(self): m2m_connectivity="flat_multiscale", mesh_layout="rectilinear", mesh_layout_kwargs=dict( - grid_spacing=3, + mesh_node_spacing=3, refinement_factor=3, max_num_refinement_levels=3, ), @@ -419,7 +419,7 @@ def test_flat_multiscale_4star_pattern(self): m2m_connectivity="flat_multiscale", mesh_layout="rectilinear", mesh_layout_kwargs=dict( - grid_spacing=3, + mesh_node_spacing=3, refinement_factor=3, max_num_refinement_levels=3, ), @@ -442,7 +442,7 @@ def test_hierarchical_with_new_api(self): m2m_connectivity="hierarchical", mesh_layout="rectilinear", mesh_layout_kwargs=dict( - grid_spacing=3, + mesh_node_spacing=3, refinement_factor=3, max_num_refinement_levels=3, ), @@ -464,7 +464,7 @@ def test_hierarchical_4star_intra(self): m2m_connectivity="hierarchical", mesh_layout="rectilinear", mesh_layout_kwargs=dict( - grid_spacing=3, + mesh_node_spacing=3, refinement_factor=3, max_num_refinement_levels=3, ), @@ -484,7 +484,7 @@ def test_hierarchical_k3_nearest(self): m2m_connectivity="hierarchical", mesh_layout="rectilinear", mesh_layout_kwargs=dict( - grid_spacing=3, + mesh_node_spacing=3, refinement_factor=3, max_num_refinement_levels=3, ), @@ -513,6 +513,7 @@ def test_old_style_flat_with_mesh_node_distance(self): graph = wmg.create.create_all_graph_components( coords=xy, m2m_connectivity="flat", + mesh_layout="rectilinear", m2m_connectivity_kwargs=dict(mesh_node_distance=3), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", @@ -532,6 +533,7 @@ def test_old_style_flat_multiscale(self): graph = wmg.create.create_all_graph_components( coords=xy, m2m_connectivity="flat_multiscale", + mesh_layout="rectilinear", m2m_connectivity_kwargs=dict( mesh_node_distance=3, level_refinement_factor=3, @@ -553,6 +555,7 @@ def test_old_style_hierarchical(self): graph = wmg.create.create_all_graph_components( coords=xy, m2m_connectivity="hierarchical", + mesh_layout="rectilinear", m2m_connectivity_kwargs=dict( mesh_node_distance=3, level_refinement_factor=3, @@ -582,6 +585,7 @@ def test_kwargs_dict_not_mutated(self): wmg.create.create_all_graph_components( coords=xy, m2m_connectivity="hierarchical", + mesh_layout="rectilinear", m2m_connectivity_kwargs=original_kwargs, g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", @@ -606,7 +610,7 @@ def test_unsupported_mesh_layout_raises(self): coords=xy, m2m_connectivity="flat", mesh_layout="triangular", - mesh_layout_kwargs=dict(grid_spacing=3), + mesh_layout_kwargs=dict(mesh_node_spacing=3), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", ) @@ -618,19 +622,19 @@ def test_unsupported_m2m_connectivity_raises(self): coords=xy, m2m_connectivity="some_unknown", mesh_layout="rectilinear", - mesh_layout_kwargs=dict(grid_spacing=3), + mesh_layout_kwargs=dict(mesh_node_spacing=3), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", ) - def test_grid_spacing_too_large_raises(self): + def test_mesh_node_spacing_too_large_raises(self): xy = test_utils.create_fake_xy(N=10) with pytest.raises(ValueError, match="too large"): wmg.create.create_all_graph_components( coords=xy, m2m_connectivity="flat", mesh_layout="rectilinear", - mesh_layout_kwargs=dict(grid_spacing=100), + mesh_layout_kwargs=dict(mesh_node_spacing=100), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", ) @@ -657,7 +661,7 @@ def test_keisler_archetype_matches_new_api(self): coords=xy, m2m_connectivity="flat", mesh_layout="rectilinear", - mesh_layout_kwargs=dict(grid_spacing=3), + mesh_layout_kwargs=dict(mesh_node_spacing=3), m2m_connectivity_kwargs=dict(pattern="8-star"), g2m_connectivity="within_radius", m2g_connectivity="nearest_neighbours", @@ -683,7 +687,7 @@ def test_graphcast_archetype_matches_new_api(self): m2m_connectivity="flat_multiscale", mesh_layout="rectilinear", mesh_layout_kwargs=dict( - grid_spacing=3, + mesh_node_spacing=3, refinement_factor=3, max_num_refinement_levels=3, ), @@ -714,7 +718,7 @@ def test_oskarsson_archetype_matches_new_api(self): m2m_connectivity="hierarchical", mesh_layout="rectilinear", mesh_layout_kwargs=dict( - grid_spacing=3, + mesh_node_spacing=3, refinement_factor=3, max_num_refinement_levels=3, ), @@ -743,7 +747,7 @@ class TestCoordinateCreationEdgeCases: def test_minimum_grid_2x2(self): """Smallest possible grid: 2x2 nodes.""" xy = test_utils.create_fake_xy(N=10) - G = create_single_level_2d_mesh_coordinates(xy, nx=2, ny=2) + G = create_single_level_2d_mesh_primitive(xy, nx=2, ny=2) assert len(G.nodes) == 4 # 2x2 grid: cardinal edges = 2*(2*1) = 4, diagonal edges = 2*(1*1) = 2 cardinal = sum( @@ -758,7 +762,7 @@ def test_minimum_grid_2x2(self): def test_single_row_grid(self): """Grid with only 1 row (nx=5, ny=1).""" xy = test_utils.create_fake_xy(N=10) - G = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=1) + G = create_single_level_2d_mesh_primitive(xy, nx=5, ny=1) assert len(G.nodes) == 5 # 5x1 grid: only horizontal cardinal edges, no diagonals cardinal = sum( @@ -773,7 +777,7 @@ def test_single_row_grid(self): def test_single_column_grid(self): """Grid with only 1 column (nx=1, ny=5).""" xy = test_utils.create_fake_xy(N=10) - G = create_single_level_2d_mesh_coordinates(xy, nx=1, ny=5) + G = create_single_level_2d_mesh_primitive(xy, nx=1, ny=5) assert len(G.nodes) == 5 cardinal = sum( 1 for _, _, d in G.edges(data=True) if d["adjacency_type"] == "cardinal" @@ -787,14 +791,14 @@ def test_single_column_grid(self): def test_1x1_grid_no_edges(self): """Grid with a single node (1x1): should have no edges.""" xy = test_utils.create_fake_xy(N=10) - G = create_single_level_2d_mesh_coordinates(xy, nx=1, ny=1) + G = create_single_level_2d_mesh_primitive(xy, nx=1, ny=1) assert len(G.nodes) == 1 assert len(G.edges) == 0 def test_large_grid(self): """Larger grid should still work correctly.""" xy = test_utils.create_fake_xy(N=50) - G = create_single_level_2d_mesh_coordinates(xy, nx=10, ny=10) + G = create_single_level_2d_mesh_primitive(xy, nx=10, ny=10) assert len(G.nodes) == 100 expected_cardinal = 2 * (10 * 9) # 180 expected_diagonal = 2 * (9 * 9) # 162 @@ -810,7 +814,7 @@ def test_large_grid(self): def test_node_positions_within_bounds(self): """Node positions should be within the xy bounds.""" xy = test_utils.create_fake_xy(N=20) - G = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=5) + G = create_single_level_2d_mesh_primitive(xy, nx=5, ny=5) x_min, y_min = np.amin(xy, axis=0) x_max, y_max = np.amax(xy, axis=0) for node in G.nodes: @@ -821,8 +825,8 @@ def test_node_positions_within_bounds(self): def test_multirange_with_max_levels_1(self): """Multi-range with max_num_levels=1 should return single-level list.""" xy = test_utils.create_fake_xy(N=30) - G_list = create_multirange_2d_mesh_coordinates( - max_num_levels=1, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_list = create_multirange_2d_mesh_primitives( + max_num_levels=1, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) assert len(G_list) == 1 assert G_list[0].graph["level"] == 0 @@ -830,8 +834,8 @@ def test_multirange_with_max_levels_1(self): def test_multirange_with_none_max_levels(self): """max_num_levels=None should auto-compute levels.""" xy = test_utils.create_fake_xy(N=30) - G_list = create_multirange_2d_mesh_coordinates( - max_num_levels=None, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_list = create_multirange_2d_mesh_primitives( + max_num_levels=None, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) assert isinstance(G_list, list) assert len(G_list) >= 1 @@ -839,8 +843,8 @@ def test_multirange_with_none_max_levels(self): def test_multirange_refinement_factor_5(self): """Test with a different refinement factor.""" xy = test_utils.create_fake_xy(N=50) - G_list = create_multirange_2d_mesh_coordinates( - max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=5 + G_list = create_multirange_2d_mesh_primitives( + max_num_levels=3, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=5 ) if len(G_list) >= 2: for i in range(len(G_list) - 1): @@ -853,7 +857,7 @@ class TestConnectivityCreationEdgeCases: def test_directed_graph_from_1x1(self): """Creating directed graph from a single-node coordinate graph.""" xy = test_utils.create_fake_xy(N=10) - G_coords = create_single_level_2d_mesh_coordinates(xy, nx=1, ny=1) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=1, ny=1) G_directed = create_directed_mesh_graph(G_coords, pattern="8-star") assert isinstance(G_directed, nx.DiGraph) assert len(G_directed.nodes) == 1 @@ -862,7 +866,7 @@ def test_directed_graph_from_1x1(self): def test_directed_graph_from_2x1(self): """Creating directed graph from a 2x1 grid.""" xy = test_utils.create_fake_xy(N=10) - G_coords = create_single_level_2d_mesh_coordinates(xy, nx=2, ny=1) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=2, ny=1) G_4star = create_directed_mesh_graph(G_coords, pattern="4-star") G_8star = create_directed_mesh_graph(G_coords, pattern="8-star") # 2x1: 1 edge, both patterns should have same (no diagonals possible) @@ -872,7 +876,7 @@ def test_directed_graph_from_2x1(self): def test_4star_is_subset_of_8star(self): """All edges in 4-star should exist in 8-star.""" xy = test_utils.create_fake_xy(N=10) - G_coords = create_single_level_2d_mesh_coordinates(xy, nx=5, ny=5) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=5, ny=5) G_4star = create_directed_mesh_graph(G_coords, pattern="4-star") G_8star = create_directed_mesh_graph(G_coords, pattern="8-star") for u, v in G_4star.edges(): @@ -881,7 +885,7 @@ def test_4star_is_subset_of_8star(self): def test_edge_lengths_are_positive(self): """All edge lengths should be positive.""" xy = test_utils.create_fake_xy(N=10) - G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=4, ny=4) G = create_directed_mesh_graph(G_coords, pattern="8-star") for u, v, d in G.edges(data=True): assert d["len"] > 0, f"Edge ({u},{v}) has non-positive length" @@ -889,7 +893,7 @@ def test_edge_lengths_are_positive(self): def test_vdiff_antisymmetric(self): """vdiff(u,v) should be -vdiff(v,u).""" xy = test_utils.create_fake_xy(N=10) - G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=4, ny=4) G = create_directed_mesh_graph(G_coords, pattern="8-star") for u, v in G.edges(): if G.has_edge(v, u): @@ -901,7 +905,7 @@ def test_vdiff_antisymmetric(self): def test_edge_len_matches_vdiff_norm(self): """Edge length should equal the norm of vdiff.""" xy = test_utils.create_fake_xy(N=10) - G_coords = create_single_level_2d_mesh_coordinates(xy, nx=4, ny=4) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=4, ny=4) G = create_directed_mesh_graph(G_coords, pattern="8-star") for u, v, d in G.edges(data=True): expected_len = np.sqrt(np.sum(d["vdiff"] ** 2)) @@ -913,8 +917,8 @@ def test_edge_len_matches_vdiff_norm(self): def test_flat_multiscale_single_level_input(self): """Flat multiscale with a single-level list should work (degenerate case).""" xy = test_utils.create_fake_xy(N=30) - G_coords_list = create_multirange_2d_mesh_coordinates( - max_num_levels=1, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_coords_list = create_multirange_2d_mesh_primitives( + max_num_levels=1, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) G = create_flat_multiscale_from_coordinates(G_coords_list, pattern="8-star") assert isinstance(G, nx.DiGraph) @@ -923,8 +927,8 @@ def test_flat_multiscale_single_level_input(self): def test_flat_multiscale_4star_vs_8star_edge_count(self): """4-star flat_multiscale should have fewer edges than 8-star.""" xy = test_utils.create_fake_xy(N=30) - G_coords_list = create_multirange_2d_mesh_coordinates( - max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_coords_list = create_multirange_2d_mesh_primitives( + max_num_levels=3, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) G_4 = create_flat_multiscale_from_coordinates(G_coords_list, pattern="4-star") G_8 = create_flat_multiscale_from_coordinates(G_coords_list, pattern="8-star") @@ -933,8 +937,8 @@ def test_flat_multiscale_4star_vs_8star_edge_count(self): def test_hierarchical_single_level_raises(self): """Hierarchical with only 1 level should raise ValueError.""" xy = test_utils.create_fake_xy(N=30) - G_coords_list = create_multirange_2d_mesh_coordinates( - max_num_levels=1, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_coords_list = create_multirange_2d_mesh_primitives( + max_num_levels=1, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) with pytest.raises(ValueError, match="At least two mesh levels"): create_hierarchical_from_coordinates(G_coords_list) @@ -942,8 +946,8 @@ def test_hierarchical_single_level_raises(self): def test_hierarchical_edge_direction_attributes(self): """Every edge in hierarchical graph must have a 'direction' attribute.""" xy = test_utils.create_fake_xy(N=30) - G_coords_list = create_multirange_2d_mesh_coordinates( - max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_coords_list = create_multirange_2d_mesh_primitives( + max_num_levels=3, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) G = create_hierarchical_from_coordinates(G_coords_list) for u, v, d in G.edges(data=True): @@ -953,8 +957,8 @@ def test_hierarchical_edge_direction_attributes(self): def test_hierarchical_up_down_symmetry(self): """For each 'down' edge (u,v), there should be an 'up' edge (v,u).""" xy = test_utils.create_fake_xy(N=30) - G_coords_list = create_multirange_2d_mesh_coordinates( - max_num_levels=3, xy=xy, grid_spacing=3, interlevel_refinement_factor=3 + G_coords_list = create_multirange_2d_mesh_primitives( + max_num_levels=3, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) G = create_hierarchical_from_coordinates(G_coords_list) for u, v, d in G.edges(data=True): @@ -973,7 +977,7 @@ def test_flat_default_pattern_is_8star(self): coords=xy, m2m_connectivity="flat", mesh_layout="rectilinear", - mesh_layout_kwargs=dict(grid_spacing=3), + mesh_layout_kwargs=dict(mesh_node_spacing=3), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", ) @@ -981,7 +985,7 @@ def test_flat_default_pattern_is_8star(self): coords=xy, m2m_connectivity="flat", mesh_layout="rectilinear", - mesh_layout_kwargs=dict(grid_spacing=3), + mesh_layout_kwargs=dict(mesh_node_spacing=3), m2m_connectivity_kwargs=dict(pattern="8-star"), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", @@ -996,7 +1000,7 @@ def test_flat_multiscale_default_pattern_is_8star(self): m2m_connectivity="flat_multiscale", mesh_layout="rectilinear", mesh_layout_kwargs=dict( - grid_spacing=3, + mesh_node_spacing=3, refinement_factor=3, max_num_refinement_levels=3, ), @@ -1008,7 +1012,7 @@ def test_flat_multiscale_default_pattern_is_8star(self): m2m_connectivity="flat_multiscale", mesh_layout="rectilinear", mesh_layout_kwargs=dict( - grid_spacing=3, + mesh_node_spacing=3, refinement_factor=3, max_num_refinement_levels=3, ), @@ -1026,7 +1030,7 @@ def test_hierarchical_default_kwargs(self): m2m_connectivity="hierarchical", mesh_layout="rectilinear", mesh_layout_kwargs=dict( - grid_spacing=3, + mesh_node_spacing=3, refinement_factor=3, max_num_refinement_levels=3, ), @@ -1038,7 +1042,7 @@ def test_hierarchical_default_kwargs(self): m2m_connectivity="hierarchical", mesh_layout="rectilinear", mesh_layout_kwargs=dict( - grid_spacing=3, + mesh_node_spacing=3, refinement_factor=3, max_num_refinement_levels=3, ), @@ -1051,17 +1055,17 @@ def test_hierarchical_default_kwargs(self): ) assert len(graph_default.edges) == len(graph_explicit.edges) - def test_mesh_layout_default_is_rectilinear(self): - """When mesh_layout not specified, it should default to rectilinear.""" + def test_mesh_layout_is_required(self): + """When mesh_layout not specified, it should raise TypeError.""" xy = test_utils.create_fake_xy(N=32) - graph = wmg.create.create_all_graph_components( - coords=xy, - m2m_connectivity="flat", - mesh_layout_kwargs=dict(grid_spacing=3), - g2m_connectivity="nearest_neighbour", - m2g_connectivity="nearest_neighbour", - ) - assert isinstance(graph, nx.DiGraph) + with pytest.raises(TypeError, match="mesh_layout"): + wmg.create.create_all_graph_components( + coords=xy, + m2m_connectivity="flat", + mesh_layout_kwargs=dict(mesh_node_spacing=3), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) def test_return_components_flat(self): """return_components=True should return dict with g2m, m2m, m2g.""" @@ -1070,7 +1074,7 @@ def test_return_components_flat(self): coords=xy, m2m_connectivity="flat", mesh_layout="rectilinear", - mesh_layout_kwargs=dict(grid_spacing=3), + mesh_layout_kwargs=dict(mesh_node_spacing=3), m2m_connectivity_kwargs=dict(pattern="8-star"), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", @@ -1091,7 +1095,7 @@ def test_return_components_hierarchical(self): m2m_connectivity="hierarchical", mesh_layout="rectilinear", mesh_layout_kwargs=dict( - grid_spacing=3, + mesh_node_spacing=3, refinement_factor=3, max_num_refinement_levels=3, ), @@ -1116,7 +1120,7 @@ def test_return_components_flat_multiscale(self): m2m_connectivity="flat_multiscale", mesh_layout="rectilinear", mesh_layout_kwargs=dict( - grid_spacing=3, + mesh_node_spacing=3, refinement_factor=3, max_num_refinement_levels=3, ), @@ -1136,7 +1140,7 @@ def test_flat_multiscale_no_sub_dicts_interface(self): m2m_connectivity="flat_multiscale", mesh_layout="rectilinear", mesh_layout_kwargs=dict( - grid_spacing=3, + mesh_node_spacing=3, refinement_factor=3, max_num_refinement_levels=3, ), @@ -1155,7 +1159,7 @@ def test_decode_mask_with_new_api(self): coords=xy, m2m_connectivity="flat", mesh_layout="rectilinear", - mesh_layout_kwargs=dict(grid_spacing=3), + mesh_layout_kwargs=dict(mesh_node_spacing=3), m2m_connectivity_kwargs=dict(pattern="8-star"), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", @@ -1176,6 +1180,7 @@ def test_old_kwargs_with_flat_multiscale_compat(self): graph = wmg.create.create_all_graph_components( coords=xy, m2m_connectivity="flat_multiscale", + mesh_layout="rectilinear", m2m_connectivity_kwargs=dict( mesh_node_distance=3, level_refinement_factor=3, @@ -1191,7 +1196,7 @@ def test_old_kwargs_with_flat_multiscale_compat(self): assert len(deprecation_warnings) >= 3 # Check the new names appear in the messages msgs = " ".join(str(x.message) for x in deprecation_warnings) - assert "grid_spacing" in msgs + assert "mesh_node_spacing" in msgs assert "refinement_factor" in msgs assert "max_num_refinement_levels" in msgs assert isinstance(graph, nx.DiGraph) @@ -1207,7 +1212,7 @@ def test_all_mesh_nodes_have_pos(self): coords=xy, m2m_connectivity="flat", mesh_layout="rectilinear", - mesh_layout_kwargs=dict(grid_spacing=3), + mesh_layout_kwargs=dict(mesh_node_spacing=3), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", ) @@ -1221,7 +1226,7 @@ def test_all_edges_have_component(self): coords=xy, m2m_connectivity="flat", mesh_layout="rectilinear", - mesh_layout_kwargs=dict(grid_spacing=3), + mesh_layout_kwargs=dict(mesh_node_spacing=3), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", ) @@ -1236,7 +1241,7 @@ def test_all_edges_have_len_and_vdiff(self): coords=xy, m2m_connectivity="flat", mesh_layout="rectilinear", - mesh_layout_kwargs=dict(grid_spacing=3), + mesh_layout_kwargs=dict(mesh_node_spacing=3), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", ) @@ -1252,7 +1257,7 @@ def test_graph_is_directed(self): coords=xy, m2m_connectivity=connectivity, mesh_layout="rectilinear", - mesh_layout_kwargs=dict(grid_spacing=3), + mesh_layout_kwargs=dict(mesh_node_spacing=3), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", ) @@ -1265,7 +1270,7 @@ def test_flat_4star_strictly_fewer_m2m_edges(self): coords=xy, m2m_connectivity="flat", mesh_layout="rectilinear", - mesh_layout_kwargs=dict(grid_spacing=3), + mesh_layout_kwargs=dict(mesh_node_spacing=3), m2m_connectivity_kwargs=dict(pattern="4-star"), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", @@ -1275,7 +1280,7 @@ def test_flat_4star_strictly_fewer_m2m_edges(self): coords=xy, m2m_connectivity="flat", mesh_layout="rectilinear", - mesh_layout_kwargs=dict(grid_spacing=3), + mesh_layout_kwargs=dict(mesh_node_spacing=3), m2m_connectivity_kwargs=dict(pattern="8-star"), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", @@ -1295,7 +1300,7 @@ def test_hierarchical_has_same_up_down_edge_count(self): m2m_connectivity="hierarchical", mesh_layout="rectilinear", mesh_layout_kwargs=dict( - grid_spacing=3, + mesh_node_spacing=3, refinement_factor=3, max_num_refinement_levels=3, ), @@ -1317,4 +1322,4 @@ def test_hierarchical_has_same_up_down_edge_count(self): assert up_count == down_count, ( f"Up edges ({up_count}) != Down edges ({down_count})" ) - assert up_count > 0, "Should have at least some up/down edges" \ No newline at end of file + assert up_count > 0, "Should have at least some up/down edges" From 21d33e8c74574893a5a7b2f97f009af5fceec321 Mon Sep 17 00:00:00 2001 From: prajwal Date: Sun, 8 Mar 2026 10:33:38 +0530 Subject: [PATCH 06/21] refactor: add coords module with two-step mesh creation (primitives + directed graph) Add type annotations and expose pattern argument in create_multirange_2d_mesh_graphs. Each mesh primitive (nx.Graph) encodes spatial adjacency with 'cardinal'/'diagonal' edge annotations; create_directed_mesh_graph converts to nx.DiGraph with pattern filtering ('4-star' or '8-star'). --- .../create/mesh/coords.py | 26 +++++++++++++------ 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/src/weather_model_graphs/create/mesh/coords.py b/src/weather_model_graphs/create/mesh/coords.py index 47dd30c..1b76709 100644 --- a/src/weather_model_graphs/create/mesh/coords.py +++ b/src/weather_model_graphs/create/mesh/coords.py @@ -159,7 +159,7 @@ def create_directed_mesh_graph( return dg -def create_single_level_2d_mesh_graph(xy, nx, ny): +def create_single_level_2d_mesh_graph(xy: np.ndarray, nx: int, ny: int): """ Create directed graph with nx * ny nodes representing a 2D grid with positions spanning the range of xy coordinate values (first dimension @@ -201,7 +201,10 @@ def create_single_level_2d_mesh_graph(xy, nx, ny): def create_multirange_2d_mesh_primitives( - max_num_levels, xy, mesh_node_spacing=3, interlevel_refinement_factor=3 + max_num_levels: int, + xy: np.ndarray, + mesh_node_spacing: float = 3, + interlevel_refinement_factor: float = 3, ): """ Create a list of undirected mesh primitive graphs (nx.Graph) representing @@ -285,7 +288,11 @@ def create_multirange_2d_mesh_primitives( def create_multirange_2d_mesh_graphs( - max_num_levels, xy, mesh_node_distance=3, level_refinement_factor=3 + max_num_levels: int, + xy: np.ndarray, + mesh_node_distance: float = 3, + level_refinement_factor: float = 3, + pattern: str = "8-star", ): """ Create a list of 2D grid mesh graphs representing different levels of edge-length @@ -297,7 +304,7 @@ def create_multirange_2d_mesh_graphs( Internally uses the two-step process: 1. create_multirange_2d_mesh_primitives (coordinate creation) - 2. create_directed_mesh_graph (connectivity creation, pattern="8-star") + 2. create_directed_mesh_graph (connectivity creation) Parameters ---------- @@ -305,15 +312,18 @@ def create_multirange_2d_mesh_graphs( Number of edge-distance levels in mesh graph xy : np.ndarray Grid point coordinates, shaped [N_grid_points, 2] - mesh_node_distance: float + mesh_node_distance : float Distance (in x- and y-direction) between created mesh nodes, in coordinate system of xy - level_refinement_factor: float + level_refinement_factor : float Refinement factor between grid points and bottom level of mesh hierarchy + pattern : str + Connectivity pattern for directed graph creation: ``"4-star"`` or + ``"8-star"`` (default: ``"8-star"``) Returns ------- - G_all_levels : list of networkx.Graph + G_all_levels : list of networkx.DiGraph List of networkx graphs for each level representing the connectivity of the mesh within each level """ @@ -326,7 +336,7 @@ def create_multirange_2d_mesh_graphs( G_all_levels = [] for g_coords in G_coords_list: - g_directed = create_directed_mesh_graph(g_coords, pattern="8-star") + g_directed = create_directed_mesh_graph(g_coords, pattern=pattern) G_all_levels.append(g_directed) return G_all_levels From 8d18530c64c513671ddf0514cf36813c42b3dc41 Mon Sep 17 00:00:00 2001 From: prajwal Date: Sun, 8 Mar 2026 10:34:04 +0530 Subject: [PATCH 07/21] refactor: use explicit keyword argument names in flat mesh functions Add create_flat_multiscale_from_coordinates and create_flat_singlescale_from_coordinates that take coordinate graphs from the first step. Use explicit kwarg names in all function calls (e.g. mesh_coords.create_single_level_2d_mesh_graph(xy=xy, nx=nx, ny=ny)). --- src/weather_model_graphs/create/mesh/kinds/flat.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/weather_model_graphs/create/mesh/kinds/flat.py b/src/weather_model_graphs/create/mesh/kinds/flat.py index 4329389..deda1ad 100644 --- a/src/weather_model_graphs/create/mesh/kinds/flat.py +++ b/src/weather_model_graphs/create/mesh/kinds/flat.py @@ -272,4 +272,4 @@ def create_flat_singlescale_mesh_graph(xy, mesh_node_distance: float): " so that the mesh nodes are spaced closer together?" ) - return mesh_coords.create_single_level_2d_mesh_graph(xy, nx, ny) + return mesh_coords.create_single_level_2d_mesh_graph(xy=xy, nx=nx, ny=ny) From 4f5a8de0d257f7f6b096fa2d7aee0a402541c8c3 Mon Sep 17 00:00:00 2001 From: prajwal Date: Sun, 8 Mar 2026 10:34:14 +0530 Subject: [PATCH 08/21] refactor: put default values in call signature for create_hierarchical_from_coordinates Instead of using None-to-default mapping pattern, set defaults directly: intra_level={'pattern': '8-star'} inter_level={'pattern': 'nearest', 'k': 1} This makes the actual defaults visible in the function signature. --- .../create/mesh/kinds/hierarchical.py | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/weather_model_graphs/create/mesh/kinds/hierarchical.py b/src/weather_model_graphs/create/mesh/kinds/hierarchical.py index d1a9ca3..19a0cd2 100644 --- a/src/weather_model_graphs/create/mesh/kinds/hierarchical.py +++ b/src/weather_model_graphs/create/mesh/kinds/hierarchical.py @@ -10,8 +10,8 @@ def create_hierarchical_from_coordinates( G_coords_list: List[networkx.Graph], - intra_level: Optional[Dict[str, object]] = None, - inter_level: Optional[Dict[str, object]] = None, + intra_level: Dict[str, object] = {"pattern": "8-star"}, + inter_level: Dict[str, object] = {"pattern": "nearest", "k": 1}, ): """ Create a hierarchical multiscale mesh graph from a list of mesh primitive @@ -35,11 +35,11 @@ def create_hierarchical_from_coordinates( - Node attributes: ``"pos"`` (np.ndarray of shape [2,]), ``"type"`` (str) - Edge attributes: ``"adjacency_type"`` (str, ``"cardinal"`` or ``"diagonal"``) Created by ``create_multirange_2d_mesh_primitives``. - intra_level : dict, optional + intra_level : dict Configuration for intra-level connectivity. Keys: - ``"pattern"`` (str): ``"4-star"`` or ``"8-star"``. Default: ``{"pattern": "8-star"}`` - inter_level : dict, optional + inter_level : dict Configuration for inter-level connectivity. Keys: - ``"pattern"`` (str): Currently only ``"nearest"`` is supported. - ``"k"`` (int): Number of nearest neighbours for inter-level connections. @@ -52,11 +52,6 @@ def create_hierarchical_from_coordinates( edges (direction="same"), inter-level down edges (direction="down"), and inter-level up edges (direction="up"). """ - if intra_level is None: - intra_level = {"pattern": "8-star"} - if inter_level is None: - inter_level = {"pattern": "nearest", "k": 1} - intra_level_pattern = intra_level.get("pattern", "8-star") inter_level_pattern = inter_level.get("pattern", "nearest") inter_level_k = inter_level.get("k", 1) @@ -222,8 +217,13 @@ def create_hierarchical_multiscale_mesh_graph( interlevel_refinement_factor=level_refinement_factor, ) + kwargs = {} + if intra_level is not None: + kwargs["intra_level"] = intra_level + if inter_level is not None: + kwargs["inter_level"] = inter_level + return create_hierarchical_from_coordinates( G_coords_list, - intra_level=intra_level, - inter_level=inter_level, + **kwargs, ) From e7e3c03e714823e2907fa698129842469a7d36fe Mon Sep 17 00:00:00 2001 From: prajwal Date: Sun, 8 Mar 2026 10:34:30 +0530 Subject: [PATCH 09/21] refactor: use loguru logger.warning() and implement two-step mesh creation in base - Use logger.warning() from loguru instead of warnings.warn() in _migrate_deprecated_kwargs - Add mesh_layout + mesh_layout_kwargs parameters for coordinate creation step - Implement two-step process (coordinate creation then connectivity creation) for all three m2m types: flat, hierarchical, flat_multiscale --- src/weather_model_graphs/create/base.py | 49 ++++++++++++------------- 1 file changed, 24 insertions(+), 25 deletions(-) diff --git a/src/weather_model_graphs/create/base.py b/src/weather_model_graphs/create/base.py index ddee05b..48bbd11 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/weather_model_graphs/create/base.py @@ -8,7 +8,6 @@ function uses `connect_nodes_across_graphs` to connect nodes across the component graphs. """ -import warnings from typing import Iterable import networkx @@ -41,12 +40,12 @@ def _migrate_deprecated_kwargs(mesh_layout_kwargs, m2m_connectivity_kwargs): """Migrate old-style kwargs to the new mesh_layout_kwargs structure. In the old API, ``mesh_node_distance``, ``level_refinement_factor``, and - ``max_num_levels`` were passed via ``m2m_connectivity_kwargs``. In the new + ``max_num_levels`` were passed via ``m2m_connectivity_kwargs``. In the new design these belong in ``mesh_layout_kwargs`` (as ``mesh_node_spacing``, ``refinement_factor``, and ``max_num_refinement_levels`` respectively). - This helper emits ``DeprecationWarning`` for each migrated key and moves - the value into *mesh_layout_kwargs*. It is intended to be removed once the + This helper emits deprecation warnings for each migrated key and moves + the value into *mesh_layout_kwargs*. It is intended to be removed once the old API is no longer supported. Parameters @@ -62,31 +61,25 @@ def _migrate_deprecated_kwargs(mesh_layout_kwargs, m2m_connectivity_kwargs): Updated (mesh_layout_kwargs, m2m_connectivity_kwargs). """ if "mesh_node_distance" in m2m_connectivity_kwargs and "mesh_node_spacing" not in mesh_layout_kwargs: - warnings.warn( + logger.warning( "Passing 'mesh_node_distance' in m2m_connectivity_kwargs is deprecated. " - "Use mesh_layout_kwargs=dict(mesh_node_spacing=...) instead.", - DeprecationWarning, - stacklevel=3, + "Use mesh_layout_kwargs=dict(mesh_node_spacing=...) instead." ) mesh_layout_kwargs["mesh_node_spacing"] = m2m_connectivity_kwargs.pop( "mesh_node_distance" ) if "level_refinement_factor" in m2m_connectivity_kwargs and "refinement_factor" not in mesh_layout_kwargs: - warnings.warn( + logger.warning( "Passing 'level_refinement_factor' in m2m_connectivity_kwargs is deprecated. " - "Use mesh_layout_kwargs=dict(refinement_factor=...) instead.", - DeprecationWarning, - stacklevel=3, + "Use mesh_layout_kwargs=dict(refinement_factor=...) instead." ) mesh_layout_kwargs["refinement_factor"] = ( m2m_connectivity_kwargs.pop("level_refinement_factor") ) if "max_num_levels" in m2m_connectivity_kwargs and "max_num_refinement_levels" not in mesh_layout_kwargs: - warnings.warn( + logger.warning( "Passing 'max_num_levels' in m2m_connectivity_kwargs is deprecated. " - "Use mesh_layout_kwargs=dict(max_num_refinement_levels=...) instead.", - DeprecationWarning, - stacklevel=3, + "Use mesh_layout_kwargs=dict(max_num_refinement_levels=...) instead." ) mesh_layout_kwargs["max_num_refinement_levels"] = m2m_connectivity_kwargs.pop( "max_num_levels" @@ -137,15 +130,15 @@ def create_all_graph_components( mesh_layout: - "rectilinear": Uniform regular grid with ``mesh_node_spacing`` resolution. - Produces an undirected mesh primitive with 4-star (cardinal) and - 8-star (cardinal + diagonal) spatial adjacency edges. + Produces an undirected mesh primitive with 4-star (cardinal) and + 8-star (cardinal + diagonal) spatial adjacency edges. mesh_layout_kwargs (for mesh_layout="rectilinear"): - mesh_node_spacing: float, distance between mesh nodes in coordinate units. - refinement_factor: int, refinement factor between levels - (for multi-level and hierarchical mesh graphs, default: 3) + (for multi-level and hierarchical mesh graphs, default: 3) - max_num_refinement_levels: int, maximum number of mesh levels - (for multi-level and hierarchical mesh graphs) + (for multi-level and hierarchical mesh graphs) Wherever the ``pattern`` argument appears below it defines the spatial neighbourhood connectivity: @@ -154,11 +147,11 @@ def create_all_graph_components( m2m_connectivity: - "flat": Create a single-level directed mesh graph. - m2m_connectivity_kwargs: pattern (default: "8-star") + m2m_connectivity_kwargs: pattern (default: "8-star") - "flat_multiscale": Create a flat multiscale mesh graph. - m2m_connectivity_kwargs: pattern (default: "8-star") + m2m_connectivity_kwargs: pattern (default: "8-star") - "hierarchical": Create a hierarchical mesh graph with up/down connections. - m2m_connectivity_kwargs: intra_level=dict(pattern=...), inter_level=dict(pattern=..., k=...) + m2m_connectivity_kwargs: intra_level=dict(pattern=...), inter_level=dict(pattern=..., k=...) m2g_connectivity: - "nearest_neighbour": Find the nearest neighbour in mesh for each node in grid @@ -308,10 +301,16 @@ def create_all_graph_components( # hierarchical mesh graph have three sub-graphs: # `m2m` (mesh-to-mesh), `mesh_up` (up edge connections) and # `mesh_down` (down edge connections) + hier_kwargs = {} + intra = m2m_connectivity_kwargs.get("intra_level") + inter = m2m_connectivity_kwargs.get("inter_level") + if intra is not None: + hier_kwargs["intra_level"] = intra + if inter is not None: + hier_kwargs["inter_level"] = inter graph_components["m2m"] = create_hierarchical_from_coordinates( G_coords_list, - intra_level=m2m_connectivity_kwargs.get("intra_level"), - inter_level=m2m_connectivity_kwargs.get("inter_level"), + **hier_kwargs, ) # Only connect grid to bottom level of hierarchy grid_connect_graph = split_graph_by_edge_attribute( From ef478a2842c414b6dd9cde17a330304b04b1e5a3 Mon Sep 17 00:00:00 2001 From: prajwal Date: Sun, 8 Mar 2026 10:34:42 +0530 Subject: [PATCH 10/21] refactor: update archetypes and tests to use mesh_layout parameter - All three archetype functions now pass mesh_layout='rectilinear' and mesh_layout_kwargs with mesh_node_spacing - Reorder arguments: coords_crs/graph_crs before connectivity args - Add mesh_layout='rectilinear' to test_create_graph_generic and test_plot --- src/weather_model_graphs/create/archetype.py | 34 ++++++-------------- 1 file changed, 10 insertions(+), 24 deletions(-) diff --git a/src/weather_model_graphs/create/archetype.py b/src/weather_model_graphs/create/archetype.py index d8ac806..f216335 100644 --- a/src/weather_model_graphs/create/archetype.py +++ b/src/weather_model_graphs/create/archetype.py @@ -61,16 +61,12 @@ def create_keisler_graph( graph_crs=graph_crs, mesh_layout="rectilinear", mesh_layout_kwargs=dict(mesh_node_spacing=mesh_node_distance), - g2m_connectivity="within_radius", - g2m_connectivity_kwargs=dict( - rel_max_dist=0.51, - ), m2m_connectivity="flat", m2m_connectivity_kwargs=dict(pattern="8-star"), + g2m_connectivity="within_radius", + g2m_connectivity_kwargs=dict(rel_max_dist=0.51), m2g_connectivity="nearest_neighbours", - m2g_connectivity_kwargs=dict( - max_num_neighbours=4, - ), + m2g_connectivity_kwargs=dict(max_num_neighbours=4), decode_mask=decode_mask, return_components=return_components, ) @@ -142,18 +138,12 @@ def create_graphcast_graph( refinement_factor=level_refinement_factor, max_num_refinement_levels=max_num_levels, ), - g2m_connectivity="within_radius", - g2m_connectivity_kwargs=dict( - rel_max_dist=0.51, - ), m2m_connectivity="flat_multiscale", - m2m_connectivity_kwargs=dict( - pattern="8-star", - ), + m2m_connectivity_kwargs=dict(pattern="8-star"), + g2m_connectivity="within_radius", + g2m_connectivity_kwargs=dict(rel_max_dist=0.51), m2g_connectivity="nearest_neighbours", - m2g_connectivity_kwargs=dict( - max_num_neighbours=4, - ), + m2g_connectivity_kwargs=dict(max_num_neighbours=4), decode_mask=decode_mask, return_components=return_components, ) @@ -230,19 +220,15 @@ def create_oskarsson_hierarchical_graph( refinement_factor=level_refinement_factor, max_num_refinement_levels=max_num_levels, ), - g2m_connectivity="within_radius", - g2m_connectivity_kwargs=dict( - rel_max_dist=0.51, - ), m2m_connectivity="hierarchical", m2m_connectivity_kwargs=dict( 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", - m2g_connectivity_kwargs=dict( - max_num_neighbours=4, - ), + m2g_connectivity_kwargs=dict(max_num_neighbours=4), decode_mask=decode_mask, return_components=return_components, ) From 305e66aa8ed165958d233c78f63cbaed3e13c19d Mon Sep 17 00:00:00 2001 From: prajwal Date: Sun, 8 Mar 2026 10:41:15 +0530 Subject: [PATCH 11/21] test: update backward-compat tests to capture loguru warnings instead of DeprecationWarning Since _migrate_deprecated_kwargs now uses loguru logger.warning() instead of warnings.warn(), the tests need to use a loguru StringIO sink to capture and verify the deprecation warning messages. --- tests/test_mesh_layout.py | 69 ++++++++++++++++++++------------------- 1 file changed, 36 insertions(+), 33 deletions(-) diff --git a/tests/test_mesh_layout.py b/tests/test_mesh_layout.py index a5c15ec..4b0875b 100644 --- a/tests/test_mesh_layout.py +++ b/tests/test_mesh_layout.py @@ -13,11 +13,13 @@ 6. Error handling for invalid inputs """ +import io import warnings import networkx as nx import numpy as np import pytest +from loguru import logger import tests.utils as test_utils import weather_model_graphs as wmg @@ -508,8 +510,9 @@ class TestBackwardCompatibility: def test_old_style_flat_with_mesh_node_distance(self): xy = test_utils.create_fake_xy(N=32) - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") + log_output = io.StringIO() + handler_id = logger.add(log_output, format="{message}", level="WARNING") + try: graph = wmg.create.create_all_graph_components( coords=xy, m2m_connectivity="flat", @@ -518,18 +521,16 @@ def test_old_style_flat_with_mesh_node_distance(self): g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", ) - # Should have deprecation warning - deprecation_warnings = [ - x for x in w if issubclass(x.category, DeprecationWarning) - ] - assert len(deprecation_warnings) >= 1 - assert "mesh_node_distance" in str(deprecation_warnings[0].message) + finally: + logger.remove(handler_id) + assert "mesh_node_distance" in log_output.getvalue() assert isinstance(graph, nx.DiGraph) def test_old_style_flat_multiscale(self): xy = test_utils.create_fake_xy(N=32) - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") + log_output = io.StringIO() + handler_id = logger.add(log_output, format="{message}", level="WARNING") + try: graph = wmg.create.create_all_graph_components( coords=xy, m2m_connectivity="flat_multiscale", @@ -542,16 +543,19 @@ def test_old_style_flat_multiscale(self): g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", ) - deprecation_warnings = [ - x for x in w if issubclass(x.category, DeprecationWarning) - ] - assert len(deprecation_warnings) >= 3 # 3 migrated kwargs + finally: + logger.remove(handler_id) + log_text = log_output.getvalue() + assert "mesh_node_distance" in log_text + assert "level_refinement_factor" in log_text + assert "max_num_levels" in log_text assert isinstance(graph, nx.DiGraph) def test_old_style_hierarchical(self): xy = test_utils.create_fake_xy(N=32) - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") + log_output = io.StringIO() + handler_id = logger.add(log_output, format="{message}", level="WARNING") + try: graph = wmg.create.create_all_graph_components( coords=xy, m2m_connectivity="hierarchical", @@ -564,10 +568,12 @@ def test_old_style_hierarchical(self): g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", ) - deprecation_warnings = [ - x for x in w if issubclass(x.category, DeprecationWarning) - ] - assert len(deprecation_warnings) >= 3 + finally: + logger.remove(handler_id) + log_text = log_output.getvalue() + assert "mesh_node_distance" in log_text + assert "level_refinement_factor" in log_text + assert "max_num_levels" in log_text assert isinstance(graph, nx.DiGraph) def test_kwargs_dict_not_mutated(self): @@ -1173,10 +1179,11 @@ class TestBackwardCompatEdgeCases: def test_old_kwargs_with_flat_multiscale_compat(self): """Old-style flat_multiscale kwargs should trigger deprecation warnings - and be migrated to the new names.""" + (via loguru) and be migrated to the new names.""" xy = test_utils.create_fake_xy(N=32) - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") + log_output = io.StringIO() + handler_id = logger.add(log_output, format="{message}", level="WARNING") + try: graph = wmg.create.create_all_graph_components( coords=xy, m2m_connectivity="flat_multiscale", @@ -1189,16 +1196,12 @@ def test_old_kwargs_with_flat_multiscale_compat(self): g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", ) - deprecation_warnings = [ - x for x in w if issubclass(x.category, DeprecationWarning) - ] - # Should have 3 deprecation warnings - assert len(deprecation_warnings) >= 3 - # Check the new names appear in the messages - msgs = " ".join(str(x.message) for x in deprecation_warnings) - assert "mesh_node_spacing" in msgs - assert "refinement_factor" in msgs - assert "max_num_refinement_levels" in msgs + finally: + logger.remove(handler_id) + log_text = log_output.getvalue() + assert "mesh_node_spacing" in log_text + assert "refinement_factor" in log_text + assert "max_num_refinement_levels" in log_text assert isinstance(graph, nx.DiGraph) From 4372a3f43b99ec20f08d1eeb1d8dfa5eaa825e7c Mon Sep 17 00:00:00 2001 From: prajwal Date: Tue, 10 Mar 2026 23:32:14 +0530 Subject: [PATCH 12/21] =?UTF-8?q?style:=20fix=20linting=20=E2=80=94=20run?= =?UTF-8?q?=20black,=20isort,=20remove=20empty=20f-strings?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/weather_model_graphs/create/base.py | 31 ++++++++++++------- .../create/mesh/coords.py | 12 ++----- .../create/mesh/kinds/flat.py | 6 +--- .../create/mesh/kinds/hierarchical.py | 12 +++---- tests/test_mesh_layout.py | 22 +++++++------ 5 files changed, 40 insertions(+), 43 deletions(-) diff --git a/src/weather_model_graphs/create/base.py b/src/weather_model_graphs/create/base.py index 48bbd11..efbd7a9 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/weather_model_graphs/create/base.py @@ -23,6 +23,10 @@ split_on_edge_attribute_existance, ) from .grid import create_grid_graph_nodes +from .mesh.coords import ( + create_multirange_2d_mesh_primitives, + create_single_level_2d_mesh_primitive, +) from .mesh.kinds.flat import ( create_flat_multiscale_from_coordinates, create_flat_singlescale_from_coordinates, @@ -30,10 +34,6 @@ from .mesh.kinds.hierarchical import ( create_hierarchical_from_coordinates, ) -from .mesh.coords import ( - create_multirange_2d_mesh_primitives, - create_single_level_2d_mesh_primitive, -) def _migrate_deprecated_kwargs(mesh_layout_kwargs, m2m_connectivity_kwargs): @@ -60,7 +60,10 @@ def _migrate_deprecated_kwargs(mesh_layout_kwargs, m2m_connectivity_kwargs): tuple[dict, dict] Updated (mesh_layout_kwargs, m2m_connectivity_kwargs). """ - if "mesh_node_distance" in m2m_connectivity_kwargs and "mesh_node_spacing" not in mesh_layout_kwargs: + if ( + "mesh_node_distance" in m2m_connectivity_kwargs + and "mesh_node_spacing" not in mesh_layout_kwargs + ): logger.warning( "Passing 'mesh_node_distance' in m2m_connectivity_kwargs is deprecated. " "Use mesh_layout_kwargs=dict(mesh_node_spacing=...) instead." @@ -68,15 +71,21 @@ def _migrate_deprecated_kwargs(mesh_layout_kwargs, m2m_connectivity_kwargs): mesh_layout_kwargs["mesh_node_spacing"] = m2m_connectivity_kwargs.pop( "mesh_node_distance" ) - if "level_refinement_factor" in m2m_connectivity_kwargs and "refinement_factor" not in mesh_layout_kwargs: + if ( + "level_refinement_factor" in m2m_connectivity_kwargs + and "refinement_factor" not in mesh_layout_kwargs + ): logger.warning( "Passing 'level_refinement_factor' in m2m_connectivity_kwargs is deprecated. " "Use mesh_layout_kwargs=dict(refinement_factor=...) instead." ) - mesh_layout_kwargs["refinement_factor"] = ( - m2m_connectivity_kwargs.pop("level_refinement_factor") + mesh_layout_kwargs["refinement_factor"] = m2m_connectivity_kwargs.pop( + "level_refinement_factor" ) - if "max_num_levels" in m2m_connectivity_kwargs and "max_num_refinement_levels" not in mesh_layout_kwargs: + if ( + "max_num_levels" in m2m_connectivity_kwargs + and "max_num_refinement_levels" not in mesh_layout_kwargs + ): logger.warning( "Passing 'max_num_levels' in m2m_connectivity_kwargs is deprecated. " "Use mesh_layout_kwargs=dict(max_num_refinement_levels=...) instead." @@ -254,9 +263,7 @@ def create_all_graph_components( "want to decrease the `mesh_node_spacing` so that the mesh nodes " "are spaced closer together?" ) - G_mesh_coords = create_single_level_2d_mesh_primitive( - xy, nx_mesh, ny_mesh - ) + G_mesh_coords = create_single_level_2d_mesh_primitive(xy, nx_mesh, ny_mesh) else: raise NotImplementedError( f"mesh_layout='{mesh_layout}' is not yet supported. " diff --git a/src/weather_model_graphs/create/mesh/coords.py b/src/weather_model_graphs/create/mesh/coords.py index 1b76709..acb4898 100644 --- a/src/weather_model_graphs/create/mesh/coords.py +++ b/src/weather_model_graphs/create/mesh/coords.py @@ -78,9 +78,7 @@ def create_single_level_2d_mesh_primitive(xy: np.ndarray, nx: int, ny: int): return g -def create_directed_mesh_graph( - G_undirected: networkx.Graph, pattern: str = "8-star" -): +def create_directed_mesh_graph(G_undirected: networkx.Graph, pattern: str = "8-star"): """ Convert an undirected mesh primitive graph with spatial adjacency edges to a directed mesh graph (nx.DiGraph) based on the specified connectivity pattern. @@ -139,9 +137,7 @@ def create_directed_mesh_graph( dg = networkx.DiGraph(g_filtered) for u, v in g_filtered.edges(): d = np.sqrt( - np.sum( - (G_undirected.nodes[u]["pos"] - G_undirected.nodes[v]["pos"]) ** 2 - ) + np.sum((G_undirected.nodes[u]["pos"] - G_undirected.nodes[v]["pos"]) ** 2) ) dg.edges[u, v]["len"] = d dg.edges[u, v]["vdiff"] = ( @@ -270,9 +266,7 @@ def create_multirange_2d_mesh_primitives( G_all_levels = [] for lev in range(mesh_levels_to_create): # 0-index mesh levels # Compute number of nodes on level separate for each direction - nodes_x, nodes_y = ( - nleaf / (interlevel_refinement_factor**lev) - ).astype(int) + nodes_x, nodes_y = (nleaf / (interlevel_refinement_factor**lev)).astype(int) g = create_single_level_2d_mesh_primitive(xy, nodes_x, nodes_y) # Add level information to nodes, edges and full graph for node in g.nodes: diff --git a/src/weather_model_graphs/create/mesh/kinds/flat.py b/src/weather_model_graphs/create/mesh/kinds/flat.py index deda1ad..c0a1b90 100644 --- a/src/weather_model_graphs/create/mesh/kinds/flat.py +++ b/src/weather_model_graphs/create/mesh/kinds/flat.py @@ -131,11 +131,7 @@ def create_flat_multiscale_from_coordinates( :, ] .reshape( - int( - num_nodes_x - * num_nodes_y - / (interlevel_refinement_factor**2) - ), + int(num_nodes_x * num_nodes_y / (interlevel_refinement_factor**2)), 2, ) ) diff --git a/src/weather_model_graphs/create/mesh/kinds/hierarchical.py b/src/weather_model_graphs/create/mesh/kinds/hierarchical.py index 19a0cd2..e870f4f 100644 --- a/src/weather_model_graphs/create/mesh/kinds/hierarchical.py +++ b/src/weather_model_graphs/create/mesh/kinds/hierarchical.py @@ -64,9 +64,7 @@ def create_hierarchical_from_coordinates( # Convert each level's coordinate graph to directed graph with chosen pattern Gs_all_levels = [ - mesh_coords.create_directed_mesh_graph( - g_coords, pattern=intra_level_pattern - ) + mesh_coords.create_directed_mesh_graph(g_coords, pattern=intra_level_pattern) for g_coords in G_coords_list ] @@ -76,8 +74,8 @@ def create_hierarchical_from_coordinates( raise ValueError( "At least two mesh levels are required for hierarchical mesh graph. " "You may need to reduce the level refinement factor " - f"or increase the max number of levels " - f"or number of grid points." + "or increase the max number of levels " + "or number of grid points." ) # Relabel nodes of each level with level index first @@ -130,9 +128,7 @@ def create_hierarchical_from_coordinates( # add edge from coarser to finer G_down.add_edge(u, v) d = np.sqrt( - np.sum( - (G_down.nodes[u]["pos"] - G_down.nodes[v]["pos"]) ** 2 - ) + np.sum((G_down.nodes[u]["pos"] - G_down.nodes[v]["pos"]) ** 2) ) G_down.edges[u, v]["len"] = d G_down.edges[u, v]["vdiff"] = ( diff --git a/tests/test_mesh_layout.py b/tests/test_mesh_layout.py index 4b0875b..56a677d 100644 --- a/tests/test_mesh_layout.py +++ b/tests/test_mesh_layout.py @@ -36,7 +36,6 @@ create_hierarchical_from_coordinates, ) - # ==================== # Step 1: Coordinate creation tests # ==================== @@ -841,7 +840,10 @@ def test_multirange_with_none_max_levels(self): """max_num_levels=None should auto-compute levels.""" xy = test_utils.create_fake_xy(N=30) G_list = create_multirange_2d_mesh_primitives( - max_num_levels=None, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 + max_num_levels=None, + xy=xy, + mesh_node_spacing=3, + interlevel_refinement_factor=3, ) assert isinstance(G_list, list) assert len(G_list) >= 1 @@ -904,8 +906,9 @@ def test_vdiff_antisymmetric(self): for u, v in G.edges(): if G.has_edge(v, u): np.testing.assert_allclose( - G.edges[u, v]["vdiff"], -G.edges[v, u]["vdiff"], - err_msg=f"vdiff not antisymmetric for ({u},{v})" + G.edges[u, v]["vdiff"], + -G.edges[v, u]["vdiff"], + err_msg=f"vdiff not antisymmetric for ({u},{v})", ) def test_edge_len_matches_vdiff_norm(self): @@ -916,8 +919,9 @@ def test_edge_len_matches_vdiff_norm(self): for u, v, d in G.edges(data=True): expected_len = np.sqrt(np.sum(d["vdiff"] ** 2)) np.testing.assert_allclose( - d["len"], expected_len, - err_msg=f"Edge ({u},{v}) len doesn't match vdiff norm" + d["len"], + expected_len, + err_msg=f"Edge ({u},{v}) len doesn't match vdiff norm", ) def test_flat_multiscale_single_level_input(self): @@ -1322,7 +1326,7 @@ def test_hierarchical_has_same_up_down_edge_count(self): down_count = sum( 1 for _, _, d in m2m.edges(data=True) if d.get("direction") == "down" ) - assert up_count == down_count, ( - f"Up edges ({up_count}) != Down edges ({down_count})" - ) + assert ( + up_count == down_count + ), f"Up edges ({up_count}) != Down edges ({down_count})" assert up_count > 0, "Should have at least some up/down edges" From 4f48ed3459fc8f43abb3ae57ca6cb7eac67e26fb Mon Sep 17 00:00:00 2001 From: prajwal Date: Sat, 21 Mar 2026 01:34:27 +0530 Subject: [PATCH 13/21] refactor: remove inline defaults and simplify m2m_connectivity_kwargs forwarding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Remove explicit m2m_connectivity_kwargs from all three archetype functions (keisler, graphcast, oskarsson_hierarchical) — let downstream from_coordinates functions define their own defaults - Remove inline pattern = kwargs.get('pattern', '8-star') in base.py flat and flat_multiscale branches — pass **m2m_connectivity_kwargs directly - Remove inline refinement_factor default of 3 — conditionally build primitives_kwargs so function signature defaults are respected - Simplify hierarchical branch: pass **m2m_connectivity_kwargs directly instead of constructing a new hier_kwargs dict Addresses review feedback from leifdenby (Mar 18, 2026) --- src/weather_model_graphs/create/archetype.py | 6 --- src/weather_model_graphs/create/base.py | 53 ++++++++++---------- 2 files changed, 27 insertions(+), 32 deletions(-) diff --git a/src/weather_model_graphs/create/archetype.py b/src/weather_model_graphs/create/archetype.py index f216335..6c94166 100644 --- a/src/weather_model_graphs/create/archetype.py +++ b/src/weather_model_graphs/create/archetype.py @@ -62,7 +62,6 @@ def create_keisler_graph( mesh_layout="rectilinear", mesh_layout_kwargs=dict(mesh_node_spacing=mesh_node_distance), m2m_connectivity="flat", - m2m_connectivity_kwargs=dict(pattern="8-star"), g2m_connectivity="within_radius", g2m_connectivity_kwargs=dict(rel_max_dist=0.51), m2g_connectivity="nearest_neighbours", @@ -139,7 +138,6 @@ def create_graphcast_graph( 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", @@ -221,10 +219,6 @@ def create_oskarsson_hierarchical_graph( max_num_refinement_levels=max_num_levels, ), m2m_connectivity="hierarchical", - m2m_connectivity_kwargs=dict( - 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", diff --git a/src/weather_model_graphs/create/base.py b/src/weather_model_graphs/create/base.py index efbd7a9..ff61e66 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/weather_model_graphs/create/base.py @@ -271,9 +271,8 @@ def create_all_graph_components( ) # --- Step 2: Connectivity creation --- - pattern = m2m_connectivity_kwargs.get("pattern", "8-star") graph_components["m2m"] = create_flat_singlescale_from_coordinates( - G_mesh_coords, pattern=pattern + G_mesh_coords, **m2m_connectivity_kwargs ) grid_connect_graph = graph_components["m2m"] @@ -283,20 +282,25 @@ def create_all_graph_components( mesh_node_spacing = mesh_layout_kwargs.get("mesh_node_spacing") if mesh_node_spacing is None: mesh_node_spacing = mesh_layout_kwargs.get("grid_spacing") - refinement_factor = mesh_layout_kwargs.get("refinement_factor", 3) - max_num_refinement_levels = mesh_layout_kwargs.get( - "max_num_refinement_levels" - ) if mesh_node_spacing is None: raise ValueError( "mesh_layout='rectilinear' with m2m_connectivity='hierarchical' " "requires 'mesh_node_spacing' in mesh_layout_kwargs." ) - G_coords_list = create_multirange_2d_mesh_primitives( - max_num_levels=max_num_refinement_levels, + primitives_kwargs = dict( xy=xy, mesh_node_spacing=mesh_node_spacing, - interlevel_refinement_factor=refinement_factor, + ) + if "refinement_factor" in mesh_layout_kwargs: + primitives_kwargs["interlevel_refinement_factor"] = ( + mesh_layout_kwargs["refinement_factor"] + ) + if "max_num_refinement_levels" in mesh_layout_kwargs: + primitives_kwargs["max_num_levels"] = mesh_layout_kwargs[ + "max_num_refinement_levels" + ] + G_coords_list = create_multirange_2d_mesh_primitives( + **primitives_kwargs, ) else: raise NotImplementedError( @@ -308,16 +312,9 @@ def create_all_graph_components( # hierarchical mesh graph have three sub-graphs: # `m2m` (mesh-to-mesh), `mesh_up` (up edge connections) and # `mesh_down` (down edge connections) - hier_kwargs = {} - intra = m2m_connectivity_kwargs.get("intra_level") - inter = m2m_connectivity_kwargs.get("inter_level") - if intra is not None: - hier_kwargs["intra_level"] = intra - if inter is not None: - hier_kwargs["inter_level"] = inter graph_components["m2m"] = create_hierarchical_from_coordinates( G_coords_list, - **hier_kwargs, + **m2m_connectivity_kwargs, ) # Only connect grid to bottom level of hierarchy grid_connect_graph = split_graph_by_edge_attribute( @@ -330,20 +327,25 @@ def create_all_graph_components( mesh_node_spacing = mesh_layout_kwargs.get("mesh_node_spacing") if mesh_node_spacing is None: mesh_node_spacing = mesh_layout_kwargs.get("grid_spacing") - refinement_factor = mesh_layout_kwargs.get("refinement_factor", 3) - max_num_refinement_levels = mesh_layout_kwargs.get( - "max_num_refinement_levels" - ) if mesh_node_spacing is None: raise ValueError( "mesh_layout='rectilinear' with m2m_connectivity='flat_multiscale' " "requires 'mesh_node_spacing' in mesh_layout_kwargs." ) - G_coords_list = create_multirange_2d_mesh_primitives( - max_num_levels=max_num_refinement_levels, + primitives_kwargs = dict( xy=xy, mesh_node_spacing=mesh_node_spacing, - interlevel_refinement_factor=refinement_factor, + ) + if "refinement_factor" in mesh_layout_kwargs: + primitives_kwargs["interlevel_refinement_factor"] = ( + mesh_layout_kwargs["refinement_factor"] + ) + if "max_num_refinement_levels" in mesh_layout_kwargs: + primitives_kwargs["max_num_levels"] = mesh_layout_kwargs[ + "max_num_refinement_levels" + ] + G_coords_list = create_multirange_2d_mesh_primitives( + **primitives_kwargs, ) else: raise NotImplementedError( @@ -352,10 +354,9 @@ def create_all_graph_components( ) # --- Step 2: Connectivity creation --- - pattern = m2m_connectivity_kwargs.get("pattern", "8-star") graph_components["m2m"] = create_flat_multiscale_from_coordinates( G_coords_list, - pattern=pattern, + **m2m_connectivity_kwargs, ) grid_connect_graph = graph_components["m2m"] else: From 0b6ff4ae3bbaea5d31f314544b0c11f4907333ba Mon Sep 17 00:00:00 2001 From: prajwal Date: Mon, 23 Mar 2026 23:08:05 +0530 Subject: [PATCH 14/21] refactor: rename mesh.kinds to mesh.connectivity, use **kwargs, update notebook - Rename wmg.create.mesh.kinds -> wmg.create.mesh.connectivity (clearer separation between coordinate creation and connectivity) - Replace explicit 'pattern' param with **kwargs in create_flat_singlescale_from_coordinates and create_flat_multiscale_from_coordinates (future-proofs for triangular and other mesh layouts) - Update imports in base.py and tests/test_mesh_layout.py - Update docs/creating_the_graph.ipynb to use new mesh_layout API - Fix black and isort formatting --- docs/creating_the_graph.ipynb | 5 ++-- src/weather_model_graphs/create/base.py | 24 +++++++-------- .../mesh/{kinds => connectivity}/__init__.py | 0 .../mesh/{kinds => connectivity}/flat.py | 29 +++++++------------ .../{kinds => connectivity}/hierarchical.py | 0 tests/test_mesh_layout.py | 14 ++++----- 6 files changed, 32 insertions(+), 40 deletions(-) rename src/weather_model_graphs/create/mesh/{kinds => connectivity}/__init__.py (100%) rename src/weather_model_graphs/create/mesh/{kinds => connectivity}/flat.py (90%) rename src/weather_model_graphs/create/mesh/{kinds => connectivity}/hierarchical.py (100%) diff --git a/docs/creating_the_graph.ipynb b/docs/creating_the_graph.ipynb index 8cbc913..4283e75 100644 --- a/docs/creating_the_graph.ipynb +++ b/docs/creating_the_graph.ipynb @@ -411,8 +411,9 @@ "graph = wmg.create.create_all_graph_components(\n", " m2m_connectivity=\"flat_multiscale\",\n", " coords=xy,\n", - " m2m_connectivity_kwargs=dict(\n", - " mesh_node_distance=2, level_refinement_factor=3, max_num_levels=None\n", + " mesh_layout=\"rectilinear\",\n", + " mesh_layout_kwargs=dict(\n", + " mesh_node_spacing=2, refinement_factor=3,\n", " ),\n", " g2m_connectivity=\"nearest_neighbour\",\n", " m2g_connectivity=\"nearest_neighbour\",\n", diff --git a/src/weather_model_graphs/create/base.py b/src/weather_model_graphs/create/base.py index ff61e66..d7f37f2 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/weather_model_graphs/create/base.py @@ -23,17 +23,17 @@ split_on_edge_attribute_existance, ) from .grid import create_grid_graph_nodes -from .mesh.coords import ( - create_multirange_2d_mesh_primitives, - create_single_level_2d_mesh_primitive, -) -from .mesh.kinds.flat import ( +from .mesh.connectivity.flat import ( create_flat_multiscale_from_coordinates, create_flat_singlescale_from_coordinates, ) -from .mesh.kinds.hierarchical import ( +from .mesh.connectivity.hierarchical import ( create_hierarchical_from_coordinates, ) +from .mesh.coords import ( + create_multirange_2d_mesh_primitives, + create_single_level_2d_mesh_primitive, +) def _migrate_deprecated_kwargs(mesh_layout_kwargs, m2m_connectivity_kwargs): @@ -292,9 +292,9 @@ def create_all_graph_components( mesh_node_spacing=mesh_node_spacing, ) if "refinement_factor" in mesh_layout_kwargs: - primitives_kwargs["interlevel_refinement_factor"] = ( - mesh_layout_kwargs["refinement_factor"] - ) + primitives_kwargs["interlevel_refinement_factor"] = mesh_layout_kwargs[ + "refinement_factor" + ] if "max_num_refinement_levels" in mesh_layout_kwargs: primitives_kwargs["max_num_levels"] = mesh_layout_kwargs[ "max_num_refinement_levels" @@ -337,9 +337,9 @@ def create_all_graph_components( mesh_node_spacing=mesh_node_spacing, ) if "refinement_factor" in mesh_layout_kwargs: - primitives_kwargs["interlevel_refinement_factor"] = ( - mesh_layout_kwargs["refinement_factor"] - ) + primitives_kwargs["interlevel_refinement_factor"] = mesh_layout_kwargs[ + "refinement_factor" + ] if "max_num_refinement_levels" in mesh_layout_kwargs: primitives_kwargs["max_num_levels"] = mesh_layout_kwargs[ "max_num_refinement_levels" diff --git a/src/weather_model_graphs/create/mesh/kinds/__init__.py b/src/weather_model_graphs/create/mesh/connectivity/__init__.py similarity index 100% rename from src/weather_model_graphs/create/mesh/kinds/__init__.py rename to src/weather_model_graphs/create/mesh/connectivity/__init__.py diff --git a/src/weather_model_graphs/create/mesh/kinds/flat.py b/src/weather_model_graphs/create/mesh/connectivity/flat.py similarity index 90% rename from src/weather_model_graphs/create/mesh/kinds/flat.py rename to src/weather_model_graphs/create/mesh/connectivity/flat.py index c0a1b90..5c192fc 100644 --- a/src/weather_model_graphs/create/mesh/kinds/flat.py +++ b/src/weather_model_graphs/create/mesh/connectivity/flat.py @@ -46,7 +46,7 @@ def _check_required_graph_attributes(G: networkx.Graph, context: str): def create_flat_multiscale_from_coordinates( G_coords_list: List[networkx.Graph], - pattern: str = "8-star", + **kwargs, ): """ Create flat multiscale mesh graph from a list of coordinate graphs. @@ -58,10 +58,6 @@ def create_flat_multiscale_from_coordinates( In a flat multiscale graph, coarser levels are merged into the finer level by coincident node positions (no separate inter-level connectivity needed). - The ``pattern`` argument defines the spatial neighbourhood connectivity: - - ``"4-star"``: only cardinal directions (horizontal and vertical neighbours) - - ``"8-star"``: cardinal directions plus diagonals (all 8 surrounding neighbours) - Parameters ---------- G_coords_list : list of networkx.Graph @@ -71,9 +67,9 @@ def create_flat_multiscale_from_coordinates( - Edge attributes: ``"adjacency_type"`` (str, ``"cardinal"`` or ``"diagonal"``) - Graph attribute: ``"interlevel_refinement_factor"`` (int) Created by ``create_multirange_2d_mesh_primitives``. - pattern : str - Connectivity pattern for intra-level edges: ``"4-star"`` or ``"8-star"`` - (default: ``"8-star"``) + **kwargs + Additional keyword arguments passed to ``create_directed_mesh_graph`` + (e.g. ``pattern="8-star"``). Returns ------- @@ -107,7 +103,7 @@ def create_flat_multiscale_from_coordinates( # Convert each level's coordinate graph to directed graph with chosen pattern G_all_levels = [ - mesh_coords.create_directed_mesh_graph(g_coords, pattern=pattern) + mesh_coords.create_directed_mesh_graph(g_coords, **kwargs) for g_coords in G_coords_list ] @@ -155,9 +151,7 @@ def create_flat_multiscale_from_coordinates( return G_tot -def create_flat_singlescale_from_coordinates( - G_coords: networkx.Graph, pattern: str = "8-star" -): +def create_flat_singlescale_from_coordinates(G_coords: networkx.Graph, **kwargs): """ Create a flat single-scale directed mesh graph from a mesh primitive graph. @@ -165,10 +159,6 @@ def create_flat_singlescale_from_coordinates( It converts an undirected mesh primitive graph to a directed mesh graph using the specified connectivity pattern. - The ``pattern`` argument defines the spatial neighbourhood connectivity: - - ``"4-star"``: only cardinal directions (horizontal and vertical neighbours) - - ``"8-star"``: cardinal directions plus diagonals (all 8 surrounding neighbours) - Parameters ---------- G_coords : networkx.Graph @@ -176,8 +166,9 @@ def create_flat_singlescale_from_coordinates( - Node attributes: ``"pos"`` (np.ndarray of shape [2,]), ``"type"`` (str) - Edge attributes: ``"adjacency_type"`` (str, ``"cardinal"`` or ``"diagonal"``) Created by ``create_single_level_2d_mesh_primitive``. - pattern : str - Connectivity pattern: ``"4-star"`` or ``"8-star"`` (default: ``"8-star"``) + **kwargs + Additional keyword arguments passed to ``create_directed_mesh_graph`` + (e.g. ``pattern="8-star"``). Returns ------- @@ -187,7 +178,7 @@ def create_flat_singlescale_from_coordinates( _check_required_graph_attributes( G_coords, "create_flat_singlescale_from_coordinates" ) - return mesh_coords.create_directed_mesh_graph(G_coords, pattern=pattern) + return mesh_coords.create_directed_mesh_graph(G_coords, **kwargs) def create_flat_multiscale_mesh_graph( diff --git a/src/weather_model_graphs/create/mesh/kinds/hierarchical.py b/src/weather_model_graphs/create/mesh/connectivity/hierarchical.py similarity index 100% rename from src/weather_model_graphs/create/mesh/kinds/hierarchical.py rename to src/weather_model_graphs/create/mesh/connectivity/hierarchical.py diff --git a/tests/test_mesh_layout.py b/tests/test_mesh_layout.py index 56a677d..845d31a 100644 --- a/tests/test_mesh_layout.py +++ b/tests/test_mesh_layout.py @@ -23,18 +23,18 @@ import tests.utils as test_utils import weather_model_graphs as wmg -from weather_model_graphs.create.mesh.coords import ( - create_directed_mesh_graph, - create_multirange_2d_mesh_primitives, - create_single_level_2d_mesh_primitive, -) -from weather_model_graphs.create.mesh.kinds.flat import ( +from weather_model_graphs.create.mesh.connectivity.flat import ( create_flat_multiscale_from_coordinates, create_flat_singlescale_from_coordinates, ) -from weather_model_graphs.create.mesh.kinds.hierarchical import ( +from weather_model_graphs.create.mesh.connectivity.hierarchical import ( create_hierarchical_from_coordinates, ) +from weather_model_graphs.create.mesh.coords import ( + create_directed_mesh_graph, + create_multirange_2d_mesh_primitives, + create_single_level_2d_mesh_primitive, +) # ==================== # Step 1: Coordinate creation tests From 51ea1891d233c32f18e81c14e6efa8ce7a07260a Mon Sep 17 00:00:00 2001 From: prajwal Date: Tue, 24 Mar 2026 01:04:07 +0530 Subject: [PATCH 15/21] fix: pre-commit linting + notebook max_num_refinement_levels - Fix isort import ordering in base.py - Fix black-jupyter formatting in creating_the_graph.ipynb - Add missing max_num_refinement_levels param to notebook cell - All pre-commit hooks pass: trailing-whitespace, end-of-file-fixer, isort, black, black-jupyter, flake8 --- docs/creating_the_graph.ipynb | 4 +++- src/weather_model_graphs/create/base.py | 4 +--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/creating_the_graph.ipynb b/docs/creating_the_graph.ipynb index 4283e75..76bcbcd 100644 --- a/docs/creating_the_graph.ipynb +++ b/docs/creating_the_graph.ipynb @@ -413,7 +413,9 @@ " coords=xy,\n", " mesh_layout=\"rectilinear\",\n", " mesh_layout_kwargs=dict(\n", - " mesh_node_spacing=2, refinement_factor=3,\n", + " mesh_node_spacing=2,\n", + " refinement_factor=3,\n", + " max_num_refinement_levels=3,\n", " ),\n", " g2m_connectivity=\"nearest_neighbour\",\n", " m2g_connectivity=\"nearest_neighbour\",\n", diff --git a/src/weather_model_graphs/create/base.py b/src/weather_model_graphs/create/base.py index d7f37f2..f907e23 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/weather_model_graphs/create/base.py @@ -27,9 +27,7 @@ create_flat_multiscale_from_coordinates, create_flat_singlescale_from_coordinates, ) -from .mesh.connectivity.hierarchical import ( - create_hierarchical_from_coordinates, -) +from .mesh.connectivity.hierarchical import create_hierarchical_from_coordinates from .mesh.coords import ( create_multirange_2d_mesh_primitives, create_single_level_2d_mesh_primitive, From edbce82b8b1ce5cdd943aafb0850913a5011468a Mon Sep 17 00:00:00 2001 From: prajwal Date: Wed, 25 Mar 2026 22:26:02 +0530 Subject: [PATCH 16/21] refactor: separate coordinate creation (Step 1) from connectivity creation (Step 2) in base.py Per leifdenby's review: reorder if-statements so that: - Step 1 (coordinate creation) is written once, branching on mesh_layout - Step 2 (connectivity creation) is written once, branching on m2m_connectivity This eliminates ~80 lines of duplicated coordinate-creation code that previously existed inside each of the flat/hierarchical/flat_multiscale branches. The G_mesh_coords variable holds either: - a single nx.Graph (for flat singlescale) - a List[nx.Graph] (for hierarchical and flat_multiscale) Also adds early validation of m2m_connectivity before Step 1 to give a clear NotImplementedError immediately for unsupported values. All 193 tests pass, pre-commit hooks pass. --- src/weather_model_graphs/create/base.py | 140 +++++++++--------------- 1 file changed, 53 insertions(+), 87 deletions(-) diff --git a/src/weather_model_graphs/create/base.py b/src/weather_model_graphs/create/base.py index f907e23..27be089 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/weather_model_graphs/create/base.py @@ -8,7 +8,7 @@ function uses `connect_nodes_across_graphs` to connect nodes across the component graphs. """ -from typing import Iterable +from typing import Iterable, List, Union import networkx import networkx as nx @@ -236,20 +236,36 @@ def create_all_graph_components( xy_tuple = coord_transformer.transform(xx=coords[:, 0], yy=coords[:, 1]) xy = np.stack(xy_tuple, axis=1) - if m2m_connectivity == "flat": - # --- Step 1: Coordinate creation based on mesh_layout --- - if mesh_layout == "rectilinear": - mesh_node_spacing = mesh_layout_kwargs.get("mesh_node_spacing") - # Backward compat: also check for old name "grid_spacing" - if mesh_node_spacing is None: - mesh_node_spacing = mesh_layout_kwargs.get("grid_spacing") - if mesh_node_spacing is None: - raise ValueError( - "mesh_layout='rectilinear' requires 'mesh_node_spacing' in " - "mesh_layout_kwargs (or 'mesh_node_distance' in " - "m2m_connectivity_kwargs for backward compatibility)." - ) - # Compute number of mesh nodes from mesh_node_spacing + # Validate m2m_connectivity early so that we raise a clear NotImplementedError + # before any coordinate creation is attempted + _supported_m2m_connectivity = {"flat", "hierarchical", "flat_multiscale"} + if m2m_connectivity not in _supported_m2m_connectivity: + raise NotImplementedError( + f"Kind {m2m_connectivity} not implemented. " + f"Supported: {sorted(_supported_m2m_connectivity)}" + ) + + # ----------------------------------------------------------------------- + # Step 1: Coordinate creation — produces the mesh primitive graph(s) + # Result type depends on m2m_connectivity: + # - flat: G_mesh_coords: nx.Graph + # - hierarchical/flat_multiscale: G_mesh_coords: List[nx.Graph] + # ----------------------------------------------------------------------- + G_mesh_coords: Union[networkx.Graph, List[networkx.Graph]] + + if mesh_layout == "rectilinear": + mesh_node_spacing = mesh_layout_kwargs.get( + "mesh_node_spacing" + ) or mesh_layout_kwargs.get("grid_spacing") + if mesh_node_spacing is None: + raise ValueError( + "mesh_layout='rectilinear' requires 'mesh_node_spacing' in " + "mesh_layout_kwargs (or 'mesh_node_distance' in " + "m2m_connectivity_kwargs for backward compatibility)." + ) + + if m2m_connectivity == "flat": + # Single-level mesh: compute nx/ny from spacing range_x, range_y = np.ptp(xy, axis=0) nx_mesh = int(range_x / mesh_node_spacing) ny_mesh = int(range_y / mesh_node_spacing) @@ -263,32 +279,8 @@ def create_all_graph_components( ) G_mesh_coords = create_single_level_2d_mesh_primitive(xy, nx_mesh, ny_mesh) else: - raise NotImplementedError( - f"mesh_layout='{mesh_layout}' is not yet supported. " - "Currently only 'rectilinear' is implemented." - ) - - # --- Step 2: Connectivity creation --- - graph_components["m2m"] = create_flat_singlescale_from_coordinates( - G_mesh_coords, **m2m_connectivity_kwargs - ) - grid_connect_graph = graph_components["m2m"] - - elif m2m_connectivity == "hierarchical": - # --- Step 1: Coordinate creation based on mesh_layout --- - if mesh_layout == "rectilinear": - mesh_node_spacing = mesh_layout_kwargs.get("mesh_node_spacing") - if mesh_node_spacing is None: - mesh_node_spacing = mesh_layout_kwargs.get("grid_spacing") - if mesh_node_spacing is None: - raise ValueError( - "mesh_layout='rectilinear' with m2m_connectivity='hierarchical' " - "requires 'mesh_node_spacing' in mesh_layout_kwargs." - ) - primitives_kwargs = dict( - xy=xy, - mesh_node_spacing=mesh_node_spacing, - ) + # Multi-level mesh: build kwargs for create_multirange_2d_mesh_primitives + primitives_kwargs = dict(xy=xy, mesh_node_spacing=mesh_node_spacing) if "refinement_factor" in mesh_layout_kwargs: primitives_kwargs["interlevel_refinement_factor"] = mesh_layout_kwargs[ "refinement_factor" @@ -297,22 +289,28 @@ def create_all_graph_components( primitives_kwargs["max_num_levels"] = mesh_layout_kwargs[ "max_num_refinement_levels" ] - G_coords_list = create_multirange_2d_mesh_primitives( - **primitives_kwargs, - ) - else: - raise NotImplementedError( - f"mesh_layout='{mesh_layout}' is not yet supported. " - "Currently only 'rectilinear' is implemented." - ) + G_mesh_coords = create_multirange_2d_mesh_primitives(**primitives_kwargs) + else: + raise NotImplementedError( + f"mesh_layout='{mesh_layout}' is not yet supported. " + "Currently only 'rectilinear' is implemented." + ) + + # ----------------------------------------------------------------------- + # Step 2: Connectivity creation — converts mesh primitives to directed graph + # ----------------------------------------------------------------------- + if m2m_connectivity == "flat": + graph_components["m2m"] = create_flat_singlescale_from_coordinates( + G_mesh_coords, **m2m_connectivity_kwargs + ) + grid_connect_graph = graph_components["m2m"] - # --- Step 2: Connectivity creation --- - # hierarchical mesh graph have three sub-graphs: + elif m2m_connectivity == "hierarchical": + # hierarchical mesh graph has three sub-graphs: # `m2m` (mesh-to-mesh), `mesh_up` (up edge connections) and # `mesh_down` (down edge connections) graph_components["m2m"] = create_hierarchical_from_coordinates( - G_coords_list, - **m2m_connectivity_kwargs, + G_mesh_coords, **m2m_connectivity_kwargs ) # Only connect grid to bottom level of hierarchy grid_connect_graph = split_graph_by_edge_attribute( @@ -320,43 +318,11 @@ def create_all_graph_components( )[0] elif m2m_connectivity == "flat_multiscale": - # --- Step 1: Coordinate creation based on mesh_layout --- - if mesh_layout == "rectilinear": - mesh_node_spacing = mesh_layout_kwargs.get("mesh_node_spacing") - if mesh_node_spacing is None: - mesh_node_spacing = mesh_layout_kwargs.get("grid_spacing") - if mesh_node_spacing is None: - raise ValueError( - "mesh_layout='rectilinear' with m2m_connectivity='flat_multiscale' " - "requires 'mesh_node_spacing' in mesh_layout_kwargs." - ) - primitives_kwargs = dict( - xy=xy, - mesh_node_spacing=mesh_node_spacing, - ) - if "refinement_factor" in mesh_layout_kwargs: - primitives_kwargs["interlevel_refinement_factor"] = mesh_layout_kwargs[ - "refinement_factor" - ] - if "max_num_refinement_levels" in mesh_layout_kwargs: - primitives_kwargs["max_num_levels"] = mesh_layout_kwargs[ - "max_num_refinement_levels" - ] - G_coords_list = create_multirange_2d_mesh_primitives( - **primitives_kwargs, - ) - else: - raise NotImplementedError( - f"mesh_layout='{mesh_layout}' is not yet supported. " - "Currently only 'rectilinear' is implemented." - ) - - # --- Step 2: Connectivity creation --- graph_components["m2m"] = create_flat_multiscale_from_coordinates( - G_coords_list, - **m2m_connectivity_kwargs, + G_mesh_coords, **m2m_connectivity_kwargs ) grid_connect_graph = graph_components["m2m"] + else: raise NotImplementedError(f"Kind {m2m_connectivity} not implemented") From eb4533f91f0164deb346ec14aba93b9cad2fa815 Mon Sep 17 00:00:00 2001 From: prajwal Date: Tue, 14 Apr 2026 01:01:18 +0530 Subject: [PATCH 17/21] refactor: move nx_mesh/ny_mesh computation into create_single_level_2d_mesh_primitive Move the definition and validation of nx_mesh and ny_mesh from base.py into create_single_level_2d_mesh_primitive() in coords.py, since these variables are related to building the coordinate primitive. The function now accepts an optional mesh_node_spacing keyword argument. When provided, nx and ny are computed internally from the coordinate extent. Existing callers passing nx/ny directly continue to work. --- src/weather_model_graphs/create/base.py | 17 ++------ .../create/mesh/coords.py | 43 ++++++++++++++++--- 2 files changed, 42 insertions(+), 18 deletions(-) diff --git a/src/weather_model_graphs/create/base.py b/src/weather_model_graphs/create/base.py index 27be089..f768968 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/weather_model_graphs/create/base.py @@ -265,19 +265,10 @@ def create_all_graph_components( ) if m2m_connectivity == "flat": - # Single-level mesh: compute nx/ny from spacing - range_x, range_y = np.ptp(xy, axis=0) - nx_mesh = int(range_x / mesh_node_spacing) - ny_mesh = int(range_y / mesh_node_spacing) - if nx_mesh == 0 or ny_mesh == 0: - raise ValueError( - "The given `mesh_node_spacing` is too large for the provided " - f"coordinates. Got mesh_node_spacing={mesh_node_spacing}, but the " - f"x-range is {range_x} and y-range is {range_y}. Maybe you " - "want to decrease the `mesh_node_spacing` so that the mesh nodes " - "are spaced closer together?" - ) - G_mesh_coords = create_single_level_2d_mesh_primitive(xy, nx_mesh, ny_mesh) + # Single-level mesh + G_mesh_coords = create_single_level_2d_mesh_primitive( + xy, mesh_node_spacing=mesh_node_spacing + ) else: # Multi-level mesh: build kwargs for create_multirange_2d_mesh_primitives primitives_kwargs = dict(xy=xy, mesh_node_spacing=mesh_node_spacing) diff --git a/src/weather_model_graphs/create/mesh/coords.py b/src/weather_model_graphs/create/mesh/coords.py index acb4898..3755fdf 100644 --- a/src/weather_model_graphs/create/mesh/coords.py +++ b/src/weather_model_graphs/create/mesh/coords.py @@ -3,7 +3,13 @@ from loguru import logger -def create_single_level_2d_mesh_primitive(xy: np.ndarray, nx: int, ny: int): +def create_single_level_2d_mesh_primitive( + xy: np.ndarray, + nx: int = None, + ny: int = None, + *, + mesh_node_spacing: float = None, +): """ Create an undirected mesh primitive graph (nx.Graph) with node positions and spatial adjacency edges, representing the coordinate creation step. @@ -24,15 +30,25 @@ def create_single_level_2d_mesh_primitive(xy: np.ndarray, nx: int, ny: int): 1. Coordinate creation (this function) -> nx.Graph with spatial adjacency 2. Connectivity creation (create_directed_mesh_graph) -> nx.DiGraph + Either provide ``nx`` and ``ny`` directly, or provide + ``mesh_node_spacing`` to have them computed automatically from the + coordinate extent of ``xy``. + Parameters ---------- xy : np.ndarray Grid point coordinates, shaped [N_grid_points, 2], with first column representing x coordinates and second column y coordinates. - nx : int - Number of nodes in x direction - ny : int - Number of nodes in y direction + nx : int, optional + Number of nodes in x direction. If not given, computed from + ``mesh_node_spacing``. + ny : int, optional + Number of nodes in y direction. If not given, computed from + ``mesh_node_spacing``. + mesh_node_spacing : float, optional + Distance between mesh nodes (in coordinate units). When provided, + ``nx`` and ``ny`` are computed as + ``int(range / mesh_node_spacing)`` and validated to be > 0. Returns ------- @@ -40,6 +56,23 @@ def create_single_level_2d_mesh_primitive(xy: np.ndarray, nx: int, ny: int): Undirected mesh primitive graph with node positions and annotated spatial adjacency edges. """ + if mesh_node_spacing is not None: + range_x, range_y = np.ptp(xy, axis=0) + nx = int(range_x / mesh_node_spacing) + ny = int(range_y / mesh_node_spacing) + if nx == 0 or ny == 0: + raise ValueError( + "The given `mesh_node_spacing` is too large for the provided " + f"coordinates. Got mesh_node_spacing={mesh_node_spacing}, but the " + f"x-range is {range_x} and y-range is {range_y}. Maybe you " + "want to decrease the `mesh_node_spacing` so that the mesh nodes " + "are spaced closer together?" + ) + elif nx is None or ny is None: + raise ValueError( + "Either provide both `nx` and `ny`, or provide " + "`mesh_node_spacing` to compute them automatically." + ) xm, xM = np.amin(xy[:, 0]), np.amax(xy[:, 0]) ym, yM = np.amin(xy[:, 1]), np.amax(xy[:, 1]) From 5335de689cf305760ccc1aef4925ff95ba9bda61 Mon Sep 17 00:00:00 2001 From: prajwal Date: Tue, 19 May 2026 21:30:12 +0530 Subject: [PATCH 18/21] refactor: move create_directed_mesh_graph to connectivity/directed and add return type annotations - Move create_directed_mesh_graph() from create/mesh/coords.py into a new create/mesh/connectivity/directed.py module, since it returns a networkx.DiGraph (connectivity) not a networkx.Graph (coordinate layout) - Keep backward-compatible re-export from coords.py (noqa: F401) and from connectivity/__init__.py so existing call sites are unaffected - Update flat.py and hierarchical.py to import directly from .directed instead of going through mesh_coords - Add return type annotations to all public functions across: create/mesh/coords.py, create/mesh/connectivity/directed.py, create/mesh/connectivity/flat.py, create/mesh/connectivity/hierarchical.py, create/base.py, create/archetype.py, create/grid/grid.py, networkx_utils.py - Add 113 tests in test_mesh_layout.py including 5 new test classes: TestDirectedMeshGraphModuleLocation, TestGraphAttributeValidation, TestDirectedMeshGraphAdditionalEdgeCases, TestReturnTypeAnnotations --- src/weather_model_graphs/create/archetype.py | 9 +- src/weather_model_graphs/create/base.py | 10 +- src/weather_model_graphs/create/grid/grid.py | 2 +- .../create/mesh/connectivity/__init__.py | 3 + .../create/mesh/connectivity/directed.py | 81 ++++ .../create/mesh/connectivity/flat.py | 22 +- .../create/mesh/connectivity/hierarchical.py | 7 +- .../create/mesh/coords.py | 89 +---- src/weather_model_graphs/networkx_utils.py | 16 +- tests/test_mesh_layout.py | 366 ++++++++++++++++++ 10 files changed, 496 insertions(+), 109 deletions(-) create mode 100644 src/weather_model_graphs/create/mesh/connectivity/directed.py diff --git a/src/weather_model_graphs/create/archetype.py b/src/weather_model_graphs/create/archetype.py index 6c94166..f5024e2 100644 --- a/src/weather_model_graphs/create/archetype.py +++ b/src/weather_model_graphs/create/archetype.py @@ -1,5 +1,6 @@ -from typing import Iterable +from typing import Dict, Iterable, Union +import networkx import pyproj from .base import create_all_graph_components @@ -12,7 +13,7 @@ def create_keisler_graph( graph_crs: pyproj.crs.CRS | None = None, decode_mask: Iterable[bool] | None = None, return_components: bool = False, -): +) -> Union[networkx.DiGraph, Dict[str, networkx.DiGraph]]: """ Create a flat LAM graph from Oskarsson et al (2023, https://arxiv.org/abs/2309.17370) This graph setup is inspired by the global graph used by Keisler (2022, https://arxiv.org/abs/2202.07575). @@ -80,7 +81,7 @@ def create_graphcast_graph( graph_crs: pyproj.crs.CRS | None = None, decode_mask: Iterable[bool] | None = None, return_components: bool = False, -): +) -> Union[networkx.DiGraph, Dict[str, networkx.DiGraph]]: """ Create a multiscale LAM graph from Oskarsson et al (2023, https://arxiv.org/abs/2309.17370) This graph setup is inspired by the global GraphCast graph used by Lam et al (2023, https://arxiv.org/abs/2212.12794) @@ -156,7 +157,7 @@ def create_oskarsson_hierarchical_graph( graph_crs: pyproj.crs.CRS | None = None, decode_mask: Iterable[bool] | None = None, return_components: bool = False, -): +) -> Union[networkx.DiGraph, Dict[str, networkx.DiGraph]]: """ Create a LAM graph following Oskarsson et al (2023, https://arxiv.org/abs/2309.17370) hierarchical architecture. diff --git a/src/weather_model_graphs/create/base.py b/src/weather_model_graphs/create/base.py index f768968..d157c3b 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/weather_model_graphs/create/base.py @@ -8,7 +8,7 @@ function uses `connect_nodes_across_graphs` to connect nodes across the component graphs. """ -from typing import Iterable, List, Union +from typing import Dict, Iterable, List, Tuple, Union import networkx import networkx as nx @@ -34,7 +34,9 @@ ) -def _migrate_deprecated_kwargs(mesh_layout_kwargs, m2m_connectivity_kwargs): +def _migrate_deprecated_kwargs( + mesh_layout_kwargs, m2m_connectivity_kwargs +) -> Tuple[dict, dict]: """Migrate old-style kwargs to the new mesh_layout_kwargs structure. In the old API, ``mesh_node_distance``, ``level_refinement_factor``, and @@ -108,7 +110,7 @@ def create_all_graph_components( graph_crs: pyproj.crs.CRS | None = None, decode_mask: Iterable[bool] | None = None, return_components: bool = False, -): +) -> Union[networkx.DiGraph, Dict[str, networkx.DiGraph]]: """ Create all graph components used in creating the message-passing graph, grid-to-mesh (g2m), mesh-to-mesh (m2m) and mesh-to-grid (m2g), @@ -382,7 +384,7 @@ def connect_nodes_across_graphs( max_dist=None, rel_max_dist=None, max_num_neighbours=None, -): +) -> networkx.DiGraph: """ Create a new graph containing the nodes in `G_source` and `G_target` and add directed edges from nodes in `G_source` to nodes in `G_target` based on the diff --git a/src/weather_model_graphs/create/grid/grid.py b/src/weather_model_graphs/create/grid/grid.py index fc5c055..bd84a24 100644 --- a/src/weather_model_graphs/create/grid/grid.py +++ b/src/weather_model_graphs/create/grid/grid.py @@ -3,7 +3,7 @@ from ...networkx_utils import prepend_node_index -def create_grid_graph_nodes(xy, level_id=-1): +def create_grid_graph_nodes(xy, level_id=-1) -> networkx.Graph: """ Create a networkx.Graph comprising only nodes for each (x,y)-point in the `xy` coordinate array (the attribute `pos` giving the (x,y)-coordinate value) and with diff --git a/src/weather_model_graphs/create/mesh/connectivity/__init__.py b/src/weather_model_graphs/create/mesh/connectivity/__init__.py index e69de29..fee08e2 100644 --- a/src/weather_model_graphs/create/mesh/connectivity/__init__.py +++ b/src/weather_model_graphs/create/mesh/connectivity/__init__.py @@ -0,0 +1,3 @@ +from .directed import create_directed_mesh_graph + +__all__ = ["create_directed_mesh_graph"] diff --git a/src/weather_model_graphs/create/mesh/connectivity/directed.py b/src/weather_model_graphs/create/mesh/connectivity/directed.py new file mode 100644 index 0000000..9a3914c --- /dev/null +++ b/src/weather_model_graphs/create/mesh/connectivity/directed.py @@ -0,0 +1,81 @@ +import networkx +import numpy as np + + +def create_directed_mesh_graph( + G_undirected: networkx.Graph, pattern: str = "8-star" +) -> networkx.DiGraph: + """ + Convert an undirected mesh primitive graph with spatial adjacency edges to a + directed mesh graph (nx.DiGraph) based on the specified connectivity pattern. + + This is the second step in the two-step mesh creation process: + 1. Coordinate creation (create_single_level_2d_mesh_primitive) -> nx.Graph + 2. Connectivity creation (this function) -> nx.DiGraph + + The ``pattern`` argument defines the spatial neighbourhood connectivity: + - ``"4-star"``: only cardinal directions (horizontal and vertical neighbours) + - ``"8-star"``: cardinal directions plus diagonals (all 8 surrounding neighbours) + + Parameters + ---------- + G_undirected : networkx.Graph + Undirected mesh primitive graph. Expected node attributes: + - ``"pos"``: np.ndarray of shape [2,], spatial coordinates. + Expected edge attributes: + - ``"adjacency_type"``: str, either ``"cardinal"`` or ``"diagonal"``. + Additional edge attributes (e.g. ``"level"``) are preserved in the + output directed graph. + pattern : str + Connectivity pattern. Options: + - ``"4-star"``: only cardinal edges (horizontal/vertical neighbours) + - ``"8-star"``: all edges (cardinal + diagonal neighbours) + + Returns + ------- + networkx.DiGraph + Directed graph with bidirectional edges, each having ``"len"`` and + ``"vdiff"`` attributes. All original edge attributes from the + primitive graph are preserved. + """ + if pattern == "4-star": + # Filter to only cardinal edges, preserving edge data + edges_to_use = [ + (u, v, d) + for u, v, d in G_undirected.edges(data=True) + if d.get("adjacency_type") == "cardinal" + ] + elif pattern == "8-star": + # Use all edges with their data + edges_to_use = list(G_undirected.edges(data=True)) + else: + raise ValueError( + f"Unknown connectivity pattern: '{pattern}'. " + "Choose '4-star' or '8-star'." + ) + + # Create filtered undirected graph with only selected edges (preserving attrs) + g_filtered = networkx.Graph() + g_filtered.add_nodes_from(G_undirected.nodes(data=True)) + g_filtered.add_edges_from(edges_to_use) + + # Convert to directed graph (creates edges in both directions) + dg = networkx.DiGraph(g_filtered) + for u, v in g_filtered.edges(): + d = np.sqrt( + np.sum((G_undirected.nodes[u]["pos"] - G_undirected.nodes[v]["pos"]) ** 2) + ) + dg.edges[u, v]["len"] = d + dg.edges[u, v]["vdiff"] = ( + G_undirected.nodes[u]["pos"] - G_undirected.nodes[v]["pos"] + ) + # Ensure reverse edge exists and has attributes + dg.edges[v, u]["len"] = d + dg.edges[v, u]["vdiff"] = ( + G_undirected.nodes[v]["pos"] - G_undirected.nodes[u]["pos"] + ) + + # Preserve graph-level attributes (dx, dy, level, etc.) + dg.graph.update(G_undirected.graph) + + return dg diff --git a/src/weather_model_graphs/create/mesh/connectivity/flat.py b/src/weather_model_graphs/create/mesh/connectivity/flat.py index 5c192fc..021aea9 100644 --- a/src/weather_model_graphs/create/mesh/connectivity/flat.py +++ b/src/weather_model_graphs/create/mesh/connectivity/flat.py @@ -5,9 +5,10 @@ from ....networkx_utils import prepend_node_index from .. import coords as mesh_coords +from .directed import create_directed_mesh_graph -def _check_required_graph_attributes(G: networkx.Graph, context: str): +def _check_required_graph_attributes(G: networkx.Graph, context: str) -> None: """Check that a coordinate graph has the required node and edge attributes. Parameters @@ -47,7 +48,7 @@ def _check_required_graph_attributes(G: networkx.Graph, context: str): def create_flat_multiscale_from_coordinates( G_coords_list: List[networkx.Graph], **kwargs, -): +) -> networkx.DiGraph: """ Create flat multiscale mesh graph from a list of coordinate graphs. @@ -103,8 +104,7 @@ def create_flat_multiscale_from_coordinates( # Convert each level's coordinate graph to directed graph with chosen pattern G_all_levels = [ - mesh_coords.create_directed_mesh_graph(g_coords, **kwargs) - for g_coords in G_coords_list + create_directed_mesh_graph(g_coords, **kwargs) for g_coords in G_coords_list ] # combine all levels to one graph @@ -151,7 +151,9 @@ def create_flat_multiscale_from_coordinates( return G_tot -def create_flat_singlescale_from_coordinates(G_coords: networkx.Graph, **kwargs): +def create_flat_singlescale_from_coordinates( + G_coords: networkx.Graph, **kwargs +) -> networkx.DiGraph: """ Create a flat single-scale directed mesh graph from a mesh primitive graph. @@ -178,12 +180,12 @@ def create_flat_singlescale_from_coordinates(G_coords: networkx.Graph, **kwargs) _check_required_graph_attributes( G_coords, "create_flat_singlescale_from_coordinates" ) - return mesh_coords.create_directed_mesh_graph(G_coords, **kwargs) + return create_directed_mesh_graph(G_coords, **kwargs) def create_flat_multiscale_mesh_graph( xy, mesh_node_distance: float, level_refinement_factor: int, max_num_levels: int -): +) -> networkx.DiGraph: """ Create flat mesh graph by merging the single-level mesh graphs across all levels in `G_all_levels`. @@ -208,7 +210,7 @@ def create_flat_multiscale_mesh_graph( Maximum number of levels in the multi-scale graph Returns ------- - G_tot : networkx.Graph + G_tot : networkx.DiGraph The merged mesh graph """ G_coords_list = mesh_coords.create_multirange_2d_mesh_primitives( @@ -224,7 +226,7 @@ def create_flat_multiscale_mesh_graph( ) -def create_flat_singlescale_mesh_graph(xy, mesh_node_distance: float): +def create_flat_singlescale_mesh_graph(xy, mesh_node_distance: float) -> networkx.DiGraph: """ Create flat mesh graph of single level @@ -243,7 +245,7 @@ def create_flat_singlescale_mesh_graph(xy, mesh_node_distance: float): in coordinate system of xy Returns ------- - G_flat : networkx.Graph + G_flat : networkx.DiGraph The flat mesh graph """ # Compute number of mesh nodes in x and y dimensions diff --git a/src/weather_model_graphs/create/mesh/connectivity/hierarchical.py b/src/weather_model_graphs/create/mesh/connectivity/hierarchical.py index e870f4f..e29b19f 100644 --- a/src/weather_model_graphs/create/mesh/connectivity/hierarchical.py +++ b/src/weather_model_graphs/create/mesh/connectivity/hierarchical.py @@ -6,13 +6,14 @@ from ....networkx_utils import prepend_node_index from .. import coords as mesh_coords +from .directed import create_directed_mesh_graph def create_hierarchical_from_coordinates( G_coords_list: List[networkx.Graph], intra_level: Dict[str, object] = {"pattern": "8-star"}, inter_level: Dict[str, object] = {"pattern": "nearest", "k": 1}, -): +) -> networkx.DiGraph: """ Create a hierarchical multiscale mesh graph from a list of mesh primitive graphs. @@ -64,7 +65,7 @@ def create_hierarchical_from_coordinates( # Convert each level's coordinate graph to directed graph with chosen pattern Gs_all_levels = [ - mesh_coords.create_directed_mesh_graph(g_coords, pattern=intra_level_pattern) + create_directed_mesh_graph(g_coords, pattern=intra_level_pattern) for g_coords in G_coords_list ] @@ -168,7 +169,7 @@ def create_hierarchical_multiscale_mesh_graph( max_num_levels: int, intra_level: Optional[Dict[str, object]] = None, inter_level: Optional[Dict[str, object]] = None, -): +) -> networkx.DiGraph: """ Create a hierarchical multiscale mesh graph with nearest neighbour connections within each level (horizontally, vertically, and diagonally), and diff --git a/src/weather_model_graphs/create/mesh/coords.py b/src/weather_model_graphs/create/mesh/coords.py index 3755fdf..3168828 100644 --- a/src/weather_model_graphs/create/mesh/coords.py +++ b/src/weather_model_graphs/create/mesh/coords.py @@ -1,7 +1,11 @@ +from typing import List + import networkx import numpy as np from loguru import logger +from .connectivity.directed import create_directed_mesh_graph # noqa: F401 + def create_single_level_2d_mesh_primitive( xy: np.ndarray, @@ -9,7 +13,7 @@ def create_single_level_2d_mesh_primitive( ny: int = None, *, mesh_node_spacing: float = None, -): +) -> networkx.Graph: """ Create an undirected mesh primitive graph (nx.Graph) with node positions and spatial adjacency edges, representing the coordinate creation step. @@ -111,84 +115,7 @@ def create_single_level_2d_mesh_primitive( return g -def create_directed_mesh_graph(G_undirected: networkx.Graph, pattern: str = "8-star"): - """ - Convert an undirected mesh primitive graph with spatial adjacency edges to a - directed mesh graph (nx.DiGraph) based on the specified connectivity pattern. - - This is the second step in the two-step mesh creation process: - 1. Coordinate creation (create_single_level_2d_mesh_primitive) -> nx.Graph - 2. Connectivity creation (this function) -> nx.DiGraph - - The ``pattern`` argument defines the spatial neighbourhood connectivity: - - ``"4-star"``: only cardinal directions (horizontal and vertical neighbours) - - ``"8-star"``: cardinal directions plus diagonals (all 8 surrounding neighbours) - - Parameters - ---------- - G_undirected : networkx.Graph - Undirected mesh primitive graph. Expected node attributes: - - ``"pos"``: np.ndarray of shape [2,], spatial coordinates. - Expected edge attributes: - - ``"adjacency_type"``: str, either ``"cardinal"`` or ``"diagonal"``. - Additional edge attributes (e.g. ``"level"``) are preserved in the - output directed graph. - pattern : str - Connectivity pattern. Options: - - ``"4-star"``: only cardinal edges (horizontal/vertical neighbours) - - ``"8-star"``: all edges (cardinal + diagonal neighbours) - - Returns - ------- - networkx.DiGraph - Directed graph with bidirectional edges, each having ``"len"`` and - ``"vdiff"`` attributes. All original edge attributes from the - primitive graph are preserved. - """ - if pattern == "4-star": - # Filter to only cardinal edges, preserving edge data - edges_to_use = [ - (u, v, d) - for u, v, d in G_undirected.edges(data=True) - if d.get("adjacency_type") == "cardinal" - ] - elif pattern == "8-star": - # Use all edges with their data - edges_to_use = list(G_undirected.edges(data=True)) - else: - raise ValueError( - f"Unknown connectivity pattern: '{pattern}'. " - "Choose '4-star' or '8-star'." - ) - - # Create filtered undirected graph with only selected edges (preserving attrs) - g_filtered = networkx.Graph() - g_filtered.add_nodes_from(G_undirected.nodes(data=True)) - g_filtered.add_edges_from(edges_to_use) - - # Convert to directed graph (creates edges in both directions) - dg = networkx.DiGraph(g_filtered) - for u, v in g_filtered.edges(): - d = np.sqrt( - np.sum((G_undirected.nodes[u]["pos"] - G_undirected.nodes[v]["pos"]) ** 2) - ) - dg.edges[u, v]["len"] = d - dg.edges[u, v]["vdiff"] = ( - G_undirected.nodes[u]["pos"] - G_undirected.nodes[v]["pos"] - ) - # Ensure reverse edge exists and has attributes - dg.edges[v, u]["len"] = d - dg.edges[v, u]["vdiff"] = ( - G_undirected.nodes[v]["pos"] - G_undirected.nodes[u]["pos"] - ) - - # Preserve graph-level attributes (dx, dy, level, etc.) - dg.graph.update(G_undirected.graph) - - return dg - - -def create_single_level_2d_mesh_graph(xy: np.ndarray, nx: int, ny: int): +def create_single_level_2d_mesh_graph(xy: np.ndarray, nx: int, ny: int) -> networkx.DiGraph: """ Create directed graph with nx * ny nodes representing a 2D grid with positions spanning the range of xy coordinate values (first dimension @@ -234,7 +161,7 @@ def create_multirange_2d_mesh_primitives( xy: np.ndarray, mesh_node_spacing: float = 3, interlevel_refinement_factor: float = 3, -): +) -> List[networkx.Graph]: """ Create a list of undirected mesh primitive graphs (nx.Graph) representing different levels of mesh resolution spanning the spatial domain of the @@ -320,7 +247,7 @@ def create_multirange_2d_mesh_graphs( mesh_node_distance: float = 3, level_refinement_factor: float = 3, pattern: str = "8-star", -): +) -> List[networkx.DiGraph]: """ Create a list of 2D grid mesh graphs representing different levels of edge-length scales spanning the spatial domain of the xy coordinates. diff --git a/src/weather_model_graphs/networkx_utils.py b/src/weather_model_graphs/networkx_utils.py index df9fbf6..99ae79e 100644 --- a/src/weather_model_graphs/networkx_utils.py +++ b/src/weather_model_graphs/networkx_utils.py @@ -1,7 +1,9 @@ import networkx +from typing import Dict, Tuple -def prepend_node_index(graph, new_index): + +def prepend_node_index(graph, new_index) -> networkx.Graph: """ Prepend node index to node tuple in graph, i.e. (i, j) -> (new_index, i, j) @@ -22,7 +24,7 @@ def prepend_node_index(graph, new_index): return networkx.relabel_nodes(graph, to_mapping, copy=True) -def sort_nodes_internally(nx_graph, node_attr=None, edge_attr=None): +def sort_nodes_internally(nx_graph, node_attr=None, edge_attr=None) -> networkx.DiGraph: # For some reason the networkx .nodes() return list can not be sorted, # but this is the ordering used by pyg when converting. # This function fixes this. @@ -47,7 +49,7 @@ class MissingEdgeAttributeError(Exception): pass -def split_graph_by_edge_attribute(graph, attr): +def split_graph_by_edge_attribute(graph, attr) -> Dict[str, networkx.Graph]: """ Split a graph into subgraphs based on an edge attribute, returning a dictionary of subgraphs keyed by the edge attribute value. @@ -101,7 +103,7 @@ def split_graph_by_edge_attribute(graph, attr): return subgraphs -def sort_nodes_in_graph(graph): +def sort_nodes_in_graph(graph) -> networkx.DiGraph: """ Creates a new networkx.DiGraph that is a copy of input, but with nodes sorted according to their label value @@ -123,7 +125,7 @@ def sort_nodes_in_graph(graph): return sorted_graph -def replace_node_labels_with_unique_ids(graph): +def replace_node_labels_with_unique_ids(graph) -> networkx.Graph: """ Rename node labels with unique id. @@ -142,7 +144,9 @@ def replace_node_labels_with_unique_ids(graph): ) -def split_on_edge_attribute_existance(graph, attr): +def split_on_edge_attribute_existance( + graph, attr +) -> Tuple[networkx.Graph, networkx.Graph]: """ Split up graph based on if edges have specific attribute. diff --git a/tests/test_mesh_layout.py b/tests/test_mesh_layout.py index 845d31a..b7863f6 100644 --- a/tests/test_mesh_layout.py +++ b/tests/test_mesh_layout.py @@ -13,6 +13,7 @@ 6. Error handling for invalid inputs """ +import inspect import io import warnings @@ -23,7 +24,14 @@ import tests.utils as test_utils import weather_model_graphs as wmg +from weather_model_graphs.create.mesh.connectivity import ( + create_directed_mesh_graph as cdmg_from_connectivity_init, +) +from weather_model_graphs.create.mesh.connectivity.directed import ( + create_directed_mesh_graph as cdmg_from_directed, +) from weather_model_graphs.create.mesh.connectivity.flat import ( + _check_required_graph_attributes, create_flat_multiscale_from_coordinates, create_flat_singlescale_from_coordinates, ) @@ -1330,3 +1338,361 @@ def test_hierarchical_has_same_up_down_edge_count(self): up_count == down_count ), f"Up edges ({up_count}) != Down edges ({down_count})" assert up_count > 0, "Should have at least some up/down edges" + + +# ==================== +# Module location and re-export tests (Leif's review: move create_directed_mesh_graph) +# ==================== + + +class TestDirectedMeshGraphModuleLocation: + """Verify create_directed_mesh_graph is defined in connectivity/directed.py + and is also accessible via backward-compat re-exports. + These tests directly verify Leif's review request to move create_directed_mesh_graph + out of create.mesh.coords and into create.mesh.connectivity. + """ + + def test_canonical_source_is_connectivity_directed(self): + """create_directed_mesh_graph should be defined in connectivity/directed.py.""" + src = inspect.getsourcefile(cdmg_from_directed) + assert "connectivity" in src.replace("\\", "/"), ( + f"Expected source in connectivity/, got: {src}" + ) + assert "directed.py" in src, f"Expected source file directed.py, got: {src}" + + def test_coords_reexport_same_function(self): + """create_directed_mesh_graph re-exported from coords should be same object.""" + assert create_directed_mesh_graph is cdmg_from_directed + + def test_connectivity_init_reexport_same_function(self): + """create_directed_mesh_graph from connectivity __init__ should be same object.""" + assert cdmg_from_connectivity_init is cdmg_from_directed + + def test_function_not_defined_in_coords_source(self): + """coords.py should not define create_directed_mesh_graph - only re-export it.""" + import weather_model_graphs.create.mesh.coords as coords_module + + src_file = inspect.getsourcefile(coords_module.create_directed_mesh_graph) + # The source file must be directed.py, NOT coords.py + assert "coords.py" not in src_file, ( + "create_directed_mesh_graph should not be defined in coords.py, " + f"but source file is: {src_file}" + ) + + def test_backward_compat_callable_from_coords(self): + """create_directed_mesh_graph should still be callable via coords namespace.""" + import weather_model_graphs.create.mesh.coords as coords_module + + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=3, ny=3) + G = coords_module.create_directed_mesh_graph(G_coords, pattern="8-star") + assert isinstance(G, nx.DiGraph) + + def test_all_connectivity_functions_return_digraph(self): + """All connectivity functions must return DiGraph (Leif's annotation request).""" + import typing + + xy = test_utils.create_fake_xy(N=10) + checks = [ + (cdmg_from_directed, "create_directed_mesh_graph"), + ] + for fn, name in checks: + ret = fn.__annotations__.get("return") + assert ret is not None, f"{name} missing return type annotation" + assert issubclass(ret, nx.DiGraph), ( + f"{name} should return DiGraph, got {ret}" + ) + + def test_all_coords_functions_have_return_annotations(self): + """All functions in coords.py should have return type annotations.""" + import weather_model_graphs.create.mesh.coords as coords_module + + functions = [ + "create_single_level_2d_mesh_primitive", + "create_single_level_2d_mesh_graph", + "create_multirange_2d_mesh_primitives", + "create_multirange_2d_mesh_graphs", + ] + for name in functions: + fn = getattr(coords_module, name) + assert "return" in fn.__annotations__, ( + f"{name} in coords.py missing return type annotation" + ) + + def test_connectivity_functions_have_return_annotations(self): + """All connectivity creation functions should have return type annotations.""" + from weather_model_graphs.create.mesh.connectivity import flat, hierarchical + + checks = [ + (flat.create_flat_multiscale_from_coordinates, "create_flat_multiscale_from_coordinates"), + (flat.create_flat_singlescale_from_coordinates, "create_flat_singlescale_from_coordinates"), + (flat.create_flat_multiscale_mesh_graph, "create_flat_multiscale_mesh_graph"), + (flat.create_flat_singlescale_mesh_graph, "create_flat_singlescale_mesh_graph"), + (hierarchical.create_hierarchical_from_coordinates, "create_hierarchical_from_coordinates"), + (hierarchical.create_hierarchical_multiscale_mesh_graph, "create_hierarchical_multiscale_mesh_graph"), + ] + for fn, name in checks: + assert "return" in fn.__annotations__, ( + f"{name} missing return type annotation" + ) + + +# ==================== +# Graph attribute validation edge cases +# ==================== + + +class TestGraphAttributeValidation: + """Tests for _check_required_graph_attributes and attribute error handling.""" + + def test_missing_pos_attribute_raises(self): + """Graph with nodes missing 'pos' attribute should raise ValueError.""" + G = nx.Graph() + G.add_node((0, 0), type="mesh") # no 'pos' + G.add_node((1, 0), type="mesh") + G.add_edge((0, 0), (1, 0), adjacency_type="cardinal") + with pytest.raises(ValueError, match="'pos' attribute"): + _check_required_graph_attributes(G, "test_context") + + def test_missing_type_attribute_raises(self): + """Graph with nodes missing 'type' attribute should raise ValueError.""" + import numpy as np + + G = nx.Graph() + G.add_node((0, 0), pos=np.array([0.0, 0.0])) # no 'type' + G.add_node((1, 0), pos=np.array([1.0, 0.0])) + G.add_edge((0, 0), (1, 0), adjacency_type="cardinal") + with pytest.raises(ValueError, match="'type' attribute"): + _check_required_graph_attributes(G, "test_context") + + def test_missing_adjacency_type_on_edge_raises(self): + """Graph with edges missing 'adjacency_type' should raise ValueError.""" + import numpy as np + + G = nx.Graph() + G.add_node((0, 0), pos=np.array([0.0, 0.0]), type="mesh") + G.add_node((1, 0), pos=np.array([1.0, 0.0]), type="mesh") + G.add_edge((0, 0), (1, 0)) # no 'adjacency_type' + with pytest.raises(ValueError, match="'adjacency_type'"): + _check_required_graph_attributes(G, "test_context") + + def test_empty_graph_passes_validation(self): + """An empty graph (no nodes, no edges) should pass validation.""" + G = nx.Graph() + # Should not raise + _check_required_graph_attributes(G, "test_context") + + def test_flat_singlescale_missing_pos_raises(self): + """create_flat_singlescale_from_coordinates should raise on bad graph.""" + import numpy as np + + G = nx.Graph() + G.add_node((0, 0), type="mesh") # missing pos + with pytest.raises(ValueError, match="'pos' attribute"): + create_flat_singlescale_from_coordinates(G, pattern="8-star") + + def test_flat_multiscale_missing_interlevel_refinement_factor_raises(self): + """create_flat_multiscale_from_coordinates should raise if + interlevel_refinement_factor is not in graph attributes.""" + xy = test_utils.create_fake_xy(N=30) + G = create_single_level_2d_mesh_primitive(xy, nx=5, ny=5) + # G is a valid coord graph but lacks interlevel_refinement_factor + with pytest.raises(ValueError, match="interlevel_refinement_factor"): + create_flat_multiscale_from_coordinates([G]) + + def test_flat_multiscale_even_refinement_factor_raises(self): + """create_flat_multiscale_from_coordinates should raise if + interlevel_refinement_factor is even (e.g. 2 or 4).""" + xy = test_utils.create_fake_xy(N=30) + G_list = create_multirange_2d_mesh_primitives( + max_num_levels=2, + xy=xy, + mesh_node_spacing=3, + interlevel_refinement_factor=3, + ) + # Manually override to an even value to trigger the error + for G in G_list: + G.graph["interlevel_refinement_factor"] = 2 + with pytest.raises(ValueError, match="odd integer"): + create_flat_multiscale_from_coordinates(G_list) + + def test_flat_multiscale_non_integer_refinement_factor_raises(self): + """Non-integer interlevel_refinement_factor should raise ValueError.""" + xy = test_utils.create_fake_xy(N=30) + G_list = create_multirange_2d_mesh_primitives( + max_num_levels=2, + xy=xy, + mesh_node_spacing=3, + interlevel_refinement_factor=3, + ) + # Manually override to a non-integer + for G in G_list: + G.graph["interlevel_refinement_factor"] = 2.5 + with pytest.raises(ValueError, match="odd integer"): + create_flat_multiscale_from_coordinates(G_list) + + +# ==================== +# Additional create_directed_mesh_graph edge cases +# ==================== + + +class TestDirectedMeshGraphAdditionalEdgeCases: + """Additional edge cases for create_directed_mesh_graph from connectivity/directed.""" + + def test_graph_attributes_preserved_for_level(self): + """Graph-level 'level' attribute from multirange should be preserved.""" + xy = test_utils.create_fake_xy(N=30) + G_list = create_multirange_2d_mesh_primitives( + max_num_levels=2, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 + ) + for G_coords in G_list: + G_directed = cdmg_from_directed(G_coords, pattern="8-star") + assert "level" in G_directed.graph, "level graph attribute not preserved" + assert "dx" in G_directed.graph, "dx graph attribute not preserved" + assert "dy" in G_directed.graph, "dy graph attribute not preserved" + + def test_4star_no_diagonal_edge_attributes(self): + """4-star graph should have no edges with adjacency_type='diagonal'.""" + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=5, ny=5) + G_4star = cdmg_from_directed(G_coords, pattern="4-star") + for u, v, d in G_4star.edges(data=True): + adj = d.get("adjacency_type") + assert adj != "diagonal", ( + f"4-star graph has diagonal edge ({u},{v})" + ) + + def test_node_count_preserved(self): + """Directed graph should have same number of nodes as undirected.""" + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=5, ny=5) + G_directed = cdmg_from_directed(G_coords, pattern="8-star") + assert len(G_directed.nodes) == len(G_coords.nodes) + + def test_node_attributes_preserved(self): + """Node pos and type should be preserved in the directed graph.""" + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=4, ny=4) + G_directed = cdmg_from_directed(G_coords, pattern="8-star") + for node in G_directed.nodes: + assert "pos" in G_directed.nodes[node] + assert "type" in G_directed.nodes[node] + + def test_invalid_pattern_message_contains_valid_options(self): + """ValueError for invalid pattern should mention the valid options.""" + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=3, ny=3) + with pytest.raises(ValueError, match="4-star"): + cdmg_from_directed(G_coords, pattern="bad-pattern") + with pytest.raises(ValueError, match="8-star"): + cdmg_from_directed(G_coords, pattern="bad-pattern") + + def test_same_result_from_coords_and_directed_import(self): + """Calling via coords re-export and via connectivity.directed should give same result.""" + import weather_model_graphs.create.mesh.coords as coords_module + + xy = test_utils.create_fake_xy(N=10) + G_coords = create_single_level_2d_mesh_primitive(xy, nx=4, ny=4) + + G_via_coords = coords_module.create_directed_mesh_graph(G_coords, pattern="8-star") + G_via_directed = cdmg_from_directed(G_coords, pattern="8-star") + + assert len(G_via_coords.nodes) == len(G_via_directed.nodes) + assert len(G_via_coords.edges) == len(G_via_directed.edges) + assert set(G_via_coords.nodes) == set(G_via_directed.nodes) + assert set(G_via_coords.edges) == set(G_via_directed.edges) + + +# ==================== +# Return type annotation tests +# ==================== + + +class TestReturnTypeAnnotations: + """Verify all public functions across the package have return type annotations. + This directly addresses Leif's review: 'ensure all functions have type annotations + for the return types. That way we can quickly scan through if nx.Graph or + nx.DiGraph is returned.' + """ + + def test_create_single_level_2d_mesh_primitive_returns_graph(self): + """create_single_level_2d_mesh_primitive should annotate -> networkx.Graph.""" + import networkx + + ret = create_single_level_2d_mesh_primitive.__annotations__["return"] + assert ret is networkx.Graph or ( + hasattr(ret, "__origin__") and issubclass(ret.__origin__, networkx.Graph) + ), f"Expected networkx.Graph, got {ret}" + + def test_create_directed_mesh_graph_returns_digraph(self): + """create_directed_mesh_graph should annotate -> networkx.DiGraph.""" + import networkx + + ret = cdmg_from_directed.__annotations__["return"] + assert ret is networkx.DiGraph or issubclass(ret, networkx.DiGraph), ( + f"Expected networkx.DiGraph, got {ret}" + ) + + def test_create_single_level_2d_mesh_graph_returns_digraph(self): + """create_single_level_2d_mesh_graph (backward compat) -> networkx.DiGraph.""" + import weather_model_graphs.create.mesh.coords as coords_module + import networkx + + fn = coords_module.create_single_level_2d_mesh_graph + ret = fn.__annotations__["return"] + assert ret is networkx.DiGraph, f"Expected networkx.DiGraph, got {ret}" + + def test_create_multirange_2d_mesh_primitives_returns_list_of_graph(self): + """create_multirange_2d_mesh_primitives should annotate -> List[networkx.Graph].""" + import typing + + ret = create_multirange_2d_mesh_primitives.__annotations__["return"] + # Should be List[networkx.Graph] + assert hasattr(ret, "__args__"), f"Expected generic list type, got {ret}" + + def test_create_flat_singlescale_from_coordinates_returns_digraph(self): + """create_flat_singlescale_from_coordinates should annotate -> networkx.DiGraph.""" + import networkx + + ret = create_flat_singlescale_from_coordinates.__annotations__["return"] + assert ret is networkx.DiGraph, f"Expected networkx.DiGraph, got {ret}" + + def test_create_flat_multiscale_from_coordinates_returns_digraph(self): + """create_flat_multiscale_from_coordinates should annotate -> networkx.DiGraph.""" + import networkx + + ret = create_flat_multiscale_from_coordinates.__annotations__["return"] + assert ret is networkx.DiGraph, f"Expected networkx.DiGraph, got {ret}" + + def test_create_hierarchical_from_coordinates_returns_digraph(self): + """create_hierarchical_from_coordinates should annotate -> networkx.DiGraph.""" + import networkx + + ret = create_hierarchical_from_coordinates.__annotations__["return"] + assert ret is networkx.DiGraph, f"Expected networkx.DiGraph, got {ret}" + + def test_create_all_graph_components_has_return_annotation(self): + """create_all_graph_components should have a return type annotation.""" + from weather_model_graphs.create.base import create_all_graph_components + + assert "return" in create_all_graph_components.__annotations__, ( + "create_all_graph_components missing return type annotation" + ) + + def test_connect_nodes_across_graphs_returns_digraph(self): + """connect_nodes_across_graphs should annotate -> networkx.DiGraph.""" + import networkx + from weather_model_graphs.create.base import connect_nodes_across_graphs + + ret = connect_nodes_across_graphs.__annotations__["return"] + assert ret is networkx.DiGraph, f"Expected networkx.DiGraph, got {ret}" + + def test_create_grid_graph_nodes_returns_graph(self): + """create_grid_graph_nodes should annotate -> networkx.Graph.""" + import networkx + from weather_model_graphs.create.grid.grid import create_grid_graph_nodes + + ret = create_grid_graph_nodes.__annotations__["return"] + assert ret is networkx.Graph, f"Expected networkx.Graph, got {ret}" + From b4f7dcf033c3bd93e9e5bd9a909abf71a0ca5a18 Mon Sep 17 00:00:00 2001 From: prajwal Date: Thu, 21 May 2026 22:52:38 +0530 Subject: [PATCH 19/21] refactor: rename connectivity/directed.py to general.py and clean up tests Per Leif's review: - Rename create/mesh/connectivity/directed.py -> general.py so the module sits at .mesh.connectivity.general alongside .flat and .hierarchical. 'general' better describes the role (a general-purpose building block consumed by both flat and hierarchical) rather than the return type. - Update all imports in __init__.py, flat.py, hierarchical.py and coords.py to reference .general instead of .directed - Remove # noqa: F401 from coords.py import (it is a real internal import, not a backward-compat re-export) - Remove TestReturnTypeAnnotations test class (static type checkers handle this, not runtime tests) - Remove TestDirectedMeshGraphModuleLocation test class (was testing re-export wiring that no longer exists) --- .../create/mesh/connectivity/__init__.py | 2 +- .../create/mesh/connectivity/flat.py | 2 +- .../connectivity/{directed.py => general.py} | 0 .../create/mesh/connectivity/hierarchical.py | 2 +- .../create/mesh/coords.py | 2 +- tests/test_mesh_layout.py | 231 ++---------------- 6 files changed, 20 insertions(+), 219 deletions(-) rename src/weather_model_graphs/create/mesh/connectivity/{directed.py => general.py} (100%) diff --git a/src/weather_model_graphs/create/mesh/connectivity/__init__.py b/src/weather_model_graphs/create/mesh/connectivity/__init__.py index fee08e2..e03c44b 100644 --- a/src/weather_model_graphs/create/mesh/connectivity/__init__.py +++ b/src/weather_model_graphs/create/mesh/connectivity/__init__.py @@ -1,3 +1,3 @@ -from .directed import create_directed_mesh_graph +from .general import create_directed_mesh_graph __all__ = ["create_directed_mesh_graph"] diff --git a/src/weather_model_graphs/create/mesh/connectivity/flat.py b/src/weather_model_graphs/create/mesh/connectivity/flat.py index 021aea9..69c375e 100644 --- a/src/weather_model_graphs/create/mesh/connectivity/flat.py +++ b/src/weather_model_graphs/create/mesh/connectivity/flat.py @@ -5,7 +5,7 @@ from ....networkx_utils import prepend_node_index from .. import coords as mesh_coords -from .directed import create_directed_mesh_graph +from .general import create_directed_mesh_graph def _check_required_graph_attributes(G: networkx.Graph, context: str) -> None: diff --git a/src/weather_model_graphs/create/mesh/connectivity/directed.py b/src/weather_model_graphs/create/mesh/connectivity/general.py similarity index 100% rename from src/weather_model_graphs/create/mesh/connectivity/directed.py rename to src/weather_model_graphs/create/mesh/connectivity/general.py diff --git a/src/weather_model_graphs/create/mesh/connectivity/hierarchical.py b/src/weather_model_graphs/create/mesh/connectivity/hierarchical.py index e29b19f..ae78da1 100644 --- a/src/weather_model_graphs/create/mesh/connectivity/hierarchical.py +++ b/src/weather_model_graphs/create/mesh/connectivity/hierarchical.py @@ -6,7 +6,7 @@ from ....networkx_utils import prepend_node_index from .. import coords as mesh_coords -from .directed import create_directed_mesh_graph +from .general import create_directed_mesh_graph def create_hierarchical_from_coordinates( diff --git a/src/weather_model_graphs/create/mesh/coords.py b/src/weather_model_graphs/create/mesh/coords.py index 3168828..0fc2c44 100644 --- a/src/weather_model_graphs/create/mesh/coords.py +++ b/src/weather_model_graphs/create/mesh/coords.py @@ -4,7 +4,7 @@ import numpy as np from loguru import logger -from .connectivity.directed import create_directed_mesh_graph # noqa: F401 +from .connectivity.general import create_directed_mesh_graph def create_single_level_2d_mesh_primitive( diff --git a/tests/test_mesh_layout.py b/tests/test_mesh_layout.py index b7863f6..035bfe3 100644 --- a/tests/test_mesh_layout.py +++ b/tests/test_mesh_layout.py @@ -13,7 +13,6 @@ 6. Error handling for invalid inputs """ -import inspect import io import warnings @@ -24,22 +23,18 @@ import tests.utils as test_utils import weather_model_graphs as wmg -from weather_model_graphs.create.mesh.connectivity import ( - create_directed_mesh_graph as cdmg_from_connectivity_init, -) -from weather_model_graphs.create.mesh.connectivity.directed import ( - create_directed_mesh_graph as cdmg_from_directed, -) from weather_model_graphs.create.mesh.connectivity.flat import ( _check_required_graph_attributes, create_flat_multiscale_from_coordinates, create_flat_singlescale_from_coordinates, ) +from weather_model_graphs.create.mesh.connectivity.general import ( + create_directed_mesh_graph, +) from weather_model_graphs.create.mesh.connectivity.hierarchical import ( create_hierarchical_from_coordinates, ) from weather_model_graphs.create.mesh.coords import ( - create_directed_mesh_graph, create_multirange_2d_mesh_primitives, create_single_level_2d_mesh_primitive, ) @@ -1340,103 +1335,6 @@ def test_hierarchical_has_same_up_down_edge_count(self): assert up_count > 0, "Should have at least some up/down edges" -# ==================== -# Module location and re-export tests (Leif's review: move create_directed_mesh_graph) -# ==================== - - -class TestDirectedMeshGraphModuleLocation: - """Verify create_directed_mesh_graph is defined in connectivity/directed.py - and is also accessible via backward-compat re-exports. - These tests directly verify Leif's review request to move create_directed_mesh_graph - out of create.mesh.coords and into create.mesh.connectivity. - """ - - def test_canonical_source_is_connectivity_directed(self): - """create_directed_mesh_graph should be defined in connectivity/directed.py.""" - src = inspect.getsourcefile(cdmg_from_directed) - assert "connectivity" in src.replace("\\", "/"), ( - f"Expected source in connectivity/, got: {src}" - ) - assert "directed.py" in src, f"Expected source file directed.py, got: {src}" - - def test_coords_reexport_same_function(self): - """create_directed_mesh_graph re-exported from coords should be same object.""" - assert create_directed_mesh_graph is cdmg_from_directed - - def test_connectivity_init_reexport_same_function(self): - """create_directed_mesh_graph from connectivity __init__ should be same object.""" - assert cdmg_from_connectivity_init is cdmg_from_directed - - def test_function_not_defined_in_coords_source(self): - """coords.py should not define create_directed_mesh_graph - only re-export it.""" - import weather_model_graphs.create.mesh.coords as coords_module - - src_file = inspect.getsourcefile(coords_module.create_directed_mesh_graph) - # The source file must be directed.py, NOT coords.py - assert "coords.py" not in src_file, ( - "create_directed_mesh_graph should not be defined in coords.py, " - f"but source file is: {src_file}" - ) - - def test_backward_compat_callable_from_coords(self): - """create_directed_mesh_graph should still be callable via coords namespace.""" - import weather_model_graphs.create.mesh.coords as coords_module - - xy = test_utils.create_fake_xy(N=10) - G_coords = create_single_level_2d_mesh_primitive(xy, nx=3, ny=3) - G = coords_module.create_directed_mesh_graph(G_coords, pattern="8-star") - assert isinstance(G, nx.DiGraph) - - def test_all_connectivity_functions_return_digraph(self): - """All connectivity functions must return DiGraph (Leif's annotation request).""" - import typing - - xy = test_utils.create_fake_xy(N=10) - checks = [ - (cdmg_from_directed, "create_directed_mesh_graph"), - ] - for fn, name in checks: - ret = fn.__annotations__.get("return") - assert ret is not None, f"{name} missing return type annotation" - assert issubclass(ret, nx.DiGraph), ( - f"{name} should return DiGraph, got {ret}" - ) - - def test_all_coords_functions_have_return_annotations(self): - """All functions in coords.py should have return type annotations.""" - import weather_model_graphs.create.mesh.coords as coords_module - - functions = [ - "create_single_level_2d_mesh_primitive", - "create_single_level_2d_mesh_graph", - "create_multirange_2d_mesh_primitives", - "create_multirange_2d_mesh_graphs", - ] - for name in functions: - fn = getattr(coords_module, name) - assert "return" in fn.__annotations__, ( - f"{name} in coords.py missing return type annotation" - ) - - def test_connectivity_functions_have_return_annotations(self): - """All connectivity creation functions should have return type annotations.""" - from weather_model_graphs.create.mesh.connectivity import flat, hierarchical - - checks = [ - (flat.create_flat_multiscale_from_coordinates, "create_flat_multiscale_from_coordinates"), - (flat.create_flat_singlescale_from_coordinates, "create_flat_singlescale_from_coordinates"), - (flat.create_flat_multiscale_mesh_graph, "create_flat_multiscale_mesh_graph"), - (flat.create_flat_singlescale_mesh_graph, "create_flat_singlescale_mesh_graph"), - (hierarchical.create_hierarchical_from_coordinates, "create_hierarchical_from_coordinates"), - (hierarchical.create_hierarchical_multiscale_mesh_graph, "create_hierarchical_multiscale_mesh_graph"), - ] - for fn, name in checks: - assert "return" in fn.__annotations__, ( - f"{name} missing return type annotation" - ) - - # ==================== # Graph attribute validation edge cases # ==================== @@ -1538,7 +1436,7 @@ def test_flat_multiscale_non_integer_refinement_factor_raises(self): class TestDirectedMeshGraphAdditionalEdgeCases: - """Additional edge cases for create_directed_mesh_graph from connectivity/directed.""" + """Additional edge cases for create_directed_mesh_graph from connectivity/general.""" def test_graph_attributes_preserved_for_level(self): """Graph-level 'level' attribute from multirange should be preserved.""" @@ -1547,7 +1445,7 @@ def test_graph_attributes_preserved_for_level(self): max_num_levels=2, xy=xy, mesh_node_spacing=3, interlevel_refinement_factor=3 ) for G_coords in G_list: - G_directed = cdmg_from_directed(G_coords, pattern="8-star") + G_directed = create_directed_mesh_graph(G_coords, pattern="8-star") assert "level" in G_directed.graph, "level graph attribute not preserved" assert "dx" in G_directed.graph, "dx graph attribute not preserved" assert "dy" in G_directed.graph, "dy graph attribute not preserved" @@ -1556,7 +1454,7 @@ def test_4star_no_diagonal_edge_attributes(self): """4-star graph should have no edges with adjacency_type='diagonal'.""" xy = test_utils.create_fake_xy(N=10) G_coords = create_single_level_2d_mesh_primitive(xy, nx=5, ny=5) - G_4star = cdmg_from_directed(G_coords, pattern="4-star") + G_4star = create_directed_mesh_graph(G_coords, pattern="4-star") for u, v, d in G_4star.edges(data=True): adj = d.get("adjacency_type") assert adj != "diagonal", ( @@ -1567,14 +1465,14 @@ def test_node_count_preserved(self): """Directed graph should have same number of nodes as undirected.""" xy = test_utils.create_fake_xy(N=10) G_coords = create_single_level_2d_mesh_primitive(xy, nx=5, ny=5) - G_directed = cdmg_from_directed(G_coords, pattern="8-star") + G_directed = create_directed_mesh_graph(G_coords, pattern="8-star") assert len(G_directed.nodes) == len(G_coords.nodes) def test_node_attributes_preserved(self): """Node pos and type should be preserved in the directed graph.""" xy = test_utils.create_fake_xy(N=10) G_coords = create_single_level_2d_mesh_primitive(xy, nx=4, ny=4) - G_directed = cdmg_from_directed(G_coords, pattern="8-star") + G_directed = create_directed_mesh_graph(G_coords, pattern="8-star") for node in G_directed.nodes: assert "pos" in G_directed.nodes[node] assert "type" in G_directed.nodes[node] @@ -1584,115 +1482,18 @@ def test_invalid_pattern_message_contains_valid_options(self): xy = test_utils.create_fake_xy(N=10) G_coords = create_single_level_2d_mesh_primitive(xy, nx=3, ny=3) with pytest.raises(ValueError, match="4-star"): - cdmg_from_directed(G_coords, pattern="bad-pattern") + create_directed_mesh_graph(G_coords, pattern="bad-pattern") with pytest.raises(ValueError, match="8-star"): - cdmg_from_directed(G_coords, pattern="bad-pattern") - - def test_same_result_from_coords_and_directed_import(self): - """Calling via coords re-export and via connectivity.directed should give same result.""" - import weather_model_graphs.create.mesh.coords as coords_module + create_directed_mesh_graph(G_coords, pattern="bad-pattern") + def test_function_is_deterministic(self): + """Calling create_directed_mesh_graph twice with same input gives same result.""" xy = test_utils.create_fake_xy(N=10) G_coords = create_single_level_2d_mesh_primitive(xy, nx=4, ny=4) - G_via_coords = coords_module.create_directed_mesh_graph(G_coords, pattern="8-star") - G_via_directed = cdmg_from_directed(G_coords, pattern="8-star") - - assert len(G_via_coords.nodes) == len(G_via_directed.nodes) - assert len(G_via_coords.edges) == len(G_via_directed.edges) - assert set(G_via_coords.nodes) == set(G_via_directed.nodes) - assert set(G_via_coords.edges) == set(G_via_directed.edges) - - -# ==================== -# Return type annotation tests -# ==================== - - -class TestReturnTypeAnnotations: - """Verify all public functions across the package have return type annotations. - This directly addresses Leif's review: 'ensure all functions have type annotations - for the return types. That way we can quickly scan through if nx.Graph or - nx.DiGraph is returned.' - """ - - def test_create_single_level_2d_mesh_primitive_returns_graph(self): - """create_single_level_2d_mesh_primitive should annotate -> networkx.Graph.""" - import networkx - - ret = create_single_level_2d_mesh_primitive.__annotations__["return"] - assert ret is networkx.Graph or ( - hasattr(ret, "__origin__") and issubclass(ret.__origin__, networkx.Graph) - ), f"Expected networkx.Graph, got {ret}" - - def test_create_directed_mesh_graph_returns_digraph(self): - """create_directed_mesh_graph should annotate -> networkx.DiGraph.""" - import networkx - - ret = cdmg_from_directed.__annotations__["return"] - assert ret is networkx.DiGraph or issubclass(ret, networkx.DiGraph), ( - f"Expected networkx.DiGraph, got {ret}" - ) - - def test_create_single_level_2d_mesh_graph_returns_digraph(self): - """create_single_level_2d_mesh_graph (backward compat) -> networkx.DiGraph.""" - import weather_model_graphs.create.mesh.coords as coords_module - import networkx - - fn = coords_module.create_single_level_2d_mesh_graph - ret = fn.__annotations__["return"] - assert ret is networkx.DiGraph, f"Expected networkx.DiGraph, got {ret}" - - def test_create_multirange_2d_mesh_primitives_returns_list_of_graph(self): - """create_multirange_2d_mesh_primitives should annotate -> List[networkx.Graph].""" - import typing - - ret = create_multirange_2d_mesh_primitives.__annotations__["return"] - # Should be List[networkx.Graph] - assert hasattr(ret, "__args__"), f"Expected generic list type, got {ret}" - - def test_create_flat_singlescale_from_coordinates_returns_digraph(self): - """create_flat_singlescale_from_coordinates should annotate -> networkx.DiGraph.""" - import networkx - - ret = create_flat_singlescale_from_coordinates.__annotations__["return"] - assert ret is networkx.DiGraph, f"Expected networkx.DiGraph, got {ret}" - - def test_create_flat_multiscale_from_coordinates_returns_digraph(self): - """create_flat_multiscale_from_coordinates should annotate -> networkx.DiGraph.""" - import networkx - - ret = create_flat_multiscale_from_coordinates.__annotations__["return"] - assert ret is networkx.DiGraph, f"Expected networkx.DiGraph, got {ret}" - - def test_create_hierarchical_from_coordinates_returns_digraph(self): - """create_hierarchical_from_coordinates should annotate -> networkx.DiGraph.""" - import networkx - - ret = create_hierarchical_from_coordinates.__annotations__["return"] - assert ret is networkx.DiGraph, f"Expected networkx.DiGraph, got {ret}" - - def test_create_all_graph_components_has_return_annotation(self): - """create_all_graph_components should have a return type annotation.""" - from weather_model_graphs.create.base import create_all_graph_components - - assert "return" in create_all_graph_components.__annotations__, ( - "create_all_graph_components missing return type annotation" - ) - - def test_connect_nodes_across_graphs_returns_digraph(self): - """connect_nodes_across_graphs should annotate -> networkx.DiGraph.""" - import networkx - from weather_model_graphs.create.base import connect_nodes_across_graphs - - ret = connect_nodes_across_graphs.__annotations__["return"] - assert ret is networkx.DiGraph, f"Expected networkx.DiGraph, got {ret}" - - def test_create_grid_graph_nodes_returns_graph(self): - """create_grid_graph_nodes should annotate -> networkx.Graph.""" - import networkx - from weather_model_graphs.create.grid.grid import create_grid_graph_nodes + G1 = create_directed_mesh_graph(G_coords, pattern="8-star") + G2 = create_directed_mesh_graph(G_coords, pattern="8-star") - ret = create_grid_graph_nodes.__annotations__["return"] - assert ret is networkx.Graph, f"Expected networkx.Graph, got {ret}" + assert set(G1.nodes) == set(G2.nodes) + assert set(G1.edges) == set(G2.edges) From 106454e79a3446ecc599d5611fa2c5d06628ecb0 Mon Sep 17 00:00:00 2001 From: prajwal Date: Thu, 28 May 2026 21:32:27 +0530 Subject: [PATCH 20/21] fix: resolve lint failures and make plot tests windows-safe - apply pre-commit formatting fixes (isort/black/flake8 cleanup) - remove unused local numpy import in test_mesh_layout - replace NamedTemporaryFile savefig usage with mkstemp pattern in plotting tests to avoid Windows PermissionError Validation: - pre-commit run -a: pass - pytest tests -q: 210 passed --- .../create/mesh/connectivity/flat.py | 4 +++- src/weather_model_graphs/create/mesh/coords.py | 4 +++- src/weather_model_graphs/networkx_utils.py | 4 ++-- tests/test_graph_creation.py | 9 +++++++-- tests/test_graph_plots.py | 9 +++++++-- tests/test_mesh_layout.py | 7 +------ 6 files changed, 23 insertions(+), 14 deletions(-) diff --git a/src/weather_model_graphs/create/mesh/connectivity/flat.py b/src/weather_model_graphs/create/mesh/connectivity/flat.py index 69c375e..4abef68 100644 --- a/src/weather_model_graphs/create/mesh/connectivity/flat.py +++ b/src/weather_model_graphs/create/mesh/connectivity/flat.py @@ -226,7 +226,9 @@ def create_flat_multiscale_mesh_graph( ) -def create_flat_singlescale_mesh_graph(xy, mesh_node_distance: float) -> networkx.DiGraph: +def create_flat_singlescale_mesh_graph( + xy, mesh_node_distance: float +) -> networkx.DiGraph: """ Create flat mesh graph of single level diff --git a/src/weather_model_graphs/create/mesh/coords.py b/src/weather_model_graphs/create/mesh/coords.py index 0fc2c44..40cec0b 100644 --- a/src/weather_model_graphs/create/mesh/coords.py +++ b/src/weather_model_graphs/create/mesh/coords.py @@ -115,7 +115,9 @@ def create_single_level_2d_mesh_primitive( return g -def create_single_level_2d_mesh_graph(xy: np.ndarray, nx: int, ny: int) -> networkx.DiGraph: +def create_single_level_2d_mesh_graph( + xy: np.ndarray, nx: int, ny: int +) -> networkx.DiGraph: """ Create directed graph with nx * ny nodes representing a 2D grid with positions spanning the range of xy coordinate values (first dimension diff --git a/src/weather_model_graphs/networkx_utils.py b/src/weather_model_graphs/networkx_utils.py index 99ae79e..ab4bf58 100644 --- a/src/weather_model_graphs/networkx_utils.py +++ b/src/weather_model_graphs/networkx_utils.py @@ -1,7 +1,7 @@ -import networkx - from typing import Dict, Tuple +import networkx + def prepend_node_index(graph, new_index) -> networkx.Graph: """ diff --git a/tests/test_graph_creation.py b/tests/test_graph_creation.py index 858299e..6e9bf0a 100644 --- a/tests/test_graph_creation.py +++ b/tests/test_graph_creation.py @@ -1,3 +1,4 @@ +import os import tempfile import cartopy.crs as ccrs @@ -21,8 +22,12 @@ def test_create_single_level_mesh_graph(): ax.scatter(xy[0, ...], xy[1, ...], color="r") ax.axison = True - with tempfile.NamedTemporaryFile(suffix=".png") as f: - fig.savefig(f.name) + fd, tmp_path = tempfile.mkstemp(suffix=".png") + os.close(fd) + try: + fig.savefig(tmp_path) + finally: + os.remove(tmp_path) @pytest.mark.parametrize("kind", ["graphcast", "keisler", "oskarsson_hierarchical"]) diff --git a/tests/test_graph_plots.py b/tests/test_graph_plots.py index eba23b5..1fe0dc3 100644 --- a/tests/test_graph_plots.py +++ b/tests/test_graph_plots.py @@ -1,3 +1,4 @@ +import os import tempfile import matplotlib.pyplot as plt @@ -63,5 +64,9 @@ def fn(): else: fn() - with tempfile.NamedTemporaryFile(suffix=".png") as f: - fig.savefig(f.name) + fd, tmp_path = tempfile.mkstemp(suffix=".png") + os.close(fd) + try: + fig.savefig(tmp_path) + finally: + os.remove(tmp_path) diff --git a/tests/test_mesh_layout.py b/tests/test_mesh_layout.py index 035bfe3..b289fcc 100644 --- a/tests/test_mesh_layout.py +++ b/tests/test_mesh_layout.py @@ -1382,8 +1382,6 @@ def test_empty_graph_passes_validation(self): def test_flat_singlescale_missing_pos_raises(self): """create_flat_singlescale_from_coordinates should raise on bad graph.""" - import numpy as np - G = nx.Graph() G.add_node((0, 0), type="mesh") # missing pos with pytest.raises(ValueError, match="'pos' attribute"): @@ -1457,9 +1455,7 @@ def test_4star_no_diagonal_edge_attributes(self): G_4star = create_directed_mesh_graph(G_coords, pattern="4-star") for u, v, d in G_4star.edges(data=True): adj = d.get("adjacency_type") - assert adj != "diagonal", ( - f"4-star graph has diagonal edge ({u},{v})" - ) + assert adj != "diagonal", f"4-star graph has diagonal edge ({u},{v})" def test_node_count_preserved(self): """Directed graph should have same number of nodes as undirected.""" @@ -1496,4 +1492,3 @@ def test_function_is_deterministic(self): assert set(G1.nodes) == set(G2.nodes) assert set(G1.edges) == set(G2.edges) - From fae6188d10cd07f93322d147bfd2382c6c4be14f Mon Sep 17 00:00:00 2001 From: prajwal Date: Thu, 28 May 2026 22:37:03 +0530 Subject: [PATCH 21/21] fix: use NamedTemporaryFile handle directly in savefig per review suggestion --- tests/test_graph_creation.py | 9 ++------- tests/test_graph_plots.py | 9 ++------- 2 files changed, 4 insertions(+), 14 deletions(-) diff --git a/tests/test_graph_creation.py b/tests/test_graph_creation.py index 6e9bf0a..7035d71 100644 --- a/tests/test_graph_creation.py +++ b/tests/test_graph_creation.py @@ -1,4 +1,3 @@ -import os import tempfile import cartopy.crs as ccrs @@ -22,12 +21,8 @@ def test_create_single_level_mesh_graph(): ax.scatter(xy[0, ...], xy[1, ...], color="r") ax.axison = True - fd, tmp_path = tempfile.mkstemp(suffix=".png") - os.close(fd) - try: - fig.savefig(tmp_path) - finally: - os.remove(tmp_path) + with tempfile.NamedTemporaryFile(suffix=".png") as fh: + fig.savefig(fh) @pytest.mark.parametrize("kind", ["graphcast", "keisler", "oskarsson_hierarchical"]) diff --git a/tests/test_graph_plots.py b/tests/test_graph_plots.py index 1fe0dc3..e096b52 100644 --- a/tests/test_graph_plots.py +++ b/tests/test_graph_plots.py @@ -1,4 +1,3 @@ -import os import tempfile import matplotlib.pyplot as plt @@ -64,9 +63,5 @@ def fn(): else: fn() - fd, tmp_path = tempfile.mkstemp(suffix=".png") - os.close(fd) - try: - fig.savefig(tmp_path) - finally: - os.remove(tmp_path) + with tempfile.NamedTemporaryFile(suffix=".png") as fh: + fig.savefig(fh)