Allow non-regularly gridded and lat-lon coordinates

CHANGELOG.md

@@ -16,6 +16,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   [\#19](https://github.com/mllam/weather-model-graphs/pull/19)
   @joeloskarsson
 
+- `save.to_pyg` can now handle any number of 1D or 2D edge or node features when
+  converting pytorch-geometric `Data` objects to `torch.Tensor` objects.
+  [\#31](https://github.com/mllam/weather-model-graphs/pull/31)
+  @maxiimilian
+
 ### Changed
 
 - Fix wrong number of mesh levels when grid is multiple of refinement factor

docs/_toc.yml

@@ -7,3 +7,4 @@ chapters:
 - file: background
 - file: design
 - file: creating_the_graph
+- file: lat_lons

docs/background.ipynb

@@ -32,15 +32,16 @@
     "\n",
     "# create some fake cartesian coordinates\n",
     "def _create_fake_xy(N=10):\n",
-    "    x = np.linspace(0.0, 1.0, N)\n",
-    "    y = np.linspace(0.0, 1.0, N)\n",
-    "    xy = np.stack(np.meshgrid(x, y), axis=0)\n",
+    "    x = np.linspace(0.0, N, N)\n",
+    "    y = np.linspace(0.0, N, N)\n",
+    "    xy_mesh = np.meshgrid(x, y)\n",
+    "    xy = np.stack([mg_coord.flatten() for mg_coord in xy_mesh], axis=1)  # Shaped (N, 2)\n",
     "    return xy\n",
     "\n",
     "\n",
-    "xy_grid = _create_fake_xy(N=10)\n",
+    "xy = _create_fake_xy(N=10)\n",
     "\n",
-    "graph = wmg.create.archetype.create_keisler_graph(xy_grid=xy_grid)\n",
+    "graph = wmg.create.archetype.create_keisler_graph(coords=xy)\n",
     "\n",
     "# remove all edges from the graph\n",
     "graph.remove_edges_from(list(graph.edges))\n",
@@ -85,9 +86,7 @@
    "source": [
     "import matplotlib.pyplot as plt\n",
     "\n",
-    "graph = wmg.create.archetype.create_keisler_graph(\n",
-    "    xy_grid=xy_grid, grid_refinement_factor=2\n",
-    ")\n",
+    "graph = wmg.create.archetype.create_keisler_graph(coords=xy, mesh_node_distance=2)\n",
     "graph_components = wmg.split_graph_by_edge_attribute(graph=graph, attr=\"component\")\n",
     "\n",
     "fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(16, 5))\n",

docs/creating_the_graph.ipynb

@@ -26,7 +26,7 @@
    "source": [
     "# The grid nodes\n",
     "\n",
-    "To get started we will create a set of fake grid nodes, which represent the geographical locations (lat/lon or x/y cartesian) where we have values for the physical fields. Here we'll just use cartesian coordinates for simplicity.\n"
+    "To get started we will create a set of fake grid nodes, which represent the geographical locations (x/y cartesian) where we have values for the physical fields."
    ]
   },
   {
@@ -47,10 +47,10 @@
    "outputs": [],
    "source": [
     "def _create_fake_xy(N=10):\n",
-    "    x = np.linspace(0, 1, N)\n",
-    "    y = np.linspace(0, 1, N)\n",
-    "    xy = np.meshgrid(x, y)\n",
-    "    xy = np.stack(xy, axis=0)\n",
+    "    x = np.linspace(0.0, N, N)\n",
+    "    y = np.linspace(0.0, N, N)\n",
+    "    xy_mesh = np.meshgrid(x, y)\n",
+    "    xy = np.stack([mg_coord.flatten() for mg_coord in xy_mesh], axis=1)  # Shaped (N, 2)\n",
     "    return xy"
    ]
   },
@@ -63,7 +63,7 @@
     "xy = _create_fake_xy(32)\n",
     "\n",
     "fig, ax = plt.subplots()\n",
-    "ax.scatter(xy[0], xy[1])\n",
+    "ax.scatter(xy[:, 0], xy[:, 1])\n",
     "ax.set_aspect(1)"
    ]
   },
@@ -130,7 +130,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "graph = wmg.create.archetype.create_keisler_graph(xy_grid=xy)\n",
+    "graph = wmg.create.archetype.create_keisler_graph(coords=xy)\n",
     "graph"
    ]
   },
@@ -200,7 +200,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "graph = wmg.create.archetype.create_graphcast_graph(xy_grid=xy)"
+    "graph = wmg.create.archetype.create_graphcast_graph(coords=xy)"
    ]
   },
   {
@@ -275,7 +275,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "graph = wmg.create.archetype.create_oskarsson_hierarchical_graph(xy_grid=xy)\n",
+    "graph = wmg.create.archetype.create_oskarsson_hierarchical_graph(coords=xy)\n",
     "graph"
    ]
   },
@@ -410,9 +410,9 @@
    "source": [
     "graph = wmg.create.create_all_graph_components(\n",
     "    m2m_connectivity=\"flat_multiscale\",\n",
-    "    xy=xy,\n",
+    "    coords=xy,\n",
     "    m2m_connectivity_kwargs=dict(\n",
-    "        grid_refinement_factor=2, level_refinement_factor=3, max_num_levels=None\n",
+    "        mesh_node_distance=2, level_refinement_factor=3, max_num_levels=None\n",
     "    ),\n",
     "    g2m_connectivity=\"nearest_neighbour\",\n",
     "    m2g_connectivity=\"nearest_neighbour\",\n",

docs/lat_lons.ipynb

@@ -0,0 +1,177 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "65f79439-0b75-483f-858b-932422f7599d",
+   "metadata": {},
+   "source": [
+    "# Working with lat-lon coordinates\n",
+    "\n",
+    "In the previous sections we have considered grid point positions `coords` given as euclidean coordinates. However, it is common that we have coordinates given as latitudes and longitudes. This notebook describes how we can constuct graphs directly using lat-lon coordinates. This is achieved by also providing a specific projection, used to project the lat-lons to a euclidean space where the graph can be construced."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8cbcd3bd-a80c-41e4-a54b-bcfb5dfbb75c",
+   "metadata": {},
+   "source": [
+    "## A motivating example"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4c710a27-7e03-4f4f-a4dd-5d1c53a1e4eb",
+   "metadata": {},
+   "source": [
+    "Let's start by defining some example lat-lons to use in our example. When using lat-lons the first column of `coords` should contain longitudes and the second column latitudes.\n",
+    "\n",
+    "In the example below we create lat-lons laid out around the geographic North Pole. These example points are equidistantly spaced, but this does not have to be the case."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9d577c2f-6dca-4b1b-8bdc-1359e6573cda",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "import cartopy.crs as ccrs\n",
+    "\n",
+    "longitudes = np.linspace(-180, 180, 40)\n",
+    "latitudes = np.linspace(65, 85, 5)  # Very close to north pole\n",
+    "\n",
+    "meshgridded_lat_lons = np.meshgrid(longitudes, latitudes)\n",
+    "coords = np.stack([mg_coord.flatten() for mg_coord in meshgridded_lat_lons], axis=1)\n",
+    "\n",
+    "fig, ax = plt.subplots(figsize=(15, 9), subplot_kw={\"projection\": ccrs.PlateCarree()})\n",
+    "ax.scatter(coords[:, 0], coords[:, 1], marker=\".\")\n",
+    "ax.coastlines()\n",
+    "ax.set_extent((-180, 180, -90, 90))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4d4f5f35-01bf-4eeb-be52-d7deabbf2039",
+   "metadata": {},
+   "source": [
+    "We first consider what happens if we directly feed these lat-lons as `coords`, treating them as if they were euclidean coordinates. In this notebook we will only create flat \"Keisler-like\" graphs, but everything works analogously for the other graph types."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8e04abfd-f3e2-4f77-908c-6e2374431e9e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import weather_model_graphs as wmg\n",
+    "\n",
+    "graph = wmg.create.archetype.create_keisler_graph(coords, mesh_node_distance=10)\n",
+    "fig, ax = plt.subplots(figsize=(15, 9), subplot_kw={\"projection\": ccrs.PlateCarree()})\n",
+    "wmg.visualise.nx_draw_with_pos_and_attr(\n",
+    "    graph, ax=ax, node_size=30, edge_color_attr=\"component\", node_color_attr=\"type\"\n",
+    ")\n",
+    "ax.coastlines()\n",
+    "ax.set_extent((-180, 180, -90, 90))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "65b39a86-d64f-46a8-a732-0e608a3fb07f",
+   "metadata": {},
+   "source": [
+    "This creates a useable mesh graph, but we can note a few problems with it:\n",
+    "\n",
+    "* There are no connections between nodes around longitude -180/180, i.e. the periodicity of longitude is not considered.\n",
+    "* All nodes at the top of the plot, close to the pole, are actually very close spatially. Yet there are no connections between them.\n",
+    "\n",
+    "These are issues both in the connection between the grid nodes and the mesh, and in the connections between mesh nodes. This points to the fact that we should probably use a different projection when building our graph. "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "103a3784-6ad4-4188-ba02-189f9cf71b2f",
+   "metadata": {},
+   "source": [
+    "## Constructing a graph within a projection\n",
+    "For our example above, let's instead try to construct the graph based on first projecting our lat-lon coordinates to within a specific projection. This can be done by giving a `projection` argument to the graph creation functions. The projection should be an instance of `cartopy.crs.CRS`. See [the cartopy documentation](https://scitools.org.uk/cartopy/docs/latest/reference/projections.html) for a list of available projections. \n",
+    "\n",
+    "We will here try the same thing as above, but using a Azimuthal equidistant projection centered at the pole:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fffd518b-f323-4e24-8544-60edeaaa8bb2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Define our projection\n",
+    "ae_projection = ccrs.AzimuthalEquidistant(central_latitude=90)\n",
+    "\n",
+    "fig, ax = plt.subplots(figsize=(15, 9), subplot_kw={\"projection\": ae_projection})\n",
+    "ax.scatter(coords[:, 0], coords[:, 1], marker=\".\", transform=ccrs.PlateCarree())\n",
+    "_ = ax.coastlines()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "afa6b933-121b-41c4-a341-9e548b92ad87",
+   "metadata": {},
+   "source": [
+    "Note that distances within projections tend to have very large magnitudes, so the distance between mesh nodes should be specified accordingly."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1ddd7ec3-0d66-4feb-a0d3-293f3194dddd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mesh_distance = (\n",
+    "    10**6\n",
+    ")  # Large euclidean distance in projection coordinates between mesh nodes\n",
+    "graph = wmg.create.archetype.create_keisler_graph(\n",
+    "    coords, mesh_node_distance=mesh_distance, projection=ae_projection\n",
+    ")  # Note that we here specify the projection argument\n",
+    "fig, ax = plt.subplots(figsize=(15, 9), subplot_kw={\"projection\": ae_projection})\n",
+    "wmg.visualise.nx_draw_with_pos_and_attr(\n",
+    "    graph, ax=ax, node_size=30, edge_color_attr=\"component\", node_color_attr=\"type\"\n",
+    ")\n",
+    "_ = ax.coastlines()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f5acf83a-df33-4925-bb32-c106e27e51a4",
+   "metadata": {},
+   "source": [
+    "Now this looks like a more reasonable graph layout, that better respects the spatial relations between the grid points. There are still things that could be tweaked further (e.g. the large number of grid nodes connected to the center mesh node), but this ends our example of defining graphs using lat-lon coordinates."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

pyproject.toml

@@ -10,6 +10,7 @@ dependencies = [
     "loguru>=0.7.2",
     "networkx>=3.3",
     "scipy>=1.13.0",
+    "cartopy>=0.24.1",
 ]
 requires-python = ">=3.10"
 readme = "README.md"