implement plots

bopen · Sep 26, 2023 · d170cdb · d170cdb
1 parent c94ad94
commit d170cdb
Showing 1 changed file with 94 additions and 50 deletions.
diff --git a/notebooks/wp4/cmip6_sea_ice_diagnostics.ipynb b/notebooks/wp4/cmip6_sea_ice_diagnostics.ipynb
@@ -497,7 +497,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Plot yearly timeseries"
+    "## Define plotting function"
    ]
   },
   {
@@ -506,59 +506,103 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "resample_freq = \"Y\"\n",
+    "def plot_timeseries(ds_era5, datasets_cmip6, func, title=None):\n",
+    "    quantiles = {\n",
+    "        experiment: func(ds)\n",
+    "        .quantile([0, 1 / 3, 1 / 2, 2 / 3, 1], dim=\"model\")\n",
+    "        .to_array()\n",
+    "        for experiment, ds in datasets_cmip6.items()\n",
+    "    }\n",
+    "    facet = func(ds_era5).to_array().plot(row=\"variable\", col=\"region\", label=\"ERA5\")\n",
+    "    for ax, sel_dict in zip(facet.axs.flatten(), facet.name_dicts.flatten()):\n",
+    "        for i, (experiment, da_quantiles) in enumerate(quantiles.items()):\n",
+    "            color = f\"C{i+1}\"\n",
+    "            da = da_quantiles.sel(sel_dict)\n",
+    "            ax.plot(\n",
+    "                da[\"time\"],\n",
+    "                da.sel(quantile=1 / 2),\n",
+    "                color=color,\n",
+    "                label=f\"CMIP6 {experiment} median\",\n",
+    "                zorder=2,\n",
+    "            )\n",
+    "            ax.fill_between(\n",
+    "                da[\"time\"],\n",
+    "                da.sel(quantile=1 / 3),\n",
+    "                da.sel(quantile=2 / 3),\n",
+    "                color=color,\n",
+    "                alpha=0.4,\n",
+    "                label=f\"CMIP6 {experiment} tertiles\",\n",
+    "                zorder=1,\n",
+    "            )\n",
+    "            ax.fill_between(\n",
+    "                da[\"time\"],\n",
+    "                da.sel(quantile=0),\n",
+    "                da.sel(quantile=1),\n",
+    "                color=color,\n",
+    "                alpha=0.2,\n",
+    "                label=f\"CMIP6 {experiment} range\",\n",
+    "                zorder=0,\n",
+    "            )\n",
+    "            ax.grid()\n",
     "\n",
-    "quantiles = {\n",
-    "    experiment: ds.resample(time=resample_freq)\n",
-    "    .mean()\n",
-    "    .quantile([0, 1 / 3, 1 / 2, 2 / 3, 1], dim=\"model\")\n",
-    "    .to_array()\n",
-    "    for experiment, ds in datasets_cmip6.items()\n",
-    "}\n",
+    "    for ax, sel_dict in zip(facet.axs[:, 0], facet.name_dicts[:, 0]):\n",
+    "        variable = sel_dict.pop(\"variable\")\n",
+    "        da = ds_era5.sel(sel_dict)[variable]\n",
+    "        ax.set_ylabel(f\"[{da.attrs['units']}]\")\n",
     "\n",
-    "facet = (\n",
-    "    ds_era5.resample(time=resample_freq)\n",
-    "    .mean()\n",
-    "    .to_array()\n",
-    "    .plot(row=\"variable\", col=\"region\", label=\"ERA5\")\n",
-    ")\n",
-    "for ax, sel_dict in zip(facet.axs.flatten(), facet.name_dicts.flatten()):\n",
-    "    for i, (experiment, da_quantiles) in enumerate(quantiles.items()):\n",
-    "        color = f\"C{i+1}\"\n",
-    "        da = da_quantiles.sel(sel_dict)\n",
-    "        ax.plot(\n",
-    "            da[\"time\"],\n",
-    "            da.sel(quantile=1 / 2),\n",
-    "            color=color,\n",
-    "            label=f\"CMIP6 {experiment} median\",\n",
-    "            zorder=2,\n",
-    "        )\n",
-    "        ax.fill_between(\n",
-    "            da[\"time\"],\n",
-    "            da.sel(quantile=1 / 3),\n",
-    "            da.sel(quantile=2 / 3),\n",
-    "            color=color,\n",
-    "            alpha=0.4,\n",
-    "            label=f\"CMIP6 {experiment} tertiles\",\n",
-    "            zorder=1,\n",
-    "        )\n",
-    "        ax.fill_between(\n",
-    "            da[\"time\"],\n",
-    "            da.sel(quantile=0),\n",
-    "            da.sel(quantile=1),\n",
-    "            color=color,\n",
-    "            alpha=0.2,\n",
-    "            label=f\"CMIP6 {experiment} range\",\n",
-    "            zorder=0,\n",
-    "        )\n",
-    "        ax.grid()\n",
+    "    facet.axs[0, -1].legend(bbox_to_anchor=(1.1, 1))\n",
+    "    if title is not None:\n",
+    "        facet.fig.suptitle(title)\n",
+    "    return facet"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Plot sliced timeseries"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def slice_dataset(ds, d_year=10):\n",
+    "    this_year = datetime.date.today().year\n",
+    "    time_slice = slice(str(this_year - d_year), str(this_year + d_year))\n",
+    "    return ds.sel(time=time_slice)\n",
     "\n",
-    "for ax, sel_dict in zip(facet.axs[:, 0], facet.name_dicts[:, 0]):\n",
-    "    variable = sel_dict.pop(\"variable\")\n",
-    "    da = ds_era5.sel(sel_dict)[variable]\n",
-    "    ax.set_ylabel(f\"[{da.attrs['units']}]\")\n",
     "\n",
-    "_ = facet.axs[0, -1].legend(bbox_to_anchor=(1.1, 1))"
+    "_ = plot_timeseries(\n",
+    "    ds_era5,\n",
+    "    datasets_cmip6,\n",
+    "    func=slice_dataset,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Plot max and min"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for reduction in (\"max\", \"min\"):\n",
+    "    plot_timeseries(\n",
+    "        ds_era5,\n",
+    "        datasets_cmip6,\n",
+    "        func=lambda ds: getattr(ds.resample(time=\"Y\"), reduction)(),\n",
+    "        title=f\"Yearly {reduction}\",\n",
+    "    )\n",
+    "    plt.show()"
    ]
   }
  ],