examples/nblibrary/atm/nmse_PSL.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2292c691-9bd9-44d2-8a3f-cb90dbe2e383",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from nmse_utils import *\n",
    "from averaging_utils import *\n",
    "import glob\n",
    "\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d67416c-a2d4-403b-85f4-647aa0a816eb",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "### Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7486e94-e493-4369-9767-90eb15c0ac3a",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "sname = \"\"\n",
    "CESM_output_dir = \"\"\n",
    "case_name = \"\"\n",
    "start_date = \"\"\n",
    "end_date = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14b554f3-40e1-4c12-9f26-c1684ac03fbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(start_date)\n",
    "print(end_date)\n",
    "print(CESM_output_dir)\n",
    "print(case_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80edad94-426a-4304-9003-78e12960613a",
   "metadata": {},
   "outputs": [],
   "source": [
    "dat = xr.open_mfdataset(\n",
    "    CESM_output_dir + \"/\" + case_name + \"/atm/proc/tseries/*.PSL*.nc\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0527e3e-cd26-46b5-8c1e-08882109e12e",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "### Read in validation data and other CMIP models for comparison (precomputed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "126e65b3-2b8c-400c-af02-2ad0b0f82e6e",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# This will need to be changed if we have a common repo\n",
    "validpath = \"/glade/campaign/cgd/cas/islas/python_savs/CUPID/PSL_validation/\"\n",
    "validpath_old = (\n",
    "    \"/glade/campaign/cgd/cas/islas/python_savs/CUPID/NMSE/validation_data_example/\"\n",
    ")\n",
    "\n",
    "# ---ERA5\n",
    "era5 = xr.open_dataset(validpath_old + \"PSL_ERA5.nc\")\n",
    "# era5 = xr.open_dataset(\"/glade/campaign/cgd/cas/islas/python_savs/CUPID/NMSE/validation_data_example/PSL_ERA5.nc\")\n",
    "era5 = era5 / 100.0  # convert to hPa\n",
    "\n",
    "# ---CESM2\n",
    "lens2 = xr.open_dataset(validpath_old + \"PSL_LENS2.nc\")\n",
    "lens2 = lens2 / 100.0  # convert to hPa\n",
    "lens2[\"lon\"] = era5.lon\n",
    "lens2[\"lat\"] = era5.lat\n",
    "\n",
    "# ---CMIP6\n",
    "modelfiles = sorted(glob.glob(validpath_old + \"CMIP6/*.nc\"))\n",
    "datcmip6 = [xr.open_dataset(ifile).mean(\"M\") for ifile in modelfiles]\n",
    "datcmip6 = xr.concat(datcmip6, dim=\"model\")\n",
    "datcmip6 = datcmip6 / 100.0\n",
    "datcmip6[\"lon\"] = era5.lon\n",
    "datcmip6[\"lat\"] = era5.lat"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab5f1441-4b6b-4386-b3aa-a0f6fa9109b8",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "### Read in the current case"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a36af8c-c8c0-42df-b162-07c6c244ab0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "CESM_output_dir = \"/glade/campaign/cesm/development/cross-wg/diagnostic_framework/CESM_output_for_testing/\"\n",
    "case_name = \"b.e23_alpha17f.BLT1850.ne30_t232.092\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b4883ef-f3a7-458b-b96e-bdb79d5abdd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "dat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09397e30-62b9-474b-a47c-4a1174137bce",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Need to change file paths here to be read in using the parameters thing\n",
    "dat = xr.open_mfdataset(\n",
    "    CESM_output_dir + \"/\" + case_name + \"/atm/proc/tseries/*.PSL*.nc\"\n",
    ")\n",
    "dat = dat.sel(time=slice(start_date, end_date))\n",
    "\n",
    "dat = dat.PSL / 100.0\n",
    "dat = dat.expand_dims(lon=era5.lon, lat=era5.lat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "073a2ad0-81e6-4817-9024-4b9b718fabb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --Compute seasonal and annual means\n",
    "dat = seasonal_climatology_weighted(dat).load()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22cc331d-413c-4a87-bd89-812ad118cf8c",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "### Compute the NMSE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6857717d-7514-45b5-ba33-a774f38b7c3e",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"moving on to compute the nmse\")\n",
    "nmse_dat = []\n",
    "nmse_cesm2 = []\n",
    "nmse_cmip6 = []\n",
    "for ivar in era5.data_vars:\n",
    "    nmse_dat.append(nmse(era5[ivar], dat[ivar]))\n",
    "    nmse_cesm2.append(nmse(era5[ivar], lens2[ivar]))\n",
    "    nmse_cmip6.append(nmse(era5[ivar], datcmip6[ivar]))\n",
    "nmse_dat = xr.merge(nmse_dat)\n",
    "nmse_cesm2 = xr.merge(nmse_cesm2)\n",
    "nmse_cmip6 = xr.merge(nmse_cmip6)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1014f119-fc3f-428b-99ca-ab9de700148d",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "### Set up the plot panel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53494900-0145-4ab2-85b8-5ed6ae347892",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def plotnmse(fig, cmip6, cesm2, cesm3, x1, x2, y1, y2, titlestr):\n",
    "    ax = fig.add_axes([x1, y1, x2 - x1, y2 - y1])\n",
    "\n",
    "    cmip6 = cmip6.sortby(cmip6, ascending=False)\n",
    "    binedges = np.arange(0, cmip6.size, 1)\n",
    "    ax.bar(\n",
    "        binedges,\n",
    "        cmip6,\n",
    "        width=1,\n",
    "        bottom=0,\n",
    "        edgecolor=\"black\",\n",
    "        color=\"gray\",\n",
    "        label=\"CMIP6\",\n",
    "    )\n",
    "\n",
    "    ax.plot(cmip6.size + 1, cesm3, \"o\", color=\"blue\", label=\"THIS RUN\")\n",
    "\n",
    "    ax.fill_between(\n",
    "        np.arange(0, cmip6.size + 3, 1) - 0.5,\n",
    "        np.arange(0, cmip6.size + 3, 1) * 0 + np.array(cesm2.min()),\n",
    "        np.arange(0, cmip6.size + 3, 1) * 0 + np.array(cesm2.max()),\n",
    "        color=\"salmon\",\n",
    "        alpha=0.5,\n",
    "        label=\"LENS2\",\n",
    "    )\n",
    "\n",
    "    ax.set_xlim(-0.5, cmip6.size + 2 - 0.5)\n",
    "    ax.set_xticks([])\n",
    "    ax.set_ylabel(\"NMSE\", fontsize=14)\n",
    "    ax.set_title(titlestr, fontsize=16)\n",
    "\n",
    "    ax.legend()\n",
    "\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56b4cd99-a27e-4f28-86c2-8013e7c7bc78",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(16, 16))\n",
    "\n",
    "ax = plotnmse(\n",
    "    fig,\n",
    "    nmse_cmip6.AM,\n",
    "    nmse_cesm2.AM,\n",
    "    nmse_dat.AM,\n",
    "    0.3,\n",
    "    0.7,\n",
    "    0.77,\n",
    "    0.93,\n",
    "    \"NMSE, SLP, AM\",\n",
    ")\n",
    "ax = plotnmse(\n",
    "    fig,\n",
    "    nmse_cmip6.DJF,\n",
    "    nmse_cesm2.DJF,\n",
    "    nmse_dat.DJF,\n",
    "    0.05,\n",
    "    0.45,\n",
    "    0.57,\n",
    "    0.72,\n",
    "    \"NMSE, SLP, DJF\",\n",
    ")\n",
    "ax = plotnmse(\n",
    "    fig,\n",
    "    nmse_cmip6.MAM,\n",
    "    nmse_cesm2.MAM,\n",
    "    nmse_dat.MAM,\n",
    "    0.55,\n",
    "    0.95,\n",
    "    0.57,\n",
    "    0.72,\n",
    "    \"NMSE, SLP, MAM\",\n",
    ")\n",
    "ax = plotnmse(\n",
    "    fig,\n",
    "    nmse_cmip6.JJA,\n",
    "    nmse_cesm2.JJA,\n",
    "    nmse_dat.JJA,\n",
    "    0.05,\n",
    "    0.45,\n",
    "    0.37,\n",
    "    0.52,\n",
    "    \"NMSE, SLP, JJA\",\n",
    ")\n",
    "ax = plotnmse(\n",
    "    fig,\n",
    "    nmse_cmip6.SON,\n",
    "    nmse_cesm2.SON,\n",
    "    nmse_dat.SON,\n",
    "    0.55,\n",
    "    0.95,\n",
    "    0.37,\n",
    "    0.52,\n",
    "    \"NMSE, SLP, SON\",\n",
    ")\n",
    "\n",
    "fig.text(\n",
    "    0.5,\n",
    "    0.99,\n",
    "    \"THIS RUN = \" + case_name + \" \" + start_date + \" to \" + end_date,\n",
    "    ha=\"center\",\n",
    "    va=\"center\",\n",
    "    fontsize=14,\n",
    "    color=\"royalblue\",\n",
    ")\n",
    "fig.text(\n",
    "    0.5,\n",
    "    0.975,\n",
    "    \"Other runs = 1979-01-01 to 2023-12-31\",\n",
    "    ha=\"center\",\n",
    "    va=\"center\",\n",
    "    fontsize=14,\n",
    ")\n",
    "fig.text(\n",
    "    0.5,\n",
    "    0.96,\n",
    "    \"Validation data = ERA5 1979-01-01 to 2023-12-31\",\n",
    "    ha=\"center\",\n",
    "    va=\"center\",\n",
    "    fontsize=14,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d01c7860-b4ca-4ab4-a05f-3d0f248a56c2",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:cupid-analysis]",
   "language": "python",
   "name": "conda-env-cupid-analysis-py"
  },
  "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}