run_calc_GWL_figures.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b435edaf-e298-4432-95c9-bef125530820",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import sys\n",
    "import os\n",
    "import pathlib\n",
    "import xarray as xr\n",
    "import palettable\n",
    "import glob\n",
    "import matplotlib.pyplot as plt\n",
    "import json\n",
    "\n",
    "sys.path.append('/g/data/mn51/users/mtb563/toolbox/modules')\n",
    "import utils\n",
    "\n",
    "sys.path.append('/g/data/mn51/users/mtb563/repos_collab/plotting_maps')\n",
    "from acs_plotting_maps import plot_acs_hazard, regions_dict, cmap_dict, tick_dict\n",
    "xr.set_options(keep_attrs=True)\n",
    "\n",
    "# User defined\n",
    "indices = ['HWAtx','HWD','HWN','TNle02','TXge35','TXge40','TXge45','TXm','TXx','HWF']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "6aa81abb-a17e-4a52-9c1a-06f95f40a62f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.colors import ListedColormap, LinearSegmentedColormap\n",
    "\n",
    "def hazard_cmaps(key):\n",
    "    cmaps = {\n",
    "        'tasmax' : ListedColormap(colors=['#E3F4FB','#C8DEE8','#91C4EA','#56B6DC','#00A2AC','#30996C','#7FC69A','#B9DA88','#DCE799','#FCE850','#EACD44','#FED98E','#F89E64','#E67754','#D24241','#AD283B','#832D57','#A2667A','#AB9487'],name='tasmax'),\n",
    "        'FFDI': palettable.colorbrewer.sequential.YlOrRd_9.mpl_colormap,\n",
    "        'EHF': ListedColormap(colors=['#FEFEFE','#FEFEBE','#FEAC00','#EC5000','#EC5000'],name='EHF') \n",
    "    }\n",
    "    return cmaps[key]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e135e609-d4c7-44d1-a32f-de0031ac7a40",
   "metadata": {},
   "outputs": [],
   "source": [
    "cmaps_ticks = {\n",
    "    'HWN':   [ palettable.colorbrewer.sequential.YlOrRd_9.mpl_colormap, [1,2,3,4,5,6,7,8,9,10] ], \n",
    "    'HWF':   [ palettable.colorbrewer.sequential.YlOrRd_9.mpl_colormap, [15,30,45,60,75,90,105,120,135,150,165] ], \n",
    "    'HWAtx': [ hazard_cmaps('tasmax'), [18,20,22,24,26,28,30,32,34,36,38,40,42,44,46,48,50,52] ],\n",
    "    'HWD':   [ palettable.colorbrewer.sequential.YlOrBr_9.mpl_colormap, [10,20,30,40,50,60,70,80,90,100] ],\n",
    "    'TXge35':  [ palettable.colorbrewer.sequential.YlOrRd_9.mpl_colormap, [1,5,10,15,20,40,80,120,160,200] ],\n",
    "    'TNle02':  [ palettable.colorbrewer.sequential.YlGnBu_9.mpl_colormap, [1,5,10,15,20,40,80,120,160,200] ]\n",
    "    }\n",
    "\n",
    "cmaps_ticks.update({'TXge40':cmaps_ticks['TXge35'], 'TXge45':cmaps_ticks['TXge35'], 'TXx':cmaps_ticks['HWAtx'],'TXm':cmaps_ticks['HWAtx']})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd75570b-03af-4c36-bad4-1c54194d0b83",
   "metadata": {},
   "outputs": [],
   "source": [
    "for index in indices:\n",
    "    for bc_method in ['bias-corrected','raw']:\n",
    "        for aggr in ['individual_models','ensemble']:\n",
    "\n",
    "            assert index in cmaps_ticks.keys(), f\"Update cmaps_ticks to include {index}\"\n",
    "            \n",
    "            ddir = f\"/g/data/ia39/ncra/heat/data/{index}/{bc_method}/{aggr}/GWL-average/\"\n",
    "            fdir = f\"/g/data/ia39/ncra/heat/figures/{index}/{bc_method}/{aggr}/GWL-average/\"\n",
    "            pathlib.Path(fdir).mkdir(parents=True, exist_ok=True)\n",
    "            \n",
    "            fnames = glob.glob(ddir+'*.nc')\n",
    "            \n",
    "            for f in fnames:\n",
    "                print(f)\n",
    "                ofile = fdir+f.rsplit('/', 1)[1].replace('.nc','.png')\n",
    "\n",
    "                if not os.path.exists(ofile):\n",
    "                    ds = utils.generalio.read_data(infiles=[f])\n",
    "                    units = ds[index].units\n",
    "                    if units == 'degC':\n",
    "                        units = u'\\N{DEGREE SIGN}C'\n",
    "                    if index in ['HWF']:\n",
    "                        units = '(days/year)'\n",
    "                    fig,ax = plot_acs_hazard(data = ds[index],\n",
    "                                             cmap = cmaps_ticks[index][0],\n",
    "                                             ticks = cmaps_ticks[index][1],\n",
    "                                             #tick_labels = cmaps_ticks[index][1],\n",
    "                                             cbar_label = units,\n",
    "                                             cbar_extend = \"both\",\n",
    "                                             title = f\"{ds[index].attrs['long_name'].capitalize()}\",\n",
    "                                             dataset_name = f\"Models: ensemble, Bias-adjustment: {ds.bias_correction_method}\" if aggr == 'ensemble' else f\"GCM: {ds.driving_model}, RCM: {ds.downscaling_model}, Bias-adjustment: {ds.bias_correction_method}\",\n",
    "                                             date_range = ds.GWL.replace('GWL','Global Warming Level').replace('degC',u'\\N{DEGREE SIGN}C'),\n",
    "                                             watermark=None,\n",
    "                                             outfile = 'temp.png',\n",
    "                                                     );\n",
    "                    metadata = ds.attrs\n",
    "                    metadata['data-history'] = metadata.pop('history')\n",
    "                    metadata['figure-history'] = utils.provenance.get_notebook_prov()\n",
    "                    \n",
    "                    plt.savefig(ofile,transparent=True,dpi=300,metadata={'provenance':json.dumps(metadata)})\n",
    "                    plt.close()\n",
    "                    del(fig,ax,ds,units,metadata,ofile)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:analysis3]",
   "language": "python",
   "name": "conda-env-analysis3-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}