Add conservative_2d: conservative regrid for grids that aren't 1D-separable

docs/notebooks/demos/demo_conservative_2d_curvilinear.ipynb

@@ -0,0 +1,215 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0",
+   "metadata": {},
+   "source": [
+    "# Conservative 2D regrid — curvilinear target\n",
+    "\n",
+    "**Conservative regridding** resamples a gridded field while preserving its\n",
+    "area-weighted integral: each output cell is the area-weighted average of\n",
+    "the source cells it overlaps. It's the right tool for fluxes and intensive\n",
+    "quantities — precipitation, sea-surface temperature, mass-balance budgets —\n",
+    "where bilinear or nearest-neighbor interpolation would bias the total.\n",
+    "\n",
+    "The fast `.regrid.conservative` accessor only works on **1D-separable**\n",
+    "grids: plain rectilinear lat/lon, where lat depends only on `y` and lon\n",
+    "only on `x`. **Curvilinear** grids store coordinates as 2D arrays\n",
+    "`lat(y, x)` / `lon(y, x)` and are common in ocean models (ORCA, tripolar)\n",
+    "or any rotated/projected setup. Their cells aren't axis-aligned, so we\n",
+    "drop to `.regrid.conservative_2d`, which builds the full 2D polygon\n",
+    "intersection.\n",
+    "\n",
+    "**In this notebook.** We regrid a smooth analytic two-bump field from a\n",
+    "regular lat/lon grid onto a 30°-rotated curvilinear target — a minimal\n",
+    "stand-in for any curvilinear ocean grid — and verify that the\n",
+    "area-weighted integral is preserved to machine precision."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import xarray as xr\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "import xarray_regrid  # noqa: F401\n",
+    "from xarray_regrid import ConservativeRegridder"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2",
+   "metadata": {},
+   "source": [
+    "## Source — regular 1° lat/lon, analytic two-bump field\n",
+    "\n",
+    "A simple synthetic field — one positive Gaussian bump in the northern\n",
+    "hemisphere, one negative in the southern — gives us something we can\n",
+    "visually track from source to target."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "lat = np.linspace(-60, 60, 121)\n",
+    "lon = np.linspace(-120, 120, 241)\n",
+    "Lo, La = np.meshgrid(lon, lat)\n",
+    "field = (\n",
+    "    np.exp(-((Lo - 40) ** 2 + (La - 20) ** 2) / 500)\n",
+    "    - np.exp(-((Lo + 60) ** 2 + (La + 15) ** 2) / 400)\n",
+    ")\n",
+    "src = xr.DataArray(\n",
+    "    field,\n",
+    "    dims=(\"latitude\", \"longitude\"),\n",
+    "    coords={\"latitude\": lat, \"longitude\": lon},\n",
+    ")\n",
+    "src.plot(figsize=(8, 3.5), cmap=\"RdBu_r\", center=0)\n",
+    "plt.title(\"source: analytic two-bump field\")\n",
+    "plt.tight_layout()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4",
+   "metadata": {},
+   "source": [
+    "## Target — rotated curvilinear grid\n",
+    "\n",
+    "To break 1D-separability we take a regular `(ny, nx)` mesh and rotate it\n",
+    "30° in the lat/lon plane. The result has the same kind of 2D coordinate\n",
+    "variables you'd find in an ORCA or rotated-pole grid: `longitude(ny, nx)`\n",
+    "and `latitude(ny, nx)` riding on a non-axis-aligned mesh."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ny, nx = 30, 50\n",
+    "xi, yi = np.meshgrid(\n",
+    "    np.linspace(-110, 110, nx),\n",
+    "    np.linspace(-45, 45, ny),\n",
+    "    indexing=\"xy\",\n",
+    ")\n",
+    "th = np.deg2rad(30)\n",
+    "lon2d = xi * np.cos(th) - yi * np.sin(th)\n",
+    "lat2d = xi * np.sin(th) + yi * np.cos(th)\n",
+    "target = xr.Dataset(coords={\n",
+    "    \"longitude\": ((\"ny\", \"nx\"), lon2d),\n",
+    "    \"latitude\":  ((\"ny\", \"nx\"), lat2d),\n",
+    "})\n",
+    "\n",
+    "fig, ax = plt.subplots(figsize=(8, 4))\n",
+    "ax.plot(lon2d, lat2d, color=\"0.3\", lw=0.4)\n",
+    "ax.plot(lon2d.T, lat2d.T, color=\"0.3\", lw=0.4)\n",
+    "ax.set_title(\"curvilinear target (30° rotation)\")\n",
+    "ax.set_xlabel(\"longitude\"); ax.set_ylabel(\"latitude\")\n",
+    "ax.set_aspect(\"equal\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6",
+   "metadata": {},
+   "source": [
+    "## Regrid and plot\n",
+    "\n",
+    "`ConservativeRegridder` builds the weight matrix once — that's the\n",
+    "expensive step, since it computes the area of every source/target polygon\n",
+    "intersection — and lets us reuse it for the apply step and for the\n",
+    "diagnostic in the next cell. The one-shot equivalent is\n",
+    "`src.regrid.conservative_2d(target, ...)`. Either way, output is on the\n",
+    "curvilinear `(ny, nx)` mesh, with NaN where target cells fall outside the\n",
+    "source domain."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "rgr = ConservativeRegridder(\n",
+    "    src, target, x_coord=\"longitude\", y_coord=\"latitude\",\n",
+    ")\n",
+    "regridded = rgr.regrid(src)\n",
+    "\n",
+    "fig, ax = plt.subplots(figsize=(8, 4))\n",
+    "pc = ax.pcolormesh(lon2d, lat2d, regridded.values, cmap=\"RdBu_r\",\n",
+    "                   shading=\"auto\", vmin=-1, vmax=1)\n",
+    "fig.colorbar(pc, ax=ax, shrink=0.8)\n",
+    "ax.set_title(\"regridded onto rotated curvilinear grid\")\n",
+    "ax.set_xlabel(\"longitude\"); ax.set_ylabel(\"latitude\")\n",
+    "ax.set_aspect(\"equal\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8",
+   "metadata": {},
+   "source": [
+    "## Conservation check\n",
+    "\n",
+    "The regridder stores its area-intersection matrix\n",
+    "`A[i, j] = area(target_i ∩ source_j)`. For target cells fully inside the\n",
+    "source domain, the area-weighted sum of outputs equals the direct A·s\n",
+    "integral to machine precision — the defining property of *conservative*\n",
+    "regridding."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "A = rgr.areas\n",
+    "src_cover = np.ravel(A.sum(axis=0).todense())\n",
+    "tgt_cover = A.sum(axis=1).todense().reshape(regridded.shape)\n",
+    "valid = np.isfinite(regridded.values)\n",
+    "\n",
+    "direct = float((src.values.ravel() * src_cover).sum())\n",
+    "via_regrid = float((regridded.values[valid] * tgt_cover[valid]).sum())\n",
+    "print(f\"direct        : {direct:.6f}\")\n",
+    "print(f\"via regrid    : {via_regrid:.6f}\")\n",
+    "print(f\"relative err  : {abs(direct - via_regrid) / max(abs(direct), 1e-12):.2e}\")\n",
+    "print(f\"coverage      : {valid.mean():.2%} of target cells inside source\")"
+   ]
+  }
+ ],
+ "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"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}