Allow non-regularly gridded and lat-lon coordinates

CHANGELOG.md

@@ -7,6 +7,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [unreleased](https://github.com/mllam/weather-model-graphs/compare/v0.2.0...HEAD)
 
+### Added
+
+- Add coords_crs and graph_crs arguments to allow for using lat-lons coordinates
+  or other CRSs as input. These are then converted to the specific CRS used when
+  constructing the graph.
+  [\#32](https://github.com/mllam/weather-model-graphs/pull/32), @joeloskarsson
+
+### Changed
+
+- Change coordinate input to array of shape [N_grid_points, 2] (was previously
+  [2, Ny, Nx]), to allow for non-regularly gridded coordinates
+  [\#32](https://github.com/mllam/weather-model-graphs/pull/32), @joeloskarsson
 
 ## [v0.2.0](https://github.com/mllam/weather-model-graphs/releases/tag/v0.2.0)
 

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 where we have values for the physical fields. We will here work with cartesian x/y coordinates. See [this page](./lat_lons.ipynb) for how to use lat/lon coordinates in weather-model-graphs."
    ]
   },
   {
@@ -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)"
    ]
   },
@@ -80,7 +80,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Lets start with a simple mesh which only has nearest neighbour connections. At the moment `weather-model-graphs` creates a square mesh that sits within the spatial domain spanned by the grid nodes. Techniques for adding non-square meshes are in development."
+    "Lets start with a simple mesh which only has nearest neighbour connections. At the moment `weather-model-graphs` creates a rectangular mesh that sits within the spatial domain spanned by the grid nodes (specifically within the axis-aligned bounding box of the grid nodes). Techniques for adding non-square meshes are in development."
    ]
   },
   {
@@ -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",
@@ -458,7 +458,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.14"
+   "version": "3.10.15"
   }
  },
  "nbformat": 4,

docs/lat_lons.ipynb

@@ -0,0 +1,180 @@
+{
+ "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 Cartesian 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 specifying the Coordinate Reference System (CRS) of `coords` and the CRS that the graph construction should be carried out in. `coords` will then be projected to this new CRS before any calculations are carried out."
+   ]
+  },
+  {
+   "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 Cartesian 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 another CRS with 2-dimensional cartesian coordinates. This can be done by giving the `coords_crs` and `graph_crs` arguments to the graph creation functions. Theses arguments should both be instances of `pyproj.crs.CRS` ([pyproj docs.](https://pyproj4.github.io/pyproj/stable/api/crs/crs.html#pyproj.crs.CRS)). Nicely, they can be `cartopy.crs.CRS`, which provides easy ways to specify such CRSs. For more advanced use cases a `pyproj.crs.CRS` can be specified directly. See [the cartopy documentation](https://scitools.org.uk/cartopy/docs/latest/reference/projections.html) for a list of readily available CRSs to use for projecting the coordinates. \n",
+    "\n",
+    "We will here try the same thing as above, but using a Azimuthal equidistant projection centered at the pole. The CRS of our lat-lon coordinates will be `cartopy.crs.PlateCarree` and we want to project this to `cartopy.crs.AzimuthalEquidistant`:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fffd518b-f323-4e24-8544-60edeaaa8bb2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Define our projection\n",
+    "coords_crs = ccrs.PlateCarree()\n",
+    "graph_crs = ccrs.AzimuthalEquidistant(central_latitude=90)\n",
+    "\n",
+    "fig, ax = plt.subplots(figsize=(15, 9), subplot_kw={\"projection\": graph_crs})\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, coords_crs=coords_crs, graph_crs=graph_crs\n",
+    ")  # Note that we here specify the projection argument\n",
+    "fig, ax = plt.subplots(figsize=(15, 9), subplot_kw={\"projection\": graph_crs})\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.\n",
+    "\n",
+    "It can be noted that this projection between different CRSs provides more general functionality than just handling lat-lon coordinates. It is entirely possible to transform from any `coords_crs` to any `graph_crs` using these arguments."
+   ]
+  }
+ ],
+ "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.15"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

pyproject.toml

@@ -10,6 +10,7 @@ dependencies = [
     "loguru>=0.7.2",
     "networkx>=3.3",
     "scipy>=1.13.0",
+    "pyproj>=3.7.0",
 ]
 requires-python = ">=3.10"
 readme = "README.md"
@@ -22,6 +23,7 @@ pytorch = [
 visualisation = [
     "matplotlib>=3.8.4",
     "ipykernel>=6.29.4",
+    "cartopy>=0.24.1",
 ]
 docs = [
     "jupyter-book>=1.0.0",