conservative regridding using `grid_indexing` and `grid_weights`

in-memory-regridding.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c405248-e3f5-4d2b-88a8-ee735a75e010",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cf_xarray  # noqa: F401\n",
    "import dask\n",
    "import geoarrow.rust.core as geoarrow\n",
    "import grid_indexing\n",
    "import grid_weights\n",
    "import h3ronpy\n",
    "import lonboard\n",
    "import numpy as np\n",
    "import shapely\n",
    "import xarray as xr\n",
    "import xdggs\n",
    "\n",
    "xr.set_options(keep_attrs=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b724aad-0d70-423d-ba59-78886a62b23b",
   "metadata": {},
   "source": [
    "## conservative regridding for \"arbitrary\" grids\n",
    "\n",
    "(as long as you can represent your grid cells as geoarrow polygons this should work)\n",
    "\n",
    "This is a work-in-progress, there's a couple of features missing:\n",
    "- support for large grids: while the center coordinates may fit into memory, the polygons duplicate a lot of coordinates, causing even medium-sized grids to be too large to fit into memory on a laptop (which is my personal criterion for this task). I've been working on a dask-based distributed tree implementation, for more on that see the bottom of the notebook.\n",
    "- cell querying on a sphere: the underlying data structure (a R*Tree) works in 2D euclidean space, so any query does not account for the poles / date line. Even worse, the tree currently does not correctly support date line-spanning grid cells"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "832eead9-3cf5-451f-9601-5bd146efbc8d",
   "metadata": {},
   "source": [
    "### setup: dependencies\n",
    "\n",
    "This requires two (WIP) libraries:\n",
    "- https://github.com/keewis/grid-indexing\n",
    "- https://github.com/keewis/grid-weights\n",
    "\n",
    "Neither of which is uploaded to PyPI, yet, so will have to compiled from source. Both include rust code, so this may take some time.\n",
    "\n",
    "My personal way of installing it is\n",
    "```sh\n",
    "$ git clone https://github.com/keewis/grid-indexing\n",
    "$ cd grid-indexing\n",
    "$ maturin develop\n",
    "```\n",
    "(and the same for `grid-weights`)\n",
    "\n",
    "Additional dependencies:\n",
    "- `xdggs`\n",
    "- `lonboard`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cabc3c25-3061-4960-81f4-8beb6db3c27b",
   "metadata": {},
   "source": [
    "### setup: define source and target grids\n",
    "\n",
    "We'll use `xarray`'s `\"air_temperature\"` tutorial dataset as that has a fairly small grid (but any of `xarray`'s gridded tutorial datasets should work).\n",
    "\n",
    "The `infer_cell_geometries` functions uses `cf_xarray` to estimate the bounds of each cell, then constructs `geoarrow` polygons from those. It supports rectilinear and curvilinear cells currently."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfa4de47-1c4a-4a36-b1bd-0b2c500793aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "source_grid = (\n",
    "    xr.tutorial.open_dataset(\"air_temperature\")\n",
    "    .assign_coords(lon=lambda ds: (ds[\"lon\"] + 180) % 360 - 180)\n",
    "    .isel(lon=slice(None, -1))\n",
    ")\n",
    "source_geoms = grid_indexing.infer_cell_geometries(source_grid)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6d1ee33-b2c4-42ee-af5f-48efb7d00892",
   "metadata": {},
   "source": [
    "For the target grid, we'll use H3 (any DGGS would work for here) because it very clearly demonstrates that we can truly use arbitrary polygons as the target grid. The point, however, is that as long as it is possible to represent the grid as polygons it should work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6781ebf5-4e53-41bc-916c-f305a9863808",
   "metadata": {},
   "outputs": [],
   "source": [
    "min_lat = source_grid[\"lat\"].min().item()\n",
    "max_lat = source_grid[\"lat\"].max().item()\n",
    "min_lon = source_grid[\"lon\"].min().item()\n",
    "max_lon = source_grid[\"lon\"].max().item()\n",
    "\n",
    "bbox = shapely.box(min_lon, min_lat, max_lon, max_lat)\n",
    "\n",
    "cell_ids = h3ronpy.vector.geometry_to_cells(bbox, resolution=3).to_numpy()\n",
    "target_grid = (\n",
    "    xr.Dataset(coords={\"cell_ids\": (\"cells\", cell_ids)})\n",
    "    .dggs.decode({\"grid_name\": \"h3\", \"level\": 3})\n",
    "    .dggs.assign_latlon_coords()\n",
    ")\n",
    "target_geoms = target_grid.dggs.grid_info.cell_boundaries(cell_ids, backend=\"geoarrow\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c92bc5b9-a2b2-4d08-85fd-7ab1533887eb",
   "metadata": {},
   "source": [
    "### step 1: search overlapping cells\n",
    "\n",
    "This is the most expensive part: from all the source cells, find the ones that overlap (not just touch!) the target cells.\n",
    "\n",
    "This currently uses the [rstar](https://github.com/georust/rstar) crate to eliminate as many cells as possible before computing the polygon intersection using the [geo](https://github.com/georust/geo) crate (which is the most expensive part). The result is then returned a sparse boolean matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2104890f-69cd-4f92-b9e8-0e2d189842fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "tree = grid_indexing.Index(source_geoms)\n",
    "tree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1f7ea19-89cd-438b-ad7c-a20c0b0c4a94",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "overlapping = tree.query_overlap(target_geoms)\n",
    "overlapping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97bd8ac6-293c-4006-bb77-31927a066ac0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _visualize_grid(geoms, color, alpha=0.8, **layer_kwargs):\n",
    "    from arro3.core import Array, Schema, Table\n",
    "\n",
    "    colors = {\n",
    "        \"red\": [255, 0, 0],\n",
    "        \"green\": [0, 255, 0],\n",
    "        \"blue\": [0, 0, 255],\n",
    "        \"yellow\": [255, 255, 0],\n",
    "    }\n",
    "\n",
    "    alpha_ = int((1 - alpha) * 255)\n",
    "    color_ = colors[color] + [alpha_]\n",
    "\n",
    "    array = Array.from_arrow(geoms)\n",
    "    arrays = {\"geometry\": array, \"value\": Array.from_numpy(np.arange(len(array)))}\n",
    "    fields = [array.field.with_name(name) for name, array in arrays.items()]\n",
    "    schema = Schema(fields)\n",
    "\n",
    "    table = Table.from_arrays(list(arrays.values()), schema=schema)\n",
    "\n",
    "    return lonboard.PolygonLayer(\n",
    "        table=table,\n",
    "        filled=True,\n",
    "        get_fill_color=color_,\n",
    "        get_line_color=\"black\",\n",
    "        auto_highlight=True,\n",
    "        wireframe=True,\n",
    "        **layer_kwargs,\n",
    "    )\n",
    "\n",
    "\n",
    "def visualize_result(source_cells, target_cells, result, index):\n",
    "    def mask_geoarrow(arr, mask):\n",
    "        shapely_ = geoarrow.to_shapely(arr)\n",
    "        return geoarrow.from_shapely(shapely_[mask])\n",
    "\n",
    "    target = target_cells[index]\n",
    "    mask = result[index, :]\n",
    "    print(\"cells found:\", np.sum(mask))\n",
    "\n",
    "    source = mask_geoarrow(source_cells, mask)\n",
    "    target_cell = _visualize_grid(target, color=\"blue\", alpha=0.8)\n",
    "    selected = _visualize_grid(source, color=\"yellow\", alpha=0.6)\n",
    "    source_grid = _visualize_grid(source_cells, color=\"red\", alpha=0.9)\n",
    "\n",
    "    return lonboard.Map([source_grid, target_cell, selected])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2c74a6b-de2b-4f68-a96b-3cff30743170",
   "metadata": {},
   "source": [
    "To verify the results, we can visualize the result (red is the source grid, blue the selected target cell, and yellow are the source cells that overlap the target cell):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcb908cf-b4d2-4a92-a243-6bcfe4329e54",
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_result(source_geoms, target_geoms, overlapping.todense(), index=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6aebd6a2-8b8d-42d9-a8bb-b6ebbcb782ff",
   "metadata": {},
   "source": [
    "### step 2: compute the regridding weights\n",
    "\n",
    "Once we know the overlapping cells, the computation of the weights is relatively straightforward: unpack the sparse matrix and compute the fractional overlap between each selected source cell and the corresponding target cell (plus normalize at the end to account for edge cells). The result is, once again, a sparse matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1ec9d77-99d8-4d69-9c26-a2dc5017c061",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "weights_ = grid_weights.conservative_regridding(source_geoms, target_geoms, overlapping)\n",
    "weights_"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4d93e24-6901-4182-87da-174df720d82f",
   "metadata": {},
   "source": [
    "### step 3: apply the regridding weights\n",
    "\n",
    "This is largely solved, but requires some wrangling to get the weights into a `DataArray`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fc4120d-5442-4859-8465-5a4df27a1410",
   "metadata": {},
   "outputs": [],
   "source": [
    "source_coords = (\n",
    "    source_grid[[\"lon\", \"lat\"]]\n",
    "    .rename_vars({\"lon\": \"source_lon\", \"lat\": \"source_lat\"})\n",
    "    .stack(source_cells=[\"source_lat\", \"source_lon\"])\n",
    "    .coords\n",
    ")\n",
    "target_coords = (\n",
    "    target_grid[[\"longitude\", \"latitude\"]]\n",
    "    .rename_vars({\"longitude\": \"target_lon\", \"latitude\": \"target_lat\"})\n",
    "    .rename_dims({\"cells\": \"target_cells\"})\n",
    "    .coords\n",
    ")\n",
    "\n",
    "coords = source_coords.assign(target_coords)\n",
    "weights = xr.DataArray(\n",
    "    weights_,\n",
    "    coords=coords,\n",
    "    dims=[\"target_cells\", \"source_cells\"],\n",
    "    attrs={\"source_dims\": [\"lon\", \"lat\"]},\n",
    ")\n",
    "weights"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e79ed68-d095-46df-9ac0-1572fe78fbda",
   "metadata": {},
   "source": [
    "With that done, we can apply the weights to the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72290e96-a7cd-4992-9460-6b1de99ca6e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def regrid(arr, weights):\n",
    "    src_dims = weights.attrs[\"source_dims\"]\n",
    "\n",
    "    to_regrid = arr.stack(source_cells=src_dims)\n",
    "\n",
    "    regridded = xr.dot(to_regrid.variable, weights.chunk(), dims=\"source_cells\")\n",
    "\n",
    "    return (\n",
    "        regridded.assign_coords(\n",
    "            weights.to_dataset(name=\"weights\")[[\"target_lat\", \"target_lon\"]]\n",
    "        )\n",
    "        #.unstack(\"target_cells\")\n",
    "        .rename({\"target_lat\": \"latitude\", \"target_lon\": \"longitude\"})\n",
    "    )\n",
    "\n",
    "\n",
    "regridded = regrid(source_grid[\"air\"], weights).compute()\n",
    "regridded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e332e1d2-f6ea-4377-9d26-6459c48861ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "regridded.isel(time=0).copy().dggs.explore(center=273.15, alpha=0.8, cmap=\"coolwarm\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbed96f5-e6d6-413e-8ec1-5de79118be99",
   "metadata": {},
   "source": [
    "### Future Work\n",
    "\n",
    "While the result already looks reasonable, this still remains to be verified: preliminary tests on regridding to another rectilinear grid show that there's a fairly large difference to what `xesmf` returns. While I suspect this has to do with the way the polygons are constructed, I don't know _for sure_.\n",
    "\n",
    "Additionally, there's the two problems stated at the top. I didn't look into a spherical rtree (or into using DGGS cells as bins), but for the memory issue there's [#22](https://github.com/keewis/grid-indexing/pull/22) that wraps the tree into a dask-based distributed tree. However, while that already works if only the target grid is chunked, chunking in the source grid will cause the sparse boolean matrix to not return the correct result (a weird offset pattern can be observed)."
   ]
  }
 ],
 "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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}