Add notebook illustrating xgcm.transform (#118)

Co-authored-by: Anderson Banihirwe <[email protected]>

docs/source/examples/index.md

@@ -11,6 +11,7 @@ lateral-fill-idealized.ipynb
 lateral-fill-model-grid.ipynb
 pop_div_curl_xr_xgcm_metrics_compare.ipynb
 CloseHeatBudget_POP2.ipynb
+xgcm-transform.ipynb
 xoak-example.ipynb
 cesm-cmip-cf-xarray.ipynb
 ```

docs/source/examples/xgcm-transform.ipynb

@@ -0,0 +1,213 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "ec5f388e-bbaf-4174-9b96-253e0b01517c",
+   "metadata": {},
+   "source": [
+    "# Using `xgcm.transform` to interpolate to isopycnal space"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fae3a14b-39b8-4746-91d4-f2392aa94199",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext watermark\n",
+    "\n",
+    "import numpy as np\n",
+    "import xarray as xr\n",
+    "import xgcm\n",
+    "\n",
+    "import pop_tools\n",
+    "\n",
+    "%watermark -iv"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "806bf32a-a0c5-4be3-9863-103f91b37c87",
+   "metadata": {},
+   "source": [
+    "## Load data\n",
+    "\n",
+    "Read an existing sample  dataset for the tropical pacific. Unfortunately, this particular subset does not have `SALT`, so for demonstration purposes we set `SALT=35`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b8b1ea07-db6b-4cff-921a-e2dadd69988c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# open sample data\n",
+    "filepath = pop_tools.DATASETS.fetch('Pac_POP0.1_JRA_IAF_1993-12-6-test.nc')\n",
+    "ds = xr.open_dataset(filepath)\n",
+    "\n",
+    "# get DZU and DZT, needed for operations later on\n",
+    "filepath_g = pop_tools.DATASETS.fetch(\"Pac_grid_pbc_1301x305x62.tx01_62l.2013-07-13.nc\")\n",
+    "ds_g = xr.open_dataset(filepath_g)\n",
+    "ds.update(ds_g[[\"DZU\", \"DZT\"]])\n",
+    "\n",
+    "# There is no salinity, so let's  create a fake SALT variable\n",
+    "ds[\"SALT\"] = 35 * xr.ones_like(ds.TEMP)\n",
+    "ds.SALT.attrs[\"grid_loc\"] = ds.TEMP.attrs[\"grid_loc\"]\n",
+    "ds"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "319bacd0-70cd-4db0-b8b3-4d6a13b1fe4a",
+   "metadata": {},
+   "source": [
+    "## Construct xgcm-compatible dataset and  xgcm grid object"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7c720875-04ae-4f72-854c-b07fc3cf0b53",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "metrics = {\n",
+    "    (\"X\",): [\"DXU\", \"DXT\"],  # X distances\n",
+    "    (\"Y\",): [\"DYU\", \"DYT\"],  # Y distances\n",
+    "    (\"Z\",): [\"DZU\", \"DZT\"],  # Z distances\n",
+    "    (\"X\", \"Y\"): [\"UAREA\", \"TAREA\"],\n",
+    "}\n",
+    "\n",
+    "# here we get the xgcm compatible dataset\n",
+    "grid, xds = pop_tools.to_xgcm_grid_dataset(\n",
+    "    ds,\n",
+    "    periodic=False,\n",
+    "    metrics=metrics,\n",
+    "    boundary={\"X\": \"extend\", \"Y\": \"extend\", \"Z\": \"extend\"},\n",
+    ")\n",
+    "\n",
+    "xds[\"rho\"] = pop_tools.eos(xds.SALT, xds.TEMP, depth=xds.z_t * 1e-2)\n",
+    "xds"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3ffb9724-7b1b-4c58-ab53-6def8ebaec75",
+   "metadata": {},
+   "source": [
+    "## Visualize surface density field"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c35b7a9c-4b16-4e66-84eb-c3d9662edfc0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "(xds.rho - 1000).isel(z_t=0).plot(robust=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "93fa9558-0fc6-4aaf-a379-b5cc5fc6a1b7",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "## Regrid to density space using `grid.transform`\n",
+    "\n",
+    "See https://xgcm.readthedocs.io/en/latest/transform.html for more"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8791868d-bc2f-4c82-bdda-7fceed8020d9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "regridded = grid.transform(\n",
+    "    xds.TEMP,\n",
+    "    axis=\"Z\",\n",
+    "    target=np.arange(25, 30, 0.5),\n",
+    "    target_data=xds.rho - 1000,\n",
+    "    method=\"linear\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5cc5018f-bfac-4ff0-b776-9499af9e1c1f",
+   "metadata": {},
+   "source": [
+    "## Visualize temperature on a single isopycnal surface"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5d14e949-3873-4d7a-b53e-1e8482fa349f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "regridded.isel(rho=5).plot(robust=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d634fc1f-7514-426a-98fd-68a212459232",
+   "metadata": {},
+   "source": [
+    "## Find depth of isopycnal surface"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7432d557-d34b-4084-9365-08b5029f1991",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "regridded.coords[\"z_t\"] = grid.transform(\n",
+    "    xds.z_t.broadcast_like(xds.rho),\n",
+    "    axis=\"Z\",\n",
+    "    target=np.arange(24, 30, 0.5),\n",
+    "    target_data=xds.rho - 1000,\n",
+    "    method=\"linear\",\n",
+    ")\n",
+    "regridded.z_t.isel(rho=5).plot(robust=True)"
+   ]
+  }
+ ],
+ "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.8.6"
+  },
+  "widgets": {
+   "application/vnd.jupyter.widget-state+json": {
+    "state": {},
+    "version_major": 2,
+    "version_minor": 0
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}