analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4d047299",
   "metadata": {
    "papermill": {
     "duration": 0.027065,
     "end_time": "2023-03-06T02:49:49.549053",
     "exception": false,
     "start_time": "2023-03-06T02:49:49.521988",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# UNSEEN analysis\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95569a41",
   "metadata": {
    "papermill": {
     "duration": 9.54623,
     "end_time": "2023-03-06T02:49:59.109857",
     "exception": false,
     "start_time": "2023-03-06T02:49:49.563627",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import calendar\n",
    "import cartopy\n",
    "import cartopy.crs as ccrs\n",
    "import cmocean\n",
    "from collections import Counter\n",
    "import datetime\n",
    "import geopandas as gp\n",
    "import glob\n",
    "from IPython.display import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "from scipy.stats import genextreme as gev\n",
    "\n",
    "# import seaborn as sns\n",
    "import xclim as xc\n",
    "\n",
    "from unseen import fileio, independence, stability, eva, general_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56e9e0ab",
   "metadata": {
    "papermill": {
     "duration": 0.024169,
     "end_time": "2023-03-06T02:49:59.159559",
     "exception": false,
     "start_time": "2023-03-06T02:49:59.135390",
     "status": "completed"
    },
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# Optional parameters\n",
    "# (This cell is tagged \"parameters\")\n",
    "dpi = 300"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9787631",
   "metadata": {
    "papermill": {
     "duration": 0.042552,
     "end_time": "2023-03-06T02:49:59.256849",
     "exception": false,
     "start_time": "2023-03-06T02:49:59.214297",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Required parameters\n",
    "assert \"metric\" in locals(), \"Must provide a metric name\"\n",
    "assert \"var\" in locals(), \"Must provide a variable name\"\n",
    "assert \"model_name\" in locals(), \"Must provide a model name\"\n",
    "assert \"region_name\" in locals(), \"Must provide a region name\"\n",
    "# assert 'min_lead' in locals(), \"Must provide a minimum lead time\"\n",
    "assert \"shape_overlap\" in locals(), \"Must provide a shape overlap fraction\"\n",
    "assert \"metric_plot_label\" in locals(), \"Must provide a plot axis label for the metric\"\n",
    "assert (\n",
    "    \"metric_plot_upper_limit\" in locals()\n",
    "), \"Must provide an upper limit for the metric value\"\n",
    "assert os.path.isfile(\n",
    "    obs_file\n",
    "), \"Must provide an observations data file (papermill option -p obs_file [filepath])\"\n",
    "assert os.path.isfile(\n",
    "    file_list\n",
    "), \"Must provide the input model files list (papermill option -p file_list [filepath])\"\n",
    "assert os.path.isfile(\n",
    "    shape_file\n",
    "), \"Must provide the shapefile (papermill option -p shape_file [filepath])\"\n",
    "assert os.path.isfile(\n",
    "    model_file\n",
    "), \"Must provide an model data file (papermill option -p model_file [filepath])\"\n",
    "assert os.path.isfile(\n",
    "    model_add_bc_file\n",
    "), \"Must provide a model additive bias corrected data file (papermill option -p model_add_bc_file [filepath])\"\n",
    "assert os.path.isfile(\n",
    "    model_mulc_bc_file\n",
    "), \"Must provide a model multiplicative bias corrected data file (papermill option -p model_mulc_bc_file [filepath])\"\n",
    "assert os.path.isfile(\n",
    "    similarity_add_bc_file\n",
    "), \"Must provide an additive bias corrected similarity test file (papermill option -p similarity_add_bias_file [filepath])\"\n",
    "assert os.path.isfile(\n",
    "    similarity_mulc_bc_file\n",
    "), \"Must provide an multiplicative bias corrected similarity test file (papermill option -p similarity_mulc_bias_file [filepath])\"\n",
    "assert os.path.isfile(\n",
    "    similarity_raw_file\n",
    "), \"Must provide an raw data similarity test file (papermill option -p similarity_raw_file [filepath])\"\n",
    "assert os.path.isfile(\n",
    "    independence_file\n",
    "), \"Must provide an independence test plot (papermill option -p independence_file [filepath])\"\n",
    "assert os.path.isfile(\n",
    "    stability_plot_empirical\n",
    "), \"Must provide an empirical stability test plot (papermill option -p stability_plot_empirical [filepath])\"\n",
    "assert os.path.isfile(\n",
    "    stability_plot_gev\n",
    "), \"Must provide a GEV stability test plot (papermill option -p stability_plot_gev [filepath])\"\n",
    "assert os.path.isfile(\n",
    "    moments_add_plot\n",
    "), \"Must provide a moments test plot for the additive bias corrected data (papermill option -p moments_add_plot [filepath])\"\n",
    "assert os.path.isfile(\n",
    "    moments_mulc_plot\n",
    "), \"Must provide a moments test plot for the multiplicative bias corrected data (papermill option -p moments_mulc_plot [filepath])\"\n",
    "assert os.path.isfile(\n",
    "    moments_raw_plot\n",
    "), \"Must provide a moments test plot for the raw data (papermill option -p moments_raw_plot [filepath])\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddfc8ca8",
   "metadata": {},
   "outputs": [],
   "source": [
    "units_dict = {\"pr\": \"mm day-1\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c8da4f2",
   "metadata": {
    "papermill": {
     "duration": 0.012656,
     "end_time": "2023-03-06T02:49:59.291041",
     "exception": false,
     "start_time": "2023-03-06T02:49:59.278385",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Observations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f40952f6",
   "metadata": {
    "papermill": {
     "duration": 0.722984,
     "end_time": "2023-03-06T02:50:00.024198",
     "exception": false,
     "start_time": "2023-03-06T02:49:59.301214",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "obs_ds = fileio.open_dataset(obs_file)\n",
    "obs_ds = obs_ds.dropna(\"time\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c701895",
   "metadata": {
    "papermill": {
     "duration": 0.20619,
     "end_time": "2023-03-06T02:50:00.356187",
     "exception": false,
     "start_time": "2023-03-06T02:50:00.149997",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "years = obs_ds[\"time\"].dt.year.values\n",
    "obs_df = pd.DataFrame(index=years)\n",
    "obs_df[var] = obs_ds[var].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8bfcbcf",
   "metadata": {
    "papermill": {
     "duration": 0.16929,
     "end_time": "2023-03-06T02:50:00.545182",
     "exception": false,
     "start_time": "2023-03-06T02:50:00.375892",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "obs_max_event = obs_ds[\"pr\"].values.max()\n",
    "print(obs_max_event)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d082e6af",
   "metadata": {
    "papermill": {
     "duration": 0.014695,
     "end_time": "2023-03-06T02:50:00.583605",
     "exception": false,
     "start_time": "2023-03-06T02:50:00.568910",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Model ensemble"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f127f972",
   "metadata": {
    "papermill": {
     "duration": 0.010278,
     "end_time": "2023-03-06T02:50:00.604350",
     "exception": false,
     "start_time": "2023-03-06T02:50:00.594072",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Region selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c5bfc47",
   "metadata": {
    "papermill": {
     "duration": 0.021806,
     "end_time": "2023-03-06T02:50:00.636845",
     "exception": false,
     "start_time": "2023-03-06T02:50:00.615039",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "with open(file_list) as f:\n",
    "    all_files = f.read()\n",
    "    first_file = all_files.split(\"\\n\", 1)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb14f729",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Plot region shapefile outline using the first file in the list\n",
    "# if shape_file and file_list:\n",
    "#     with open(file_list) as f:\n",
    "#         all_files = f.read()\n",
    "#         first_file = all_files.split(\"\\n\", 1)[0]\n",
    "\n",
    "#     shapes = gp.read_file(shape_file)\n",
    "#     isel_dict = {}\n",
    "#     if model_name == \"CAFE\":\n",
    "#         isel_dict[\"ensemble\"] = 0\n",
    "\n",
    "#     region_ds = fileio.open_dataset(\n",
    "#         first_file,\n",
    "#         metadata_file=metadata_file,\n",
    "#         variables=[var],\n",
    "#         lat_bnds=[-44, -10],\n",
    "#         lon_bnds=[113, 155],\n",
    "#         shapefile=shapefile,\n",
    "#         shape_overlap=shape_overlap,\n",
    "#         units={var: units_dict[var]},\n",
    "#         isel=isel_dict,\n",
    "#     )\n",
    "\n",
    "#     fig = plt.figure(figsize=[12, 8])\n",
    "#     ax = fig.add_subplot(111, projection=ccrs.PlateCarree())\n",
    "#     region_ds[var].mean(\"time\", keep_attrs=True).plot(\n",
    "#         ax=ax,\n",
    "#         transform=ccrs.PlateCarree(),\n",
    "#         cmap=\"viridis_r\",\n",
    "#     )\n",
    "#     ax.coastlines()\n",
    "#     ax.add_feature(cartopy.feature.STATES)\n",
    "#     ax.add_geometries(\n",
    "#         shapes.geometry, ccrs.PlateCarree(), facecolor=\"none\", edgecolor=\"orange\"\n",
    "#     )\n",
    "#     plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9de3847d",
   "metadata": {
    "papermill": {
     "duration": 0.026146,
     "end_time": "2023-03-06T02:50:00.676691",
     "exception": false,
     "start_time": "2023-03-06T02:50:00.650545",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "first_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d218574b",
   "metadata": {
    "papermill": {
     "duration": 0.097571,
     "end_time": "2023-03-06T02:50:00.797733",
     "exception": false,
     "start_time": "2023-03-06T02:50:00.700162",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "shapes = gp.read_file(shape_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b9e6200",
   "metadata": {
    "papermill": {
     "duration": 0.018812,
     "end_time": "2023-03-06T02:50:00.828503",
     "exception": false,
     "start_time": "2023-03-06T02:50:00.809691",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "isel_dict = {}\n",
    "if model_name == \"CAFE\":\n",
    "    isel_dict[\"ensemble\"] = 0\n",
    "    metadata_file = (\n",
    "        \"/g/data/xv83/unseen-projects/code/dataset_config/dataset_cafe_daily.yml\"\n",
    "    )\n",
    "else:\n",
    "    metadata_file = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c232f49",
   "metadata": {
    "papermill": {
     "duration": 1.716976,
     "end_time": "2023-03-06T02:50:02.557802",
     "exception": false,
     "start_time": "2023-03-06T02:50:00.840826",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "region_ds = fileio.open_dataset(\n",
    "    first_file,\n",
    "    metadata_file=metadata_file,\n",
    "    variables=[var],\n",
    "    lat_bnds=[-44, -10],\n",
    "    lon_bnds=[113, 155],\n",
    "    shapefile=shape_file,\n",
    "    shape_overlap=shape_overlap,\n",
    "    units={var: units_dict[var]},\n",
    "    isel=isel_dict,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0eba243a",
   "metadata": {
    "papermill": {
     "duration": 9.343738,
     "end_time": "2023-03-06T02:50:11.928011",
     "exception": false,
     "start_time": "2023-03-06T02:50:02.584273",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=[12, 8])\n",
    "ax = fig.add_subplot(111, projection=ccrs.PlateCarree())\n",
    "region_ds[var].mean(\"time\", keep_attrs=True).plot(\n",
    "    ax=ax,\n",
    "    transform=ccrs.PlateCarree(),\n",
    "    cmap=\"viridis_r\",\n",
    ")\n",
    "ax.coastlines()\n",
    "ax.add_feature(cartopy.feature.STATES)\n",
    "ax.add_geometries(\n",
    "    shapes.geometry, ccrs.PlateCarree(), facecolor=\"none\", edgecolor=\"orange\"\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "146854fe",
   "metadata": {
    "papermill": {
     "duration": 0.036872,
     "end_time": "2023-03-06T02:50:12.012408",
     "exception": false,
     "start_time": "2023-03-06T02:50:11.975536",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Model data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1620b927",
   "metadata": {
    "papermill": {
     "duration": 0.093555,
     "end_time": "2023-03-06T02:50:12.140348",
     "exception": false,
     "start_time": "2023-03-06T02:50:12.046793",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_ds = fileio.open_dataset(model_file)\n",
    "model_add_bc_ds = fileio.open_dataset(model_add_bc_file)\n",
    "model_mulc_bc_ds = fileio.open_dataset(model_mulc_bc_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3d4100a",
   "metadata": {
    "papermill": {
     "duration": 0.084387,
     "end_time": "2023-03-06T02:50:12.265816",
     "exception": false,
     "start_time": "2023-03-06T02:50:12.181429",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b25bd815",
   "metadata": {
    "papermill": {
     "duration": 0.034589,
     "end_time": "2023-03-06T02:50:12.345116",
     "exception": false,
     "start_time": "2023-03-06T02:50:12.310527",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Independence testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b819d698",
   "metadata": {
    "papermill": {
     "duration": 0.058675,
     "end_time": "2023-03-06T02:50:12.448582",
     "exception": false,
     "start_time": "2023-03-06T02:50:12.389907",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Drop dependent lead times based on the minimum independent lead time\n",
    "\n",
    "independence = fileio.open_dataset(independence_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76db7fd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "independence.point_plot(ds_independence, independence_plot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5ccabe2",
   "metadata": {
    "papermill": {
     "duration": 0.04419,
     "end_time": "2023-03-06T02:50:12.529595",
     "exception": false,
     "start_time": "2023-03-06T02:50:12.485405",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "min_lead = min_lead_ds[\"min_lead\"].load()\n",
    "print(min_lead)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b40107c",
   "metadata": {
    "papermill": {
     "duration": 1.838431,
     "end_time": "2023-03-06T02:50:14.460463",
     "exception": false,
     "start_time": "2023-03-06T02:50:12.622032",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_da = (\n",
    "    model_ds[var].groupby(f\"init_date.month\").where(model_ds[\"lead_time\"] >= min_lead)\n",
    ").dropna(\"lead_time\")\n",
    "model_da_stacked = model_da.stack({\"sample\": [\"ensemble\", \"init_date\", \"lead_time\"]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10aaf4a0",
   "metadata": {
    "papermill": {
     "duration": 0.025289,
     "end_time": "2023-03-06T02:50:14.506502",
     "exception": false,
     "start_time": "2023-03-06T02:50:14.481213",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_da_stacked.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0b0019b",
   "metadata": {
    "papermill": {
     "duration": 1.808975,
     "end_time": "2023-03-06T02:50:16.331816",
     "exception": false,
     "start_time": "2023-03-06T02:50:14.522841",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_add_bc_da = (\n",
    "    model_add_bc_ds[var]\n",
    "    .groupby(f\"init_date.month\")\n",
    "    .where(model_add_bc_ds[\"lead_time\"] >= min_lead)\n",
    "    .dropna(\"lead_time\")\n",
    ")\n",
    "model_add_bc_da_stacked = model_add_bc_da.stack(\n",
    "    {\"sample\": [\"ensemble\", \"init_date\", \"lead_time\"]}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d4fdc9e",
   "metadata": {
    "papermill": {
     "duration": 0.028043,
     "end_time": "2023-03-06T02:50:16.377205",
     "exception": false,
     "start_time": "2023-03-06T02:50:16.349162",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_add_bc_da_stacked.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a7afb26-20cf-43e6-b647-9f40e7586682",
   "metadata": {
    "papermill": {
     "duration": 1.648828,
     "end_time": "2023-03-06T02:50:18.044617",
     "exception": false,
     "start_time": "2023-03-06T02:50:16.395789",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_mulc_bc_da = (\n",
    "    model_mulc_bc_ds[var]\n",
    "    .groupby(f\"init_date.month\")\n",
    "    .where(model_mulc_bc_ds[\"lead_time\"] >= min_lead)\n",
    "    .dropna(\"lead_time\")\n",
    ")\n",
    "model_mulc_bc_da_stacked = model_mulc_bc_da.stack(\n",
    "    {\"sample\": [\"ensemble\", \"init_date\", \"lead_time\"]}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80b2f212-366a-4dcd-bfbe-1f5dc232cb64",
   "metadata": {
    "papermill": {
     "duration": 0.026066,
     "end_time": "2023-03-06T02:50:18.100254",
     "exception": false,
     "start_time": "2023-03-06T02:50:18.074188",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_mulc_bc_da_stacked.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e91cbd6",
   "metadata": {},
   "source": [
    "### Stability testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00c55fd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "Image(filename=stability_plot_empirical)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cdc1506",
   "metadata": {},
   "outputs": [],
   "source": [
    "Image(filename=stability_plot_gev)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e039e61",
   "metadata": {
    "papermill": {
     "duration": 0.013056,
     "end_time": "2023-03-06T02:50:18.132496",
     "exception": false,
     "start_time": "2023-03-06T02:50:18.119440",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Bias correction and similarity testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf650636",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_gev_shape, model_gev_loc, model_gev_scale = eva.fit_gev(model_da_stacked.values)\n",
    "model_add_bc_gev_shape, model_add_bc_gev_loc, model_add_bc_gev_scale = eva.fit_gev(\n",
    "    model_add_bc_da_stacked.values\n",
    ")\n",
    "model_mulc_bc_gev_shape, model_mulc_bc_gev_loc, model_mulc_bc_gev_scale = eva.fit_gev(\n",
    "    model_mulc_bc_da_stacked.values\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "110c554e",
   "metadata": {
    "papermill": {
     "duration": 11.219097,
     "end_time": "2023-03-06T02:50:29.373856",
     "exception": false,
     "start_time": "2023-03-06T02:50:18.154759",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def plot_distributions(include_additive=False):\n",
    "    \"\"\"Plot various distributions\"\"\"\n",
    "\n",
    "    fig = plt.figure(figsize=[10, 6])\n",
    "    gev_xvals = np.arange(0, metric_plot_upper_limit)\n",
    "\n",
    "    model_da.plot.hist(\n",
    "        bins=50,\n",
    "        density=True,\n",
    "        label=model_name,\n",
    "        alpha=0.7,\n",
    "        facecolor=\"tab:blue\",\n",
    "    )\n",
    "    model_gev_pdf = gev.pdf(gev_xvals, model_gev_shape, model_gev_loc, model_gev_scale)\n",
    "    plt.plot(gev_xvals, model_gev_pdf, color=\"tab:blue\", linewidth=4.0)\n",
    "\n",
    "    if include_additive:\n",
    "        model_add_bc_ds[var].plot.hist(\n",
    "            bins=50,\n",
    "            density=True,\n",
    "            label=f\"{model_name} additive bias corrected\",\n",
    "            facecolor=\"tab:red\",\n",
    "            alpha=0.7,\n",
    "        )\n",
    "        model_add_bc_gev_pdf = gev.pdf(\n",
    "            gev_xvals,\n",
    "            model_add_bc_gev_shape,\n",
    "            model_add_bc_gev_loc,\n",
    "            model_add_bc_gev_scale,\n",
    "        )\n",
    "        plt.plot(gev_xvals, model_add_bc_gev_pdf, color=\"tab:red\", linewidth=4.0)\n",
    "\n",
    "    model_mulc_bc_ds[var].plot.hist(\n",
    "        bins=50,\n",
    "        density=True,\n",
    "        label=f\"{model_name} multiplicative bias corrected\",\n",
    "        facecolor=\"tab:orange\",\n",
    "        alpha=0.7,\n",
    "    )\n",
    "    model_mulc_bc_gev_pdf = gev.pdf(\n",
    "        gev_xvals,\n",
    "        model_mulc_bc_gev_shape,\n",
    "        model_mulc_bc_gev_loc,\n",
    "        model_mulc_bc_gev_scale,\n",
    "    )\n",
    "    plt.plot(gev_xvals, model_mulc_bc_gev_pdf, color=\"tab:orange\", linewidth=4.0)\n",
    "\n",
    "    obs_ds[var].plot.hist(\n",
    "        bins=50, density=True, label=\"observations\", facecolor=\"tab:gray\", alpha=0.7\n",
    "    )\n",
    "    obs_gev_shape, obs_gev_loc, obs_gev_scale = eva.fit_gev(obs_ds[var].values)\n",
    "    obs_gev_pdf = gev.pdf(gev_xvals, obs_gev_shape, obs_gev_loc, obs_gev_scale)\n",
    "    plt.plot(gev_xvals, obs_gev_pdf, color=\"tab:gray\", linewidth=4.0)\n",
    "\n",
    "    plt.xlabel(metric_plot_label)\n",
    "    plt.ylabel(\"probability\")\n",
    "    plt.title(f\"{metric} for {region_name} ({model_name})\")\n",
    "    plt.xlim(0, metric_plot_upper_limit)\n",
    "    plt.legend()\n",
    "    # if not include_additive:\n",
    "    #    outfile = f'/g/data/xv83/unseen-projects/outputs/tc-jasper/figures/{metric}_histogram_{model_name}_{region_name}.png'\n",
    "    #    print(outfile)\n",
    "    #    plt.savefig(outfile, bbox_inches='tight', facecolor='white', dpi=dpi)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ab62780",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_distributions()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d58c8965",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_distributions(include_additive=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0865704a-5962-4cf3-ae1a-72ea501aeac8",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"raw data shape parameter: {model_gev_shape:.3f}\")\n",
    "print(f\"raw data location parameter: {model_gev_loc:.3f}\")\n",
    "print(f\"raw data scale parameter: {model_gev_scale:.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ee68ad7-da09-4cb8-88c3-5a2c3b830cd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"additive corrected shape parameter: {model_add_bc_gev_shape:.3f}\")\n",
    "print(f\"additive corrected location parameter: {model_add_bc_gev_loc:.3f}\")\n",
    "print(f\"additive corrected scale parameter: {model_add_bc_gev_scale:.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51ebbd48-8972-4dbf-9fb5-604f81110fb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"multiplicative corrected shape parameter: {model_mulc_bc_gev_shape:.3f}\")\n",
    "print(f\"multiplicative corrected location parameter: {model_mulc_bc_gev_loc:.3f}\")\n",
    "print(f\"multiplicative corrected scale parameter: {model_mulc_bc_gev_scale:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a8210c0",
   "metadata": {
    "papermill": {
     "duration": 0.013525,
     "end_time": "2023-03-06T02:50:29.419562",
     "exception": false,
     "start_time": "2023-03-06T02:50:29.406037",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We can look at p-values for the KS-test and Anderson-Darling test for each lead time.\n",
    "\n",
    "p > 0.05 means the null hypothesis (that the two samples are from the same population) can't be rejected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a21af1f7",
   "metadata": {
    "papermill": {
     "duration": 0.036933,
     "end_time": "2023-03-06T02:50:29.470129",
     "exception": false,
     "start_time": "2023-03-06T02:50:29.433196",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "similarity_raw_ds = fileio.open_dataset(similarity_raw_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e937c26",
   "metadata": {},
   "outputs": [],
   "source": [
    "similarity_raw_ds[\"ks_statistic\"].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a599ddf",
   "metadata": {
    "papermill": {
     "duration": 0.037587,
     "end_time": "2023-03-06T02:50:29.526637",
     "exception": false,
     "start_time": "2023-03-06T02:50:29.489050",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "similarity_raw_ds[\"ks_pval\"].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63c1fcb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "similarity_raw_ds[\"ad_statistic\"].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05eb014d",
   "metadata": {
    "papermill": {
     "duration": 0.039515,
     "end_time": "2023-03-06T02:50:29.591203",
     "exception": false,
     "start_time": "2023-03-06T02:50:29.551688",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "similarity_raw_ds[\"ad_pval\"].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6b89403",
   "metadata": {},
   "outputs": [],
   "source": [
    "Image(filename=moments_raw_plot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b3f7de0",
   "metadata": {
    "papermill": {
     "duration": 0.035255,
     "end_time": "2023-03-06T02:50:29.644538",
     "exception": false,
     "start_time": "2023-03-06T02:50:29.609283",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "similarity_add_bias_ds = fileio.open_dataset(similarity_add_bc_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1a4eb0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "similarity_add_bias_ds[\"ks_statistic\"].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d388e95d",
   "metadata": {
    "papermill": {
     "duration": 0.037209,
     "end_time": "2023-03-06T02:50:29.707192",
     "exception": false,
     "start_time": "2023-03-06T02:50:29.669983",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "similarity_add_bias_ds[\"ks_pval\"].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb7f39c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "similarity_add_bias_ds[\"ad_statistic\"].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7cee743",
   "metadata": {
    "papermill": {
     "duration": 0.038861,
     "end_time": "2023-03-06T02:50:29.772718",
     "exception": false,
     "start_time": "2023-03-06T02:50:29.733857",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "similarity_add_bias_ds[\"ad_pval\"].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8c3009f",
   "metadata": {},
   "outputs": [],
   "source": [
    "Image(filename=moments_add_plot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e52d7e00-1510-46bc-ba47-ca75c162b03c",
   "metadata": {
    "papermill": {
     "duration": 0.035393,
     "end_time": "2023-03-06T02:50:29.826306",
     "exception": false,
     "start_time": "2023-03-06T02:50:29.790913",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "similarity_mulc_bias_ds = fileio.open_dataset(similarity_mulc_bc_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a8a6348",
   "metadata": {},
   "outputs": [],
   "source": [
    "similarity_mulc_bias_ds[\"ks_statistic\"].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0de30b63-1886-4ed8-9307-dabf49c5d484",
   "metadata": {
    "papermill": {
     "duration": 0.038219,
     "end_time": "2023-03-06T02:50:29.889688",
     "exception": false,
     "start_time": "2023-03-06T02:50:29.851469",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "similarity_mulc_bias_ds[\"ks_pval\"].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77233634",
   "metadata": {},
   "outputs": [],
   "source": [
    "similarity_mulc_bias_ds[\"ad_statistic\"].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "666ff4e6-2437-42ab-87d3-59231280b122",
   "metadata": {
    "papermill": {
     "duration": 0.04243,
     "end_time": "2023-03-06T02:50:29.952508",
     "exception": false,
     "start_time": "2023-03-06T02:50:29.910078",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "similarity_mulc_bias_ds[\"ad_pval\"].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ea34514",
   "metadata": {},
   "outputs": [],
   "source": [
    "Image(filename=moments_mulc_plot)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "692c7b58",
   "metadata": {
    "papermill": {
     "duration": 0.015205,
     "end_time": "2023-03-06T02:50:29.989940",
     "exception": false,
     "start_time": "2023-03-06T02:50:29.974735",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Return curves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff0f8f73",
   "metadata": {
    "papermill": {
     "duration": 0.027172,
     "end_time": "2023-03-06T02:50:30.078823",
     "exception": false,
     "start_time": "2023-03-06T02:50:30.051651",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def gev_event_likelihood(event, shape, loc, scale):\n",
    "    \"\"\"Get likelihood stats for an event from a GEV fit\"\"\"\n",
    "\n",
    "    probability = gev.sf(event, shape, loc=loc, scale=scale)\n",
    "    return_period = 1.0 / probability\n",
    "    percentile = (1 - probability) * 100\n",
    "    print(f\"{percentile:.2f}% percentile\")\n",
    "    print(f\"{return_period:.0f} year return period\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6a40828",
   "metadata": {
    "papermill": {
     "duration": 0.014216,
     "end_time": "2023-03-06T02:50:30.170868",
     "exception": false,
     "start_time": "2023-03-06T02:50:30.156652",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "#### Additive bias corrected data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39e4acf0",
   "metadata": {
    "papermill": {
     "duration": 1.828571,
     "end_time": "2023-03-06T02:50:33.858615",
     "exception": false,
     "start_time": "2023-03-06T02:50:32.030044",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "n_events_add_bc, n_population_add_bc, return_period_add_bc, percentile_add_bc = (\n",
    "    eva.event_in_context(\n",
    "        model_add_bc_da_stacked.values,\n",
    "        obs_max_event,\n",
    "        \"above\",\n",
    "    )\n",
    ")\n",
    "print(\"ADDITIVE BIAS CORRECTED DATA\")\n",
    "print(f\"{n_events_add_bc} events in {n_population_add_bc} samples\")\n",
    "print(f\"{percentile_add_bc:.2f}% percentile\")\n",
    "print(f\"{return_period_add_bc:.0f} year return period\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e0a064d",
   "metadata": {
    "papermill": {
     "duration": 0.029851,
     "end_time": "2023-03-06T02:50:33.919514",
     "exception": false,
     "start_time": "2023-03-06T02:50:33.889663",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"ADDITIVE BIAS CORRECTED GEV\")\n",
    "gev_event_likelihood(\n",
    "    obs_max_event, model_add_bc_gev_shape, model_add_bc_gev_loc, model_add_bc_gev_scale\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "378e3988-31c5-4312-96ef-72d7820ef8e4",
   "metadata": {
    "papermill": {
     "duration": 0.027586,
     "end_time": "2023-03-06T02:50:53.797469",
     "exception": false,
     "start_time": "2023-03-06T02:50:53.769883",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "#### Multiplicative bias corrected data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17d6421d-1097-4775-a8da-1072eae0e792",
   "metadata": {
    "papermill": {
     "duration": 2.002438,
     "end_time": "2023-03-06T02:50:55.816763",
     "exception": false,
     "start_time": "2023-03-06T02:50:53.814325",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=[8, 6])\n",
    "ax = fig.add_subplot()\n",
    "eva.plot_gev_return_curve(\n",
    "    ax,\n",
    "    model_mulc_bc_da_stacked,\n",
    "    obs_max_event,\n",
    "    bootstrap_method=\"non-parametric\",\n",
    "    n_bootstraps=1000,\n",
    "    ylabel=metric_plot_label,\n",
    "    ylim=[50, metric_plot_upper_limit],\n",
    "    text=True,\n",
    ")\n",
    "# outfile = f'/g/data/xv83/unseen-projects/outputs/tc-jasper/figures/{metric}_return-curve_{model_name}_{region_name}.png'\n",
    "# print(outfile)\n",
    "# plt.savefig(outfile, bbox_inches='tight', facecolor='white', dpi=dpi)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0175a01-4152-4ff2-b197-f00c153e1389",
   "metadata": {
    "papermill": {
     "duration": 1.754092,
     "end_time": "2023-03-06T02:50:57.604112",
     "exception": false,
     "start_time": "2023-03-06T02:50:55.850020",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "n_events_mulc_bc, n_population_mulc_bc, return_period_mulc_bc, percentile_mulc_bc = (\n",
    "    eva.event_in_context(model_add_bc_da_stacked.values, obs_max_event, \"above\")\n",
    ")\n",
    "print(\"MULTIPLICATIVE BIAS CORRECTED DATA\")\n",
    "print(f\"{n_events_mulc_bc} events in {n_population_mulc_bc} samples\")\n",
    "print(f\"{percentile_mulc_bc:.2f}% percentile\")\n",
    "print(f\"{return_period_mulc_bc:.0f} year return period\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd2f4347-4051-45a1-be36-e1a69617439a",
   "metadata": {
    "papermill": {
     "duration": 0.030973,
     "end_time": "2023-03-06T02:50:57.672728",
     "exception": false,
     "start_time": "2023-03-06T02:50:57.641755",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"MULTIPLICATIVE BIAS CORRECTED GEV\")\n",
    "gev_event_likelihood(\n",
    "    obs_max_event,\n",
    "    model_mulc_bc_gev_shape,\n",
    "    model_mulc_bc_gev_loc,\n",
    "    model_mulc_bc_gev_scale,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6a2c14c",
   "metadata": {
    "papermill": {
     "duration": 0.029166,
     "end_time": "2023-03-06T02:51:17.677692",
     "exception": false,
     "start_time": "2023-03-06T02:51:17.648526",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "#### Raw data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89138fa6",
   "metadata": {
    "papermill": {
     "duration": 1.543689,
     "end_time": "2023-03-06T02:51:20.960995",
     "exception": false,
     "start_time": "2023-03-06T02:51:19.417306",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "n_events, n_population, return_period, percentile = eva.event_in_context(\n",
    "    model_da_stacked.values, obs_max_event, \"above\"\n",
    ")\n",
    "print(\"RAW DATA\")\n",
    "print(f\"{n_events} events in {n_population} samples\")\n",
    "print(f\"{percentile:.2f}% percentile\")\n",
    "print(f\"{return_period:.0f} year return period\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55ca0308",
   "metadata": {
    "papermill": {
     "duration": 0.049123,
     "end_time": "2023-03-06T02:51:21.048673",
     "exception": false,
     "start_time": "2023-03-06T02:51:20.999550",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"RAW GEV\")\n",
    "gev_event_likelihood(obs_max_event, model_gev_shape, model_gev_loc, model_gev_scale)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bcb7462",
   "metadata": {
    "papermill": {
     "duration": 0.020545,
     "end_time": "2023-03-06T02:51:21.096682",
     "exception": false,
     "start_time": "2023-03-06T02:51:21.076137",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Seasonality"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e28f48c",
   "metadata": {
    "papermill": {
     "duration": 1.576099,
     "end_time": "2023-03-06T02:51:22.693705",
     "exception": false,
     "start_time": "2023-03-06T02:51:21.117606",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_df = model_da_stacked.to_dataframe()\n",
    "model_df = model_df.dropna()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1a16c7e",
   "metadata": {
    "papermill": {
     "duration": 0.046031,
     "end_time": "2023-03-06T02:51:22.781975",
     "exception": false,
     "start_time": "2023-03-06T02:51:22.735944",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "event_months = [int(date[5:7]) for date in model_df[\"event_time\"].values]\n",
    "event_years = [int(date[0:4]) for date in model_df[\"event_time\"].values]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8d46fba",
   "metadata": {
    "papermill": {
     "duration": 0.166342,
     "end_time": "2023-03-06T02:51:22.974104",
     "exception": false,
     "start_time": "2023-03-06T02:51:22.807762",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "month_counts = Counter(event_months)\n",
    "months = np.arange(1, 13)\n",
    "counts = [month_counts[month] for month in months]\n",
    "\n",
    "plt.bar(months, counts)\n",
    "plt.title(f\"{metric} timing for {region_name} ({model_name})\")\n",
    "plt.ylabel(\"number of events\")\n",
    "plt.xlabel(\"month\")\n",
    "xlabels = [calendar.month_abbr[i] for i in months]\n",
    "plt.xticks(months, xlabels)\n",
    "# outfile = f'/g/data/xv83/unseen-projects/outputs/tc-jasper/figures/{metric}_timing_{model_name}_{region_name}.png'\n",
    "# print(outfile)\n",
    "# plt.savefig(outfile, bbox_inches='tight', facecolor='white', dpi=dpi)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9405e65",
   "metadata": {
    "papermill": {
     "duration": 0.023664,
     "end_time": "2023-03-06T02:53:06.797372",
     "exception": false,
     "start_time": "2023-03-06T02:53:06.773708",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Meteorology"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ed38598",
   "metadata": {
    "papermill": {
     "duration": 0.048508,
     "end_time": "2023-03-06T02:53:06.869915",
     "exception": false,
     "start_time": "2023-03-06T02:53:06.821407",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def plot_average(\n",
    "    start_date,\n",
    "    end_date,\n",
    "    color_da=None,\n",
    "    color_var=None,\n",
    "    contour_da=None,\n",
    "    contour_var=None,\n",
    "):\n",
    "    \"\"\"Plot average\"\"\"\n",
    "\n",
    "    fig = plt.figure(figsize=[10, 5])\n",
    "    map_proj = ccrs.PlateCarree(central_longitude=180)\n",
    "    ax1 = fig.add_subplot(1, 1, 1, projection=map_proj)\n",
    "\n",
    "    if color_var:\n",
    "        if color_var == \"pr\":\n",
    "            levels = [0, 100, 200, 300, 400, 500, 600, 700, 800]\n",
    "            label = \"total precipitation (mm)\"\n",
    "            cmap = cmocean.cm.rain\n",
    "            extend = \"max\"\n",
    "        elif color_var == \"ua300\":\n",
    "            levels = [-30, -25, -20, -15, -10, -5, 0, 5, 10, 15, 20, 25, 30]\n",
    "            label = \"300hPa zonal wind\"\n",
    "            cmap = \"RdBu_r\"\n",
    "            extend = \"both\"\n",
    "        else:\n",
    "            raise ValueError(\"Invalid color variable\")\n",
    "        color_da.plot(\n",
    "            ax=ax1,\n",
    "            transform=ccrs.PlateCarree(),\n",
    "            cmap=cmap,\n",
    "            levels=levels,\n",
    "            extend=extend,\n",
    "            cbar_kwargs={\"label\": label},\n",
    "        )\n",
    "\n",
    "    if contour_var:\n",
    "        if contour_var == \"z500\":\n",
    "            levels = np.arange(5000, 6300, 50)\n",
    "        elif contour_var == \"psl\":\n",
    "            levels = np.arange(900, 1100, 2.5)\n",
    "        elif contour_var == \"ua300\":\n",
    "            levels = np.arange(15, 60, 5)\n",
    "        else:\n",
    "            raise ValueError(\"Invalid contour variable\")\n",
    "        lines = contour_da.plot.contour(\n",
    "            ax=ax1, transform=ccrs.PlateCarree(), levels=levels, colors=[\"0.1\"]\n",
    "        )\n",
    "        ax1.clabel(lines, colors=[\"0.1\"], manual=False, inline=True)\n",
    "\n",
    "    ax1.coastlines()\n",
    "    ax1.set_extent([90, 205, -55, 10], crs=ccrs.PlateCarree())\n",
    "    ax1.gridlines(linestyle=\"--\", draw_labels=True)\n",
    "    if contour_var:\n",
    "        ax1.set_title(\n",
    "            f\"Average {contour_var} ({contour_da.units}), {start_date} to {end_date}\"\n",
    "        )\n",
    "    else:\n",
    "        ax1.set_title(f\"{start_date} to {end_date}\")\n",
    "    #    plt.savefig(\n",
    "    #        f'/home/599/dbi599/east-coast-rain/{color_var}-{contour_var}_mean_BARRA2_{start_date}_{end_date}.png',\n",
    "    #        bbox_inches='tight',\n",
    "    #        facecolor='white',\n",
    "    #        dpi=300\n",
    "    #    )\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac05f6f2",
   "metadata": {
    "papermill": {
     "duration": 0.053741,
     "end_time": "2023-03-06T02:53:06.948685",
     "exception": false,
     "start_time": "2023-03-06T02:53:06.894944",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "ranked_events = model_df.sort_values(by=[\"pr\"], ascending=False)\n",
    "ranked_events.head(n=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaa93741",
   "metadata": {
    "papermill": {
     "duration": 1.25781,
     "end_time": "2023-03-06T02:53:08.243572",
     "exception": false,
     "start_time": "2023-03-06T02:53:06.985762",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_add_bc_df = model_add_bc_ds[\"pr\"].to_dataframe()\n",
    "model_add_bc_df = model_add_bc_df.dropna()\n",
    "model_add_bc_df = model_add_bc_df.reset_index()\n",
    "ranked_add_bc_events = model_add_bc_df.sort_values(by=[\"pr\"], ascending=False)\n",
    "ranked_add_bc_events.head(n=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fe3506d-7703-4271-960f-1cc7b3f7aa19",
   "metadata": {
    "papermill": {
     "duration": 1.490974,
     "end_time": "2023-03-06T02:53:09.781177",
     "exception": false,
     "start_time": "2023-03-06T02:53:08.290203",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_mulc_bc_df = model_mulc_bc_ds[\"pr\"].to_dataframe()\n",
    "model_mulc_bc_df = model_mulc_bc_df.dropna()\n",
    "model_mulc_bc_df = model_mulc_bc_df.reset_index()\n",
    "ranked_mulc_bc_events = model_mulc_bc_df.sort_values(by=[\"pr\"], ascending=False)\n",
    "ranked_mulc_bc_events.head(n=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29d990f5",
   "metadata": {
    "papermill": {
     "duration": 2.213593,
     "end_time": "2023-03-06T02:53:12.040544",
     "exception": true,
     "start_time": "2023-03-06T02:53:09.826951",
     "status": "failed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_dcpp_member(ensemble, model_name):\n",
    "    \"\"\"Get the DCPP member information (e.g. r1i1p1f1)\"\"\"\n",
    "\n",
    "    r = ensemble\n",
    "    i = 1\n",
    "    p = 1\n",
    "    f = 1\n",
    "    if model_name == \"CanESM5\":\n",
    "        p = 2\n",
    "    elif model_name == \"EC-Earth3\":\n",
    "        if ensemble > 10:\n",
    "            r = ensemble - 5\n",
    "            i = 2\n",
    "    elif model_name == \"HadGEM3-GC31-MM\":\n",
    "        f = 2\n",
    "    elif model_name == \"NorCPM1\":\n",
    "        if ensemble > 10:\n",
    "            r = ensemble - 10\n",
    "            i = 2\n",
    "\n",
    "    member = f\"r{r}i{i}p{p}f{f}\"\n",
    "\n",
    "    return member\n",
    "\n",
    "\n",
    "def plot_dcpp_meteorology(ranked_events, model_name):\n",
    "    \"\"\"Meteorology plot for DCPP data\"\"\"\n",
    "\n",
    "    for index, row in ranked_events.head(n=4).iterrows():\n",
    "        init_date = int(row[\"init_date\"].strftime(\"%Y\"))\n",
    "        if model_name in [\"CanESM5\", \"IPSL-CM6A-LR\"]:\n",
    "            init_date = init_date - 1\n",
    "        ensemble = int(row[\"ensemble\"]) + 1\n",
    "        member = get_dcpp_member(ensemble, model_name)\n",
    "        end_date = row[\"event_time\"]\n",
    "        rx_value = row[\"pr\"]\n",
    "        print(f\"{init_date} ({end_date}), member: {member}, {rx_value}mm\")\n",
    "\n",
    "        start_datetime = datetime.datetime.strptime(\n",
    "            end_date, \"%Y-%m-%d\"\n",
    "        ) - datetime.timedelta(days=14)\n",
    "        start_date = start_datetime.strftime(\"%Y-%m-%d\")\n",
    "\n",
    "        pr_files = glob.glob(\n",
    "            f\"/g/data/oi10/replicas/CMIP6/DCPP/*/{model_name}/dcppA-hindcast/s{init_date}-{member}/day/pr/*/*/*.nc\"\n",
    "        )\n",
    "        pr_files.sort()\n",
    "        if pr_files:\n",
    "            ds_pr = fileio.open_dataset(\n",
    "                pr_files,\n",
    "                variables=[\"pr\"],\n",
    "            )\n",
    "            pr_da = ds_pr.sel({\"time\": slice(start_date, end_date)})[\"pr\"]\n",
    "            pr_da = xc.units.convert_units_to(pr_da, \"mm d-1\")\n",
    "            pr_da_sum = pr_da.sum(\"time\", keep_attrs=True)\n",
    "\n",
    "        psl_files = glob.glob(\n",
    "            f\"/g/data/oi10/replicas/CMIP6/DCPP/*/{model_name}/dcppA-hindcast/s{init_date}-{member}/day/psl/*/*/*.nc\"\n",
    "        )\n",
    "        psl_files.sort()\n",
    "        if psl_files:\n",
    "            ds_psl = fileio.open_dataset(\n",
    "                psl_files,\n",
    "                variables=[\"psl\"],\n",
    "            )\n",
    "            psl_da = ds_psl.sel({\"time\": slice(start_date, end_date)})[\"psl\"]\n",
    "            psl_da = xc.units.convert_units_to(psl_da, \"hPa\")\n",
    "            psl_da_mean = psl_da.mean(\"time\", keep_attrs=True)\n",
    "\n",
    "        zg_files = glob.glob(\n",
    "            f\"/g/data/oi10/replicas/CMIP6/DCPP/*/{model_name}/dcppA-hindcast/s{init_date}-{member}/day/zg/*/*/*.nc\"\n",
    "        )\n",
    "        zg_files.sort()\n",
    "        if zg_files:\n",
    "            ds_zg = fileio.open_dataset(\n",
    "                zg_files,\n",
    "                variables=[\"zg\"],\n",
    "            )\n",
    "            zg_da = ds_zg.sel({\"time\": slice(start_date, end_date)})[\"zg\"]\n",
    "            z500_da = zg_da.sel({\"plev\": 50000}, method=\"nearest\")\n",
    "            z500_da_mean = z500_da.mean(\"time\", keep_attrs=True)\n",
    "\n",
    "        ua_files = glob.glob(\n",
    "            f\"/g/data/oi10/replicas/CMIP6/DCPP/*/{model_name}/dcppA-hindcast/s{init_date}-{member}/day/ua/*/*/*.nc\"\n",
    "        )\n",
    "        ua_files.sort()\n",
    "        if ua_files:\n",
    "            ds_ua = fileio.open_dataset(\n",
    "                ua_files,\n",
    "                variables=[\"ua\"],\n",
    "            )\n",
    "            ua_da = ds_ua.sel({\"time\": slice(start_date, end_date)})[\"ua\"]\n",
    "            ua300_da = ua_da.sel({\"plev\": 30000}, method=\"nearest\")\n",
    "            ua300_da_mean = ua300_da.mean(\"time\", keep_attrs=True)\n",
    "\n",
    "        if pr_files and psl_files:\n",
    "            plot_average(start_date, end_date, pr_da_sum, \"pr\", psl_da_mean, \"psl\")\n",
    "        if pr_files and zg_files:\n",
    "            plot_average(start_date, end_date, pr_da_sum, \"pr\", z500_da_mean, \"z500\")\n",
    "        if ua_files and zg_files:\n",
    "            plot_average(\n",
    "                start_date, end_date, ua300_da_mean, \"ua300\", z500_da_mean, \"z500\"\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "677c7d4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_cafe_meteorology(ranked_events):\n",
    "    \"\"\"Meteorology plot for DCPP data\"\"\"\n",
    "\n",
    "    for index, row in ranked_events.head(n=4).iterrows():\n",
    "        init_date = row[\"init_date\"].strftime(\"%Y%m%d\")\n",
    "        ensemble = row[\"ensemble\"]\n",
    "        end_date = row[\"event_time\"]\n",
    "        rx_value = row[\"pr\"]\n",
    "        print(f\"{init_date} ({end_date}), ensemble:{ensemble}, {rx_value}mm\")\n",
    "\n",
    "        atmos_file = f\"/g/data/xv83/dcfp/CAFE-f6/c5-d60-pX-f6-{init_date}/atmos_isobaric_daily.zarr.zip\"\n",
    "        atmos_ds = fileio.open_dataset(\n",
    "            atmos_file,\n",
    "            metadata_file=\"/g/data/xv83/unseen-projects/code/dataset_config/dataset_cafe_daily.yml\",\n",
    "            variables=[\"h500\", \"slp\", \"pr\", \"ucomp\"],\n",
    "        )\n",
    "        start_datetime = datetime.datetime.strptime(\n",
    "            end_date, \"%Y-%m-%d\"\n",
    "        ) - datetime.timedelta(days=14)\n",
    "        start_date = start_datetime.strftime(\"%Y-%m-%d\")\n",
    "        z500_da = atmos_ds.sel(\n",
    "            {\"ensemble\": ensemble, \"time\": slice(start_date, end_date)}\n",
    "        )[\"h500\"]\n",
    "        pr_da = atmos_ds.sel(\n",
    "            {\"ensemble\": ensemble, \"time\": slice(start_date, end_date)}\n",
    "        )[\"pr\"]\n",
    "        pr_da = general_utils.convert_units(pr_da, \"mm d-1\")\n",
    "        psl_da = atmos_ds.sel(\n",
    "            {\"ensemble\": ensemble, \"time\": slice(start_date, end_date)}\n",
    "        )[\"slp\"]\n",
    "        ua300_da = atmos_ds.sel(\n",
    "            {\"ensemble\": ensemble, \"level\": 300, \"time\": slice(start_date, end_date)}\n",
    "        )[\"ucomp\"]\n",
    "\n",
    "        z500_da_mean = z500_da.mean(\"time\", keep_attrs=True)\n",
    "        pr_da_sum = pr_da.sum(\"time\", keep_attrs=True)\n",
    "        psl_da_mean = psl_da.mean(\"time\", keep_attrs=True)\n",
    "        ua300_da_mean = ua300_da.mean(\"time\", keep_attrs=True)\n",
    "\n",
    "        plot_average(start_date, end_date, pr_da_sum, \"pr\", z500_da_mean, \"z500\")\n",
    "        plot_average(start_date, end_date, pr_da_sum, \"pr\", psl_da_mean, \"psl\")\n",
    "        plot_average(start_date, end_date, ua300_da_mean, \"ua300\", z500_da_mean, \"z500\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d1d2b07",
   "metadata": {},
   "outputs": [],
   "source": [
    "if model_name == \"CAFE\":\n",
    "    plot_cafe_meteorology(ranked_events)\n",
    "else:\n",
    "    plot_dcpp_meteorology(ranked_events, model_name)"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "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.10.9"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 208.255213,
   "end_time": "2023-03-06T02:53:13.009470",
   "environment_variables": {},
   "exception": true,
   "input_path": "analysis.ipynb",
   "output_path": "analysis_CanESM5.ipynb",
   "parameters": {
    "agcd_file": "/g/data/xv83/dbi599/east-coast-rain/data/Rx15day_AGCD-CSIRO_r005_1900-2022_annual-aug-to-sep_east-coast-flood-region.zarr.zip",
    "file_list": "/home/599/dbi599/east-coast-rain/file_lists/CanESM5_dcppA-hindcast_pr_files.txt",
    "independence_plot": "/g/data/xv83/dbi599/east-coast-rain/figures/independence-test_Rx15day_CanESM5-dcppA-hindcast_196101-201701_annual-aug-to-sep_east-coast-flood-region.png",
    "min_lead": 1,
    "model_add_bc_file": "/g/data/xv83/dbi599/east-coast-rain/data/Rx15day_CanESM5-dcppA-hindcast_196101-201701_annual-aug-to-sep_east-coast-flood-region_bias-corrected-AGCD-CSIRO-additive.zarr.zip",
    "model_file": "/g/data/xv83/dbi599/east-coast-rain/data/Rx15day_CanESM5-dcppA-hindcast_196101-201701_annual-aug-to-sep_east-coast-flood-region.zarr.zip",
    "model_mulc_bc_file": "/g/data/xv83/dbi599/east-coast-rain/data/Rx15day_CanESM5-dcppA-hindcast_196101-201701_annual-aug-to-sep_east-coast-flood-region_bias-corrected-AGCD-CSIRO-multiplicative.zarr.zip",
    "model_name": "CanESM5",
    "nino_file": "/g/data/xv83/dbi599/east-coast-rain/data/nino34-anomaly_CanESM5-dcppA-hindcast_196101-201701_base-196101-201701.nc",
    "region_name": "east-coast-flood-region",
    "shape_file": "/g/data/xv83/dbi599/east-coast-rain/shapefiles/east-coast-flood.shp",
    "similarity_add_bc_file": "/g/data/xv83/dbi599/east-coast-rain/data/similarity-test_Rx15day_CanESM5-dcppA-hindcast_196101-201701_annual-aug-to-sep_east-coast-flood-region_bias-corrected-AGCD-CSIRO-additive.zarr.zip",
    "similarity_mulc_bc_file": "/g/data/xv83/dbi599/east-coast-rain/data/similarity-test_Rx15day_CanESM5-dcppA-hindcast_196101-201701_annual-aug-to-sep_east-coast-flood-region_bias-corrected-AGCD-CSIRO-multiplicative.zarr.zip",
    "similarity_raw_file": "/g/data/xv83/dbi599/east-coast-rain/data/similarity-test_Rx15day_CanESM5-dcppA-hindcast_196101-201701_annual-aug-to-sep_east-coast-flood-region_AGCD-CSIRO.zarr.zip"
   },
   "start_time": "2023-03-06T02:49:44.754257",
   "version": "2.3.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}