diff --git a/CHANGELOG.md b/CHANGELOG.md index ef84eba..89fa7fb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,24 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Add `mesh_layout="prebuilt"` support to `create_all_graph_components` for + user-provided mesh node positions (e.g. ICON grid vertices or an + observation-station network), given as a nodes-only `networkx.Graph` (or a + bare `[N, 2]` coordinate array) via `mesh_layout_kwargs=dict(mesh_graph=...)`. + Mesh edges are built in the connectivity step directly from the node + positions (`method="delaunay"`); hierarchical meshes are declared with an + integer `level` node attribute. New module `create/mesh/layout/prebuilt.py` + contains the input validation and primitive creation; the connectivity step + now also validates that an explicit `pattern` matches the adjacency types + present in the mesh primitive instead of silently producing an empty mesh. + [\#79](https://github.com/mllam/weather-model-graphs/issues/79), @prajwal-tech07 +- Add `mesh_layout="triangular"` support to `create_all_graph_components`, using + `networkx.triangular_lattice_graph` to produce an equilateral-triangle lattice + with 6-connectivity. Supports all three `m2m_connectivity` modes: `flat`, + `hierarchical`, and `flat_multiscale`. New module + `create/mesh/connectivity/triangular.py` contains the coordinate and + connectivity creation functions for triangular meshes. + [\#80](https://github.com/mllam/weather-model-graphs/issues/80), @prajwal-tech07 - 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). diff --git a/docs/_toc.yml b/docs/_toc.yml index 1943552..9dc4146 100644 --- a/docs/_toc.yml +++ b/docs/_toc.yml @@ -7,5 +7,7 @@ chapters: - file: background - file: design - file: creating_the_graph +- file: mesh_layout +- file: prebuilt_mesh - file: lat_lons - file: decoding_mask diff --git a/docs/mesh_layout.ipynb b/docs/mesh_layout.ipynb new file mode 100644 index 0000000..be78843 --- /dev/null +++ b/docs/mesh_layout.ipynb @@ -0,0 +1,248 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "073d0b26", + "metadata": {}, + "source": [ + "# Changing the mesh layout\n", + "\n", + "The mesh layout controls the topology of the `m2m` (mesh-to-mesh) component of the graph.\n", + "By default, `weather-model-graphs` uses a **rectilinear** mesh, where nodes sit on a regular\n", + "rectangular grid and edges connect each node to its 8 nearest neighbours (cardinal + diagonal).\n", + "\n", + "As of v0.5.0, a **triangular** mesh layout is also supported. This places nodes on an equilateral-\n", + "triangle lattice, giving each interior node exactly 6 neighbours instead of 8. The 6-connectivity\n", + "is more isotropic and is expected to improve message-passing in graph neural network weather models.\n", + "\n", + "In this notebook we use the [Keisler 2022](https://arxiv.org/abs/2202.07575) graph archetype to\n", + "contrast three variants:\n", + "\n", + "1. **Default rectilinear** mesh (the archetype's built-in default)\n", + "2. **Rectilinear with finer mesh spacing** (more mesh nodes, denser connectivity)\n", + "3. **Triangular mesh** at the same spacing as variant 1\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2595994f", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "import weather_model_graphs as wmg" + ] + }, + { + "cell_type": "markdown", + "id": "d3edb0f9", + "metadata": {}, + "source": [ + "## Set up a fake grid\n", + "\n", + "We start from a regular 32 × 32 grid of Cartesian (x, y) coordinates. These represent the\n", + "locations of the input/output data (grid nodes in g2m / m2g)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7eb67af8", + "metadata": {}, + "outputs": [], + "source": [ + "xs, ys = np.meshgrid(np.linspace(0, 10, 32), np.linspace(0, 10, 32))\n", + "xy = np.stack([xs.flatten(), ys.flatten()], axis=-1)\n", + "\n", + "fig, ax = plt.subplots(figsize=(4, 4))\n", + "ax.scatter(xy[:, 0], xy[:, 1], s=2)\n", + "ax.set_aspect(1)\n", + "ax.set_title(\"Grid nodes\")" + ] + }, + { + "cell_type": "markdown", + "id": "686ee7f2", + "metadata": {}, + "source": [ + "## Example 1 — Rectilinear mesh (default spacing)\n", + "\n", + "`create_keisler_graph` uses `mesh_layout='rectilinear'` with `mesh_node_distance=3` by default.\n", + "Each interior mesh node connects to its 8 neighbours (4-star cardinal + 4 diagonals)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c7661b63", + "metadata": {}, + "outputs": [], + "source": [ + "graph_rectilinear = wmg.create.archetype.create_keisler_graph(\n", + " coords=xy,\n", + " mesh_node_distance=3,\n", + ")\n", + "\n", + "m2m_rectilinear = wmg.split_graph_by_edge_attribute(\n", + " graph_rectilinear, attr=\"component\"\n", + ")[\"m2m\"]\n", + "\n", + "print(f\"Mesh nodes : {m2m_rectilinear.number_of_nodes()}\")\n", + "print(f\"Mesh edges : {m2m_rectilinear.number_of_edges()}\")\n", + "\n", + "fig, ax = plt.subplots(figsize=(5, 5))\n", + "wmg.visualise.nx_draw_with_pos_and_attr(m2m_rectilinear, ax=ax, node_size=30)\n", + "ax.set_title(\"Rectilinear mesh — default spacing (mesh_node_distance=3)\")" + ] + }, + { + "cell_type": "markdown", + "id": "e34fb561", + "metadata": {}, + "source": [ + "## Example 2 — Rectilinear mesh with finer spacing\n", + "\n", + "Halving `mesh_node_distance` roughly quadruples the number of mesh nodes and gives a denser\n", + "rectilinear mesh over the same domain." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7ebc4f15", + "metadata": {}, + "outputs": [], + "source": [ + "graph_fine = wmg.create.archetype.create_keisler_graph(\n", + " coords=xy,\n", + " mesh_node_distance=1.5,\n", + ")\n", + "\n", + "m2m_fine = wmg.split_graph_by_edge_attribute(graph_fine, attr=\"component\")[\"m2m\"]\n", + "\n", + "print(f\"Mesh nodes : {m2m_fine.number_of_nodes()}\")\n", + "print(f\"Mesh edges : {m2m_fine.number_of_edges()}\")\n", + "\n", + "fig, ax = plt.subplots(figsize=(5, 5))\n", + "wmg.visualise.nx_draw_with_pos_and_attr(m2m_fine, ax=ax, node_size=10)\n", + "ax.set_title(\"Rectilinear mesh — finer spacing (mesh_node_distance=1.5)\")" + ] + }, + { + "cell_type": "markdown", + "id": "1acbf174", + "metadata": {}, + "source": [ + "## Example 3 — Triangular mesh\n", + "\n", + "Setting `mesh_layout='triangular'` places nodes on an equilateral-triangle lattice.\n", + "Each interior node has exactly **6 neighbours** (vs. 8 for rectilinear), which provides\n", + "more isotropic spatial connectivity.\n", + "\n", + "We keep the same g2m / m2g connectivity settings as the Keisler archetype (within-radius\n", + "encoding, 4-nearest-neighbour decoding) and use the same mesh spacing as Example 1." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "086f9fb1", + "metadata": {}, + "outputs": [], + "source": [ + "graph_triangular = wmg.create.create_all_graph_components(\n", + " coords=xy,\n", + " mesh_layout=\"triangular\",\n", + " mesh_layout_kwargs=dict(mesh_node_spacing=3),\n", + " m2m_connectivity=\"flat\",\n", + " g2m_connectivity=\"within_radius\",\n", + " g2m_connectivity_kwargs=dict(rel_max_dist=0.51),\n", + " m2g_connectivity=\"nearest_neighbours\",\n", + " m2g_connectivity_kwargs=dict(max_num_neighbours=4),\n", + ")\n", + "\n", + "m2m_triangular = wmg.split_graph_by_edge_attribute(graph_triangular, attr=\"component\")[\n", + " \"m2m\"\n", + "]\n", + "\n", + "print(f\"Mesh nodes : {m2m_triangular.number_of_nodes()}\")\n", + "print(f\"Mesh edges : {m2m_triangular.number_of_edges()}\")\n", + "\n", + "fig, ax = plt.subplots(figsize=(5, 5))\n", + "wmg.visualise.nx_draw_with_pos_and_attr(m2m_triangular, ax=ax, node_size=30)\n", + "ax.set_title(\"Triangular mesh (mesh_node_spacing=3)\")" + ] + }, + { + "cell_type": "markdown", + "id": "9bd01e4c", + "metadata": {}, + "source": [ + "## Side-by-side comparison\n", + "\n", + "Plotting the `m2m` component of all three graphs side by side makes the difference in\n", + "topology clear: rectilinear nodes form a square grid with 8-connectivity, while triangular\n", + "nodes form a hexagonal-offset grid with 6-connectivity." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fdd4e7aa", + "metadata": {}, + "outputs": [], + "source": [ + "fig, axes = plt.subplots(1, 3, figsize=(15, 5))\n", + "\n", + "configs = [\n", + " (m2m_rectilinear, \"Rectilinear\\n(default spacing)\", 30),\n", + " (m2m_fine, \"Rectilinear\\n(finer spacing)\", 10),\n", + " (m2m_triangular, \"Triangular\\n(same spacing as default)\", 30),\n", + "]\n", + "\n", + "# Compute shared axis limits from node positions across all graphs\n", + "all_pos = np.concatenate(\n", + " [\n", + " np.array([data[\"pos\"] for _, data in graph.nodes(data=True)])\n", + " for graph, _, _ in configs\n", + " ]\n", + ")\n", + "x_min, y_min = all_pos.min(axis=0)\n", + "x_max, y_max = all_pos.max(axis=0)\n", + "pad = max(x_max - x_min, y_max - y_min) * 0.05\n", + "\n", + "for ax, (graph, title, ns) in zip(axes, configs):\n", + " wmg.visualise.nx_draw_with_pos_and_attr(graph, ax=ax, node_size=ns)\n", + " ax.set_title(title)\n", + " ax.set_xlim(x_min - pad, x_max + pad)\n", + " ax.set_ylim(y_min - pad, y_max + pad)\n", + " ax.set_aspect(1.0)\n", + "\n", + "fig.tight_layout()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/prebuilt_mesh.ipynb b/docs/prebuilt_mesh.ipynb new file mode 100644 index 0000000..1613aea --- /dev/null +++ b/docs/prebuilt_mesh.ipynb @@ -0,0 +1,311 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Bring your own mesh: the prebuilt mesh layout\n", + "\n", + "The generated mesh layouts (`rectilinear`, `triangular`) place mesh nodes for you.\n", + "With `mesh_layout=\"prebuilt\"` you instead supply **your own mesh node positions** --\n", + "for example ICON grid vertices, MPAS cell centres, or an observation-station network --\n", + "and `weather-model-graphs` builds the encode-process-decode graph around them.\n", + "\n", + "The two-step mesh creation process still applies, with a twist:\n", + "\n", + "1. **Coordinate creation** (`mesh_layout=\"prebuilt\"`): your nodes are validated and\n", + " passed through as an *edge-less node cloud* -- no adjacency is invented here.\n", + "2. **Connectivity creation** (`m2m_connectivity`): mesh edges are built directly from\n", + " the node positions (`method=\"delaunay\"`, the default) and directed with\n", + " `len`/`vdiff` edge features.\n", + "\n", + "```{note}\n", + "Currently the prebuilt layout supports **nodes-only** input: the mesh graph you\n", + "provide must not contain any edges. Support for user-provided edges is planned --\n", + "see the design discussion in\n", + "[issue #79](https://github.com/mllam/weather-model-graphs/issues/79).\n", + "```\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import networkx as nx\n", + "import numpy as np\n", + "\n", + "import weather_model_graphs as wmg" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## The input contract\n", + "\n", + "Provide a `networkx.Graph` where **every node** has:\n", + "\n", + "- `pos`: `np.ndarray` of shape `(2,)` -- the node position. **These must be in the\n", + " same coordinate system as the grid `coords`** you pass to\n", + " `create_all_graph_components` (the library cannot check this for you!).\n", + "- `type`: the string `\"mesh\"`.\n", + "- `level`: an integer, **only** for hierarchical meshes (lowest value = finest\n", + " level); either all nodes have one or none do.\n", + "\n", + "Node positions must be unique. For the simplest case you can also pass a bare\n", + "`np.ndarray` of shape `[N, 2]` instead of a graph.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Set up a fake grid and some mesh nodes\n", + "\n", + "We create a regular grid of (x, y) coordinates (the locations of the input/output\n", + "data) and a set of irregular \"station-like\" mesh node positions inside the same\n", + "domain.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "xs, ys = np.meshgrid(np.linspace(0, 10, 32), np.linspace(0, 10, 32))\n", + "xy = np.stack([xs.flatten(), ys.flatten()], axis=-1)\n", + "\n", + "rng = np.random.default_rng(seed=42)\n", + "mesh_xy = rng.random((40, 2)) * 10\n", + "\n", + "my_mesh = nx.Graph()\n", + "for i, pos in enumerate(mesh_xy):\n", + " my_mesh.add_node(i, pos=pos, type=\"mesh\")\n", + "\n", + "fig, ax = plt.subplots(figsize=(5, 5))\n", + "ax.scatter(xy[:, 0], xy[:, 1], s=2, alpha=0.3, label=\"grid points\")\n", + "ax.scatter(mesh_xy[:, 0], mesh_xy[:, 1], s=40, marker=\"^\", label=\"my mesh nodes\")\n", + "ax.legend()\n", + "ax.set_title(\"User-provided mesh nodes over the data grid\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Example 1 -- flat mesh from nodes only\n", + "\n", + "With nodes-only input the connectivity step builds the mesh edges by Delaunay\n", + "triangulation of the node positions (`method=\"delaunay\"` is the default, shown\n", + "here explicitly).\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "graph_flat = wmg.create.create_all_graph_components(\n", + " coords=xy,\n", + " mesh_layout=\"prebuilt\",\n", + " mesh_layout_kwargs=dict(mesh_graph=my_mesh),\n", + " m2m_connectivity=\"flat\",\n", + " m2m_connectivity_kwargs=dict(method=\"delaunay\"),\n", + " g2m_connectivity=\"nearest_neighbour\",\n", + " m2g_connectivity=\"nearest_neighbour\",\n", + ")\n", + "\n", + "m2m_flat = wmg.split_graph_by_edge_attribute(graph_flat, attr=\"component\")[\"m2m\"]\n", + "\n", + "print(f\"Mesh nodes : {m2m_flat.number_of_nodes()}\")\n", + "print(f\"Mesh edges : {m2m_flat.number_of_edges()}\")\n", + "\n", + "fig, ax = plt.subplots(figsize=(5, 5))\n", + "wmg.visualise.nx_draw_with_pos_and_attr(m2m_flat, ax=ax, node_size=30)\n", + "ax.set_title(\"Flat prebuilt mesh (Delaunay connectivity)\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The convenience form skips building the graph yourself -- pass the coordinate\n", + "array directly:\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "graph_from_array = wmg.create.create_all_graph_components(\n", + " coords=xy,\n", + " mesh_layout=\"prebuilt\",\n", + " mesh_layout_kwargs=dict(mesh_graph=mesh_xy),\n", + " m2m_connectivity=\"flat\",\n", + " g2m_connectivity=\"nearest_neighbour\",\n", + " m2g_connectivity=\"nearest_neighbour\",\n", + ")\n", + "graph_from_array.number_of_nodes()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Example 2 -- hierarchical mesh from `level` attributes\n", + "\n", + "For a hierarchical mesh, give every node an integer `level` attribute (lowest\n", + "value = finest). Within each level the mesh edges are built per level\n", + "(`intra_level=dict(method=\"delaunay\")`); between levels, `mesh_up`/`mesh_down`\n", + "edges are created by nearest-neighbour search (`inter_level`), exactly as for the\n", + "generated layouts.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "my_mesh_hier = nx.Graph()\n", + "for i, pos in enumerate(mesh_xy):\n", + " my_mesh_hier.add_node((\"fine\", i), pos=pos, type=\"mesh\", level=1)\n", + "coarse_xy = rng.random((8, 2)) * 10\n", + "for i, pos in enumerate(coarse_xy):\n", + " my_mesh_hier.add_node((\"coarse\", i), pos=pos, type=\"mesh\", level=2)\n", + "\n", + "components = wmg.create.create_all_graph_components(\n", + " coords=xy,\n", + " mesh_layout=\"prebuilt\",\n", + " mesh_layout_kwargs=dict(mesh_graph=my_mesh_hier),\n", + " m2m_connectivity=\"hierarchical\",\n", + " m2m_connectivity_kwargs=dict(\n", + " intra_level=dict(method=\"delaunay\"),\n", + " inter_level=dict(pattern=\"nearest\", k=1),\n", + " ),\n", + " g2m_connectivity=\"nearest_neighbour\",\n", + " m2g_connectivity=\"nearest_neighbour\",\n", + " return_components=True,\n", + ")\n", + "\n", + "m2m_hier = components[\"m2m\"]\n", + "for direction in (\"same\", \"up\", \"down\"):\n", + " n = sum(\n", + " 1 for _, _, d in m2m_hier.edges(data=True) if d.get(\"direction\") == direction\n", + " )\n", + " print(f\"{direction:>5} edges: {n}\")\n", + "\n", + "fig, ax = plt.subplots(figsize=(5, 5))\n", + "wmg.visualise.nx_draw_with_pos_and_attr(m2m_hier, ax=ax, node_size=30)\n", + "ax.set_title(\"Hierarchical prebuilt mesh (2 levels)\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Helpful validation errors\n", + "\n", + "The input contract is checked up front, with errors that say what to fix:\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "bad_mesh = nx.Graph()\n", + "bad_mesh.add_node(0, pos=np.array([1.0, 1.0]), type=\"mesh\")\n", + "bad_mesh.add_node(1, pos=np.array([1.0, 1.0]), type=\"mesh\") # duplicate!\n", + "\n", + "try:\n", + " wmg.create.create_all_graph_components(\n", + " coords=xy,\n", + " mesh_layout=\"prebuilt\",\n", + " mesh_layout_kwargs=dict(mesh_graph=bad_mesh),\n", + " m2m_connectivity=\"flat\",\n", + " g2m_connectivity=\"nearest_neighbour\",\n", + " m2g_connectivity=\"nearest_neighbour\",\n", + " )\n", + "except ValueError as e:\n", + " print(f\"ValueError: {e}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "with_edges = nx.Graph()\n", + "with_edges.add_node(0, pos=np.array([0.0, 0.0]), type=\"mesh\")\n", + "with_edges.add_node(1, pos=np.array([5.0, 5.0]), type=\"mesh\")\n", + "with_edges.add_edge(0, 1) # user-provided edges: not yet supported\n", + "\n", + "try:\n", + " wmg.create.create_all_graph_components(\n", + " coords=xy,\n", + " mesh_layout=\"prebuilt\",\n", + " mesh_layout_kwargs=dict(mesh_graph=with_edges),\n", + " m2m_connectivity=\"flat\",\n", + " g2m_connectivity=\"nearest_neighbour\",\n", + " m2g_connectivity=\"nearest_neighbour\",\n", + " )\n", + "except NotImplementedError as e:\n", + " print(f\"NotImplementedError: {e}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Loading real mesh sources\n", + "\n", + "The library deliberately has no file-format support here -- whatever the source,\n", + "you load it into the node-cloud graph in a few lines. For example, station\n", + "locations from a CSV with columns `x`, `y` (already in the grid's coordinate\n", + "system):\n", + "\n", + "```python\n", + "import pandas as pd\n", + "\n", + "df = pd.read_csv(\"stations.csv\")\n", + "my_mesh = nx.Graph()\n", + "for i, row in df.iterrows():\n", + " my_mesh.add_node(i, pos=np.array([row.x, row.y]), type=\"mesh\")\n", + "```\n", + "\n", + "or ICON grid vertices from its NetCDF grid file (remember to transform lon/lat to\n", + "the grid's projected coordinate system first, e.g. with `pyproj`).\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/src/weather_model_graphs/create/base.py b/src/weather_model_graphs/create/base.py index d157c3b..9f70c86 100644 --- a/src/weather_model_graphs/create/base.py +++ b/src/weather_model_graphs/create/base.py @@ -8,6 +8,7 @@ function uses `connect_nodes_across_graphs` to connect nodes across the component graphs. """ +import warnings from typing import Dict, Iterable, List, Tuple, Union import networkx @@ -28,9 +29,21 @@ create_flat_singlescale_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 .mesh.layout.prebuilt import ( + create_multi_level_prebuilt_mesh_primitives, + create_single_level_prebuilt_mesh_primitive, +) +from .mesh.layout.rectilinear import ( + create_multirange_2d_mesh_primitives as create_multirange_2d_rectilinear_mesh_primitives, +) +from .mesh.layout.rectilinear import ( + create_single_level_2d_mesh_primitive as create_single_level_2d_rectilinear_mesh_primitive, +) +from .mesh.layout.triangular import ( + create_multirange_2d_mesh_primitives as create_multirange_2d_triangular_mesh_primitives, +) +from .mesh.layout.triangular import ( + create_single_level_2d_mesh_primitive as create_single_level_2d_triangular_mesh_primitive, ) @@ -141,26 +154,53 @@ def create_all_graph_components( - "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"): + - "triangular": Regular triangular lattice with ``mesh_node_spacing`` + resolution. Uses ``networkx.triangular_lattice_graph`` to produce + equilateral triangles with 6-connectivity. A CRS warning is emitted + if ``graph_crs`` is geographic (lat/lon). + - "prebuilt": User-provided mesh node positions (e.g. ICON grid vertices + or an observation-station network). The mesh nodes are taken from + ``mesh_layout_kwargs["mesh_graph"]`` and mesh edges are built in the + connectivity step (``method``, Delaunay triangulation by default). + Currently nodes-only input is supported (the given graph must not + contain edges) and ``m2m_connectivity`` must be "flat" or + "hierarchical". See ``create.mesh.layout.prebuilt`` for the input + contract. + + mesh_layout_kwargs (for mesh_layout="rectilinear" or "triangular"): - 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: + mesh_layout_kwargs (for mesh_layout="prebuilt"): + - mesh_graph: networkx.Graph with node attributes ``pos`` (np.ndarray of + shape (2,), same coordinate system as ``coords``), ``type`` ("mesh"), + and -- for hierarchical meshes -- integer ``level`` (lowest value = + finest level); or an np.ndarray of shape [N_mesh_nodes, 2] with node + positions. No ``mesh_node_spacing`` is needed (spacing is implied by + the node positions). + + Wherever the ``pattern`` argument appears below it selects a subset of + the spatial adjacency edges created by the (generated) mesh layout: + - not given (default): use every edge the layout produced - ``"4-star"``: only cardinal directions (horizontal and vertical neighbours) - ``"8-star"``: cardinal plus diagonal neighbours (all 8 surrounding nodes) + For mesh_layout="prebuilt" the primitives are edge-less node clouds, so + ``pattern`` does not apply; the ``method`` argument (default: + ``"delaunay"``) selects how mesh edges are built from the node positions. m2m_connectivity: - "flat": Create a single-level directed mesh graph. - m2m_connectivity_kwargs: pattern (default: "8-star") + m2m_connectivity_kwargs: pattern (generated layouts) or + method (mesh_layout="prebuilt", default: "delaunay") - "flat_multiscale": Create a flat multiscale mesh graph. - m2m_connectivity_kwargs: pattern (default: "8-star") + m2m_connectivity_kwargs: pattern + (not yet supported for mesh_layout="prebuilt") - "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=... or method=...), + inter_level=dict(pattern=..., k=...) m2g_connectivity: - "nearest_neighbour": Find the nearest neighbour in mesh for each node in grid @@ -255,46 +295,116 @@ def create_all_graph_components( # ----------------------------------------------------------------------- G_mesh_coords: Union[networkx.Graph, List[networkx.Graph]] - if mesh_layout == "rectilinear": + # CRS warning for triangular layout in geographic coordinates + if ( + mesh_layout == "triangular" + and graph_crs is not None + and graph_crs.is_geographic + ): + warnings.warn( + "mesh_layout='triangular' produces non-uniform physical spacing in " + "geographic coordinates. Consider mesh_layout='icosahedral' for " + "uniform coverage on a sphere.", + UserWarning, + stacklevel=2, + ) + + # Validate mesh_layout and resolve the requested mesh node spacing once + # (shared by all m2m_connectivity modes and layouts). + if mesh_layout not in ("rectilinear", "triangular", "prebuilt"): + raise NotImplementedError( + f"mesh_layout='{mesh_layout}' is not yet supported. " + "Currently supported: 'rectilinear', 'triangular', 'prebuilt'." + ) + + if mesh_layout == "prebuilt": + # Prebuilt meshes carry their own node positions, so no + # mesh_node_spacing is needed (spacing is implied by the positions) + mesh_graph = mesh_layout_kwargs.get("mesh_graph") + if mesh_graph is None: + raise ValueError( + "mesh_layout='prebuilt' requires 'mesh_graph' in " + "mesh_layout_kwargs: a networkx.Graph of mesh nodes (or an " + "np.ndarray of node positions with shape [N_mesh_nodes, 2])." + ) + else: 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 " + f"mesh_layout='{mesh_layout}' 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 - G_mesh_coords = create_single_level_2d_mesh_primitive( + # Pick the coordinate-creation function based on the mesh_layout value, + # nested inside the m2m_connectivity branch (single-level vs multi-level). + if m2m_connectivity == "flat": + # Single-level mesh primitive + if mesh_layout == "rectilinear": + G_mesh_coords = create_single_level_2d_rectilinear_mesh_primitive( + xy, mesh_node_spacing=mesh_node_spacing + ) + elif mesh_layout == "triangular": + G_mesh_coords = create_single_level_2d_triangular_mesh_primitive( xy, mesh_node_spacing=mesh_node_spacing ) + elif mesh_layout == "prebuilt": + G_mesh_coords = create_single_level_prebuilt_mesh_primitive(mesh_graph) else: - # 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" - ] - if "max_num_refinement_levels" in mesh_layout_kwargs: - primitives_kwargs["max_num_levels"] = mesh_layout_kwargs[ - "max_num_refinement_levels" - ] - G_mesh_coords = create_multirange_2d_mesh_primitives(**primitives_kwargs) + raise NotImplementedError( + f"mesh_layout='{mesh_layout}' is not implemented. " + "Supported layouts: 'rectilinear', 'triangular', 'prebuilt'." + ) + elif mesh_layout == "prebuilt": + # Multi-level prebuilt primitives, split by the nodes' `level` attribute + if m2m_connectivity == "flat_multiscale": + raise NotImplementedError( + "m2m_connectivity='flat_multiscale' is not yet supported for " + "mesh_layout='prebuilt'. Use 'flat' or 'hierarchical'." + ) + G_mesh_coords = create_multi_level_prebuilt_mesh_primitives(mesh_graph) else: - raise NotImplementedError( - f"mesh_layout='{mesh_layout}' is not yet supported. " - "Currently only 'rectilinear' is implemented." - ) + # Multi-level mesh primitives (flat_multiscale or hierarchical) + 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" + ] + if mesh_layout == "rectilinear": + G_mesh_coords = create_multirange_2d_rectilinear_mesh_primitives( + **primitives_kwargs + ) + elif mesh_layout == "triangular": + G_mesh_coords = create_multirange_2d_triangular_mesh_primitives( + **primitives_kwargs + ) + else: + raise NotImplementedError( + f"mesh_layout='{mesh_layout}' is not implemented. " + "Supported layouts: 'rectilinear', 'triangular', 'prebuilt'." + ) # ----------------------------------------------------------------------- # Step 2: Connectivity creation — converts mesh primitives to directed graph # ----------------------------------------------------------------------- if m2m_connectivity == "flat": + # `pattern` selects adjacency edges of generated layouts; `method` + # builds edges from node positions for edge-less (prebuilt) + # primitives. Both default to "use what the layout implies" when not + # given (all layout edges, and Delaunay triangulation respectively). + conn_kwargs = { + key: m2m_connectivity_kwargs[key] + for key in ("pattern", "method") + if key in m2m_connectivity_kwargs + } graph_components["m2m"] = create_flat_singlescale_from_coordinates( - G_mesh_coords, **m2m_connectivity_kwargs + G_mesh_coords, **conn_kwargs ) grid_connect_graph = graph_components["m2m"] @@ -302,8 +412,15 @@ def create_all_graph_components( # hierarchical mesh graph has three sub-graphs: # `m2m` (mesh-to-mesh), `mesh_up` (up edge connections) and # `mesh_down` (down edge connections) + hierarchical_kwargs = {} + intra_level = m2m_connectivity_kwargs.get("intra_level") + if intra_level is not None: + hierarchical_kwargs["intra_level"] = intra_level + inter_level = m2m_connectivity_kwargs.get("inter_level") + if inter_level is not None: + hierarchical_kwargs["inter_level"] = inter_level graph_components["m2m"] = create_hierarchical_from_coordinates( - G_mesh_coords, **m2m_connectivity_kwargs + G_mesh_coords, **hierarchical_kwargs ) # Only connect grid to bottom level of hierarchy grid_connect_graph = split_graph_by_edge_attribute( @@ -311,8 +428,13 @@ def create_all_graph_components( )[0] elif m2m_connectivity == "flat_multiscale": + conn_kwargs = { + key: m2m_connectivity_kwargs[key] + for key in ("pattern",) + if key in m2m_connectivity_kwargs + } graph_components["m2m"] = create_flat_multiscale_from_coordinates( - G_mesh_coords, **m2m_connectivity_kwargs + G_mesh_coords, **conn_kwargs ) grid_connect_graph = graph_components["m2m"] diff --git a/src/weather_model_graphs/create/mesh/__init__.py b/src/weather_model_graphs/create/mesh/__init__.py index 2d3d7cc..56aa1dd 100644 --- a/src/weather_model_graphs/create/mesh/__init__.py +++ b/src/weather_model_graphs/create/mesh/__init__.py @@ -1,6 +1,19 @@ -from .coords import ( - create_directed_mesh_graph, +from .connectivity.flat import create_flat_multiscale_from_coordinates +from .connectivity.general import create_directed_mesh_graph +from .layout.prebuilt import ( + create_multi_level_prebuilt_mesh_primitives, + create_single_level_prebuilt_mesh_primitive, + validate_prebuilt_mesh_nodes, +) +from .layout.rectilinear import ( + create_multirange_2d_mesh_graphs, create_multirange_2d_mesh_primitives, create_single_level_2d_mesh_graph, create_single_level_2d_mesh_primitive, ) +from .layout.triangular import ( + create_multirange_2d_mesh_primitives as create_multirange_2d_triangular_mesh_primitives, +) +from .layout.triangular import ( + create_single_level_2d_mesh_primitive as create_single_level_2d_triangular_mesh_primitive, +) diff --git a/src/weather_model_graphs/create/mesh/connectivity/flat.py b/src/weather_model_graphs/create/mesh/connectivity/flat.py index 4abef68..da57b06 100644 --- a/src/weather_model_graphs/create/mesh/connectivity/flat.py +++ b/src/weather_model_graphs/create/mesh/connectivity/flat.py @@ -2,9 +2,10 @@ import networkx import numpy as np +import scipy.spatial from ....networkx_utils import prepend_node_index -from .. import coords as mesh_coords +from ..layout import rectilinear as mesh_layout from .general import create_directed_mesh_graph @@ -59,6 +60,12 @@ 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). + This works for any mesh layout. When the mesh primitives carry the integer + ``(i, j)`` grid-index node labels of a complete rectangular lattice (the + rectilinear layout) the merge uses fast index arithmetic. Otherwise (e.g. a + triangular lattice, whose nodes do not form a full rectangle) it falls back + to layout-agnostic position-based matching with a KD-tree. + Parameters ---------- G_coords_list : list of networkx.Graph @@ -67,7 +74,8 @@ def create_flat_multiscale_from_coordinates( - 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``. + Created by ``create_multirange_2d_mesh_primitives`` (rectilinear) or + ``create_multirange_2d_mesh_primitives`` (triangular). **kwargs Additional keyword arguments passed to ``create_directed_mesh_graph`` (e.g. ``pattern="8-star"``). @@ -82,32 +90,69 @@ def create_flat_multiscale_from_coordinates( 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 - or interlevel_refinement_factor % 2 != 1 - ): - raise ValueError( - "The `interlevel_refinement_factor` must be an odd integer. " - f"Given value: {interlevel_refinement_factor}." - ) - # Convert each level's coordinate graph to directed graph with chosen pattern G_all_levels = [ create_directed_mesh_graph(g_coords, **kwargs) for g_coords in G_coords_list ] - # combine all levels to one graph + # Decide how to merge coincident nodes across levels. The fast index path is + # only valid for the rectilinear ``grid_2d`` layout, where coarse-level grid + # indices coincide in *position* with finer nodes. Its structural signature + # is the presence of ``"diagonal"`` adjacency edges (8-star lattice). Other + # layouts -- e.g. the triangular lattice, which has only ``"cardinal"`` edges + # and offset rows -- fall back to layout-agnostic position matching (KD-tree). + # The check is done on the undirected coordinate graph so it is independent + # of the connectivity ``pattern`` (which may filter diagonals out later). + grid_indexed = any( + d.get("adjacency_type") == "diagonal" + for _, _, d in G_coords_list[0].edges(data=True) + ) + + if grid_indexed: + # The index-arithmetic merge relies on a known, odd refinement factor + # between levels; this requirement is specific to the grid-index path. + # The position-based fallback below works for any refinement factor. + 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 + or interlevel_refinement_factor % 2 != 1 + ): + raise ValueError( + "The `interlevel_refinement_factor` must be an odd integer. " + f"Given value: {interlevel_refinement_factor}." + ) + G_tot = _merge_levels_by_grid_index(G_all_levels, interlevel_refinement_factor) + else: + G_tot = _merge_levels_by_position(G_all_levels) + + # Relabel mesh nodes to start with 0 + G_tot = prepend_node_index(G_tot, 0) + + # add dx and dy to graph + G_tot.graph["dx"] = {i: g.graph["dx"] for i, g in enumerate(G_all_levels)} + G_tot.graph["dy"] = {i: g.graph["dy"] for i, g in enumerate(G_all_levels)} + + return G_tot + + +def _merge_levels_by_grid_index( + G_all_levels: List[networkx.DiGraph], interlevel_refinement_factor: int +) -> networkx.DiGraph: + """Merge multiscale levels using integer ``(i, j)`` grid-index arithmetic. + + This is the original rectilinear merge: coarser-level nodes are matched to + their coincident finer-level nodes purely from the grid indices, so the + result is unchanged for the rectilinear layout. + """ + G_all_levels = list(G_all_levels) G_tot = G_all_levels[0] # First node at level l+1 share position with node (offset, offset) at level l level_offset = interlevel_refinement_factor // 2 @@ -141,12 +186,46 @@ def create_flat_multiscale_from_coordinates( 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) + return G_tot - # add dx and dy to graph - G_tot.graph["dx"] = {i: g.graph["dx"] for i, g in enumerate(G_all_levels)} - G_tot.graph["dy"] = {i: g.graph["dy"] for i, g in enumerate(G_all_levels)} + +def _merge_levels_by_position( + G_all_levels: List[networkx.DiGraph], +) -> networkx.DiGraph: + """Merge multiscale levels by coincident node positions (KD-tree). + + Layout-agnostic fallback used when node labels do not form a complete + ``(i, j)`` grid (e.g. the triangular layout). For each coarser level, any + node whose position coincides (within floating-point tolerance) with an + existing finer-level node is merged with it, so multi-resolution edges + share the same node identity. + """ + # Prepend level index so labels are unique across levels before merging + G_levels = [ + prepend_node_index(g, level_i) for level_i, g in enumerate(G_all_levels) + ] + + G_tot = G_levels[0] + for lev in range(1, len(G_levels)): + G_coarse = G_levels[lev] + + # KDTree of existing (finer) nodes for position matching + fine_nodes = list(G_tot.nodes()) + fine_positions = np.array([G_tot.nodes[n]["pos"] for n in fine_nodes]) + kdt = scipy.spatial.KDTree(fine_positions) + + # Find which coarse nodes coincide with existing fine nodes + relabel_map = {} + for node in G_coarse.nodes(): + pos = G_coarse.nodes[node]["pos"] + dist, idx = kdt.query(pos) + if dist < 1e-8: + relabel_map[node] = fine_nodes[idx] + + if relabel_map: + G_coarse = networkx.relabel_nodes(G_coarse, relabel_map) + + G_tot = networkx.compose(G_tot, G_coarse) return G_tot @@ -213,7 +292,7 @@ def create_flat_multiscale_mesh_graph( G_tot : networkx.DiGraph The merged mesh graph """ - G_coords_list = mesh_coords.create_multirange_2d_mesh_primitives( + G_coords_list = mesh_layout.create_multirange_2d_mesh_primitives( max_num_levels=max_num_levels, xy=xy, mesh_node_spacing=mesh_node_distance, @@ -263,4 +342,4 @@ def create_flat_singlescale_mesh_graph( " so that the mesh nodes are spaced closer together?" ) - return mesh_coords.create_single_level_2d_mesh_graph(xy=xy, nx=nx, ny=ny) + return mesh_layout.create_single_level_2d_mesh_graph(xy=xy, nx=nx, ny=ny) diff --git a/src/weather_model_graphs/create/mesh/connectivity/general.py b/src/weather_model_graphs/create/mesh/connectivity/general.py index 9a3914c..be879e8 100644 --- a/src/weather_model_graphs/create/mesh/connectivity/general.py +++ b/src/weather_model_graphs/create/mesh/connectivity/general.py @@ -1,57 +1,116 @@ import networkx import numpy as np +import scipy.spatial def create_directed_mesh_graph( - G_undirected: networkx.Graph, pattern: str = "8-star" + G_undirected: networkx.Graph, pattern: str = None, method: str = None ) -> 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. + Convert an undirected mesh primitive graph to a directed mesh graph + (nx.DiGraph). This is the second step in the two-step mesh creation process: - 1. Coordinate creation (create_single_level_2d_mesh_primitive) -> nx.Graph + 1. Coordinate creation (mesh layout module) -> 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) + Two kinds of mesh primitive are supported: + + - **Primitives with adjacency edges** (the generated layouts, + ``rectilinear``/``triangular``): the edges carry an ``adjacency_type`` + attribute and the optional ``pattern`` argument selects a subset of + them. When no ``pattern`` is given, every edge the layout produced is + used. When a ``pattern`` is given it must match the adjacency types + present in the primitive, otherwise a ``ValueError`` is raised listing + what is available (rather than silently producing an empty mesh). + - **Edge-less primitives** (the ``prebuilt`` layout's node clouds): there + is no adjacency to select from, so the directed edges are built + directly from the node positions using ``method`` (currently only + ``"delaunay"``, the default: Delaunay triangulation of the node + positions). ``pattern`` does not apply to node clouds. 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"``. + Expected edge attributes (only when the primitive has edges): + - ``"adjacency_type"``: str, e.g. ``"cardinal"`` or ``"diagonal"``. Additional edge attributes (e.g. ``"level"``) are preserved in the output directed graph. - pattern : str - Connectivity pattern. Options: + pattern : str, optional + Connectivity pattern for primitives with adjacency edges. Options: + - ``None`` (default): use every edge the layout produced - ``"4-star"``: only cardinal edges (horizontal/vertical neighbours) - ``"8-star"``: all edges (cardinal + diagonal neighbours) + method : str, optional + Edge construction method for edge-less primitives (node clouds). + Options: + - ``None`` (default): resolves to ``"delaunay"`` for node clouds + - ``"delaunay"``: Delaunay triangulation of the node positions Returns ------- networkx.DiGraph Directed graph with bidirectional edges, each having ``"len"`` and - ``"vdiff"`` attributes. All original edge attributes from the - primitive graph are preserved. + ``"vdiff"`` attributes. All original node, edge and graph attributes + from the primitive graph are preserved. + + Raises + ------ + ValueError + If ``pattern`` does not match the adjacency types present in the + primitive, if ``pattern`` is given for an edge-less primitive, if + ``method`` is given for a primitive that already has adjacency + edges, or if the node positions are degenerate (e.g. all collinear) + so that no triangulation exists. + NotImplementedError + If an unknown ``method`` is requested. """ - if pattern == "4-star": + if G_undirected.number_of_edges() == 0 and G_undirected.number_of_nodes() > 0: + return _create_directed_mesh_graph_from_node_cloud( + G_undirected, pattern=pattern, method=method + ) + + if method is not None: + raise ValueError( + f"method='{method}' was given, but the mesh primitive already " + "has adjacency edges (created by the mesh layout). The 'method' " + "argument only applies to edge-less primitives (node clouds " + "from mesh_layout='prebuilt')." + ) + + if pattern is None: + # Use every edge the layout produced + edges_to_use = list(G_undirected.edges(data=True)) + elif 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" ] + if len(edges_to_use) == 0: + available = sorted( + { + str(d.get("adjacency_type")) + for _, _, d in G_undirected.edges(data=True) + } + ) + raise ValueError( + "pattern='4-star' selects edges with " + "adjacency_type='cardinal', but the mesh primitive has no " + f"such edges (available adjacency types: {available}). " + "Omit 'pattern' to use every edge the layout produced." + ) 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'." + "Choose '4-star', '8-star', or omit 'pattern' to use every " + "edge the layout produced." ) # Create filtered undirected graph with only selected edges (preserving attrs) @@ -79,3 +138,70 @@ def create_directed_mesh_graph( dg.graph.update(G_undirected.graph) return dg + + +def _create_directed_mesh_graph_from_node_cloud( + G_nodes: networkx.Graph, pattern: str = None, method: str = None +) -> networkx.DiGraph: + """Build a directed mesh graph directly from an edge-less node cloud. + + The directed edges are constructed straight from the node positions + (no intermediate undirected adjacency graph is built). + """ + # A single-node primitive has no edges under any semantics; don't reject + # a 'pattern' that a generated-layout code path may have passed along. + if pattern is not None and G_nodes.number_of_nodes() > 1: + raise ValueError( + f"pattern='{pattern}' was given, but the mesh primitive has no " + "adjacency edges to select from (it is a node cloud from " + "mesh_layout='prebuilt'). Use method='delaunay' (the default) " + "to control how edges are constructed from the node positions." + ) + if method is None: + method = "delaunay" + if method != "delaunay": + raise NotImplementedError( + f"method='{method}' is not implemented for building mesh edges " + "from node positions. Currently supported: 'delaunay'." + ) + + dg = networkx.DiGraph() + dg.add_nodes_from(G_nodes.nodes(data=True)) + dg.graph.update(G_nodes.graph) + + nodes = list(G_nodes.nodes) + positions = np.array( + [np.asarray(G_nodes.nodes[n]["pos"], dtype=float) for n in nodes] + ) + n_nodes = len(nodes) + + # Delaunay triangulation needs >= 3 non-collinear points; smaller node + # clouds get the only sensible connectivity directly. + if n_nodes == 1: + return dg + if n_nodes == 2: + undirected_pairs = {(0, 1)} + else: + try: + triangulation = scipy.spatial.Delaunay(positions) + except scipy.spatial.QhullError as exc: + raise ValueError( + "Delaunay triangulation of the prebuilt mesh nodes failed " + f"({n_nodes} nodes). This typically means the node positions " + "are degenerate (e.g. all collinear). Provide at least 3 " + "non-collinear mesh node positions." + ) from exc + undirected_pairs = set() + for simplex in triangulation.simplices: + for i in range(3): + a, b = int(simplex[i]), int(simplex[(i + 1) % 3]) + undirected_pairs.add((min(a, b), max(a, b))) + + for ia, ib in sorted(undirected_pairs): + u, v = nodes[ia], nodes[ib] + vdiff = positions[ia] - positions[ib] + d = float(np.sqrt(np.sum(vdiff**2))) + dg.add_edge(u, v, len=d, vdiff=vdiff) + dg.add_edge(v, u, len=d, vdiff=-vdiff) + + return dg diff --git a/src/weather_model_graphs/create/mesh/connectivity/hierarchical.py b/src/weather_model_graphs/create/mesh/connectivity/hierarchical.py index ae78da1..bdb0571 100644 --- a/src/weather_model_graphs/create/mesh/connectivity/hierarchical.py +++ b/src/weather_model_graphs/create/mesh/connectivity/hierarchical.py @@ -5,14 +5,14 @@ import scipy from ....networkx_utils import prepend_node_index -from .. import coords as mesh_coords +from ..layout import rectilinear as mesh_layout from .general 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}, + intra_level: Optional[Dict[str, object]] = None, + inter_level: Optional[Dict[str, object]] = None, ) -> networkx.DiGraph: """ Create a hierarchical multiscale mesh graph from a list of mesh primitive @@ -23,10 +23,14 @@ def create_hierarchical_from_coordinates( 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) + Intra-level connectivity is controlled by ``intra_level``: + - For primitives with adjacency edges (generated layouts), the optional + ``intra_level["pattern"]`` selects a subset of them (``"4-star"``: + only cardinal neighbours; ``"8-star"``: cardinal plus diagonal). When + no pattern is given, every edge the layout produced is used. + - For edge-less primitives (``mesh_layout="prebuilt"`` node clouds), + ``intra_level["method"]`` selects how edges are built from the node + positions per level (currently only ``"delaunay"``, the default). Parameters ---------- @@ -34,13 +38,18 @@ def create_hierarchical_from_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"``) - Created by ``create_multirange_2d_mesh_primitives``. - intra_level : dict + - Edge attributes (when the primitive has edges): ``"adjacency_type"`` + (str, ``"cardinal"`` or ``"diagonal"``) + Created by ``create_multirange_2d_mesh_primitives`` (generated + layouts) or ``create_multi_level_prebuilt_mesh_primitives`` + (prebuilt). + intra_level : dict, optional Configuration for intra-level connectivity. Keys: - - ``"pattern"`` (str): ``"4-star"`` or ``"8-star"``. - Default: ``{"pattern": "8-star"}`` - inter_level : dict + - ``"pattern"`` (str): ``"4-star"`` or ``"8-star"`` (primitives with + adjacency edges only; default: use every edge). + - ``"method"`` (str): edge construction method for edge-less + primitives (default: ``"delaunay"``). + 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. @@ -53,7 +62,12 @@ def create_hierarchical_from_coordinates( edges (direction="same"), inter-level down edges (direction="down"), and inter-level up edges (direction="up"). """ - intra_level_pattern = intra_level.get("pattern", "8-star") + if intra_level is None: + intra_level = {} + if inter_level is None: + inter_level = {} + intra_level_pattern = intra_level.get("pattern") + intra_level_method = intra_level.get("method") inter_level_pattern = inter_level.get("pattern", "nearest") inter_level_k = inter_level.get("k", 1) @@ -63,9 +77,12 @@ def create_hierarchical_from_coordinates( "for hierarchical graphs. Only 'nearest' is currently implemented." ) - # Convert each level's coordinate graph to directed graph with chosen pattern + # Convert each level's coordinate graph to directed graph with chosen + # pattern (adjacency-edge primitives) or method (edge-less primitives) Gs_all_levels = [ - create_directed_mesh_graph(g_coords, pattern=intra_level_pattern) + create_directed_mesh_graph( + g_coords, pattern=intra_level_pattern, method=intra_level_method + ) for g_coords in G_coords_list ] @@ -207,7 +224,7 @@ def create_hierarchical_multiscale_mesh_graph( A directed graph containing the hierarchical mesh with intra-level, up, and down edges. """ - G_coords_list = mesh_coords.create_multirange_2d_mesh_primitives( + G_coords_list = mesh_layout.create_multirange_2d_mesh_primitives( max_num_levels=max_num_levels, xy=xy, mesh_node_spacing=mesh_node_distance, diff --git a/src/weather_model_graphs/create/mesh/coords.py b/src/weather_model_graphs/create/mesh/coords.py index 40cec0b..7b4ad7f 100644 --- a/src/weather_model_graphs/create/mesh/coords.py +++ b/src/weather_model_graphs/create/mesh/coords.py @@ -1,298 +1,23 @@ -from typing import List +""" +Backward-compatibility re-exports from ``layout.rectilinear``. -import networkx -import numpy as np -from loguru import logger +All coordinate creation functions have been moved to +``wmg.create.mesh.layout.rectilinear``. This module re-exports them +so that existing imports continue to work. +""" from .connectivity.general import create_directed_mesh_graph - - -def create_single_level_2d_mesh_primitive( - xy: np.ndarray, - nx: int = None, - 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. - - 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 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 - 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, 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 - ------- - networkx.Graph - 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]) - - # avoid nodes on border - dx = (xM - xm) / nx - dy = (yM - ym) / ny - lx = np.linspace(xm + dx / 2, xM - dx / 2, nx) - ly = np.linspace(ym + dy / 2, yM - dy / 2, ny) - - mg = np.meshgrid(lx, ly) - g = networkx.grid_2d_graph(len(lx), len(ly)) - - # Node name and `pos` attribute takes form (x, y) - for node in g.nodes: - node_xi, node_yi = node # Extract x and y index from node to index mx - g.nodes[node]["pos"] = np.array( - [mg[0][node_yi, node_xi], mg[1][node_yi, node_xi]] - ) - g.nodes[node]["type"] = "mesh" - - # Mark existing grid_2d_graph edges as cardinal (4-star adjacency) - for u, v in g.edges(): - 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_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 - 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_primitive (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_primitive(xy, nx, ny) - return create_directed_mesh_graph(G_coords, pattern="8-star") - - -def create_multirange_2d_mesh_primitives( - max_num_levels: int, - 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 - xy coordinates. - - 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. - - 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_spacing : float - Distance (in x- and y-direction) between created mesh nodes, - in coordinate system of xy - 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 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. - """ - # 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 / 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 - # 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 = 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 - # direction, create mesh levels corresponding to the minimum of the two - mesh_levels_to_create = max_mesh_levels.min() - - if max_num_levels: - # Limit the levels in mesh graph - mesh_levels_to_create = min(mesh_levels_to_create, max_num_levels) - - logger.debug(f"mesh_levels: {mesh_levels_to_create}, nleaf: {nleaf}") - - # multi resolution tree levels - 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) - 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 - 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: int, - xy: np.ndarray, - 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. - 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_primitives (coordinate creation) - 2. create_directed_mesh_graph (connectivity creation) - - 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 - pattern : str - Connectivity pattern for directed graph creation: ``"4-star"`` or - ``"8-star"`` (default: ``"8-star"``) - - Returns - ------- - G_all_levels : list of networkx.DiGraph - List of networkx graphs for each level representing the connectivity - of the mesh within each level - """ - G_coords_list = create_multirange_2d_mesh_primitives( - max_num_levels=max_num_levels, - xy=xy, - mesh_node_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=pattern) - G_all_levels.append(g_directed) - - return G_all_levels +from .layout.rectilinear import ( + create_multirange_2d_mesh_graphs, + create_multirange_2d_mesh_primitives, + create_single_level_2d_mesh_graph, + create_single_level_2d_mesh_primitive, +) + +__all__ = [ + "create_directed_mesh_graph", + "create_multirange_2d_mesh_graphs", + "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/layout/__init__.py b/src/weather_model_graphs/create/mesh/layout/__init__.py new file mode 100644 index 0000000..4f07a55 --- /dev/null +++ b/src/weather_model_graphs/create/mesh/layout/__init__.py @@ -0,0 +1,18 @@ +""" +Mesh layout modules. + +Each layout module defines how mesh node coordinates are placed in space +(coordinate creation step). The resulting undirected primitive graphs are +then consumed by the connectivity modules to produce directed mesh graphs. + +Available layouts: + +- ``rectilinear``: nodes placed on a uniform rectangular grid. +- ``triangular``: nodes placed on a regular (equilateral) triangular lattice. +- ``prebuilt``: nodes taken from a user-provided mesh graph (edge-less node + clouds; mesh edges are built in the connectivity step). +""" + +from . import prebuilt, rectilinear, triangular + +__all__ = ["prebuilt", "rectilinear", "triangular"] diff --git a/src/weather_model_graphs/create/mesh/layout/prebuilt.py b/src/weather_model_graphs/create/mesh/layout/prebuilt.py new file mode 100644 index 0000000..e6aba36 --- /dev/null +++ b/src/weather_model_graphs/create/mesh/layout/prebuilt.py @@ -0,0 +1,308 @@ +""" +Prebuilt mesh layout: coordinate creation from user-provided mesh nodes. + +Unlike the generated layouts (``rectilinear``, ``triangular``), the prebuilt +layout does not place mesh nodes itself -- the user supplies their own node +positions (e.g. ICON grid vertices, an observation-station network, or any +custom point set) and the library builds the graph around them. + +This is the coordinate creation step in the two-step mesh creation process: + +1. **Coordinate creation** (this module) -> edge-less ``nx.Graph`` (a "node + cloud") with ``pos`` and ``type`` node attributes. No adjacency edges are + created here: how a point cloud gets connected is a *connectivity* + decision, so edge construction happens in the connectivity step (see + ``create_directed_mesh_graph`` in ``connectivity.general``, which builds + directed edges directly from the node positions with + ``method="delaunay"``). +2. **Connectivity creation** -> ``nx.DiGraph`` with ``len`` and ``vdiff`` + edge attributes. + +Currently only *nodes-only* input is supported: the input graph must not +contain any edges. Support for user-provided edges (using them as the mesh +adjacency) is planned as a follow-up -- see the design discussion in +https://github.com/mllam/weather-model-graphs/issues/79. + +The user input contract: + +- an ``nx.Graph`` (or edge-less ``nx.DiGraph``) whose nodes carry: + + - ``pos``: ``np.ndarray`` of shape ``(2,)`` -- the node position, **in the + same coordinate system as the grid coordinates** passed to + ``create_all_graph_components`` + - ``type``: ``str``, must be ``"mesh"`` + - ``level``: ``int``, only for hierarchical meshes (lowest value = finest + level); must be present on either all nodes or none + +- or, for convenience, a bare ``np.ndarray`` of shape ``[N, 2]`` with node + positions (a nodes-only, single-level mesh). +""" + +from typing import List, Union + +import networkx +import numpy as np +import scipy.spatial +from loguru import logger + + +def validate_prebuilt_mesh_nodes( + mesh_graph: networkx.Graph, require_levels: bool = False +) -> None: + """ + Validate that a user-provided mesh graph satisfies the prebuilt nodes-only + input contract. + + Every node must have a ``pos`` attribute (``np.ndarray`` of shape ``(2,)`` + with finite values) and a ``type`` attribute equal to ``"mesh"``. Node + positions must be unique. The graph must not contain any edges + (user-provided edges are not yet supported, see issue #79). + + Parameters + ---------- + mesh_graph : networkx.Graph + User-provided graph to validate. + require_levels : bool + If True, additionally require an integer ``level`` attribute on every + node, with at least two distinct level values (needed for hierarchical + meshes). + + Raises + ------ + ValueError + If the graph is empty, a node attribute is missing or malformed, + positions are duplicated, or level attributes are inconsistent. + NotImplementedError + If the graph contains edges (nodes+edges input is not yet supported). + """ + if mesh_graph.number_of_nodes() == 0: + raise ValueError( + "mesh_layout='prebuilt' requires a mesh_graph with at least one node." + ) + + if mesh_graph.number_of_edges() > 0: + raise NotImplementedError( + "mesh_layout='prebuilt' currently only supports nodes-only input, " + f"but the given mesh_graph has {mesh_graph.number_of_edges()} " + "edge(s). Mesh connectivity is built in the connectivity step " + "(method='delaunay' by default). Support for user-provided edges " + "is planned -- see " + "https://github.com/mllam/weather-model-graphs/issues/79." + ) + + n_with_level = 0 + positions = [] + for node, data in mesh_graph.nodes(data=True): + if "pos" not in data: + raise ValueError( + f"Node {node!r} is missing the required 'pos' attribute. All " + "nodes in a prebuilt mesh must have 'pos' as an np.ndarray of " + "shape (2,)." + ) + pos = np.asarray(data["pos"]) + if pos.shape != (2,): + raise ValueError( + f"Node {node!r} has 'pos' with shape {pos.shape}, expected " + "(2,). All nodes in a prebuilt mesh must have 'pos' as an " + "np.ndarray of shape (2,)." + ) + if not np.all(np.isfinite(pos.astype(float))): + raise ValueError( + f"Node {node!r} has a non-finite 'pos' value ({pos}). Node " + "positions must be finite numbers." + ) + if data.get("type") != "mesh": + raise ValueError( + f"Node {node!r} has type={data.get('type')!r}, expected " + "'mesh'. All nodes in a prebuilt mesh must have the 'type' " + "attribute set to 'mesh' (the 'grid' node type is reserved " + "for the grid nodes created from the `coords` argument)." + ) + if "level" in data: + n_with_level += 1 + if not isinstance(data["level"], (int, np.integer)): + raise ValueError( + f"Node {node!r} has a non-integer 'level' attribute " + f"({data['level']!r}). Mesh levels must be integers " + "(lowest value = finest level)." + ) + positions.append(pos.astype(float)) + + n_nodes = mesh_graph.number_of_nodes() + if 0 < n_with_level < n_nodes: + raise ValueError( + f"Only {n_with_level} of {n_nodes} nodes have a 'level' " + "attribute. For a hierarchical prebuilt mesh every node must " + "have a 'level'; for a flat mesh no node should have one." + ) + + positions_arr = np.stack(positions) + n_unique = np.unique(positions_arr, axis=0).shape[0] + if n_unique < n_nodes: + raise ValueError( + f"The mesh_graph contains duplicate node positions ({n_nodes} " + f"nodes but only {n_unique} unique positions). Duplicate " + "positions would produce zero-length mesh edges." + ) + + if require_levels: + if n_with_level == 0: + raise ValueError( + "Hierarchical prebuilt meshes require an integer 'level' " + "attribute on every node (lowest value = finest level), but " + "no node has one." + ) + levels = {int(data["level"]) for _, data in mesh_graph.nodes(data=True)} + if len(levels) < 2: + raise ValueError( + "At least two distinct mesh levels are required for a " + f"hierarchical prebuilt mesh, but only level(s) " + f"{sorted(levels)} were found." + ) + + +def _as_node_cloud_graph( + mesh_graph: Union[networkx.Graph, np.ndarray] +) -> networkx.Graph: + """Normalize prebuilt-mesh input to an undirected node-cloud graph. + + Accepts either a graph (undirected or directed -- direction is the + library's to assign, so an edge-less DiGraph is treated as its + undirected node set) or a bare ``[N, 2]`` coordinate array. + """ + if isinstance(mesh_graph, np.ndarray): + xy = np.asarray(mesh_graph, dtype=float) + if xy.ndim != 2 or xy.shape[1] != 2: + raise ValueError( + "A prebuilt mesh given as an array must have shape " + f"[N_mesh_nodes, 2], got {xy.shape}." + ) + g = networkx.Graph() + for i, pos in enumerate(xy): + g.add_node(i, pos=pos, type="mesh") + return g + if isinstance(mesh_graph, networkx.Graph): # includes DiGraph + return networkx.Graph(mesh_graph) + raise TypeError( + "mesh_graph must be a networkx.Graph (or edge-less DiGraph) or an " + f"np.ndarray of shape [N, 2], got {type(mesh_graph).__name__}." + ) + + +def _estimate_node_spacing(positions: np.ndarray) -> float: + """Median nearest-neighbour distance -- the characteristic node spacing. + + Used to fill the ``dx``/``dy`` graph attributes that generated layouts + derive from their lattice spacing (needed e.g. by the hierarchical + connectivity step and relative-distance grid connection methods). + """ + if positions.shape[0] < 2: + return 0.0 + kdt = scipy.spatial.KDTree(positions) + # k=2: the nearest neighbour that isn't the node itself + dists, _ = kdt.query(positions, k=2) + return float(np.median(dists[:, 1])) + + +def _node_cloud_primitive(g_cloud: networkx.Graph) -> networkx.Graph: + """Build one edge-less mesh primitive from a validated node cloud. + + Node labels are replaced by ``(i,)`` integer tuples (insertion order) so + they sort against the grid node labels and support the level-index + prepending used by hierarchical connectivity. Only the contract + attributes (``pos``, ``type``) are carried over. + """ + g = networkx.Graph() + positions = [] + for i, (_, data) in enumerate(g_cloud.nodes(data=True)): + pos = np.asarray(data["pos"], dtype=float) + g.add_node((i,), pos=pos, type="mesh") + positions.append(pos) + spacing = _estimate_node_spacing(np.stack(positions)) + g.graph["dx"] = spacing + g.graph["dy"] = spacing + return g + + +def create_single_level_prebuilt_mesh_primitive( + mesh_graph: Union[networkx.Graph, np.ndarray] +) -> networkx.Graph: + """ + Create a single-level mesh primitive from user-provided mesh nodes. + + This is the coordinate creation step for ``mesh_layout="prebuilt"`` with + flat connectivity. The result is an *edge-less* undirected graph (a node + cloud): mesh adjacency for a point cloud is built in the connectivity + step (``method="delaunay"`` by default). + + Parameters + ---------- + mesh_graph : networkx.Graph or np.ndarray + User-provided mesh nodes (see the module docstring for the input + contract). If nodes carry a ``level`` attribute it is ignored (with + a warning) -- use ``m2m_connectivity="hierarchical"`` to build a + hierarchical mesh from the levels. + + Returns + ------- + networkx.Graph + Edge-less mesh primitive. Node attributes: ``pos`` + (np.ndarray of shape ``(2,)``), ``type`` (``"mesh"``). Graph + attributes: ``dx``, ``dy`` (median nearest-neighbour node spacing). + """ + g_cloud = _as_node_cloud_graph(mesh_graph) + validate_prebuilt_mesh_nodes(g_cloud) + if any("level" in d for _, d in g_cloud.nodes(data=True)): + logger.warning( + "The prebuilt mesh_graph nodes carry 'level' attributes but a " + "single-level (flat) mesh was requested; the levels are ignored. " + "Use m2m_connectivity='hierarchical' to build a hierarchical " + "mesh from them." + ) + return _node_cloud_primitive(g_cloud) + + +def create_multi_level_prebuilt_mesh_primitives( + mesh_graph: Union[networkx.Graph, np.ndarray] +) -> List[networkx.Graph]: + """ + Create per-level mesh primitives from user-provided mesh nodes with + ``level`` attributes. + + This is the coordinate creation step for ``mesh_layout="prebuilt"`` with + hierarchical connectivity. Nodes are split by their integer ``level`` + attribute (lowest value = finest level) into one *edge-less* primitive + per level; intra-level adjacency is built per level in the connectivity + step (``intra_level=dict(method="delaunay")`` by default) and inter-level + up/down edges by nearest-neighbour search (``inter_level``). + + Parameters + ---------- + mesh_graph : networkx.Graph + User-provided mesh nodes with ``pos``, ``type`` and ``level`` + attributes on every node (see the module docstring). + + Returns + ------- + list[networkx.Graph] + Edge-less mesh primitives, one per level, ordered finest first. + Each carries the graph attributes ``level`` (0-based level index), + ``dx`` and ``dy`` (median nearest-neighbour spacing of that level). + """ + g_cloud = _as_node_cloud_graph(mesh_graph) + validate_prebuilt_mesh_nodes(g_cloud, require_levels=True) + + user_levels = sorted({int(d["level"]) for _, d in g_cloud.nodes(data=True)}) + + primitives = [] + for level_index, user_level in enumerate(user_levels): + level_nodes = [ + n for n, d in g_cloud.nodes(data=True) if int(d["level"]) == user_level + ] + g_level = _node_cloud_primitive(g_cloud.subgraph(level_nodes)) + for node in g_level.nodes: + g_level.nodes[node]["level"] = level_index + g_level.graph["level"] = level_index + primitives.append(g_level) + + return primitives diff --git a/src/weather_model_graphs/create/mesh/layout/rectilinear.py b/src/weather_model_graphs/create/mesh/layout/rectilinear.py new file mode 100644 index 0000000..eb42755 --- /dev/null +++ b/src/weather_model_graphs/create/mesh/layout/rectilinear.py @@ -0,0 +1,313 @@ +""" +Rectilinear mesh layout: coordinate creation for uniform rectangular grids. + +Uses ``networkx.grid_2d_graph`` to produce a rectilinear lattice with +4-star (cardinal) and 8-star (cardinal + diagonal) spatial adjacency edges. + +This is the coordinate creation step in the two-step mesh creation process: + +1. **Coordinate creation** (this module) -> ``nx.Graph`` with ``pos``, ``type``, + and ``adjacency_type`` attributes. +2. **Connectivity creation** (``create_directed_mesh_graph`` in + ``connectivity.general``) -> ``nx.DiGraph`` with ``len`` and ``vdiff`` + edge attributes. +""" + +from typing import List + +import networkx +import numpy as np +from loguru import logger + +from ..connectivity.general import create_directed_mesh_graph + + +def create_single_level_2d_mesh_primitive( + xy: np.ndarray, + nx: int = None, + 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. + + 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 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 + 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, 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 + ------- + networkx.Graph + 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]) + + # avoid nodes on border + dx = (xM - xm) / nx + dy = (yM - ym) / ny + lx = np.linspace(xm + dx / 2, xM - dx / 2, nx) + ly = np.linspace(ym + dy / 2, yM - dy / 2, ny) + + mg = np.meshgrid(lx, ly) + g = networkx.grid_2d_graph(len(lx), len(ly)) + + # Node name and `pos` attribute takes form (x, y) + for node in g.nodes: + node_xi, node_yi = node # Extract x and y index from node to index mx + g.nodes[node]["pos"] = np.array( + [mg[0][node_yi, node_xi], mg[1][node_yi, node_xi]] + ) + g.nodes[node]["type"] = "mesh" + + # Mark existing grid_2d_graph edges as cardinal (4-star adjacency) + for u, v in g.edges(): + 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_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 + 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_primitive (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_primitive(xy, nx, ny) + return create_directed_mesh_graph(G_coords, pattern="8-star") + + +def create_multirange_2d_mesh_primitives( + max_num_levels: int, + 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 + xy coordinates. + + 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. + + 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_spacing : float + Distance (in x- and y-direction) between created mesh nodes, + in coordinate system of xy + 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 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. + """ + # 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 / 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 + # 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 = 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 + # direction, create mesh levels corresponding to the minimum of the two + mesh_levels_to_create = max_mesh_levels.min() + + if max_num_levels: + # Limit the levels in mesh graph + mesh_levels_to_create = min(mesh_levels_to_create, max_num_levels) + + logger.debug(f"mesh_levels: {mesh_levels_to_create}, nleaf: {nleaf}") + + # multi resolution tree levels + 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) + 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 + 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: int, + xy: np.ndarray, + 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. + 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_primitives (coordinate creation) + 2. create_directed_mesh_graph (connectivity creation) + + 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 + pattern : str + Connectivity pattern for directed graph creation: ``"4-star"`` or + ``"8-star"`` (default: ``"8-star"``) + + Returns + ------- + G_all_levels : list of networkx.DiGraph + List of networkx graphs for each level representing the connectivity + of the mesh within each level + """ + G_coords_list = create_multirange_2d_mesh_primitives( + max_num_levels=max_num_levels, + xy=xy, + mesh_node_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=pattern) + G_all_levels.append(g_directed) + + return G_all_levels diff --git a/src/weather_model_graphs/create/mesh/layout/triangular.py b/src/weather_model_graphs/create/mesh/layout/triangular.py new file mode 100644 index 0000000..9042cae --- /dev/null +++ b/src/weather_model_graphs/create/mesh/layout/triangular.py @@ -0,0 +1,217 @@ +""" +Triangular mesh layout: coordinate creation for regular triangular lattices. + +Uses ``networkx.triangular_lattice_graph`` to produce an equilateral-triangle +lattice with 6-connectivity (each interior node has 6 neighbours). This +mirrors the rectilinear layout (which uses ``networkx.grid_2d_graph`` +with 8-connectivity) and plugs into the same two-step process: + +1. **Coordinate creation** (this module) -> ``nx.Graph`` with ``pos``, ``type``, + and ``adjacency_type`` attributes. +2. **Connectivity creation** (``create_directed_mesh_graph`` in + ``connectivity.general``) -> ``nx.DiGraph`` with ``len`` and ``vdiff`` + edge attributes. +""" + +from typing import List + +import networkx +import numpy as np +from loguru import logger + + +def create_single_level_2d_mesh_primitive( + xy: np.ndarray, + nx: int = None, + ny: int = None, + *, + mesh_node_spacing: float = None, +) -> networkx.Graph: + """ + Create an undirected triangular mesh primitive graph (``nx.Graph``) with + node positions and spatial adjacency edges. + + This is analogous to ``create_single_level_2d_mesh_primitive`` in the + rectilinear layout but uses ``networkx.triangular_lattice_graph`` instead + of ``grid_2d_graph``. + + In a triangular lattice, each interior node has 6 neighbours (vs. 8 for + the rectilinear lattice with diagonals), providing more isotropic message + passing. + + The nodes form a grid of ``(ny + 1)`` rows and ``(nx + 1) // 2`` columns, + with odd-row nodes shifted horizontally. Positions are scaled and offset + so that the mesh spans the coordinate domain of *xy* (with nodes inset + from the border by half a cell width in each direction). + + Either provide ``nx`` and ``ny`` directly, or provide ``mesh_node_spacing`` + to have them computed automatically from the coordinate extent of ``xy`` + (mirroring ``create_single_level_2d_mesh_primitive``). + + Parameters + ---------- + xy : np.ndarray + Grid point coordinates, shaped ``[N_grid_points, 2]``. + nx : int, optional + Number of triangle columns (passed as *n* to + ``triangular_lattice_graph``). If not given, computed from + ``mesh_node_spacing``. + ny : int, optional + Number of triangle rows (passed as *m* to + ``triangular_lattice_graph``). 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 from the coordinate extent of ``xy`` + (``ny`` accounts for the ``sqrt(3)/2`` triangular row spacing) and + validated to be > 0. + + Returns + ------- + networkx.Graph + Undirected mesh primitive graph. Node attributes: ``pos`` + (np.ndarray[2,]), ``type`` (``"mesh"``). Edge attributes: + ``adjacency_type`` (always ``"cardinal"`` -- triangular lattices have + only one class of edge). Graph attributes: ``dx``, ``dy``. + """ + 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 * np.sqrt(3) / 2)) + 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]) + + # Create the raw triangular lattice + g_raw = networkx.triangular_lattice_graph(ny, nx, with_positions=True) + + if g_raw.number_of_nodes() == 0: + raise ValueError( + f"triangular_lattice_graph({ny}, {nx}) produced 0 nodes. " + "Increase nx/ny or decrease mesh_node_spacing." + ) + + # Gather raw positions to compute extent + raw_positions = np.array([g_raw.nodes[n]["pos"] for n in g_raw.nodes()]) + raw_xmin, raw_ymin = raw_positions.min(axis=0) + raw_xmax, raw_ymax = raw_positions.max(axis=0) + raw_extent_x = raw_xmax - raw_xmin + raw_extent_y = raw_ymax - raw_ymin + + # Domain extent with half-cell inset + domain_x = xM - xm + domain_y = yM - ym + + # Scale factors -- map raw lattice extent to domain extent (inset by half + # a cell in each direction, mirroring the rectilinear approach) + if raw_extent_x > 0: + scale_x = domain_x / (raw_extent_x + 1.0) # +1 for inset + else: + scale_x = domain_x # single column + if raw_extent_y > 0: + scale_y = domain_y / (raw_extent_y + np.sqrt(3) / 2) # +row_h for inset + else: + scale_y = domain_y # single row + + # Effective dx/dy for graph attributes + dx = scale_x + dy = scale_y * (np.sqrt(3) / 2) + + # Offset so mesh is centred within domain + offset_x = xm + (domain_x - raw_extent_x * scale_x) / 2 + offset_y = ym + (domain_y - raw_extent_y * scale_y) / 2 + + # Build output graph with scaled positions + g = networkx.Graph() + for node in g_raw.nodes(): + raw_pos = g_raw.nodes[node]["pos"] + pos = np.array( + [ + offset_x + (raw_pos[0] - raw_xmin) * scale_x, + offset_y + (raw_pos[1] - raw_ymin) * scale_y, + ] + ) + g.add_node(node, pos=pos, type="mesh") + + for u, v in g_raw.edges(): + g.add_edge(u, v, adjacency_type="cardinal") + + g.graph["dx"] = dx + g.graph["dy"] = dy + + return g + + +def create_multirange_2d_mesh_primitives( + max_num_levels, + xy: np.ndarray, + mesh_node_spacing: float = 3, + interlevel_refinement_factor: int = 3, +) -> List[networkx.Graph]: + """ + Create a list of undirected triangular mesh primitive graphs representing + different levels of mesh resolution. + + Mirrors ``create_multirange_2d_mesh_primitives`` in the rectilinear layout + but uses triangular lattice topology at each level. + + Parameters + ---------- + max_num_levels : int + Maximum number of levels in the multi-scale graph. + xy : np.ndarray + Grid point coordinates, shaped ``[N_grid_points, 2]``. + mesh_node_spacing : float + Distance between mesh nodes at the finest level, in coordinate units. + interlevel_refinement_factor : int + Factor by which mesh node count decreases per level. + + Returns + ------- + list[networkx.Graph] + Triangular mesh primitive graphs, one per level. + """ + coord_extent = np.ptp(xy, axis=0) + # For triangular lattice, ny accounts for row spacing of sqrt(3)/2 + max_nx = int(coord_extent[0] / mesh_node_spacing) + max_ny = int(coord_extent[1] / (mesh_node_spacing * np.sqrt(3) / 2)) + + max_nodes_bottom = np.array([max_nx, max_ny]) + + max_mesh_levels_float = np.log(max_nodes_bottom) / np.log( + interlevel_refinement_factor + ) + max_mesh_levels = max_mesh_levels_float.astype(int) + nleaf = interlevel_refinement_factor**max_mesh_levels + + mesh_levels_to_create = max_mesh_levels.min() + if max_num_levels: + mesh_levels_to_create = min(mesh_levels_to_create, max_num_levels) + + logger.debug(f"triangular mesh_levels: {mesh_levels_to_create}, nleaf: {nleaf}") + + G_all_levels = [] + for lev in range(mesh_levels_to_create): + nodes_x, nodes_y = (nleaf / (interlevel_refinement_factor**lev)).astype(int) + g = create_single_level_2d_mesh_primitive(xy, nodes_x, nodes_y) + for node in g.nodes: + g.nodes[node]["level"] = lev + for edge in g.edges: + g.edges[edge]["level"] = lev + g.graph["level"] = lev + g.graph["interlevel_refinement_factor"] = interlevel_refinement_factor + G_all_levels.append(g) + + return G_all_levels diff --git a/tests/test_mesh_layout.py b/tests/test_mesh_layout.py index b289fcc..fbeb805 100644 --- a/tests/test_mesh_layout.py +++ b/tests/test_mesh_layout.py @@ -34,7 +34,7 @@ from weather_model_graphs.create.mesh.connectivity.hierarchical import ( create_hierarchical_from_coordinates, ) -from weather_model_graphs.create.mesh.coords import ( +from weather_model_graphs.create.mesh.layout.rectilinear import ( create_multirange_2d_mesh_primitives, create_single_level_2d_mesh_primitive, ) @@ -617,7 +617,7 @@ def test_unsupported_mesh_layout_raises(self): wmg.create.create_all_graph_components( coords=xy, m2m_connectivity="flat", - mesh_layout="triangular", + mesh_layout="nonexistent_layout", mesh_layout_kwargs=dict(mesh_node_spacing=3), g2m_connectivity="nearest_neighbour", m2g_connectivity="nearest_neighbour", diff --git a/tests/test_prebuilt_mesh.py b/tests/test_prebuilt_mesh.py new file mode 100644 index 0000000..9de3bf7 --- /dev/null +++ b/tests/test_prebuilt_mesh.py @@ -0,0 +1,464 @@ +""" +Tests for mesh_layout="prebuilt" support (Issue #79). + +Tests verify: +1. Input validation (nodes-only contract: pos/type/level attributes, + duplicate positions, edges rejected) +2. Single- and multi-level primitive creation (edge-less node clouds, + tuple relabelling, dx/dy spacing estimate, level splitting) +3. Directed mesh graph construction from node clouds + (method="delaunay": Delaunay edges, bidirectional len/vdiff, + small/degenerate node clouds, pattern-vs-method argument validation) +4. Integration through create_all_graph_components for flat and + hierarchical connectivity (including the np.ndarray convenience input) +5. Generated-layout behaviour is unchanged (pattern default equivalence) +""" + +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.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.layout.prebuilt import ( + create_multi_level_prebuilt_mesh_primitives, + create_single_level_prebuilt_mesh_primitive, + validate_prebuilt_mesh_nodes, +) + +# =========================== +# Fixtures +# =========================== + + +@pytest.fixture +def xy_grid(): + """Grid point coordinates covering [0, 10]^2.""" + return test_utils.create_fake_xy(N=10) * 10 / 10 + + +@pytest.fixture +def mesh_xy(): + """Irregular mesh node positions inside the grid domain.""" + rng = np.random.default_rng(seed=7) + return rng.random((25, 2)) * 10 + + +def _nodes_only_graph(positions, level=None, label=lambda i: i): + """Build a nodes-only mesh graph from an [N, 2] position array.""" + g = nx.Graph() + for i, pos in enumerate(positions): + attrs = dict(pos=np.asarray(pos, dtype=float), type="mesh") + if level is not None: + attrs["level"] = level + g.add_node(label(i), **attrs) + return g + + +@pytest.fixture +def mesh_graph(mesh_xy): + """Nodes-only mesh graph with string labels.""" + return _nodes_only_graph(mesh_xy, label=lambda i: f"station_{i}") + + +@pytest.fixture +def mesh_graph_two_levels(mesh_xy): + """Nodes-only mesh graph with two levels (1 = fine, 2 = coarse).""" + rng = np.random.default_rng(seed=11) + g = _nodes_only_graph(mesh_xy, level=1, label=lambda i: ("f", i)) + for i, pos in enumerate(rng.random((6, 2)) * 10): + g.add_node(("c", i), pos=pos, type="mesh", level=2) + return g + + +# =========================== +# 1. Input validation +# =========================== + + +class TestValidatePrebuiltMeshNodes: + def test_valid_nodes_pass(self, mesh_graph): + validate_prebuilt_mesh_nodes(mesh_graph) + + def test_empty_graph_raises(self): + with pytest.raises(ValueError, match="at least one node"): + validate_prebuilt_mesh_nodes(nx.Graph()) + + def test_missing_pos_raises(self): + g = nx.Graph() + g.add_node(0, type="mesh") + with pytest.raises(ValueError, match="missing the required 'pos'"): + validate_prebuilt_mesh_nodes(g) + + def test_wrong_pos_shape_raises(self): + g = nx.Graph() + g.add_node(0, pos=np.array([1.0, 2.0, 3.0]), type="mesh") + with pytest.raises(ValueError, match="shape"): + validate_prebuilt_mesh_nodes(g) + + def test_non_finite_pos_raises(self): + g = nx.Graph() + g.add_node(0, pos=np.array([np.nan, 0.0]), type="mesh") + with pytest.raises(ValueError, match="non-finite"): + validate_prebuilt_mesh_nodes(g) + + def test_missing_type_raises(self): + g = nx.Graph() + g.add_node(0, pos=np.array([0.0, 0.0])) + with pytest.raises(ValueError, match="type"): + validate_prebuilt_mesh_nodes(g) + + def test_wrong_type_value_raises(self): + g = nx.Graph() + g.add_node(0, pos=np.array([0.0, 0.0]), type="grid") + with pytest.raises(ValueError, match="expected 'mesh'"): + validate_prebuilt_mesh_nodes(g) + + def test_duplicate_positions_raise(self): + g = _nodes_only_graph([[0.0, 0.0], [1.0, 1.0], [0.0, 0.0]]) + with pytest.raises(ValueError, match="duplicate node positions"): + validate_prebuilt_mesh_nodes(g) + + def test_edges_not_yet_supported(self, mesh_graph): + mesh_graph.add_edge("station_0", "station_1") + with pytest.raises(NotImplementedError, match="nodes-only"): + validate_prebuilt_mesh_nodes(mesh_graph) + + def test_mixed_level_presence_raises(self, mesh_xy): + g = _nodes_only_graph(mesh_xy) + g.nodes[0]["level"] = 1 + with pytest.raises(ValueError, match="'level'"): + validate_prebuilt_mesh_nodes(g) + + def test_non_integer_level_raises(self, mesh_xy): + g = _nodes_only_graph(mesh_xy, level=1) + g.nodes[0]["level"] = "fine" + with pytest.raises(ValueError, match="non-integer 'level'"): + validate_prebuilt_mesh_nodes(g) + + def test_require_levels_missing_raises(self, mesh_graph): + with pytest.raises(ValueError, match="no node has one"): + validate_prebuilt_mesh_nodes(mesh_graph, require_levels=True) + + def test_require_levels_single_level_raises(self, mesh_xy): + g = _nodes_only_graph(mesh_xy, level=1) + with pytest.raises(ValueError, match="two distinct mesh levels"): + validate_prebuilt_mesh_nodes(g, require_levels=True) + + +# =========================== +# 2. Primitive creation (coordinate creation step) +# =========================== + + +class TestSingleLevelPrimitive: + def test_is_edge_less_node_cloud(self, mesh_graph, mesh_xy): + g = create_single_level_prebuilt_mesh_primitive(mesh_graph) + assert g.number_of_nodes() == len(mesh_xy) + assert g.number_of_edges() == 0 + + def test_nodes_relabelled_to_tuples(self, mesh_graph): + g = create_single_level_prebuilt_mesh_primitive(mesh_graph) + assert all(isinstance(n, tuple) for n in g.nodes) + # tuple labels must sort against the (level_id, i) grid node labels + assert sorted(g.nodes) == list(g.nodes) + + def test_positions_preserved(self, mesh_graph, mesh_xy): + g = create_single_level_prebuilt_mesh_primitive(mesh_graph) + positions = np.stack([g.nodes[n]["pos"] for n in g.nodes]) + assert np.allclose(np.sort(positions, axis=0), np.sort(mesh_xy, axis=0)) + + def test_spacing_estimate_set(self, mesh_graph): + g = create_single_level_prebuilt_mesh_primitive(mesh_graph) + assert g.graph["dx"] > 0 + assert g.graph["dx"] == g.graph["dy"] + + def test_ndarray_input(self, mesh_xy): + g = create_single_level_prebuilt_mesh_primitive(mesh_xy) + assert g.number_of_nodes() == len(mesh_xy) + assert g.number_of_edges() == 0 + + def test_bad_ndarray_shape_raises(self): + with pytest.raises(ValueError, match=r"\[N_mesh_nodes, 2\]"): + create_single_level_prebuilt_mesh_primitive(np.zeros((3, 4))) + + def test_bad_input_type_raises(self): + with pytest.raises(TypeError, match="mesh_graph must be"): + create_single_level_prebuilt_mesh_primitive([[0, 0], [1, 1]]) + + def test_edge_less_digraph_accepted(self, mesh_xy): + dg = nx.DiGraph() + for i, pos in enumerate(mesh_xy): + dg.add_node(i, pos=pos, type="mesh") + g = create_single_level_prebuilt_mesh_primitive(dg) + assert g.number_of_nodes() == len(mesh_xy) + + +class TestMultiLevelPrimitives: + def test_split_by_level_finest_first(self, mesh_graph_two_levels, mesh_xy): + primitives = create_multi_level_prebuilt_mesh_primitives(mesh_graph_two_levels) + assert len(primitives) == 2 + # level 1 (finest, 25 nodes) must come first as level index 0 + assert primitives[0].number_of_nodes() == len(mesh_xy) + assert primitives[0].graph["level"] == 0 + assert primitives[1].number_of_nodes() == 6 + assert primitives[1].graph["level"] == 1 + + def test_primitives_are_edge_less(self, mesh_graph_two_levels): + primitives = create_multi_level_prebuilt_mesh_primitives(mesh_graph_two_levels) + assert all(g.number_of_edges() == 0 for g in primitives) + + def test_level_values_need_not_be_contiguous(self, mesh_xy): + g = _nodes_only_graph(mesh_xy[:10], level=3) + for i, pos in enumerate(mesh_xy[10:15]): + g.add_node(("coarse", i), pos=pos, type="mesh", level=7) + primitives = create_multi_level_prebuilt_mesh_primitives(g) + assert [p.graph["level"] for p in primitives] == [0, 1] + assert primitives[0].number_of_nodes() == 10 + + def test_per_level_spacing_estimates(self, mesh_graph_two_levels): + primitives = create_multi_level_prebuilt_mesh_primitives(mesh_graph_two_levels) + assert all(g.graph["dx"] > 0 for g in primitives) + + +# =========================== +# 3. Directed mesh graph from node clouds (connectivity step) +# =========================== + + +class TestNodeCloudDirectedGraph: + def test_delaunay_bidirectional_len_vdiff(self, mesh_graph): + g_prim = create_single_level_prebuilt_mesh_primitive(mesh_graph) + dg = create_directed_mesh_graph(g_prim) + assert isinstance(dg, nx.DiGraph) + assert dg.number_of_edges() > 0 + for u, v, d in dg.edges(data=True): + assert dg.has_edge(v, u) + assert d["len"] > 0 + assert np.allclose(d["vdiff"], -dg.edges[v, u]["vdiff"]) + assert np.isclose(d["len"], np.linalg.norm(d["vdiff"])) + + def test_delaunay_edges_match_scipy(self, mesh_xy): + import scipy.spatial + + g_prim = create_single_level_prebuilt_mesh_primitive(mesh_xy) + dg = create_directed_mesh_graph(g_prim) + tri = scipy.spatial.Delaunay(mesh_xy) + expected_pairs = set() + for simplex in tri.simplices: + for i in range(3): + a, b = sorted((simplex[i], simplex[(i + 1) % 3])) + expected_pairs.add((a, b)) + assert dg.number_of_edges() == 2 * len(expected_pairs) + + def test_single_node_no_edges(self): + g_prim = create_single_level_prebuilt_mesh_primitive(np.array([[1.0, 2.0]])) + dg = create_directed_mesh_graph(g_prim) + assert dg.number_of_nodes() == 1 + assert dg.number_of_edges() == 0 + + def test_two_nodes_bidirectional_pair(self): + g_prim = create_single_level_prebuilt_mesh_primitive( + np.array([[0.0, 0.0], [3.0, 4.0]]) + ) + dg = create_directed_mesh_graph(g_prim) + assert dg.number_of_edges() == 2 + (d,) = [d for _, _, d in dg.edges(data=True) if d["vdiff"][0] < 0] + assert np.isclose(d["len"], 5.0) + + def test_three_nodes_triangle(self): + g_prim = create_single_level_prebuilt_mesh_primitive( + np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]]) + ) + dg = create_directed_mesh_graph(g_prim) + assert dg.number_of_edges() == 6 + + def test_collinear_nodes_raise(self): + g_prim = create_single_level_prebuilt_mesh_primitive( + np.array([[0.0, 0.0], [1.0, 0.0], [2.0, 0.0], [3.0, 0.0]]) + ) + with pytest.raises(ValueError, match="collinear"): + create_directed_mesh_graph(g_prim) + + def test_pattern_on_node_cloud_raises(self, mesh_graph): + g_prim = create_single_level_prebuilt_mesh_primitive(mesh_graph) + with pytest.raises(ValueError, match="method='delaunay'"): + create_directed_mesh_graph(g_prim, pattern="8-star") + + def test_unknown_method_raises(self, mesh_graph): + g_prim = create_single_level_prebuilt_mesh_primitive(mesh_graph) + with pytest.raises(NotImplementedError, match="'delaunay'"): + create_directed_mesh_graph(g_prim, method="knn") + + def test_method_on_lattice_primitive_raises(self, xy_grid): + from weather_model_graphs.create.mesh.layout.rectilinear import ( + create_single_level_2d_mesh_primitive, + ) + + g_prim = create_single_level_2d_mesh_primitive(xy_grid, nx=4, ny=4) + with pytest.raises(ValueError, match="already"): + create_directed_mesh_graph(g_prim, method="delaunay") + + +# =========================== +# 4. Integration through create_all_graph_components +# =========================== + + +class TestFlatEndToEnd: + def test_flat_components(self, xy_grid, mesh_graph, mesh_xy): + components = wmg.create.create_all_graph_components( + coords=xy_grid, + mesh_layout="prebuilt", + mesh_layout_kwargs=dict(mesh_graph=mesh_graph), + m2m_connectivity="flat", + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + return_components=True, + ) + assert set(components.keys()) == {"m2m", "g2m", "m2g"} + m2m = components["m2m"] + mesh_nodes = [n for n, d in m2m.nodes(data=True) if d.get("type") == "mesh"] + assert len(mesh_nodes) == len(mesh_xy) + assert m2m.number_of_edges() > 0 + assert components["g2m"].number_of_edges() == len(mesh_xy) + + def test_explicit_delaunay_method_matches_default(self, xy_grid, mesh_graph): + kwargs = dict( + coords=xy_grid, + mesh_layout="prebuilt", + mesh_layout_kwargs=dict(mesh_graph=mesh_graph), + m2m_connectivity="flat", + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + return_components=True, + ) + default = wmg.create.create_all_graph_components(**kwargs) + explicit = wmg.create.create_all_graph_components( + m2m_connectivity_kwargs=dict(method="delaunay"), **kwargs + ) + assert default["m2m"].number_of_edges() == explicit["m2m"].number_of_edges() + + def test_merged_single_graph(self, xy_grid, mesh_xy): + graph = wmg.create.create_all_graph_components( + coords=xy_grid, + mesh_layout="prebuilt", + mesh_layout_kwargs=dict(mesh_graph=mesh_xy), + m2m_connectivity="flat", + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + assert graph.number_of_nodes() == len(xy_grid) + len(mesh_xy) + + def test_missing_mesh_graph_raises(self, xy_grid): + with pytest.raises(ValueError, match="mesh_graph"): + wmg.create.create_all_graph_components( + coords=xy_grid, + mesh_layout="prebuilt", + m2m_connectivity="flat", + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + + def test_flat_multiscale_not_supported(self, xy_grid, mesh_graph): + with pytest.raises(NotImplementedError, match="flat_multiscale"): + wmg.create.create_all_graph_components( + coords=xy_grid, + mesh_layout="prebuilt", + mesh_layout_kwargs=dict(mesh_graph=mesh_graph), + m2m_connectivity="flat_multiscale", + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + + +class TestHierarchicalEndToEnd: + def test_hierarchical_components(self, xy_grid, mesh_graph_two_levels, mesh_xy): + components = wmg.create.create_all_graph_components( + coords=xy_grid, + mesh_layout="prebuilt", + mesh_layout_kwargs=dict(mesh_graph=mesh_graph_two_levels), + m2m_connectivity="hierarchical", + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + return_components=True, + ) + m2m = components["m2m"] + directions = { + d["direction"] for _, _, d in m2m.edges(data=True) if "direction" in d + } + assert directions == {"same", "up", "down"} + n_up = sum(1 for _, _, d in m2m.edges(data=True) if d.get("direction") == "up") + # nearest with k=1: one up edge per fine node + assert n_up == len(mesh_xy) + # the grid connects only to the finest level + assert components["g2m"].number_of_edges() == len(mesh_xy) + + def test_intra_level_method_kwarg(self, xy_grid, mesh_graph_two_levels): + components = wmg.create.create_all_graph_components( + coords=xy_grid, + mesh_layout="prebuilt", + mesh_layout_kwargs=dict(mesh_graph=mesh_graph_two_levels), + m2m_connectivity="hierarchical", + m2m_connectivity_kwargs=dict( + intra_level=dict(method="delaunay"), + inter_level=dict(pattern="nearest", k=2), + ), + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + return_components=True, + ) + n_up = sum( + 1 + for _, _, d in components["m2m"].edges(data=True) + if d.get("direction") == "up" + ) + # k=2: two up edges per fine node + assert n_up == 2 * mesh_graph_two_levels.number_of_nodes() - 2 * 6 + + def test_hierarchical_without_levels_raises(self, xy_grid, mesh_graph): + with pytest.raises(ValueError, match="level"): + wmg.create.create_all_graph_components( + coords=xy_grid, + mesh_layout="prebuilt", + mesh_layout_kwargs=dict(mesh_graph=mesh_graph), + m2m_connectivity="hierarchical", + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + ) + + def test_direct_hierarchical_from_prebuilt_primitives(self, mesh_graph_two_levels): + primitives = create_multi_level_prebuilt_mesh_primitives(mesh_graph_two_levels) + m2m = create_hierarchical_from_coordinates(primitives) + assert m2m.number_of_edges() > 0 + assert set(m2m.graph["dx"].keys()) == {0, 1} + + +# =========================== +# 5. Generated layouts unchanged (pattern default equivalence) +# =========================== + + +class TestGeneratedLayoutsUnchanged: + @pytest.mark.parametrize("mesh_layout", ["rectilinear", "triangular"]) + def test_no_pattern_equals_8_star(self, xy_grid, mesh_layout): + kwargs = dict( + coords=xy_grid, + mesh_layout=mesh_layout, + mesh_layout_kwargs=dict(mesh_node_spacing=2), + m2m_connectivity="flat", + g2m_connectivity="nearest_neighbour", + m2g_connectivity="nearest_neighbour", + return_components=True, + ) + default = wmg.create.create_all_graph_components(**kwargs) + explicit = wmg.create.create_all_graph_components( + m2m_connectivity_kwargs=dict(pattern="8-star"), **kwargs + ) + assert default["m2m"].number_of_edges() == explicit["m2m"].number_of_edges() diff --git a/tests/test_triangular_mesh.py b/tests/test_triangular_mesh.py new file mode 100644 index 0000000..c0587f3 --- /dev/null +++ b/tests/test_triangular_mesh.py @@ -0,0 +1,1000 @@ +""" +Tests for mesh_layout="triangular" support (Issue #80). + +Tests verify: +1. Primitive creation (node count, positions, adjacency_type, type attrs) +2. Single-level directed graph (bidirectional edges, len/vdiff attrs, 6-connectivity) +3. Multirange primitive creation +4. Flat multiscale mesh graph (position-based merging via two-step API) +5. Hierarchical mesh graph (triangular primitives + generic hierarchical connectivity) +6. Integration through create_all_graph_components for all m2m_connectivity modes +7. Edge cases (zero nodes, single-level hierarchical) +8. Numerical correctness (len symmetry, vdiff reciprocity) +""" + +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.connectivity.flat import ( + create_flat_multiscale_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.layout.triangular import ( + create_multirange_2d_mesh_primitives as create_multirange_2d_triangular_mesh_primitives, +) +from weather_model_graphs.create.mesh.layout.triangular import ( + create_single_level_2d_mesh_primitive as create_single_level_2d_triangular_mesh_primitive, +) + +# =========================== +# Fixtures +# =========================== + + +@pytest.fixture +def xy_small(): + """Small 10x10 domain with 4 corner grid points.""" + return np.array([[0, 0], [10, 0], [0, 10], [10, 10]], dtype=float) + + +@pytest.fixture +def xy_medium(): + """Medium domain with many grid points.""" + return test_utils.create_fake_xy(N=20) + + +@pytest.fixture +def xy_rectangular(): + """Non-square domain.""" + return test_utils.create_rectangular_fake_xy(Nx=15, Ny=10) + + +@pytest.fixture +def xy_offset(): + """Domain not starting at origin.""" + return np.array([[5, 3], [15, 3], [5, 13], [15, 13]], dtype=float) + + +@pytest.fixture +def xy_large(): + """Larger domain with many grid points.""" + return test_utils.create_fake_xy(N=50) + + +@pytest.fixture +def xy_wide(): + """Very wide, short domain.""" + return test_utils.create_rectangular_fake_xy(Nx=30, Ny=5) + + +# =========================== +# Step 1: Triangular Primitive (Coordinate Creation) +# =========================== + + +class TestTriangularPrimitive: + """Tests for create_single_level_2d_triangular_mesh_primitive.""" + + def test_returns_undirected_graph(self, xy_small): + G = create_single_level_2d_triangular_mesh_primitive(xy_small, nx=5, ny=5) + assert isinstance(G, nx.Graph) + assert not isinstance(G, nx.DiGraph) + + def test_nodes_have_pos_and_type(self, xy_small): + G = create_single_level_2d_triangular_mesh_primitive(xy_small, 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" + pos = G.nodes[node]["pos"] + assert len(pos) == 2 + assert np.isfinite(pos).all() + + def test_nonzero_node_count(self, xy_small): + G = create_single_level_2d_triangular_mesh_primitive(xy_small, nx=6, ny=6) + assert G.number_of_nodes() > 0 + + def test_has_edges(self, xy_small): + G = create_single_level_2d_triangular_mesh_primitive(xy_small, nx=6, ny=6) + assert G.number_of_edges() > 0 + + def test_all_edges_are_cardinal(self, xy_small): + """Triangular lattice has only cardinal edges (no diagonal distinction).""" + G = create_single_level_2d_triangular_mesh_primitive(xy_small, nx=5, ny=5) + 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"] == "cardinal" + + def test_graph_has_dx_dy(self, xy_small): + G = create_single_level_2d_triangular_mesh_primitive(xy_small, nx=5, ny=5) + assert "dx" in G.graph + assert "dy" in G.graph + assert G.graph["dx"] > 0 + assert G.graph["dy"] > 0 + + def test_positions_within_domain(self, xy_small): + """Mesh node positions should lie within the coordinate domain.""" + xm, xM = xy_small[:, 0].min(), xy_small[:, 0].max() + ym, yM = xy_small[:, 1].min(), xy_small[:, 1].max() + G = create_single_level_2d_triangular_mesh_primitive(xy_small, nx=6, ny=6) + for node in G.nodes: + pos = G.nodes[node]["pos"] + assert xm <= pos[0] <= xM, f"x={pos[0]} out of [{xm}, {xM}]" + assert ym <= pos[1] <= yM, f"y={pos[1]} out of [{ym}, {yM}]" + + def test_raises_on_zero_nodes(self): + """nx=0 or ny=0 should produce 0 nodes and raise.""" + xy = np.array([[0, 0], [1, 0], [0, 1], [1, 1]], dtype=float) + with pytest.raises(ValueError, match="produced 0 nodes"): + create_single_level_2d_triangular_mesh_primitive(xy, nx=0, ny=0) + + def test_rectangular_domain(self, xy_rectangular): + """Works with non-square domains.""" + G = create_single_level_2d_triangular_mesh_primitive(xy_rectangular, nx=8, ny=5) + assert G.number_of_nodes() > 0 + assert G.number_of_edges() > 0 + + def test_minimal_lattice(self, xy_small): + """Smallest valid lattice (nx=1, ny=1) should produce nodes and edges.""" + G = create_single_level_2d_triangular_mesh_primitive(xy_small, nx=1, ny=1) + assert G.number_of_nodes() >= 2 + assert G.number_of_edges() >= 1 + + def test_large_lattice(self, xy_small): + """Large nx/ny values should produce many nodes.""" + G = create_single_level_2d_triangular_mesh_primitive(xy_small, nx=20, ny=20) + assert G.number_of_nodes() > 100 + + def test_asymmetric_nx_ny(self, xy_small): + """Very different nx and ny should still work.""" + G = create_single_level_2d_triangular_mesh_primitive(xy_small, nx=15, ny=3) + assert G.number_of_nodes() > 0 + assert G.number_of_edges() > 0 + + def test_offset_domain(self, xy_offset): + """Domain not starting at origin: positions should still be within bounds.""" + G = create_single_level_2d_triangular_mesh_primitive(xy_offset, nx=5, ny=5) + xm, xM = 5.0, 15.0 + ym, yM = 3.0, 13.0 + for node in G.nodes: + pos = G.nodes[node]["pos"] + assert xm <= pos[0] <= xM + assert ym <= pos[1] <= yM + + def test_positions_are_numpy_arrays(self, xy_small): + """Node positions should be numpy arrays.""" + G = create_single_level_2d_triangular_mesh_primitive(xy_small, nx=4, ny=4) + for node in G.nodes: + assert isinstance(G.nodes[node]["pos"], np.ndarray) + + def test_no_self_loops(self, xy_small): + """Primitive graph should have no self-loops.""" + G = create_single_level_2d_triangular_mesh_primitive(xy_small, nx=6, ny=6) + for u, v in G.edges(): + assert u != v + + def test_wide_domain(self, xy_wide): + """Very wide, short domain should still produce valid mesh.""" + G = create_single_level_2d_triangular_mesh_primitive(xy_wide, nx=10, ny=3) + assert G.number_of_nodes() > 0 + for node in G.nodes: + pos = G.nodes[node]["pos"] + assert np.isfinite(pos).all() + + +# =========================== +# Step 2: Directed Mesh Graph (Connectivity Creation) +# =========================== + + +class TestTriangularDirectedGraph: + """Tests for directed graph creation from triangular primitives.""" + + def test_returns_digraph(self, xy_small): + G_coords = create_single_level_2d_triangular_mesh_primitive( + xy_small, nx=5, ny=5 + ) + G = create_directed_mesh_graph(G_coords, pattern="4-star") + assert isinstance(G, nx.DiGraph) + + def test_edges_are_bidirectional(self, xy_small): + G_coords = create_single_level_2d_triangular_mesh_primitive( + xy_small, nx=5, ny=5 + ) + G = create_directed_mesh_graph(G_coords, pattern="4-star") + for u, v in G.edges(): + assert G.has_edge(v, u), f"Edge ({u}, {v}) missing reverse" + + def test_edges_have_len_and_vdiff(self, xy_small): + G_coords = create_single_level_2d_triangular_mesh_primitive( + xy_small, nx=5, ny=5 + ) + G = create_directed_mesh_graph(G_coords, pattern="4-star") + for u, v, d in G.edges(data=True): + assert "len" in d, f"Edge ({u}, {v}) missing 'len'" + assert "vdiff" in d, f"Edge ({u}, {v}) missing 'vdiff'" + assert d["len"] > 0 + assert len(d["vdiff"]) == 2 + + def test_len_symmetry(self, xy_small): + """Edge length should be the same in both directions.""" + G_coords = create_single_level_2d_triangular_mesh_primitive( + xy_small, nx=5, ny=5 + ) + G = create_directed_mesh_graph(G_coords, pattern="4-star") + for u, v in G.edges(): + if G.has_edge(v, u): + np.testing.assert_allclose(G[u][v]["len"], G[v][u]["len"], atol=1e-10) + + def test_vdiff_reciprocity(self, xy_small): + """vdiff(u->v) should equal -vdiff(v->u).""" + G_coords = create_single_level_2d_triangular_mesh_primitive( + xy_small, nx=5, ny=5 + ) + G = create_directed_mesh_graph(G_coords, pattern="4-star") + for u, v in G.edges(): + if G.has_edge(v, u): + np.testing.assert_allclose( + G[u][v]["vdiff"], -G[v][u]["vdiff"], atol=1e-10 + ) + + def test_node_count_preserved(self, xy_small): + """Directed graph should have same number of nodes as primitive.""" + G_coords = create_single_level_2d_triangular_mesh_primitive( + xy_small, nx=5, ny=5 + ) + G = create_directed_mesh_graph(G_coords, pattern="4-star") + assert G.number_of_nodes() == G_coords.number_of_nodes() + + def test_edge_count_is_twice_undirected(self, xy_small): + """Directed graph should have exactly 2x the undirected edge count.""" + G_coords = create_single_level_2d_triangular_mesh_primitive( + xy_small, nx=5, ny=5 + ) + G = create_directed_mesh_graph(G_coords, pattern="4-star") + assert G.number_of_edges() == 2 * G_coords.number_of_edges() + + def test_no_self_loops_directed(self, xy_small): + """Directed graph should have no self-loops.""" + G_coords = create_single_level_2d_triangular_mesh_primitive( + xy_small, nx=5, ny=5 + ) + G = create_directed_mesh_graph(G_coords, pattern="4-star") + for u, v in G.edges(): + assert u != v + + def test_pos_preserved_after_direction(self, xy_small): + """Node positions should be preserved after converting to directed.""" + G_coords = create_single_level_2d_triangular_mesh_primitive( + xy_small, nx=5, ny=5 + ) + G = create_directed_mesh_graph(G_coords, pattern="4-star") + for node in G.nodes: + assert "pos" in G.nodes[node] + assert len(G.nodes[node]["pos"]) == 2 + + def test_interior_node_degree_six(self, xy_small): + """Interior nodes of a triangular lattice should have degree 6 + (6 in-edges + 6 out-edges = 12 total in directed graph).""" + G_coords = create_single_level_2d_triangular_mesh_primitive( + xy_small, nx=8, ny=8 + ) + G = create_directed_mesh_graph(G_coords, pattern="4-star") + # At least one interior node should have degree 12 (6 in + 6 out) + max_deg = max(dict(G.degree()).values()) + assert max_deg == 12 + + def test_minimal_lattice_directed(self, xy_small): + """Minimal lattice (nx=1, ny=1) should still produce a valid directed graph.""" + G = create_directed_mesh_graph( + create_single_level_2d_triangular_mesh_primitive(xy_small, nx=1, ny=1) + ) + assert isinstance(G, nx.DiGraph) + assert G.number_of_nodes() >= 2 + assert G.number_of_edges() >= 2 # at least one bidirectional edge + + +# =========================== +# Multirange primitives +# =========================== + + +class TestMultirangeTriangularPrimitives: + """Tests for create_multirange_2d_triangular_mesh_primitives.""" + + def test_returns_list(self, xy_medium): + G_list = create_multirange_2d_triangular_mesh_primitives( + max_num_levels=3, + xy=xy_medium, + mesh_node_spacing=2, + interlevel_refinement_factor=3, + ) + assert isinstance(G_list, list) + assert len(G_list) >= 1 + + def test_each_level_is_undirected(self, xy_medium): + G_list = create_multirange_2d_triangular_mesh_primitives( + max_num_levels=3, + xy=xy_medium, + mesh_node_spacing=2, + interlevel_refinement_factor=3, + ) + for G in G_list: + assert isinstance(G, nx.Graph) + assert not isinstance(G, nx.DiGraph) + + def test_level_attributes_set(self, xy_medium): + G_list = create_multirange_2d_triangular_mesh_primitives( + max_num_levels=3, + xy=xy_medium, + mesh_node_spacing=2, + interlevel_refinement_factor=3, + ) + for lev, G in enumerate(G_list): + assert G.graph["level"] == lev + for node in G.nodes: + assert G.nodes[node]["level"] == lev + for u, v in G.edges(): + assert G.edges[u, v]["level"] == lev + + def test_finer_level_has_more_nodes(self, xy_medium): + G_list = create_multirange_2d_triangular_mesh_primitives( + max_num_levels=3, + xy=xy_medium, + mesh_node_spacing=2, + interlevel_refinement_factor=3, + ) + if len(G_list) > 1: + assert G_list[0].number_of_nodes() > G_list[1].number_of_nodes() + + def test_max_num_levels_respected(self, xy_medium): + G_list = create_multirange_2d_triangular_mesh_primitives( + max_num_levels=2, + xy=xy_medium, + mesh_node_spacing=2, + interlevel_refinement_factor=3, + ) + assert len(G_list) <= 2 + + def test_single_level(self, xy_medium): + """max_num_levels=1 should produce exactly 1 level.""" + G_list = create_multirange_2d_triangular_mesh_primitives( + max_num_levels=1, + xy=xy_medium, + mesh_node_spacing=2, + interlevel_refinement_factor=3, + ) + assert len(G_list) == 1 + assert G_list[0].graph["level"] == 0 + + def test_refinement_factor_2(self, xy_medium): + """Different refinement factor should still produce valid graphs.""" + G_list = create_multirange_2d_triangular_mesh_primitives( + max_num_levels=3, + xy=xy_medium, + mesh_node_spacing=2, + interlevel_refinement_factor=2, + ) + assert len(G_list) >= 1 + for G in G_list: + assert G.number_of_nodes() > 0 + + def test_all_levels_cover_same_domain(self, xy_medium): + """All levels should span approximately the same coordinate domain.""" + G_list = create_multirange_2d_triangular_mesh_primitives( + max_num_levels=3, + xy=xy_medium, + mesh_node_spacing=2, + interlevel_refinement_factor=3, + ) + if len(G_list) < 2: + pytest.skip("Only one level created") + # Check centers are roughly the same across levels + centers = [] + for G in G_list: + positions = np.array([G.nodes[n]["pos"] for n in G.nodes]) + centers.append(positions.mean(axis=0)) + for c in centers[1:]: + np.testing.assert_allclose(c, centers[0], atol=2.0) + + def test_interlevel_refinement_factor_preserved(self, xy_medium): + """Each level should have the refinement factor as a graph attribute.""" + G_list = create_multirange_2d_triangular_mesh_primitives( + max_num_levels=3, + xy=xy_medium, + mesh_node_spacing=2, + interlevel_refinement_factor=3, + ) + for G in G_list: + assert G.graph["interlevel_refinement_factor"] == 3 + + def test_all_levels_have_edges(self, xy_medium): + """Every level should have at least some edges.""" + G_list = create_multirange_2d_triangular_mesh_primitives( + max_num_levels=3, + xy=xy_medium, + mesh_node_spacing=2, + interlevel_refinement_factor=3, + ) + for G in G_list: + assert G.number_of_edges() > 0 + + +# =========================== +# Convenience wrapper: single-level +# =========================== + + +class TestSingleLevelTriangularGraph: + """Tests for the triangular single-level directed mesh graph.""" + + def test_returns_digraph(self, xy_small): + G = create_directed_mesh_graph( + create_single_level_2d_triangular_mesh_primitive(xy_small, nx=5, ny=5) + ) + assert isinstance(G, nx.DiGraph) + + def test_has_bidirectional_edges(self, xy_small): + G = create_directed_mesh_graph( + create_single_level_2d_triangular_mesh_primitive(xy_small, nx=5, ny=5) + ) + for u, v in G.edges(): + assert G.has_edge(v, u) + + def test_edges_have_attributes(self, xy_small): + """Directed graph edges should have len and vdiff.""" + G = create_directed_mesh_graph( + create_single_level_2d_triangular_mesh_primitive(xy_small, nx=5, ny=5) + ) + for u, v, d in G.edges(data=True): + assert "len" in d + assert "vdiff" in d + + def test_with_rectangular_domain(self, xy_rectangular): + """Should work correctly on non-square domains.""" + G = create_directed_mesh_graph( + create_single_level_2d_triangular_mesh_primitive(xy_rectangular, nx=8, ny=5) + ) + assert isinstance(G, nx.DiGraph) + assert G.number_of_edges() > 0 + + def test_minimal_grid(self, xy_small): + """Minimal grid (nx=1, ny=1) should produce a valid graph.""" + G = create_directed_mesh_graph( + create_single_level_2d_triangular_mesh_primitive(xy_small, nx=1, ny=1) + ) + assert G.number_of_nodes() >= 2 + + +# =========================== +# Flat multiscale (triangular-specific merging) +# =========================== + + +class TestFlatMultiscaleTriangular: + """Tests for flat multiscale triangular mesh graph using two-step API. + + Uses ``create_multirange_2d_triangular_mesh_primitives`` (coordinate + creation) followed by ``create_flat_multiscale_from_coordinates`` + (connectivity creation with position-based merging). + """ + + def _create_multiscale( + self, + xy, + mesh_node_spacing=2.0, + interlevel_refinement_factor=3, + max_num_levels=3, + ): + """Helper: create primitives then apply position-based merging.""" + G_coords_list = create_multirange_2d_triangular_mesh_primitives( + max_num_levels=max_num_levels, + xy=xy, + mesh_node_spacing=mesh_node_spacing, + interlevel_refinement_factor=interlevel_refinement_factor, + ) + return create_flat_multiscale_from_coordinates(G_coords_list) + + def test_returns_digraph(self, xy_medium): + G = self._create_multiscale(xy_medium) + assert isinstance(G, nx.DiGraph) + + def test_has_edges(self, xy_medium): + G = self._create_multiscale(xy_medium) + assert G.number_of_edges() > 0 + + def test_edges_have_len_and_vdiff(self, xy_medium): + G = self._create_multiscale(xy_medium) + for u, v, d in G.edges(data=True): + assert "len" in d + assert "vdiff" in d + + def test_fewer_nodes_than_sum_of_levels(self, xy_medium): + """Position-based merging should produce fewer nodes than the raw + sum of all levels (coincident nodes get merged).""" + G_coords_list = create_multirange_2d_triangular_mesh_primitives( + max_num_levels=3, + xy=xy_medium, + mesh_node_spacing=2, + interlevel_refinement_factor=3, + ) + total_raw = sum(g.number_of_nodes() for g in G_coords_list) + G = create_flat_multiscale_from_coordinates(G_coords_list) + # Merged graph has at most as many nodes (usually fewer) + assert G.number_of_nodes() <= total_raw + + def test_graph_has_dx_dy_dicts(self, xy_medium): + G = self._create_multiscale(xy_medium) + assert isinstance(G.graph.get("dx"), dict) + assert isinstance(G.graph.get("dy"), dict) + + def test_bidirectional_edges(self, xy_medium): + """All edges should have a reverse.""" + G = self._create_multiscale(xy_medium) + for u, v in G.edges(): + assert G.has_edge(v, u), f"Edge ({u},{v}) no reverse" + + def test_nodes_have_pos(self, xy_medium): + """All nodes should have pos attribute.""" + G = self._create_multiscale(xy_medium) + for node in G.nodes: + assert "pos" in G.nodes[node] + assert np.isfinite(G.nodes[node]["pos"]).all() + + def test_single_level_multiscale(self): + """When domain only supports 1 level, flat_multiscale should still work.""" + xy = np.array([[0, 0], [3, 0], [0, 3], [3, 3]], dtype=float) + G = self._create_multiscale(xy, mesh_node_spacing=1.0) + assert isinstance(G, nx.DiGraph) + assert G.number_of_nodes() > 0 + + def test_refinement_factor_2(self, xy_medium): + """Refinement factor of 2 should work.""" + G = self._create_multiscale(xy_medium, interlevel_refinement_factor=2) + assert isinstance(G, nx.DiGraph) + assert G.number_of_edges() > 0 + + def test_no_self_loops(self, xy_medium): + """No self-loops in flat multiscale graph.""" + G = self._create_multiscale(xy_medium) + for u, v in G.edges(): + assert u != v + + def test_more_nodes_than_coarsest_level(self, xy_large): + """Multiscale should have more nodes than the coarsest single level.""" + G_coords_list = create_multirange_2d_triangular_mesh_primitives( + max_num_levels=3, + xy=xy_large, + mesh_node_spacing=2, + interlevel_refinement_factor=3, + ) + if len(G_coords_list) < 2: + pytest.skip("Only one level created") + coarsest_nodes = G_coords_list[-1].number_of_nodes() + G_multi = create_flat_multiscale_from_coordinates(G_coords_list) + assert G_multi.number_of_nodes() > coarsest_nodes + + +# =========================== +# Hierarchical +# =========================== + + +class TestHierarchicalTriangular: + """Tests for hierarchical mesh graph from triangular primitives. + + Uses ``create_multirange_2d_triangular_mesh_primitives`` (coordinate + creation) followed by ``create_hierarchical_from_coordinates`` (generic + hierarchical connectivity creation). + """ + + def _create_hierarchical( + self, + xy, + mesh_node_spacing=2.0, + interlevel_refinement_factor=3, + max_num_levels=3, + **kwargs, + ): + """Helper: create primitives then apply hierarchical connectivity.""" + G_coords_list = create_multirange_2d_triangular_mesh_primitives( + max_num_levels=max_num_levels, + xy=xy, + mesh_node_spacing=mesh_node_spacing, + interlevel_refinement_factor=interlevel_refinement_factor, + ) + return create_hierarchical_from_coordinates(G_coords_list, **kwargs) + + def test_returns_digraph(self, xy_medium): + G = self._create_hierarchical(xy_medium) + assert isinstance(G, nx.DiGraph) + + def test_has_edges(self, xy_medium): + G = self._create_hierarchical(xy_medium) + assert G.number_of_edges() > 0 + + def test_edges_have_level_attribute(self, xy_medium): + G = self._create_hierarchical(xy_medium) + for u, v, d in G.edges(data=True): + # Intra-level edges have 'level', inter-level have 'levels' + assert "level" in d or "levels" in d + + def test_multiple_levels_present(self, xy_medium): + G = self._create_hierarchical(xy_medium) + levels = set() + for u, v, d in G.edges(data=True): + if "level" in d: + levels.add(d["level"]) + elif "levels" in d: + # Inter-level edges like '0>1' + parts = d["levels"].split(">") + levels.update(int(p) for p in parts) + assert len(levels) >= 2, "Expected multiple levels in hierarchical graph" + + def test_single_level_raises(self, xy_small): + """Hierarchical requires >= 2 levels; single level should raise.""" + G_coords_list = create_multirange_2d_triangular_mesh_primitives( + max_num_levels=1, + xy=xy_small, + mesh_node_spacing=2.0, + interlevel_refinement_factor=3, + ) + with pytest.raises(ValueError): + create_hierarchical_from_coordinates(G_coords_list) + + def test_nodes_have_pos(self, xy_medium): + """All nodes should have pos attribute.""" + G = self._create_hierarchical(xy_medium) + for node in G.nodes: + assert "pos" in G.nodes[node] + assert np.isfinite(G.nodes[node]["pos"]).all() + + def test_custom_intra_level(self, xy_medium): + """Custom intra_level pattern should be accepted.""" + G = self._create_hierarchical(xy_medium, intra_level={"pattern": "8-star"}) + assert isinstance(G, nx.DiGraph) + assert G.number_of_edges() > 0 + + def test_custom_inter_level(self, xy_medium): + """Custom inter_level config should be accepted.""" + G = self._create_hierarchical( + xy_medium, inter_level={"pattern": "nearest", "k": 3} + ) + assert isinstance(G, nx.DiGraph) + assert G.number_of_edges() > 0 + + def test_no_self_loops(self, xy_medium): + """Hierarchical graph should have no self-loops.""" + G = self._create_hierarchical(xy_medium) + for u, v in G.edges(): + assert u != v + + def test_has_inter_level_edges(self, xy_medium): + """Should have inter-level edges connecting different levels.""" + G = self._create_hierarchical(xy_medium) + inter_level_count = sum(1 for _, _, d in G.edges(data=True) if "levels" in d) + assert inter_level_count > 0 + + def test_inter_level_edges_have_direction(self, xy_medium): + """Inter-level edges should have 'direction' attribute (up/down).""" + G = self._create_hierarchical(xy_medium) + for u, v, d in G.edges(data=True): + if "levels" in d: + assert "direction" in d + assert d["direction"] in ("up", "down") + + +# =========================== +# Integration: create_all_graph_components +# =========================== + + +class TestIntegrationTriangular: + """Full integration tests through create_all_graph_components.""" + + COMMON_KW = dict( + m2g_connectivity="nearest_neighbours", + g2m_connectivity="nearest_neighbours", + m2g_connectivity_kwargs=dict(max_num_neighbours=4), + g2m_connectivity_kwargs=dict(max_num_neighbours=4), + return_components=True, + ) + + def test_flat_triangular(self, xy_medium): + comps = wmg.create.create_all_graph_components( + coords=xy_medium, + m2m_connectivity="flat", + mesh_layout="triangular", + mesh_layout_kwargs=dict(mesh_node_spacing=2.0), + **self.COMMON_KW, + ) + m2m = comps["m2m"] + assert isinstance(m2m, nx.DiGraph) + assert m2m.number_of_nodes() > 0 + assert m2m.number_of_edges() > 0 + # Should also have g2m and m2g + assert comps["g2m"].number_of_edges() > 0 + assert comps["m2g"].number_of_edges() > 0 + + def test_hierarchical_triangular(self, xy_medium): + comps = wmg.create.create_all_graph_components( + coords=xy_medium, + m2m_connectivity="hierarchical", + mesh_layout="triangular", + mesh_layout_kwargs=dict(mesh_node_spacing=2.0, max_num_refinement_levels=3), + **self.COMMON_KW, + ) + m2m = comps["m2m"] + assert isinstance(m2m, nx.DiGraph) + assert m2m.number_of_nodes() > 0 + + def test_flat_multiscale_triangular(self, xy_medium): + comps = wmg.create.create_all_graph_components( + coords=xy_medium, + m2m_connectivity="flat_multiscale", + mesh_layout="triangular", + mesh_layout_kwargs=dict(mesh_node_spacing=2.0, max_num_refinement_levels=3), + **self.COMMON_KW, + ) + m2m = comps["m2m"] + assert isinstance(m2m, nx.DiGraph) + assert m2m.number_of_nodes() > 0 + + def test_unsupported_layout_raises(self, xy_small): + with pytest.raises(NotImplementedError, match="not yet supported"): + wmg.create.create_all_graph_components( + coords=xy_small, + m2m_connectivity="flat", + mesh_layout="nonexistent_layout", + mesh_layout_kwargs=dict(mesh_node_spacing=1.0), + **self.COMMON_KW, + ) + + def test_flat_triangular_return_combined(self, xy_medium): + """With return_components=False, returns a single composed graph.""" + kw = dict(self.COMMON_KW) + kw["return_components"] = False + G = wmg.create.create_all_graph_components( + coords=xy_medium, + m2m_connectivity="flat", + mesh_layout="triangular", + mesh_layout_kwargs=dict(mesh_node_spacing=2.0), + **kw, + ) + assert isinstance(G, nx.DiGraph) + assert G.number_of_nodes() > 0 + + def test_flat_pattern_kwarg_forwarded(self, xy_medium): + """m2m_connectivity_kwargs={'pattern': ...} should be forwarded.""" + comps = wmg.create.create_all_graph_components( + coords=xy_medium, + m2m_connectivity="flat", + mesh_layout="triangular", + mesh_layout_kwargs=dict(mesh_node_spacing=2.0), + m2m_connectivity_kwargs=dict(pattern="8-star"), + **self.COMMON_KW, + ) + assert comps["m2m"].number_of_edges() > 0 + + def test_rectilinear_still_works(self, xy_medium): + """Regression: rectilinear layout should not be broken.""" + comps = wmg.create.create_all_graph_components( + coords=xy_medium, + m2m_connectivity="flat", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(mesh_node_spacing=2.0), + **self.COMMON_KW, + ) + assert comps["m2m"].number_of_nodes() > 0 + + def test_rectilinear_flat_multiscale_still_works(self, xy_medium): + """Regression: rectilinear flat_multiscale should not be broken.""" + comps = wmg.create.create_all_graph_components( + coords=xy_medium, + m2m_connectivity="flat_multiscale", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(mesh_node_spacing=2.0, max_num_refinement_levels=3), + **self.COMMON_KW, + ) + assert comps["m2m"].number_of_nodes() > 0 + + def test_rectilinear_hierarchical_still_works(self, xy_medium): + """Regression: rectilinear hierarchical should not be broken.""" + comps = wmg.create.create_all_graph_components( + coords=xy_medium, + m2m_connectivity="hierarchical", + mesh_layout="rectilinear", + mesh_layout_kwargs=dict(mesh_node_spacing=2.0, max_num_refinement_levels=3), + **self.COMMON_KW, + ) + assert comps["m2m"].number_of_nodes() > 0 + + def test_flat_triangular_with_within_radius(self, xy_medium): + """Triangular flat with within_radius g2m/m2g connectivity.""" + comps = wmg.create.create_all_graph_components( + coords=xy_medium, + m2m_connectivity="flat", + mesh_layout="triangular", + mesh_layout_kwargs=dict(mesh_node_spacing=2.0), + m2g_connectivity="within_radius", + g2m_connectivity="within_radius", + m2g_connectivity_kwargs=dict(max_dist=5.0), + g2m_connectivity_kwargs=dict(max_dist=5.0), + return_components=True, + ) + assert comps["m2m"].number_of_nodes() > 0 + assert comps["g2m"].number_of_edges() > 0 + assert comps["m2g"].number_of_edges() > 0 + + def test_flat_triangular_with_nearest_neighbour(self, xy_medium): + """Triangular flat with nearest_neighbour (singular) connectivity.""" + comps = wmg.create.create_all_graph_components( + coords=xy_medium, + m2m_connectivity="flat", + mesh_layout="triangular", + mesh_layout_kwargs=dict(mesh_node_spacing=2.0), + m2g_connectivity="nearest_neighbour", + g2m_connectivity="nearest_neighbour", + return_components=True, + ) + assert comps["m2m"].number_of_nodes() > 0 + assert comps["g2m"].number_of_edges() > 0 + + def test_flat_no_mesh_node_spacing_raises(self, xy_small): + """Missing mesh_node_spacing should raise ValueError.""" + with pytest.raises(ValueError, match="mesh_node_spacing"): + wmg.create.create_all_graph_components( + coords=xy_small, + m2m_connectivity="flat", + mesh_layout="triangular", + mesh_layout_kwargs=dict(), + **self.COMMON_KW, + ) + + def test_flat_multiscale_no_mesh_node_spacing_raises(self, xy_small): + """Missing mesh_node_spacing in flat_multiscale should raise ValueError.""" + with pytest.raises(ValueError, match="mesh_node_spacing"): + wmg.create.create_all_graph_components( + coords=xy_small, + m2m_connectivity="flat_multiscale", + mesh_layout="triangular", + mesh_layout_kwargs=dict(max_num_refinement_levels=3), + **self.COMMON_KW, + ) + + def test_hierarchical_no_mesh_node_spacing_raises(self, xy_small): + """Missing mesh_node_spacing in hierarchical should raise ValueError.""" + with pytest.raises(ValueError, match="mesh_node_spacing"): + wmg.create.create_all_graph_components( + coords=xy_small, + m2m_connectivity="hierarchical", + mesh_layout="triangular", + mesh_layout_kwargs=dict(max_num_refinement_levels=3), + **self.COMMON_KW, + ) + + def test_all_components_have_nodes(self, xy_medium): + """All three components (g2m, m2m, m2g) should have nodes.""" + comps = wmg.create.create_all_graph_components( + coords=xy_medium, + m2m_connectivity="flat", + mesh_layout="triangular", + mesh_layout_kwargs=dict(mesh_node_spacing=2.0), + **self.COMMON_KW, + ) + for key in ("g2m", "m2m", "m2g"): + assert comps[key].number_of_nodes() > 0 + assert comps[key].number_of_edges() > 0 + + def test_large_domain_triangular(self, xy_large): + """Large domain with small spacing should produce a big graph.""" + comps = wmg.create.create_all_graph_components( + coords=xy_large, + m2m_connectivity="flat", + mesh_layout="triangular", + mesh_layout_kwargs=dict(mesh_node_spacing=3.0), + **self.COMMON_KW, + ) + assert comps["m2m"].number_of_nodes() > 50 + + +# =========================== +# Numerical correctness +# =========================== + + +class TestNumericalCorrectness: + """Test numerical properties of the triangular mesh graph.""" + + def test_edge_lengths_positive(self, xy_medium): + G = create_directed_mesh_graph( + create_single_level_2d_triangular_mesh_primitive(xy_medium, nx=5, ny=5) + ) + for u, v, d in G.edges(data=True): + assert d["len"] > 0 + + def test_vdiff_consistent_with_pos(self, xy_small): + """vdiff should equal pos(u) - pos(v).""" + G = create_directed_mesh_graph( + create_single_level_2d_triangular_mesh_primitive(xy_small, nx=5, ny=5) + ) + for u, v, d in G.edges(data=True): + pos_u = G.nodes[u]["pos"] + pos_v = G.nodes[v]["pos"] + expected_vdiff = pos_u - pos_v + np.testing.assert_allclose(d["vdiff"], expected_vdiff, atol=1e-10) + + def test_len_consistent_with_vdiff(self, xy_small): + """len should equal the L2 norm of vdiff.""" + G = create_directed_mesh_graph( + create_single_level_2d_triangular_mesh_primitive(xy_small, nx=5, ny=5) + ) + for u, v, d in G.edges(data=True): + expected_len = np.linalg.norm(d["vdiff"]) + np.testing.assert_allclose(d["len"], expected_len, atol=1e-10) + + def test_no_nan_in_edge_attrs(self, xy_small): + """Edge attributes should contain no NaN or Inf.""" + G = create_directed_mesh_graph( + create_single_level_2d_triangular_mesh_primitive(xy_small, nx=6, ny=6) + ) + for u, v, d in G.edges(data=True): + assert np.isfinite(d["len"]) + assert np.isfinite(d["vdiff"]).all() + + def test_no_zero_length_edges(self, xy_small): + """All edges should have strictly positive length.""" + G = create_directed_mesh_graph( + create_single_level_2d_triangular_mesh_primitive(xy_small, nx=6, ny=6) + ) + for u, v, d in G.edges(data=True): + assert d["len"] > 1e-12 + + def test_edge_lengths_roughly_uniform_for_interior(self, xy_small): + """For a uniform triangular lattice, all edges should have similar + length (within a narrow tolerance, accounting for scaling).""" + G = create_directed_mesh_graph( + create_single_level_2d_triangular_mesh_primitive(xy_small, nx=8, ny=8) + ) + lengths = [d["len"] for _, _, d in G.edges(data=True)] + # In a uniformly scaled equilateral mesh, all edges should be + # within ~50% of each other (accounting for aspect ratio scaling) + max_len = max(lengths) + min_len = min(lengths) + assert min_len > 0 + ratio = max_len / min_len + # For equilateral triangles with potentially different x/y scaling, + # the ratio should still be reasonable + assert ratio < 3.0, f"Edge length ratio {ratio} too large" + + def test_scaled_domain_produces_scaled_lengths(self): + """Doubling the domain should roughly double edge lengths.""" + xy1 = np.array([[0, 0], [10, 0], [0, 10], [10, 10]], dtype=float) + xy2 = np.array([[0, 0], [20, 0], [0, 20], [20, 20]], dtype=float) + G1 = create_directed_mesh_graph( + create_single_level_2d_triangular_mesh_primitive(xy1, nx=5, ny=5) + ) + G2 = create_directed_mesh_graph( + create_single_level_2d_triangular_mesh_primitive(xy2, nx=5, ny=5) + ) + avg_len1 = np.mean([d["len"] for _, _, d in G1.edges(data=True)]) + avg_len2 = np.mean([d["len"] for _, _, d in G2.edges(data=True)]) + np.testing.assert_allclose(avg_len2 / avg_len1, 2.0, rtol=0.1) + + def test_no_nan_in_positions(self, xy_medium): + """No node should have NaN in positions.""" + G = create_directed_mesh_graph( + create_single_level_2d_triangular_mesh_primitive(xy_medium, nx=5, ny=5) + ) + for node in G.nodes: + pos = G.nodes[node]["pos"] + assert isinstance(pos, np.ndarray) + assert np.isfinite(pos).all()