4a.evaluation-checkpoint.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NCI WeatherBench-4a: Evaluate forecasts (Calculate)\n",
    "\n",
    "In this notebook, we will compute all the scores for the baseline models and create the figures for the paper. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import warnings\n",
    "import os\n",
    "warnings.filterwarnings('ignore')\n",
    "os.environ[\"TF_ENABLE_ONEDNN_OPTS\"] = \"0\"\n",
    "from score import *\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "from datetime import datetime\n",
    "from dask.distributed import Client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.transition-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.scheduler.events-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/mlenv/lib/python3.9/site-packages/dask/config.py:742: FutureWarning: Dask configuration key 'distributed.comm.recent-messages-log-length' has been deprecated; please use 'distributed.admin.low-level-log-length' instead\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\"> </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px;\">Client</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Client-70745116-eb83-11ee-817e-000003f9fe80</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "\n",
       "        <tr>\n",
       "        \n",
       "            <td style=\"text-align: left;\"><strong>Connection method:</strong> Cluster object</td>\n",
       "            <td style=\"text-align: left;\"><strong>Cluster type:</strong> distributed.LocalCluster</td>\n",
       "        \n",
       "        </tr>\n",
       "\n",
       "        \n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard: </strong> <a href=\"/proxy/8787/status\" target=\"_blank\">/proxy/8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\"></td>\n",
       "            </tr>\n",
       "        \n",
       "\n",
       "        </table>\n",
       "\n",
       "        \n",
       "            <button style=\"margin-bottom: 12px;\" data-commandlinker-command=\"dask:populate-and-launch-layout\" data-commandlinker-args='{\"url\": \"/proxy/8787/status\" }'>\n",
       "                Launch dashboard in JupyterLab\n",
       "            </button>\n",
       "        \n",
       "\n",
       "        \n",
       "            <details>\n",
       "            <summary style=\"margin-bottom: 20px;\"><h3 style=\"display: inline;\">Cluster Info</h3></summary>\n",
       "            <div class=\"jp-RenderedHTMLCommon jp-RenderedHTML jp-mod-trusted jp-OutputArea-output\">\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\">\n",
       "    </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px; margin-top: 0px;\">LocalCluster</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">ecf64e7a</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard:</strong> <a href=\"/proxy/8787/status\" target=\"_blank\">/proxy/8787/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Workers:</strong> 12\n",
       "                </td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total threads:</strong> 12\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total memory:</strong> 95.00 GiB\n",
       "                </td>\n",
       "            </tr>\n",
       "            \n",
       "            <tr>\n",
       "    <td style=\"text-align: left;\"><strong>Status:</strong> running</td>\n",
       "    <td style=\"text-align: left;\"><strong>Using processes:</strong> True</td>\n",
       "</tr>\n",
       "\n",
       "            \n",
       "        </table>\n",
       "\n",
       "        <details>\n",
       "            <summary style=\"margin-bottom: 20px;\">\n",
       "                <h3 style=\"display: inline;\">Scheduler Info</h3>\n",
       "            </summary>\n",
       "\n",
       "            <div style=\"\">\n",
       "    <div>\n",
       "        <div style=\"width: 24px; height: 24px; background-color: #FFF7E5; border: 3px solid #FF6132; border-radius: 5px; position: absolute;\"> </div>\n",
       "        <div style=\"margin-left: 48px;\">\n",
       "            <h3 style=\"margin-bottom: 0px;\">Scheduler</h3>\n",
       "            <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Scheduler-574457bc-bd33-49b5-a953-5aa77360e3d6</p>\n",
       "            <table style=\"width: 100%; text-align: left;\">\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Comm:</strong> tcp://127.0.0.1:41565\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Workers:</strong> 12\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Dashboard:</strong> <a href=\"/proxy/8787/status\" target=\"_blank\">/proxy/8787/status</a>\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total threads:</strong> 12\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Started:</strong> Just now\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total memory:</strong> 95.00 GiB\n",
       "                    </td>\n",
       "                </tr>\n",
       "            </table>\n",
       "        </div>\n",
       "    </div>\n",
       "\n",
       "    <details style=\"margin-left: 48px;\">\n",
       "        <summary style=\"margin-bottom: 20px;\">\n",
       "            <h3 style=\"display: inline;\">Workers</h3>\n",
       "        </summary>\n",
       "\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 0</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:43443\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"/proxy/36149/status\" target=\"_blank\">/proxy/36149/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.92 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:39369\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /jobfs/111819797.gadi-pbs/dask-scratch-space/worker-9q0sp5j0\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 32.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 1</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:41745\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"/proxy/43599/status\" target=\"_blank\">/proxy/43599/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.92 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:46715\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /jobfs/111819797.gadi-pbs/dask-scratch-space/worker-b435rdrq\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 32.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 2</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:35551\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"/proxy/32869/status\" target=\"_blank\">/proxy/32869/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.92 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:34357\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /jobfs/111819797.gadi-pbs/dask-scratch-space/worker-3jugqutc\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 32.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 3</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:45847\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"/proxy/44811/status\" target=\"_blank\">/proxy/44811/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.92 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:35175\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /jobfs/111819797.gadi-pbs/dask-scratch-space/worker-_fvxbntj\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 32.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 4</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:43987\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"/proxy/38395/status\" target=\"_blank\">/proxy/38395/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.92 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:33087\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /jobfs/111819797.gadi-pbs/dask-scratch-space/worker-zqz5m1yq\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 32.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 5</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:46441\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"/proxy/42703/status\" target=\"_blank\">/proxy/42703/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.92 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:45517\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /jobfs/111819797.gadi-pbs/dask-scratch-space/worker-bdqflym5\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 32.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 6</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:35819\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"/proxy/46615/status\" target=\"_blank\">/proxy/46615/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.92 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:36933\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /jobfs/111819797.gadi-pbs/dask-scratch-space/worker-bjhudeq7\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 32.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 7</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:40171\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"/proxy/33685/status\" target=\"_blank\">/proxy/33685/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.92 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:34183\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /jobfs/111819797.gadi-pbs/dask-scratch-space/worker-50y_zaga\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 32.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 8</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:39225\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"/proxy/40495/status\" target=\"_blank\">/proxy/40495/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.92 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:44595\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /jobfs/111819797.gadi-pbs/dask-scratch-space/worker-c1rjcv5a\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 32.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 9</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:45715\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"/proxy/41089/status\" target=\"_blank\">/proxy/41089/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.92 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:44947\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /jobfs/111819797.gadi-pbs/dask-scratch-space/worker-r7xzu6mp\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 32.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 10</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:41887\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"/proxy/34955/status\" target=\"_blank\">/proxy/34955/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.92 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:46323\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /jobfs/111819797.gadi-pbs/dask-scratch-space/worker-4ejsoz4r\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 32.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 11</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:36263\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 1\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"/proxy/36937/status\" target=\"_blank\">/proxy/36937/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.92 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:43313\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /jobfs/111819797.gadi-pbs/dask-scratch-space/worker-7oy8tybk\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU: </strong>Tesla V100-SXM2-32GB\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>GPU memory: </strong> 32.00 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "\n",
       "    </details>\n",
       "</div>\n",
       "\n",
       "        </details>\n",
       "    </div>\n",
       "</div>\n",
       "            </details>\n",
       "        \n",
       "\n",
       "    </div>\n",
       "</div>"
      ],
      "text/plain": [
       "<Client: 'tcp://127.0.0.1:41565' processes=12 threads=12, memory=95.00 GiB>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client = Client(n_workers=12, threads_per_worker=1)  \n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import xarray as xr\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from score import *\n",
    "from collections import OrderedDict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "def to_pickle(obj, fn):\n",
    "    with open(fn, 'wb') as f:\n",
    "        pickle.dump(obj, f)\n",
    "def read_pickle(fn):\n",
    "    with open(fn, 'rb') as f:\n",
    "        return pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.set_style('darkgrid')\n",
    "sns.set_context('notebook')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "res: 5.625\n",
      "DATADIR: /g/data/wb00/NCI-Weatherbench/5.625deg\n",
      "PREDDIR: /scratch/vp91/mah900/NCI-Weatherbench/pred_dir\n"
     ]
    }
   ],
   "source": [
    "res = '5.625'\n",
    "print ('res:', res)\n",
    "DATADIR = f'/g/data/wb00/NCI-Weatherbench/{res}deg' \n",
    "PREDDIR = f\"/scratch/vp91/{os.environ['USER']}/NCI-Weatherbench/pred_dir\"\n",
    "print('DATADIR:', DATADIR)\n",
    "print('PREDDIR:', PREDDIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validation data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-03-27 02:13:49]\n",
      "/g/data/wb00/NCI-Weatherbench/5.625deg/geopotential_500/\n",
      "[2021, 2022]\n",
      "load_test_data, var: z\n",
      "load_test_data, var: t\n",
      "load_test_data, var: tp\n",
      "load_test_data, var: t2m\n",
      "CPU times: user 6.35 s, sys: 876 ms, total: 7.22 s\n",
      "Wall time: 8.9 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "\n",
    "print (f'{DATADIR}/geopotential_500/')\n",
    "valid_years = list(range( 2021, 2022+1))\n",
    "print (valid_years)\n",
    "z500_valid_files = [ file for year in valid_years for file in glob.glob (fr'{DATADIR}/geopotential/*{year}*') ] \n",
    "t850_valid_files = [ file for year in valid_years for file in glob.glob (fr'{DATADIR}/temperature/*{year}*')  ] \n",
    "z500_valid = load_test_data(z500_valid_files, 'z', slice('2021', '2022')).load()  \n",
    "t850_valid = load_test_data(t850_valid_files, 't', slice('2021', '2022')).load() \n",
    "\n",
    "# for precipitation we are taking 6 hourly accumulations\n",
    "tp_valid_files  = [ file for year in valid_years for file in glob.glob (fr'{DATADIR}/total_precipitation/*{year}*') ] \n",
    "t2m_valid_files = [ file for year in valid_years for file in glob.glob (fr'{DATADIR}/2m_temperature/*{year}*')  ] \n",
    "tp_valid = load_test_data(tp_valid_files,   'tp', slice('2021', '2022')).rolling(time=6).sum()\n",
    "tp_valid.name = 'tp'\n",
    "t2m_valid = load_test_data(t2m_valid_files, 't2m', slice('2021', '2022'))\n",
    "\n",
    "valid_data = xr.merge([z500_valid, t850_valid, tp_valid, t2m_valid])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load data\n",
    "\n",
    "First let's load all the prediction files from the previous notebooks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# to speed things up, let's only evaluate every 6th hour\n",
    "valid_data = valid_data.isel(time=slice(0, None, 6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-03-27 02:13:58]\n",
      "/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/NCI_tutorial_persistence_5.625.nc\n",
      "/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/NCI_tutorial_climatology_5.625.nc\n",
      "CPU times: user 1min 13s, sys: 19.3 s, total: 1min 32s\n",
      "Wall time: 3min 29s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "save_prefix = 'NCI_tutorial' \n",
    "print (f'{PREDDIR}/{save_prefix}_persistence_{res}.nc')\n",
    "print (f'{PREDDIR}/{save_prefix}_climatology_{res}.nc')\n",
    "persistence = xr.open_dataset(f'{PREDDIR}/{save_prefix}_persistence_{res}.nc', chunks={'time': 10}).load()\n",
    "climatology = xr.open_dataset(f'{PREDDIR}/{save_prefix}_climatology_{res}.nc', chunks={'time': 10}).load()\n",
    "weekly_climatology = xr.open_dataset(f'{PREDDIR}/{save_prefix}_weekly_climatology_{res}.nc', chunks={'time': 10}).load()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load linear regression checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-03-27 02:17:27]\n",
      "Linear regression - 3d - Merging data ...\n",
      "/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/NCI_tutorial_lr_3d_z_z.nc\n",
      "/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/NCI_tutorial_lr_3d_t_t.nc\n",
      "/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/NCI_tutorial_lr_3d_tp_tp.nc\n",
      "/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/NCI_tutorial_lr_3d_t2m_t2m.nc\n",
      "Linear regression - 5d - Merging data ...\n",
      "/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/NCI_tutorial_lr_5d_z_z.nc\n",
      "/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/NCI_tutorial_lr_5d_t_t.nc\n",
      "/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/NCI_tutorial_lr_5d_tp_tp.nc\n",
      "/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/NCI_tutorial_lr_5d_t2m_t2m.nc\n",
      "Linear regression - 6h, iterative loading ...\n",
      "/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/NCI_tutorial_lr_6h_iter.nc\n",
      "CPU times: user 1min 1s, sys: 19.2 s, total: 1min 20s\n",
      "Wall time: 2min 38s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "variables = ['z', 't', 'tp', 't2m']\n",
    "print('Linear regression - 3d - Merging data ...')\n",
    "[ print (f'{PREDDIR}/{save_prefix}_lr_3d_{v}_{v}.nc') for v in variables]\n",
    "lr_3d = xr.merge([xr.open_dataset(f'{PREDDIR}/{save_prefix}_lr_3d_{v}_{v}.nc', chunks={'time': 10}) for v in variables]).load()\n",
    "print('Linear regression - 5d - Merging data ...')\n",
    "[ print (f'{PREDDIR}/{save_prefix}_lr_5d_{v}_{v}.nc') for v in variables]\n",
    "lr_5d = xr.merge([xr.open_dataset(f'{PREDDIR}/{save_prefix}_lr_5d_{v}_{v}.nc', chunks={'time': 10}) for v in variables]).load()\n",
    "print('Linear regression - 6h, iterative loading ...')\n",
    "print(f'{PREDDIR}/{save_prefix}_lr_6h_iter.nc')\n",
    "lr_6h_iter = xr.open_dataset(f'{PREDDIR}/{save_prefix}_lr_6h_iter.nc', chunks={'time': 10}).load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-03-27 02:20:06]\n",
      "save_prefix: NCI_tutorial\n",
      "CPU times: user 90.5 ms, sys: 29.4 ms, total: 120 ms\n",
      "Wall time: 147 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "print ('save_prefix:', save_prefix)\n",
    "cnn_3d = xr.open_dataset(f'{PREDDIR}/{save_prefix}_cnn_3d.nc', chunks={'time': 10})\n",
    "cnn_5d = xr.open_dataset(f'{PREDDIR}/{save_prefix}_cnn_5d.nc', chunks={'time': 10})\n",
    "cnn_6h_iter = xr.open_dataset(f'{PREDDIR}/{save_prefix}_fccnn_6h_iter.nc', chunks={'time': 10})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "#cnn_6h_iter "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "variables = ['2m_temperature', 'total_precipitation']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute RMSE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-03-27 02:20:06]\n",
      "Merging ...\n",
      "Merging ...\n",
      "Merging ...\n",
      "CPU times: user 22min 16s, sys: 1min 6s, total: 23min 23s\n",
      "Wall time: 23min 24s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "\n",
    "func = compute_weighted_rmse\n",
    "rmse = OrderedDict({\n",
    "    'Persistence': evaluate_iterative_forecast(persistence, valid_data, func).load(),\n",
    "    'Climatology': func(climatology, valid_data).load(),\n",
    "    'Weekly clim.': func(weekly_climatology, valid_data).load(),\n",
    "    'LR (iterative)': evaluate_iterative_forecast(lr_6h_iter, valid_data, func).load(),\n",
    "    'CNN (iterative)': evaluate_iterative_forecast(cnn_6h_iter, valid_data, func).load(),\n",
    "    'LR (direct)': xr.concat(\n",
    "        [\n",
    "            func(lr_3d, valid_data),\n",
    "            func(lr_5d, valid_data)\n",
    "        ],\n",
    "        dim=pd.Index([72, 120], name='lead_time')\n",
    "    ).load(),\n",
    "    'CNN (direct)': xr.concat(\n",
    "        [\n",
    "            compute_weighted_rmse(cnn_3d, valid_data),\n",
    "            compute_weighted_rmse(cnn_5d, valid_data)\n",
    "        ],\n",
    "        dim=pd.Index([72, 120], name='lead_time')\n",
    "    ).load(),\n",
    "\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute ACC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-03-27 02:43:31]\n",
      "Merging ...\n",
      "Merging ...\n",
      "Merging ...\n",
      "CPU times: user 42min 22s, sys: 1min 53s, total: 44min 15s\n",
      "Wall time: 44min 10s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "\n",
    "func = compute_weighted_acc\n",
    "acc = OrderedDict({\n",
    "    'Persistence': evaluate_iterative_forecast(persistence, valid_data, func).load(),\n",
    "    'Climatology': func(climatology, valid_data).load(),\n",
    "    'Weekly clim.': func(weekly_climatology, valid_data).load(),\n",
    "    'LR (iterative)': evaluate_iterative_forecast(lr_6h_iter, valid_data, func).load(),\n",
    "    'CNN (iterative)': evaluate_iterative_forecast(cnn_6h_iter, valid_data, func).load(),\n",
    "    'LR (direct)': xr.concat(\n",
    "        [\n",
    "            func(lr_3d, valid_data),\n",
    "            func(lr_5d, valid_data)\n",
    "        ],\n",
    "        dim=pd.Index([72, 120], name='lead_time')\n",
    "    ).load(),\n",
    "    'CNN (direct)': xr.concat(\n",
    "        [\n",
    "            func(cnn_3d, valid_data),\n",
    "            func(cnn_5d, valid_data)\n",
    "        ],\n",
    "        dim=pd.Index([72, 120], name='lead_time')\n",
    "    ).load(),\n",
    "\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute MAE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-03-27 03:27:42]\n",
      "Merging ...\n",
      "Merging ...\n",
      "Merging ...\n",
      "CPU times: user 25min, sys: 1min 7s, total: 26min 7s\n",
      "Wall time: 26min 9s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "\n",
    "func = compute_weighted_mae\n",
    "mae = OrderedDict({\n",
    "    'Persistence': evaluate_iterative_forecast(persistence, valid_data, func).load(),\n",
    "    'Climatology': func(climatology, valid_data).load(),\n",
    "    'Weekly clim.': func(weekly_climatology, valid_data).load(),\n",
    "    'LR (iterative)': evaluate_iterative_forecast(lr_6h_iter, valid_data, func).load(),\n",
    "    'CNN (iterative)': evaluate_iterative_forecast(cnn_6h_iter, valid_data, func).load(),\n",
    "    'LR (direct)': xr.concat(\n",
    "        [\n",
    "            func(lr_3d, valid_data),\n",
    "            func(lr_5d, valid_data)\n",
    "        ],\n",
    "        dim=pd.Index([72, 120], name='lead_time')\n",
    "    ).load(),\n",
    "    'CNN (direct)': xr.concat(\n",
    "        [\n",
    "            func(cnn_3d, valid_data),\n",
    "            func(cnn_5d, valid_data)\n",
    "        ],\n",
    "        dim=pd.Index([72, 120], name='lead_time')\n",
    "    ).load(),\n",
    "\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-03-27 03:53:53]\n",
      "save_prefix: NCI_tutorial\n",
      "/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/NCI_tutorial_rmse.pkl\n",
      "/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/NCI_tutorial_acc.pkl\n",
      "/scratch/vp91/mah900/NCI-Weatherbench/pred_dir/NCI_tutorial_mae.pkl\n",
      "CPU times: user 11.5 ms, sys: 2 ms, total: 13.5 ms\n",
      "Wall time: 46.3 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "print ('save_prefix:', save_prefix)\n",
    "print (f'{PREDDIR}/{save_prefix}_rmse.pkl')\n",
    "print (f'{PREDDIR}/{save_prefix}_acc.pkl')\n",
    "print (f'{PREDDIR}/{save_prefix}_mae.pkl')    \n",
    "to_pickle(rmse, f'{PREDDIR}/{save_prefix}_rmse.pkl')\n",
    "to_pickle(acc,  f'{PREDDIR}/{save_prefix}_acc.pkl')\n",
    "to_pickle(mae,  f'{PREDDIR}/{save_prefix}_mae.pkl')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The same procedure for higher resolutions\n",
    "\n",
    "To see if the results are consistant for higher resolutions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/g/data/wb00/NCI-Weatherbench/2.8125deg\n",
      "/scratch/vp91/mah900/NCI-Weatherbench/pred_dir\n"
     ]
    }
   ],
   "source": [
    "res = '2.8125'\n",
    "DATADIR = f'/g/data/wb00/NCI-Weatherbench/{res}deg' \n",
    "print(DATADIR)\n",
    "PREDDIR = f\"/scratch/vp91/{os.environ['USER']}/NCI-Weatherbench/pred_dir\"\n",
    "print(PREDDIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-03-27 03:53:53]\n",
      "[2021, 2022]\n",
      "load_test_data, var: z\n",
      "load_test_data, var: t\n",
      "CPU times: user 23.2 s, sys: 3.48 s, total: 26.7 s\n",
      "Wall time: 1min 4s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "\n",
    "valid_years = list(range( 2021, 2022+1))\n",
    "print (valid_years)\n",
    "z500_valid_files = [ file for year in valid_years for file in glob.glob (fr'{DATADIR}/geopotential/*{year}*') ] \n",
    "t850_valid_files = [ file for year in valid_years for file in glob.glob (fr'{DATADIR}/temperature/*{year}*')  ] \n",
    "z500_valid = load_test_data(z500_valid_files, 'z', slice('2021', '2022')).load()  \n",
    "t850_valid = load_test_data(t850_valid_files, 't', slice('2021', '2022')).load() \n",
    "valid_data = xr.merge([z500_valid, t850_valid])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-03-27 03:54:58]\n",
      "save_prefix: NCI_tutorial\n",
      "CPU times: user 1min 4s, sys: 30.2 s, total: 1min 34s\n",
      "Wall time: 3min 29s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "\n",
    "save_prefix = 'NCI_tutorial' \n",
    "print ('save_prefix:', save_prefix)\n",
    "persistence = xr.open_dataset(f'{PREDDIR}/{save_prefix}_persistence_{res}.nc', chunks={'time': 10}).load()#)\n",
    "climatology = xr.open_dataset(f'{PREDDIR}/{save_prefix}_climatology_{res}.nc', chunks={'time': 10}).load()#)\n",
    "weekly_climatology = xr.open_dataset(f'{PREDDIR}/{save_prefix}_weekly_climatology_{res}.nc', chunks={'time': 10}).load()#)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-03-27 03:58:28]\n",
      "Merging ...\n",
      "CPU times: user 54.8 s, sys: 39.4 s, total: 1min 34s\n",
      "Wall time: 1min 16s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "\n",
    "rmse_persistence = evaluate_iterative_forecast(persistence, valid_data, compute_weighted_rmse).load()\n",
    "rmse_climatology = compute_weighted_rmse(climatology, valid_data).load()\n",
    "rmse_weekly_climatology = compute_weighted_rmse(weekly_climatology, valid_data).load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-03-27 03:59:44]\n",
      "CPU times: user 859 µs, sys: 79 µs, total: 938 µs\n",
      "Wall time: 899 µs\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "\n",
    "rmse = xr.merge(\n",
    "    [r.rename({'z': f'z_{n}', 't': f't_{n}'}) for r, n in zip(\n",
    "        [rmse_persistence, rmse_climatology, rmse_weekly_climatology #, \n",
    "        ],\n",
    "        ['persistence', 'climatology', 'weekly_climatology' #,\n",
    "        ]\n",
    "    )],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "save_prefix: NCI_tutorial\n"
     ]
    }
   ],
   "source": [
    "print ('save_prefix:', save_prefix)\n",
    "rmse.to_netcdf(f'{PREDDIR}/{save_prefix}_rmse_{res}.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-03-27 03:59:44]\n",
      "z_persistence [ 951.00437868 1053.50191451]\n",
      "t_persistence [4.34957183 4.64705966]\n",
      "CPU times: user 2.29 ms, sys: 0 ns, total: 2.29 ms\n",
      "Wall time: 2.04 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "\n",
    "for var in rmse:\n",
    "    if 'lead_time' in rmse[var].coords:\n",
    "        print(var, rmse[var].sel(lead_time=[3*24, 5*24]).values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-03-27 03:59:44]\n",
      "z_climatology 1086.1617434942157\n",
      "t_climatology 5.57893994332815\n",
      "z_weekly_climatology 823.7249845019733\n",
      "t_weekly_climatology 3.550191546286565\n",
      "CPU times: user 369 µs, sys: 44 µs, total: 413 µs\n",
      "Wall time: 361 µs\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "\n",
    "for var in rmse:\n",
    "    if 'lead_time' not in rmse[var].coords:\n",
    "        print(var, rmse[var].values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-03-27 03:59:44]\n"
     ]
    }
   ],
   "source": [
    "print( f'[{datetime.now().replace(microsecond=0)}]' )\n",
    "client.shutdown()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The End"
   ]
  }
 ],
 "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}