Add mesh_layout='prebuilt' for user-provided mesh graphs

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).

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

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
+}