diff --git a/src/esnb/templates/oni_nino34.ipynb b/src/esnb/templates/oni_nino34.ipynb new file mode 100644 index 0000000..e4af267 --- /dev/null +++ b/src/esnb/templates/oni_nino34.ipynb @@ -0,0 +1,633 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "e4a43ed7-7624-4da9-b213-32f892574a65", + "metadata": {}, + "source": [ + "# ONI / NINO-3.4 Variability Notebook\n", + "\n", + "J. Krasting -- NOAA/GFDL" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e67185bc-4731-4c89-b63a-ca8953428f07", + "metadata": {}, + "outputs": [], + "source": [ + "# Development mode: constantly refreshes module code\n", + "%load_ext autoreload\n", + "%autoreload 2" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2b41c926-e3dc-4234-95d7-fcbf834af0a0", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "os.environ[\"ESNB_LOG_LEVEL\"] = \"INFO\"" + ] + }, + { + "cell_type": "markdown", + "id": "6898886b-b5e1-4de0-939a-f715748be915", + "metadata": {}, + "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\n", + "from esnb.sites.gfdl import call_dmget" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8fb540b1-c6c4-45dc-b77e-d55386a0e816", + "metadata": {}, + "outputs": [], + "source": [ + "diag_name = \"ONI Variability\"\n", + "diag_desc = \"ONI and Nino34 Variability Diagnostics\"\n", + "variables = [RequestedVariable(\"tos\", \"ocean_month\")]\n", + "user_options = {\"enso_region\": [\"nino12\", \"nino3\", \"nino34\", \"nino4\"]}\n", + "workdir = \"/vftmp/John.Krasting/cm5-enso-20250904\"\n", + "diag = NotebookDiagnostic(\n", + " diag_name, diag_desc, variables=variables, workdir=workdir, **user_options\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4de00b2c-1f62-41ce-bedf-a60f9bc8f6c3", + "metadata": {}, + "outputs": [], + "source": [ + "groups = [\n", + " CaseGroup2(\"895\", date_range=(\"0001-01-01\", \"0500-12-31\"), name=\"CM4.0 (895)\", plot_color=\"blue\"),\n", + " CaseGroup2(\"2916\", date_range=(\"0001-01-01\", \"0150-12-31\"), name=\"OM4_D5 (2916)\", plot_color=\"orange\"),\n", + " CaseGroup2(\"3031\", date_range=(\"0001-01-01\", \"0150-12-31\"), name=\"B11_D5 (3031)\", plot_color=\"green\"),\n", + " CaseGroup2(\"esm45-109\", date_range=(\"0001-01-01\", \"0250-12-31\"), name=\"ESM4.5 (esm45-109)\", plot_color=\"purple\"),\n", + "]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "353fff81-2c32-4163-964c-ae6f0e18fdb5", + "metadata": {}, + "outputs": [], + "source": [ + "diag.resolve(groups)" + ] + }, + { + "cell_type": "raw", + "id": "76bae4f6-e52f-48f4-87f8-d42f5f08cfb8", + "metadata": {}, + "source": [ + "call_dmget(diag.files)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b3d5a1c7-6606-4f81-b4df-ce5b281d6a6b", + "metadata": {}, + "outputs": [], + "source": [ + "diag.open(use_cache=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5d7d5c1b-879a-43d0-8c08-54f2575268cb", + "metadata": {}, + "outputs": [], + "source": [ + "diag.write_cache()" + ] + }, + { + "cell_type": "markdown", + "id": "b621a130-a8ec-4f9f-9b1c-e462180fc22b", + "metadata": {}, + "source": [ + "## Begin the User Diagnostic Code" + ] + }, + { + "cell_type": "markdown", + "id": "644824d1-2dcf-4369-b468-e2830b563ab0", + "metadata": {}, + "source": [ + "#### Load Modules" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aafa7a0a-61ed-4106-a5c7-a352529f65c6", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "import cftime\n", + "import matplotlib.pyplot as plt\n", + "import momlevel as ml\n", + "import numpy as np\n", + "import xarray as xr\n", + "from momgrid.geoslice import geoslice" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "47b5005d-93a3-4a8e-84b1-782624756128", + "metadata": {}, + "outputs": [], + "source": [ + "esnb.sites.gfdl.convert_to_momgrid(diag)" + ] + }, + { + "cell_type": "markdown", + "id": "e4be64f7-54b1-42bb-8da6-ab2577a52966", + "metadata": {}, + "source": [ + "#### Custom functions" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b8e30b20-78b7-41f5-b024-5935358b4bca", + "metadata": {}, + "outputs": [], + "source": [ + "# get the custom variable from the diagnostic settings\n", + "enso_region = diag.diag_vars.get(\"enso_region\", None)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "48e986f2-6d83-451d-bafb-93e719d423a2", + "metadata": {}, + "outputs": [], + "source": [ + "varname = \"sst\"\n", + "xdim = \"lon\"\n", + "ydim = \"lat\"\n", + "\n", + "time_coder = xr.coders.CFDatetimeCoder(use_cftime=True)\n", + "dsobs = xr.open_dataset(\"/home/jpk/NOAA.ER.v2.sst.nc\", decode_times=time_coder)\n", + "dsobs[\"area\"] = ml.util.standard_grid_cell_area(dsobs[ydim], dsobs[xdim])\n", + "\n", + "dsobs = dsobs.sel(time=slice(\"1891-01-01\", None))\n", + "\n", + "obs_enso_ts = {}\n", + "for region in enso_region:\n", + " ds = dsobs\n", + " if region == \"nino12\":\n", + " tos = ds[varname].sel({xdim: slice(270, 280), ydim: slice(-10, 0)})\n", + " area = ds[\"area\"].sel({xdim: slice(270, 280), ydim: slice(-10, 0)})\n", + " elif region == \"nino3\":\n", + " tos = ds[varname].sel({xdim: slice(210, 270), ydim: slice(-5, 5)})\n", + " area = ds[\"area\"].sel({xdim: slice(210, 270), ydim: slice(-5, 5)})\n", + " elif region == \"nino34\":\n", + " tos = ds[varname].sel({xdim: slice(190, 240), ydim: slice(-5, 5)})\n", + " area = ds[\"area\"].sel({xdim: slice(190, 240), ydim: slice(-5, 5)})\n", + " elif region == \"nino4\":\n", + " tos = ds[varname].sel({xdim: slice(170, 210), ydim: slice(-5, 5)})\n", + " area = ds[\"area\"].sel({xdim: slice(170, 210), ydim: slice(-5, 5)})\n", + " else:\n", + " print(f\"Unknown region: {region}\")\n", + " tos = tos.weighted(area).mean((xdim, ydim))\n", + " obs_enso_ts[region] = tos.load()" + ] + }, + { + "cell_type": "markdown", + "id": "f25eba56-c14d-4264-b000-3b99bd16c70a", + "metadata": {}, + "source": [ + "### Timeseries plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ae675e94-6d92-446e-9552-0333180f9c72", + "metadata": {}, + "outputs": [], + "source": [ + "# abstract out the dimension names here\n", + "xdim = \"xh\"\n", + "ydim = \"yh\"\n", + "tvar = \"tos\"\n", + "areavar = \"areacello\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e6583696-7f2f-40e0-b7f2-92df88e2a5cf", + "metadata": {}, + "outputs": [], + "source": [ + "# Loop over groups and regions to extract timeseries\n", + "\n", + "varname = \"tos\"\n", + "varobj = diag.varmap[varname]\n", + "\n", + "enso_ts = {}\n", + "\n", + "for region in enso_region:\n", + " enso_ts[region] = {}\n", + " for group in diag.groups:\n", + " ds = group.datasets[varobj]\n", + " if region == \"nino12\":\n", + " tos = geoslice(ds[varname], x=(-90, -80), y=(-10, 0))\n", + " area = geoslice(ds[areavar], x=(-90, -80), y=(-10, 0))\n", + " elif region == \"nino3\":\n", + " tos = geoslice(ds[varname], x=(-150, -90), y=(-5, 5))\n", + " area = geoslice(ds[areavar], x=(-150, -90), y=(-5, 5))\n", + " elif region == \"nino34\":\n", + " tos = geoslice(ds[varname], x=(-170, -120), y=(-5, 5))\n", + " area = geoslice(ds[areavar], x=(-170, -120), y=(-5, 5))\n", + " elif region == \"nino4\":\n", + " tos = geoslice(ds[varname], x=(-190, -150), y=(-5, 5))\n", + " area = geoslice(ds[areavar], x=(-190, -150), y=(-5, 5))\n", + " else:\n", + " print(f\"Unknown region: {region}\")\n", + " tos = tos.weighted(area).mean((xdim, ydim))\n", + " enso_ts[region][group] = tos.load()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a16699ca-8af3-417b-bca8-716400f7fee1", + "metadata": {}, + "outputs": [], + "source": [ + "def plot_oni(ts, detrend=True, ax=None):\n", + "\n", + " # detrend, deseaon, and calculate the monthly anomalies and their stddev\n", + " arr = ml.trend.linear_detrend(ts, mode=\"correct\")\n", + " ac = ml.util.annual_cycle(arr)\n", + " ac = list(ac.values) * int(len(arr) / 12)\n", + "\n", + " # 3-month triangle filter\n", + " arr = arr.rolling(time=3, center=True).mean()\n", + " anom = arr - ac\n", + "\n", + " # Create masks for significant anomalies that persist 3+ months\n", + " pos_mask = anom >= 0.5\n", + " neg_mask = anom <= -0.5\n", + "\n", + " # Find consecutive runs of 5+ months\n", + " pos_persistent = pos_mask & (pos_mask.rolling(time=5, center=True).sum() >= 5)\n", + " neg_persistent = neg_mask & (neg_mask.rolling(time=5, center=True).sum() >= 5)\n", + "\n", + " # Count events (transitions from False to True mark new events)\n", + " pos_events = (\n", + " ((pos_persistent == True) & (pos_persistent.shift(time=1) == False))\n", + " .sum()\n", + " .values\n", + " )\n", + " neg_events = (\n", + " ((neg_persistent == True) & (neg_persistent.shift(time=1) == False))\n", + " .sum()\n", + " .values\n", + " )\n", + "\n", + " tax = anom.time.values\n", + "\n", + " if ax is None:\n", + " fig, ax = plt.subplots(figsize=(18, 4))\n", + "\n", + " ax.plot(tax, anom, color=\"black\", linewidth=0.5)\n", + "\n", + " ax.fill_between(tax, 0.5, anom, where=pos_persistent, color=\"red\", alpha=0.7)\n", + " ax.fill_between(tax, -0.5, anom, where=neg_persistent, color=\"blue\", alpha=0.7)\n", + "\n", + " plt.axhline(0, color=\"black\", linewidth=0.5)\n", + " ax.set_xlim(tax[0], tax[-1])\n", + "\n", + " return (ax, tax, pos_events, neg_events)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bab6e349-68aa-4f67-bbf4-945ba8d2c706", + "metadata": {}, + "outputs": [], + "source": [ + "esnb.nbtools.setup_plots()\n", + "\n", + "ngroups = len(diag.groups) + 1\n", + "figsize = (esnb.nbtools.FULL_PAGE, esnb.nbtools.SINGLE_COLUMN * ngroups)\n", + "subplots = (ngroups, 1)\n", + "\n", + "fig = plt.figure(figsize=figsize, dpi=150)\n", + "for n in range(0, len(diag.groups) + 1):\n", + " if n == 0:\n", + " arr = obs_enso_ts[\"nino34\"]\n", + " name = \"NOAA_ERSST v2\"\n", + " group = None\n", + " else:\n", + " group = diag.groups[n - 1]\n", + " arr = enso_ts[\"nino34\"][group]\n", + " name = group.name\n", + "\n", + " ax = plt.subplot(*subplots, n + 1)\n", + " ax, tax, pos_events, neg_events = plot_oni(arr, ax=ax)\n", + " ax.text(0, 1.03, f\"ONI (Nino3.4) - {name}\", transform=ax.transAxes, fontsize=9)\n", + "\n", + " pos_per_dec = round((pos_events / (len(tax) / 12.0)) * 10, 2)\n", + " neg_per_dec = round((neg_events / (len(tax) / 12.0)) * 10, 2)\n", + "\n", + " ax.text(\n", + " 1.0,\n", + " 1.06,\n", + " f\"Positive Events: {pos_events} ({pos_per_dec} / decade)\",\n", + " transform=ax.transAxes,\n", + " fontsize=7,\n", + " ha=\"right\",\n", + " )\n", + " ax.text(\n", + " 1.0,\n", + " 1.015,\n", + " f\"Negative Events: {neg_events} ({neg_per_dec} / decade)\",\n", + " transform=ax.transAxes,\n", + " fontsize=7,\n", + " ha=\"right\",\n", + " )\n", + "\n", + " if group is not None:\n", + " group.add_metric(f\"oni_events\", (\"positive\", float(pos_events)))\n", + " group.add_metric(f\"oni_events\", (\"negative\", float(neg_events)))\n", + " group.add_metric(f\"oni_events\", (\"positive_per_dec\", float(pos_per_dec)))\n", + " group.add_metric(f\"oni_events\", (\"negative_per_dec\", float(neg_per_dec)))\n", + " group.add_metric(f\"oni_events\", (\"nmonths\", int(len(tax))))" + ] + }, + { + "cell_type": "markdown", + "id": "7bbdd7a4-6c8f-491e-bd22-826b2cf66bd1", + "metadata": {}, + "source": [ + "### Wavelet Analysis" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "452e8b8c-e7c2-4b40-af1f-f1c84ed278aa", + "metadata": {}, + "outputs": [], + "source": [ + "import xwavelet as xw" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e72f4cd3-faf6-4952-ab17-a259a5239178", + "metadata": {}, + "outputs": [], + "source": [ + "def scale_line_widths(ax, scale_factor):\n", + " \"\"\"\n", + " Scale all line widths in a matplotlib axis by a given factor.\n", + "\n", + " Parameters:\n", + " ax : matplotlib.axes.Axes\n", + " The axis object containing the plots\n", + " scale_factor : float\n", + " Factor to multiply all line widths by\n", + " \"\"\"\n", + " # Scale line widths for regular line plots\n", + " for line in ax.get_lines():\n", + " current_width = line.get_linewidth()\n", + " line.set_linewidth(current_width * scale_factor)\n", + "\n", + " # Scale line widths for contour plots\n", + " for collection in ax.collections:\n", + " # Handle LineCollection objects (used by contour plots)\n", + " if hasattr(collection, \"get_linewidths\"):\n", + " current_widths = collection.get_linewidths()\n", + " if current_widths is not None:\n", + " # LineCollection can have array of widths or single width\n", + " if hasattr(current_widths, \"__iter__\") and len(current_widths) > 1:\n", + " new_widths = [w * scale_factor for w in current_widths]\n", + " else:\n", + " new_widths = current_widths * scale_factor\n", + " collection.set_linewidths(new_widths)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "86ae71eb-d978-401d-a00c-1e13bedd8c7a", + "metadata": {}, + "outputs": [], + "source": [ + "figsize = (esnb.nbtools.FULL_PAGE, np.ceil((ngroups + 1) / 2) * 2.5)\n", + "subplots = (int(np.ceil((ngroups + 2) / 2)), 2)\n", + "\n", + "axes = []\n", + "fig = plt.figure(figsize=figsize, dpi=200)\n", + "for n in range(0, len(diag.groups) + 1):\n", + " ax = plt.subplot(*subplots, n + 1)\n", + " if n == 0:\n", + " arr = obs_enso_ts[\"nino34\"]\n", + " name = \"NOAA_ERSST v2\"\n", + " group = None\n", + " else:\n", + " group = groups[n - 1]\n", + " arr = enso_ts[\"nino34\"][group]\n", + " name = group\n", + "\n", + " result = xw.Wavelet(arr, scaled=True)\n", + " _ = result.density(ax=ax)\n", + " scale_line_widths(ax, 0.5)\n", + " ax.text(\n", + " 0,\n", + " 1.03,\n", + " f\"Wavelet Density (Nino3.4) - {name}\",\n", + " transform=ax.transAxes,\n", + " fontsize=6,\n", + " )\n", + " ax.set_xlabel(None)\n", + "\n", + " plt.subplots_adjust(hspace=0.4, wspace=0.4)\n", + " axes.append(ax)\n", + "\n", + "esnb.nbtools.panel_letters(axes)" + ] + }, + { + "cell_type": "markdown", + "id": "887fea56-b4ab-4a0a-9104-d5411d6413f7", + "metadata": {}, + "source": [ + "### Frequency Spectra" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "253817cc-8595-4ea3-88cc-7ac46c6f7743", + "metadata": {}, + "outputs": [], + "source": [ + "def adjust_recent_line(ax, color=None, linewidth=None, label=None):\n", + " \"\"\"\n", + " Adjust the color, linewidth, and label of the most recently added line plot.\n", + "\n", + " Parameters:\n", + " -----------\n", + " ax : matplotlib.axes.Axes\n", + " The axis handle containing the line plots\n", + " color : str or tuple, optional\n", + " New color for the line (e.g., 'red', 'blue', '#FF5733', (0.2, 0.4, 0.6))\n", + " linewidth : float, optional\n", + " New line width\n", + " label : str, optional\n", + " New label for the line\n", + "\n", + " Returns:\n", + " --------\n", + " matplotlib.lines.Line2D\n", + " The modified line object\n", + " \"\"\"\n", + " # Get all line objects from the axis\n", + " lines = ax.get_lines()\n", + "\n", + " if not lines:\n", + " raise ValueError(\"No line plots found on the axis\")\n", + "\n", + " # Get the most recent line (last in the list)\n", + " recent_line = lines[-1]\n", + "\n", + " # Apply modifications if parameters are provided\n", + " if color is not None:\n", + " recent_line.set_color(color)\n", + "\n", + " if linewidth is not None:\n", + " recent_line.set_linewidth(linewidth)\n", + "\n", + " if label is not None:\n", + " recent_line.set_label(label)\n", + "\n", + " # Refresh the plot to show changes\n", + " ax.figure.canvas.draw_idle()\n", + "\n", + " return recent_line" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0d77cfa9-e4d7-4e3a-b53e-deb626bfcdcb", + "metadata": {}, + "outputs": [], + "source": [ + "figsize = (esnb.nbtools.SINGLE_COLUMN, esnb.nbtools.SINGLE_COLUMN * 1.7)\n", + "\n", + "fig = plt.figure(figsize=figsize, dpi=200)\n", + "ax = plt.subplot(1, 1, 1)\n", + "\n", + "for n in range(0, len(diag.groups) + 1):\n", + " if n == 0:\n", + " arr = obs_enso_ts[\"nino34\"]\n", + " name = \"NOAA_ERSST v2\"\n", + " group = None\n", + " color = \"black\"\n", + " linewidth = 1.5\n", + " else:\n", + " group = groups[n - 1]\n", + " arr = enso_ts[\"nino34\"][group]\n", + " name = group\n", + " color = group.plot_color\n", + " linewidth = 1.0\n", + "\n", + " result = xw.Wavelet(arr, scaled=True)\n", + " result.spectrum(ax=ax)\n", + " adjust_recent_line(ax, color=color, linewidth=linewidth, label=name)\n", + "\n", + " ax.grid(True, linewidth=0.2)\n", + " plt.legend(loc=4, fontsize=6)\n", + "\n", + " ax.text(\n", + " 0.0, 1.02, \"Frequency Spectra - Nino3.4 SST\", fontsize=8, transform=ax.transAxes\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "7d61f10f-5fde-4733-a31a-90fd399667ad", + "metadata": {}, + "source": [ + "### Part 4: Write metrics file" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8db94dba-389f-4410-89eb-9d4c8761a09f", + "metadata": {}, + "outputs": [], + "source": [ + "diag.write_metrics()" + ] + } + ], + "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 +}