diff --git a/docs/source/examples/CloseHeatBudget_POP2.ipynb b/docs/source/examples/CloseHeatBudget_POP2.ipynb new file mode 100644 index 00000000..0b03b77e --- /dev/null +++ b/docs/source/examples/CloseHeatBudget_POP2.ipynb @@ -0,0 +1,401 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Calculate POP2 heat budget using xgcm\n", + "\n", + "Contributed by Anna-Lena Deppenmeier\n", + "\n", + "Using xgcm with metric to demonstrate budget closure. \n", + "\n", + "This is an image of the POP output structure on the horizontal B-grid courtesy Yassir Eddebbar\n", + "\n", + "\"Drawing\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Import packages and define functions" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%matplotlib inline\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import xarray as xr\n", + "from tqdm import tqdm\n", + "\n", + "import pop_tools\n", + "\n", + "\n", + "def pop_find_lat_ind(loc, LATDAT):\n", + " return np.abs(LATDAT[:, 0].values - loc).argmin()\n", + "\n", + "\n", + "def pop_find_lon_ind(loc, LONDAT, direction=\"w\"):\n", + " if direction.lower() in [\"east\", \"e\"]:\n", + " value = loc\n", + " elif direction.lower() in [\"west\", \"w\"]:\n", + " value = 360 - loc\n", + " else:\n", + " print(\"I do not know which direction.\")\n", + " return np.nanargmin(np.abs(LONDAT[152, :].values - value))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Load Dataset" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "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)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ds" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# get lola inds from somewhere for indexing later on\n", + "lola_inds = {}\n", + "inds_lat = [-8, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 8]\n", + "for j in inds_lat:\n", + " if j < 0:\n", + " lola_inds[\"j_\" + str(j)[1:] + \"s\"] = pop_find_lat_ind(j, ds_g.TLAT)\n", + " else:\n", + " lola_inds[\"j_\" + str(j) + \"n\"] = pop_find_lat_ind(j, ds_g.TLAT)\n", + "\n", + "inds_lon = range(95, 185, 5)\n", + "for i in inds_lon:\n", + " lola_inds[\"i_\" + str(i) + \"_w\"] = pop_find_lon_ind(i, ds_g.TLONG)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# just making sure everything works\n", + "ds.TEMP.isel(z_t=0).mean(dim=\"time\").plot(levels=np.arange(20, 30.5, 0.5), cmap=\"RdYlBu_r\")\n", + "plt.scatter(lola_inds[\"i_140_w\"], lola_inds[\"j_0n\"], marker=\"*\");" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Set up vertical thickness and volume for scaling" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ds[\"DZT\"] = ds_g.DZT\n", + "ds[\"DZU\"] = ds_g.DZU\n", + "\n", + "ds.DZT.attrs[\"long_name\"] = \"Thickness of T cells\"\n", + "ds.DZT.attrs[\"units\"] = \"centimeter\"\n", + "ds.DZT.attrs[\"grid_loc\"] = \"3111\"\n", + "ds.DZU.attrs[\"long_name\"] = \"Thickness of U cells\"\n", + "ds.DZU.attrs[\"units\"] = \"centimeter\"\n", + "ds.DZU.attrs[\"grid_loc\"] = \"3221\"\n", + "\n", + "# make sure we have the cell volumne for calculations\n", + "VOL = (ds.DZT * ds.DXT * ds.DYT).compute()\n", + "KMT = ds_g.KMT.compute()\n", + "\n", + "for j in tqdm(range(len(KMT.nlat))):\n", + " for i in range(len(KMT.nlon)):\n", + " k = KMT.values[j, i].astype(int)\n", + " VOL.values[k:, j, i] = 0.0\n", + "\n", + "ds[\"VOL\"] = VOL\n", + "\n", + "ds.VOL.attrs[\"long_name\"] = \"volume of T cells\"\n", + "ds.VOL.attrs[\"units\"] = \"centimeter^3\"\n", + "\n", + "ds.VOL.attrs[\"grid_loc\"] = \"3111\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Set up dataset to gather budget terms" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "budget = xr.Dataset()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Set grid and data set for xgcm with metrics" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "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", + "gridxgcm, dsxgcm = 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", + "for coord in [\"nlat\", \"nlon\"]:\n", + " if coord in dsxgcm.coords:\n", + " dsxgcm = dsxgcm.drop_vars(coord)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 0) Tendency" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "budget['TEND_TEMP'] = dsxgcm.TEND_TEMP" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### i) Total heat advection" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "budget[\"UET\"] = -(gridxgcm.diff(dsxgcm.UET * dsxgcm.VOL.values, axis=\"X\") / dsxgcm.VOL)\n", + "budget[\"VNT\"] = -(gridxgcm.diff(dsxgcm.VNT * dsxgcm.VOL.values, axis=\"Y\") / dsxgcm.VOL)\n", + "budget[\"WTT\"] = (\n", + " gridxgcm.diff(dsxgcm.WTT.fillna(0) * (dsxgcm.dz * dsxgcm.DXT * dsxgcm.DYT).values, axis=\"Z\")\n", + " / dsxgcm.VOL\n", + ")\n", + "budget[\"TOT_ADV\"] = budget[\"UET\"] + budget[\"VNT\"] + budget[\"WTT\"]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### ii) Heat flux due to vertical mixing:\n", + "includes surface flux at the 0th layer" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "budget[\"DIA_IMPVF_TEMP\"] = -(\n", + " gridxgcm.diff(dsxgcm.DIA_IMPVF_TEMP * dsxgcm.TAREA, axis=\"Z\") / dsxgcm.VOL\n", + ")\n", + "\n", + "# set surface flux at 0th layer\n", + "SRF_TEMP_FLUX = (dsxgcm.SHF - dsxgcm.SHF_QSW) * dsxgcm.hflux_factor\n", + "\n", + "budget[\"DIA_IMPVF_TEMP\"][:, 0, :, :] = (\n", + " SRF_TEMP_FLUX * dsxgcm.TAREA - dsxgcm.DIA_IMPVF_TEMP.isel(z_w_bot=0) * dsxgcm.TAREA\n", + ") / dsxgcm.VOL.values[0, :, :]\n", + "\n", + "budget[\"KPP_SRC_TMP\"] = dsxgcm.KPP_SRC_TEMP\n", + "budget[\"VDIF\"] = budget[\"DIA_IMPVF_TEMP\"] + budget[\"KPP_SRC_TMP\"]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### iii) Heat flux due to horizontal diffusion" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "budget[\"HDIFE_TEMP\"] = gridxgcm.diff(dsxgcm.HDIFE_TEMP * dsxgcm.VOL.values, axis=\"X\") / dsxgcm.VOL\n", + "budget[\"HDIFN_TEMP\"] = gridxgcm.diff(dsxgcm.HDIFN_TEMP * dsxgcm.VOL.values, axis=\"Y\") / dsxgcm.VOL\n", + "budget[\"HDIF\"] = budget[\"HDIFE_TEMP\"] + budget[\"HDIFN_TEMP\"]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### iv) Solar Penetration" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "budget[\"QSW_3D\"] = -gridxgcm.diff((dsxgcm.QSW_3D * dsxgcm.hflux_factor), axis=\"Z\") / dsxgcm.DZT" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### plot to make sure it closes at (0N, 140W)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "subset = budget.isel(nlon_t=lola_inds[\"i_140_w\"], nlat_t=lola_inds[\"j_0n\"], time=0)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%%time\n", + "\n", + "fig, ax = plt.subplots(1, 2, figsize=(10, 3), sharey=True)\n", + "\n", + "# plot individual components\n", + "subset.VDIF.plot(y=\"z_t\", ylim=(300e2, 0), label=\"VDIF\", ax=ax[0])\n", + "subset.HDIF.plot(y=\"z_t\", ylim=(300e2, 0), label=\"HDIF\", ax=ax[0])\n", + "subset.QSW_3D.plot(y=\"z_t\", ylim=(300e2, 0), label=\"QSW_3D\", ax=ax[0])\n", + "subset.TOT_ADV.plot(y=\"z_t\", ylim=(300e2, 0), label=\"DIV\", ax=ax[0])\n", + "\n", + "# plot sum\n", + "(subset.QSW_3D + subset.HDIF + subset.VDIF + subset.TOT_ADV).plot(\n", + " y=\"z_t\", ylim=(300e2, 0), label=\"SUM\", ls=\"--\", ax=ax[0]\n", + ")\n", + "# plot tendency\n", + "subset.TEND_TEMP.plot(y=\"z_t\", ylim=(300e2, 0), label=\"TEND_TEMP\", ax=ax[0])\n", + "\n", + "ax[0].legend()\n", + "\n", + "# plot sum\n", + "(subset.QSW_3D + subset.HDIF + subset.VDIF + subset.TOT_ADV).plot(\n", + " y=\"z_t\", ylim=(300e2, 0), label=\"SUM\", ls=\"--\", ax=ax[1]\n", + ")\n", + "# plot tendency\n", + "subset.TEND_TEMP.plot(y=\"z_t\", ylim=(300e2, 0), label=\"TEND_TEMP\", ax=ax[1])\n", + "\n", + "ax[1].legend();" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# You may need to install watermark (conda install -c conda-forge watermark)\n", + "import xgcm # just to display the version\n", + "\n", + "%load_ext watermark\n", + "%watermark -d -iv -m -g -h\n", + "print(\"Above are the versions of the packages this works with.\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "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.8" + }, + "widgets": { + "application/vnd.jupyter.widget-state+json": { + "state": {}, + "version_major": 2, + "version_minor": 0 + } + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/docs/source/examples/POP_Grid.png b/docs/source/examples/POP_Grid.png new file mode 100644 index 00000000..94f5cd89 Binary files /dev/null and b/docs/source/examples/POP_Grid.png differ diff --git a/docs/source/examples/index.md b/docs/source/examples/index.md index b535ce73..5cd18d9d 100644 --- a/docs/source/examples/index.md +++ b/docs/source/examples/index.md @@ -10,4 +10,5 @@ region-mask.ipynb lateral-fill-idealized.ipynb lateral-fill-model-grid.ipynb pop_div_curl_xr_xgcm_metrics_compare.ipynb +CloseHeatBudget_POP2.ipynb ``` diff --git a/pop_tools/data_registry.txt b/pop_tools/data_registry.txt index 2c9ddac4..a673c8b4 100644 --- a/pop_tools/data_registry.txt +++ b/pop_tools/data_registry.txt @@ -11,6 +11,6 @@ lateral_fill_np_array_tripole_filled_SOR_ref.20200820.npz f6d803bd05cd6a16857683 POP_gx3v7.nc 7586ba68ee5b4d8de8ffe5a11974bff81a412c62af9de230088575c8016e4b84 g.e20.G.TL319_t13.control.001_hfreq.nc 439eb1abf14737341ead088bfd9a3c1b795dc1f79bc9052c09db08fef9ab83e5 g.e20.G.TL319_t13.control.001_hfreq-coarsen.nc 145659813daf1a607d0857c10699c721693333e39aeb270e116103185b7236ae -Pac_POP0.1_JRA_IAF_1993-12-6-test.nc fc9c649428c1e62108fad00644805130f58e4063e0219ed677a357c077886f8e +Pac_POP0.1_JRA_IAF_1993-12-6-test.nc f401c55771adfe6fa0b4646f9465b28cebc02b82641ed13c3c8a47fe6e24cb7f Pac_grid_pbc_1301x305x62.tx01_62l.2013-07-13.nc 9ead77ee8b1e352c9b0316664b3bfe9b84e322ca785bfd9a024da54a2b6dc60e comp-grid.tx9.1v3.20170718.zarr.zip 2c1cd41c1c803c0565bc6120268aac728757207963386b727f00b822b81155da