diff --git a/src/esnb/core/NotebookDiagnostic.py b/src/esnb/core/NotebookDiagnostic.py index 915521f..6e2e65a 100644 --- a/src/esnb/core/NotebookDiagnostic.py +++ b/src/esnb/core/NotebookDiagnostic.py @@ -456,6 +456,10 @@ def datasets(self): def access_dataset(self, id=0): return self.datasets[id] + @property + def varmap(self): + return {x.varname: x for x in self.variables} + def resolve(self, groups=None): """ Resolve datasets for the provided groups and assign them to the diff --git a/src/esnb/core/util_xr.py b/src/esnb/core/util_xr.py index 22347f5..b05f4ba 100644 --- a/src/esnb/core/util_xr.py +++ b/src/esnb/core/util_xr.py @@ -30,6 +30,10 @@ def open_paths(files, varname=None): logger.debug("Found `z_i` in dataset; associating it as a coordinate.") ds["z_i"] = _ds["z_i"] + if "deptho" in _ds.keys(): + logger.debug("Found `deptho` in dataset; associating it as a coordinate.") + ds["deptho"] = _ds["deptho"] + ds.attrs = dict(_ds.attrs) else: ds = _ds diff --git a/src/esnb/templates/MOC_z-space.ipynb b/src/esnb/templates/MOC_z-space.ipynb new file mode 100644 index 0000000..89a7456 --- /dev/null +++ b/src/esnb/templates/MOC_z-space.ipynb @@ -0,0 +1,509 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "e4a43ed7-7624-4da9-b213-32f892574a65", + "metadata": {}, + "source": [ + "# Meridional Overturning Circulation Notebook (z-space)\n", + "\n", + "This notebook plots the global MOC in depth and isopycnal (rho2) space as well as a break down of the Indo-Pacific and Atlantic basin.\n", + "\n", + "Authors: Jan-Erik Tesdal and John Krasting" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e67185bc-4731-4c89-b63a-ca8953428f07", + "metadata": { + "editable": true, + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "# Development mode: constantly refreshes module code\n", + "%load_ext autoreload\n", + "%autoreload 2" + ] + }, + { + "cell_type": "markdown", + "id": "6898886b-b5e1-4de0-939a-f715748be915", + "metadata": { + "editable": true, + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "source": [ + "## Framework Code and Diagnostic Setup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4168ea13-0a51-4f50-b0ae-8cccbd75161b", + "metadata": { + "editable": true, + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "import esnb\n", + "from esnb import NotebookDiagnostic, RequestedVariable, CaseGroup2, nbtools\n", + "from esnb.sites.gfdl import call_dmget" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a418bf87-ce1d-4161-933d-ce33023853ec", + "metadata": {}, + "outputs": [], + "source": [ + "# Define a mode (leave \"prod\" for now)\n", + "mode = \"prod\"\n", + "\n", + "# Verbosity\n", + "verbose = True\n", + "\n", + "# Give your diagnostic a name and a short description\n", + "diag_name = \"MOC z-space\"\n", + "diag_desc = \"Meridional Overturning Circulation in z-space\"\n", + "\n", + "# Define what variables you would like to analyze. The first entry is the\n", + "# variable name and the second entry is the realm (post-processing dir).\n", + "# (By default, monthly timeseries data will be loaded. TODO: add documentation\n", + "# on how to select different frequencies, multiple realms to search, etc.)\n", + "variables = [\n", + " RequestedVariable(\"vmo\", \"ocean_annual_z\", frequency=\"yearly\"),\n", + "]\n", + "\n", + "# Optional: define runtime settings or options for your diagnostic\n", + "user_options = {\"regions\": [\"global\", \"atlantic\"]}\n", + "\n", + "# Initialize the diagnostic with its name, description, vars, and options\n", + "diag = NotebookDiagnostic(diag_name, diag_desc, variables=variables, **user_options)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c5ec6425-305d-45b2-9a2e-76be3a8cb129", + "metadata": {}, + "outputs": [], + "source": [ + "# Define the groups of experiments to analyze. Provide a single dora id for one experiment\n", + "# or a list of IDs to aggregate multiple experiments into one; e.g. historical+future runs\n", + "groups = [\n", + " CaseGroup2(\"odiv-516\", date_range=(\"1993-01-01\", \"2017-12-31\"), name=\"OM5 B11 NB\"),\n", + " CaseGroup2(\"odiv-290\", date_range=(\"1993-01-01\", \"2017-12-31\"), name=\"OM5 B01 (OM4-like)\"),\n", + " CaseGroup2(\"odiv-554\", date_range=(\"0041-01-01\", \"0060-12-31\"), name=\"CM4.0 + OM5 B11\"),\n", + " CaseGroup2(\"odiv-558\", date_range=(\"0041-01-01\", \"0060-12-31\"), name=\"CM4.0\"),\n", + "]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aa1a552a-083f-4019-8d41-815bb2611c61", + "metadata": { + "editable": true, + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "# Combine the experiments with the diag request and determine what files need to be loaded:\n", + "diag.resolve(groups)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "375fbdd0-422c-4b93-b6e8-644a282d63b7", + "metadata": { + "editable": true, + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "# Print a list of file paths\n", + "# This cell and the markdown cell that follows are necessary to run this notebook\n", + "# Interactively on Dora\n", + "_ = [print(x) for x in diag.files]" + ] + }, + { + "cell_type": "markdown", + "id": "37ecd8b9-d851-4b9b-8032-ac8b77deb511", + "metadata": { + "editable": true, + "slideshow": { + "slide_type": "" + }, + "tags": [ + "stop_here" + ] + }, + "source": [ + "(The files above are necessary to run the diagnostic.)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "337fb74f-f81c-49e4-9014-69c1c33bc220", + "metadata": {}, + "outputs": [], + "source": [ + "# Check to see the dmget status before calling \"open\"\n", + "call_dmget(diag.files,status=True)\n", + "call_dmget(diag.files)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0e173eb4-5c69-40b7-87a6-2c62d1111162", + "metadata": { + "editable": true, + "slideshow": { + "slide_type": "" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "# Load the data as xarray datasets\n", + "diag.open()" + ] + }, + { + "cell_type": "markdown", + "id": "2014a755-2efc-4cd5-a7e3-76cb93d22ca2", + "metadata": {}, + "source": [ + "## Main Diagnostic" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "980c3699-e785-46bc-82e4-4a7a3308ed58", + "metadata": {}, + "outputs": [], + "source": [ + "import xarray as xr\n", + "import cmip_basins\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "import momgrid as mg" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f24d6003-cec9-40d2-9b4d-6a5b9ae40f3f", + "metadata": {}, + "outputs": [], + "source": [ + "esnb.sites.gfdl.convert_to_momgrid(diag)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fd3b31c1-dbf6-49c5-9807-6298ddf4bf25", + "metadata": {}, + "outputs": [], + "source": [ + "def calc_moc_z(\n", + " ds,\n", + " varname,\n", + " mask=1.0,\n", + " xdim=\"xh\",\n", + " tdim=\"time\",\n", + " ydim=\"yq\",\n", + " zdim=\"z_l\",\n", + " interfaces=\"z_i\",\n", + "):\n", + " arr = ds[varname]\n", + " arr = arr.mean(tdim) * mask\n", + " integ_layers = arr.sum(xdim).reindex({zdim: arr[zdim][::-1]}).cumsum(zdim)\n", + " # The result of the integration over layers is evaluated at the interfaces\n", + " # with psi = 0 as the bottom boundary condition for the integration\n", + " bottom_condition = xr.zeros_like(integ_layers.isel({zdim: 0}))\n", + " # combine bottom condition with data array\n", + " psi_raw = xr.concat(\n", + " [integ_layers.reindex({zdim: integ_layers[zdim][::-1]}), bottom_condition],\n", + " dim=zdim,\n", + " )\n", + " # rename to correct dimension and add correct vertical coordinate\n", + " psi = -psi_raw.rename({zdim: interfaces}).transpose(interfaces, ydim)\n", + " psi[interfaces] = ds[interfaces]\n", + " psi.name = \"psi\"\n", + " # # Convert kg.s-1 to Sv (1e6 m3.s-1)\n", + " moc_z = psi / rho0 / 1.0e6\n", + " return moc_z" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3acc943e-c1d3-4cb8-86ef-1a7cfd0369f8", + "metadata": {}, + "outputs": [], + "source": [ + "def plot_panel(ax, y, z, psi, depth, domain):\n", + " if domain == \"atlantic\":\n", + " levels = np.arange(-28, 30, 2)\n", + " elif domain == \"global\":\n", + " levels = np.arange(-40, 42, 2)\n", + "\n", + " cb = ax.contourf(y, z, psi, levels=levels, cmap=\"RdBu_r\")\n", + "\n", + " # Setting y-axis limits to avoid singular transformation warning for future calls\n", + " ax.set_ylim([np.min(z) - 500, np.max(z) + 500])\n", + "\n", + " cs = ax.contour(y, z, psi, levels=levels, colors=\"k\", linewidths=0.3)\n", + " zero_contours = ax.contour(y, z, psi, levels=[0], colors=\"blue\", linewidths=0)\n", + "\n", + " ax.set_yscale(\"splitscale\", zval=[6500.0, 2000.0, 0.0])\n", + " plt.axhline(y=2000, color=\"k\", linestyle=\"dashed\", linewidth=0.8)\n", + "\n", + " _ = ax.fill_between(y, depth, 6750, color=\"gray\", zorder=2)\n", + " # plt.colorbar(cb)\n", + "\n", + " stats = {}\n", + "\n", + " if domain == \"atlantic\":\n", + " Y, Z = np.meshgrid(y, z)\n", + " mask = (Y >= 20) & (Y <= 80) & (Z >= 500) & (Z <= 2000)\n", + " _psi = psi.values\n", + " max_psi = np.max(_psi[mask])\n", + " max_index = np.unravel_index(np.argmax(_psi[mask]), _psi[mask].shape)\n", + " max_y = Y[mask][max_index]\n", + " max_z = Z[mask][max_index]\n", + "\n", + " stats = {\n", + " \"max\": round(float(max_psi), 2),\n", + " \"max_lat\": round(float(max_y), 2),\n", + " \"max_depth\": round(float(max_z), 2),\n", + " }\n", + "\n", + " square_size = 10\n", + " ax.plot(\n", + " max_y,\n", + " max_z,\n", + " marker=\"s\",\n", + " color=\"yellow\",\n", + " markersize=square_size + 5,\n", + " markeredgewidth=2,\n", + " markeredgecolor=\"yellow\",\n", + " markerfacecolor=\"none\",\n", + " )\n", + " ax.annotate(\n", + " f\"{max_psi:.1f} Sv\",\n", + " xy=(max_y, max_z),\n", + " xytext=(10, 10),\n", + " textcoords=\"offset points\",\n", + " arrowprops=dict(arrowstyle=\"->\", color=\"black\"),\n", + " fontsize=5,\n", + " bbox=dict(\n", + " facecolor=\"white\",\n", + " edgecolor=\"black\",\n", + " boxstyle=\"round,pad=0.2\",\n", + " linewidth=0.5,\n", + " alpha=0.7,\n", + " ),\n", + " )\n", + "\n", + " z_values = []\n", + " for path in zero_contours.get_paths(): # Use get_paths() instead of allsegs\n", + " vertices = path.vertices\n", + " x = vertices[:, 0]\n", + " z_val = vertices[:, 1]\n", + " mask = (20 <= x) & (x <= 55) & (2500 <= z_val) & (z_val <= 4000)\n", + " z_values.extend(z_val[mask])\n", + "\n", + " mean_z = np.mean(z_values) if z_values else None\n", + "\n", + " if mean_z:\n", + " stats[\"zero_depth\"] = round(float(mean_z), 1)\n", + " ax.annotate(\n", + " f\"Mean Depth:\\n{mean_z:.0f} m\",\n", + " xy=(40, mean_z),\n", + " xytext=(-10, -30),\n", + " textcoords=\"offset points\",\n", + " arrowprops=dict(arrowstyle=\"->\", color=\"black\"),\n", + " fontsize=7,\n", + " bbox=dict(\n", + " facecolor=\"white\",\n", + " edgecolor=\"black\",\n", + " boxstyle=\"round,pad=0.2\",\n", + " linewidth=0.5,\n", + " alpha=0.7,\n", + " ),\n", + " )\n", + " return (stats, cb)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6b639275-e96d-43af-a999-5eb799aca5a8", + "metadata": {}, + "outputs": [], + "source": [ + "varname = \"vmo\"\n", + "varobj = diag.varmap[varname]\n", + "\n", + "xdim = \"xh\"\n", + "ydim = \"yq\"\n", + "zdim = \"z_l\"\n", + "interfaces = \"z_i\"\n", + "tdim = \"time\"\n", + "\n", + "rho0 = 1035.0\n", + "\n", + "moc_z_data = {}\n", + "\n", + "for region in diag.settings[\"diag_vars\"][\"regions\"]:\n", + " moc_z_data[region] = {}\n", + " for group in diag.groups:\n", + " print(f\"Calculating {region} moc z for {group.name}\")\n", + " ds = group.datasets[varobj]\n", + " deptho = mg.MOMgrid(ds.model, warn=False).to_xarray()[\"deptho\"]\n", + " deptho = xr.DataArray(deptho.values, dims=(\"yq\", \"xh\"))\n", + "\n", + " if len(ds[ydim]) != len(deptho[ydim]):\n", + " ds = ds.isel({ydim: slice(1, None)})\n", + "\n", + " if region == \"atlantic\":\n", + " # Calculate the basin mask on the v-grid\n", + " basin = cmip_basins.generate_basin_codes(ds, lon=\"geolon_v\", lat=\"geolat_v\")\n", + " mask = xr.where(\n", + " (basin == 2)\n", + " | (basin == 4)\n", + " | (basin == 6)\n", + " | (basin == 7)\n", + " | (basin == 8)\n", + " | (basin == 9),\n", + " 1.0,\n", + " np.nan,\n", + " )\n", + " else:\n", + " mask = 1.0\n", + "\n", + " y = (ds.geolat_v * mask).mean(xdim)\n", + " y = xr.DataArray(y.values, dims=(\"y\"), coords={\"y\": y.values})\n", + " yloc = xr.where(y.isnull(), False, True)\n", + "\n", + " z = ds[interfaces]\n", + "\n", + " psi = calc_moc_z(ds, varname, mask=mask)\n", + " psi = xr.DataArray(\n", + " psi.values, dims=(\"z\", \"y\"), coords={\"z\": z.values, \"y\": y.values}\n", + " )\n", + " psi = psi[:, yloc]\n", + "\n", + " depth = xr.DataArray(\n", + " (deptho * mask).max(\"xh\").values, dims=(\"y\"), coords={\"y\": y.values}\n", + " )\n", + " depth = depth[yloc]\n", + "\n", + " y = y[yloc]\n", + "\n", + " moc_z_data[region][group] = (y, z, psi, depth, region)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "491784a2-d3d0-4252-9a51-5d346890c006", + "metadata": {}, + "outputs": [], + "source": [ + "all_figs = []\n", + "\n", + "esnb.nbtools.setup_plots()\n", + "\n", + "for region in diag.settings[\"diag_vars\"][\"regions\"]:\n", + " nexps = len(diag.groups)\n", + " figsize, subplot = esnb.nbtools.get_figsize_subplots(nexps)\n", + " figsize = (5.41 * 1.25, 3.35 * 0.75 * 2)\n", + " fig = plt.figure(figsize=figsize, dpi=200)\n", + "\n", + " axes = []\n", + " for n, group in enumerate(diag.groups):\n", + " ax = plt.subplot(*subplot, n + 1)\n", + " stats, cb = plot_panel(ax, *moc_z_data[region][group])\n", + " ax.set_title(group.name)\n", + " axes.append(ax)\n", + "\n", + " if len(stats) > 0:\n", + " for k, v in stats.items():\n", + " group.add_metric(region, (k, float(v)))\n", + "\n", + " if region == \"global\":\n", + " for ax in axes:\n", + " ax.set_xlim(-78, None)\n", + "\n", + " plt.subplots_adjust(wspace=0.3)\n", + "\n", + " cbar = esnb.nbtools.bottom_colorbar(\n", + " fig, cb, orientation=\"horizontal\", extend=\"both\"\n", + " )\n", + " cbar.set_label(f\"{str(region).capitalize()} Meridional Streamfunction [Sv]\")\n", + "\n", + " # add letter labels for each panel\n", + " esnb.nbtools.panel_letters(axes, -0.12, 1.05)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14a1cf16-e012-49ae-adc5-6e3dc5eb4db2", + "metadata": {}, + "outputs": [], + "source": [ + "diag.write_metrics(\"MOC_metrics.json\")" + ] + } + ], + "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.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}