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 os\n",
    "import datetime\n",
    "import glob\n",
    "\n",
    "from collections import Counter\n",
    "import calendar\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import cartopy\n",
    "import cartopy.crs as ccrs\n",
    "import geopandas as gp\n",
    "from IPython.display import Image\n",
    "from scipy.stats import genextreme as gev\n",
    "import xclim as xc\n",
    "import cmocean\n",
    "import seaborn as sns\n",
    "\n",
    "from unseen import fileio\n",
    "from unseen import stability\n",
    "from unseen import eva\n",
    "from unseen import 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 'metric_plot_upper_limit' in locals(), \"Must provide an upper limit for the metric value\"\n",
    "assert os.path.isfile(obs_file), \"Must provide an observations data file (papermill option -p obs_file [filepath])\"\n",
    "assert os.path.isfile(file_list), \"Must provide the input model files list (papermill option -p file_list [filepath])\"\n",
    "assert os.path.isfile(shape_file), \"Must provide the shapefile (papermill option -p shape_file [filepath])\"\n",
    "assert os.path.isfile(model_file), \"Must provide an model data file (papermill option -p model_file [filepath])\"\n",
    "assert os.path.isfile(model_add_bc_file), \"Must provide a model additive bias corrected data file (papermill option -p model_add_bc_file [filepath])\"\n",
    "assert os.path.isfile(model_mulc_bc_file), \"Must provide a model multiplicative bias corrected data file (papermill option -p model_mulc_bc_file [filepath])\"\n",
    "assert os.path.isfile(similarity_add_bc_file), \"Must provide an additive bias corrected similarity test file (papermill option -p similarity_add_bias_file [filepath])\"\n",
    "assert os.path.isfile(similarity_mulc_bc_file), \"Must provide an multiplicative bias corrected similarity test file (papermill option -p similarity_mulc_bias_file [filepath])\"\n",
    "assert os.path.isfile(similarity_raw_file), \"Must provide an raw data similarity test file (papermill option -p similarity_raw_file [filepath])\"\n",
    "assert os.path.isfile(independence_plot), \"Must provide an independence test plot (papermill option -p independence_plot [filepath])\"\n",
    "assert os.path.isfile(stability_plot_empirical), \"Must provide an empirical stability test plot (papermill option -p stability_plot_empirical [filepath])\"\n",
    "assert os.path.isfile(stability_plot_gev), \"Must provide a GEV stability test plot (papermill option -p stability_plot_gev [filepath])\"\n",
    "assert os.path.isfile(moments_add_plot), \"Must provide a moments test plot for the additive bias corrected data (papermill option -p moments_add_plot [filepath])\"\n",
    "assert os.path.isfile(moments_mulc_plot), \"Must provide a moments test plot for the multiplicative bias corrected data (papermill option -p moments_mulc_plot [filepath])\"\n",
    "assert os.path.isfile(moments_raw_plot), \"Must provide a moments test plot for the raw data (papermill option -p moments_raw_plot [filepath])\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddfc8ca8",
   "metadata": {},
   "outputs": [],
   "source": [
    "units_dict = {\n",
    "    'pr': 'mm day-1'\n",
    "}"
   ]
  },
  {
   "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": "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 = '/g/data/xv83/unseen-projects/code/dataset_config/dataset_cafe_daily.yml'\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(shapes.geometry, ccrs.PlateCarree(), facecolor='none', edgecolor='orange')\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": [
    "Image(filename=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": [
    "print(min_lead)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e055048c",
   "metadata": {
    "papermill": {
     "duration": 0.038365,
     "end_time": "2023-03-06T02:50:12.601760",
     "exception": false,
     "start_time": "2023-03-06T02:50:12.563395",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_da = model_ds[var].where(model_ds['lead_time'] >= 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_stacked = model_da.dropna('lead_time').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_stacked = model_add_bc_ds.dropna('lead_time')[var].stack({'sample': ['ensemble', 'init_date', 'lead_time']})"
   ]
  },
  {
   "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_stacked = model_mulc_bc_ds.dropna('lead_time')[var].stack({'sample': ['ensemble', 'init_date', 'lead_time']})"
   ]
  },
  {
   "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(model_add_bc_da_stacked.values)\n",
    "model_mulc_bc_gev_shape, model_mulc_bc_gev_loc, model_mulc_bc_gev_scale = eva.fit_gev(model_mulc_bc_da_stacked.values)"
   ]
  },
  {
   "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, density=True, label=model_name, alpha=0.7, 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, density=True, label=f'{model_name} additive bias corrected', facecolor='tab:red', alpha=0.7\n",
    "        )\n",
    "        model_add_bc_gev_pdf = gev.pdf(gev_xvals, model_add_bc_gev_shape, model_add_bc_gev_loc, model_add_bc_gev_scale)\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, density=True, label=f'{model_name} multiplicative bias corrected', facecolor='tab:orange', alpha=0.7\n",
    "    )\n",
    "    model_mulc_bc_gev_pdf = gev.pdf(gev_xvals, model_mulc_bc_gev_shape, model_mulc_bc_gev_loc, model_mulc_bc_gev_scale)\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. / 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 = eva.event_in_context(\n",
    "    model_add_bc_da_stacked.values,\n",
    "    obs_max_event,\n",
    "    'above',\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(obs_max_event, model_add_bc_gev_shape, model_add_bc_gev_loc, model_add_bc_gev_scale)"
   ]
  },
  {
   "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 = eva.event_in_context(model_add_bc_da_stacked.values, obs_max_event, 'above')\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(obs_max_event, model_mulc_bc_gev_shape, model_mulc_bc_gev_loc, model_mulc_bc_gev_scale)"
   ]
  },
  {
   "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(model_da_stacked.values, obs_max_event, 'above')\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(start_date, end_date, color_da=None, color_var=None, contour_da=None, contour_var=None):\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,\n",
    "            transform=ccrs.PlateCarree(),\n",
    "            levels=levels,\n",
    "            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(f'Average {contour_var} ({contour_da.units}), {start_date} to {end_date}')\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(end_date, \"%Y-%m-%d\") - datetime.timedelta(days=14)\n",
    "        start_date = start_datetime.strftime(\"%Y-%m-%d\")\n",
    "\n",
    "        pr_files = glob.glob(f'/g/data/oi10/replicas/CMIP6/DCPP/*/{model_name}/dcppA-hindcast/s{init_date}-{member}/day/pr/*/*/*.nc')\n",
    "        pr_files.sort()\n",
    "        if pr_files:\n",
    "            ds_pr = fileio.open_dataset(pr_files, variables=['pr'],)\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(f'/g/data/oi10/replicas/CMIP6/DCPP/*/{model_name}/dcppA-hindcast/s{init_date}-{member}/day/psl/*/*/*.nc')\n",
    "        psl_files.sort()\n",
    "        if psl_files:\n",
    "            ds_psl = fileio.open_dataset(psl_files, variables=['psl'],)\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(f'/g/data/oi10/replicas/CMIP6/DCPP/*/{model_name}/dcppA-hindcast/s{init_date}-{member}/day/zg/*/*/*.nc')\n",
    "        zg_files.sort()\n",
    "        if zg_files:\n",
    "            ds_zg = fileio.open_dataset(zg_files, variables=['zg'],)\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(f'/g/data/oi10/replicas/CMIP6/DCPP/*/{model_name}/dcppA-hindcast/s{init_date}-{member}/day/ua/*/*/*.nc')\n",
    "        ua_files.sort()\n",
    "        if ua_files:\n",
    "            ds_ua = fileio.open_dataset(ua_files, variables=['ua'],)\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(start_date, end_date, ua300_da_mean, 'ua300', z500_da_mean, 'z500')"
   ]
  },
  {
   "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(end_date, \"%Y-%m-%d\") - datetime.timedelta(days=14)\n",
    "        start_date = start_datetime.strftime(\"%Y-%m-%d\")\n",
    "        z500_da = atmos_ds.sel({'ensemble': ensemble, 'time': slice(start_date, end_date)})['h500']\n",
    "        pr_da = atmos_ds.sel({'ensemble': ensemble, 'time': slice(start_date, end_date)})['pr']\n",
    "        pr_da = general_utils.convert_units(pr_da, 'mm d-1')\n",
    "        psl_da = atmos_ds.sel({'ensemble': ensemble, 'time': slice(start_date, end_date)})['slp']\n",
    "        ua300_da = atmos_ds.sel({'ensemble': ensemble, 'level': 300, 'time': slice(start_date, end_date)})['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
}