plot_mme_vs_sme.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f54eb26",
   "metadata": {},
   "outputs": [],
   "source": [
    "##!/usr/bin/env python\n",
    "\"\"\"plot_mme_vs_sme.py\n",
    "\n",
    "This script compares the PDFs of the EHI values for the CMIP Multi-Model Ensembles (MME)\n",
    "and the Single-Model Ensembles (SME) against Berkeley Earth for different regions\n",
    "\n",
    "Using the Excess Heat Index \n",
    "Baseline Period: 1950-2014, 90th Percentile Threshold\n",
    "\n",
    "Author: Annette L Hirsch @ CLEX, UNSW. Sydney (Australia)\n",
    "email: a.hirsch@unsw.edu.au\n",
    "Created: Wed Nov 11 10:54:39 AEDT 2020\n",
    "\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6449b8ec",
   "metadata": {},
   "source": [
    "## Load Python Packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d15ba67a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "import numpy as np\n",
    "import netCDF4 as nc\n",
    "import matplotlib.pyplot as plt\n",
    "import xarray as xr\n",
    "import xesmf as xe\n",
    "import regionmask\n",
    "import seaborn as sns; sns.set(style=\"white\", color_codes=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9ea7dbc",
   "metadata": {},
   "source": [
    "### Updated IPCC regions\n",
    "\n",
    "This is a python package containing the new regions: import regionmask\n",
    "\n",
    "Iturbide, M., Guti�errez, J.M., Alves, L.M., Bedia, J., Cimadevilla, E., Cofi~no, A.S., Cerezo-Mota, R., Di Luca, A., Faria, S.H., Gorodetskaya, I., Hauser, M., Herrera, S., Hewitt, H.T., Hennessy, K.J., Jones, R.G., Krakovska, S., Manzanas, R., Marínez-Castro, D., Narisma, G.T., Nurhati, I.S., Pinto, I., Seneviratne, S.I., van den Hurk, B., Vera, C.S., 2020. An update of IPCC climate reference regions for subcontinental analysis of climate model data: definition and aggregated datasets. Earth Syst. Sci. Data Discuss. https://doi.org/10.5194/essd-2019-258\n",
    "\n",
    "Git repo with the new regions: https://github.com/SantanderMetGroup/ATLAS/tree/master/reference-regions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06185906",
   "metadata": {},
   "outputs": [],
   "source": [
    "rlabels = ['GIC', 'NWN', 'NEN', 'WNA', 'CNA', 'ENA', 'NCA', 'SCA', 'CAR', 'NWS',\n",
    "                  'NSA', 'NES', 'SAM', 'SWS', 'SES', 'SSA', 'NEU', 'WCE', 'EEU', 'MED',\n",
    "                  'SAH', 'WAF', 'CAF', 'NEAF', 'SEAF', 'WSAF', 'ESAF', 'MDG', 'RAR', 'WSB',\n",
    "                  'ESB', 'RFE', 'WCA', 'ECA', 'TIB', 'EAS', 'ARP', 'SAS', 'SEA', 'NAU',\n",
    "                  'CAU', 'EAU', 'SAU', 'NZ']#, 'EAN', 'WAN']\n",
    "nreg = len(rlabels)\n",
    "\n",
    "ar6_land = regionmask.defined_regions.ar6.land"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d443b77",
   "metadata": {},
   "source": [
    "### Berkeley Earth\n",
    "\n",
    "This data has nlat = 180, nlon = 360 so it is 1deg data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89b4d6e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "berkdir = '/g/data/w97/azh561/cmip_heatwave/BerkleyEarth/'\n",
    "\n",
    "# Define the filename convention:\n",
    "bprefix = 'EHF_'\n",
    "bendfix = '_summer_BerkleyEarth_1950-2014.nc'\n",
    "\n",
    "# for the cumulative heat\n",
    "hwcfile = 'EHF_HWC_summer_BerkleyEarth_1950-2014.nc'\n",
    "ehifile = 'BerkleyEarth_HW_daily_EHI_value_1950_2017_bp_1950_2014.nc'\n",
    "spellfile = 'BerkleyEarth_HW_daily_EHF_1950_2017_bp_1950_2014.nc'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "361823dc",
   "metadata": {},
   "source": [
    "### CMIP6 data\n",
    "\n",
    "They are all on different resolutions! So will require regridding to plot and evaluate:\n",
    "\n",
    "    - finest resolution: nlat = 256, nlon = 512 -> 0.7 deg x 0.7 deg\n",
    "    - coarsest resolution: nlat = 96, nlon = 144 -> 1.875 deg x 2.5 deg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da0b2b00",
   "metadata": {},
   "outputs": [],
   "source": [
    "cmip6dir = '/g/data/w97/azh561/cmip_heatwave/cmip6_ehf/'\n",
    "\n",
    "# Define the models\n",
    "models6 = ['ACCESS-CM2','AWI-CM-1-1-MR','AWI-ESM-1-1-LR','BCC-CSM2-MR','BCC-ESM1','CanESM5','CMCC-CM2-SR5',\n",
    "'EC-Earth3-AerChem','EC-Earth3','EC-Earth3-Veg','EC-Earth3-Veg-LR','FGOALS-f3-L','FGOALS-g3','GFDL-CM4',\n",
    "'GFDL-ESM4','GISS-E2-1-G','INM-CM4-8','INM-CM5-0','IPSL-CM6A-LR','KACE-1-0-G','KIOST-ESM','MIROC6',\n",
    "'MPI-ESM-1-2-HAM','MPI-ESM1-2-HR','MPI-ESM1-2-LR','MRI-ESM2-0','NESM3','NorCPM1','NorESM2-LM','NorESM2-MM','TaiESM1']\n",
    "nmod6 = len(models6)\n",
    "\n",
    "# Define the filename convention:\n",
    "mendfix6 = '_historical_r1i1p1f1_yearly_summer.nc'\n",
    "\n",
    "# The EHI files\n",
    "ehiefix6 = '_historical_r1i1p1f1_daily.nc'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a21cd420",
   "metadata": {},
   "source": [
    "### CMIP6 Models with 10 members"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "768fa7d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "mmdir = '/g/data/w97/azh561/cmip_heatwave/cmip6_sme_ehf/'\n",
    "mmendfix = '_yearly_summer.nc'\n",
    "\n",
    "mm1 = ['IPSL-CM6A-LR_historical_r2i1p1f1','IPSL-CM6A-LR_historical_r3i1p1f1','IPSL-CM6A-LR_historical_r4i1p1f1',\n",
    "'IPSL-CM6A-LR_historical_r5i1p1f1','IPSL-CM6A-LR_historical_r6i1p1f1','IPSL-CM6A-LR_historical_r7i1p1f1',\n",
    "'IPSL-CM6A-LR_historical_r8i1p1f1','IPSL-CM6A-LR_historical_r9i1p1f1','IPSL-CM6A-LR_historical_r10i1p1f1']\n",
    "\n",
    "mm2 = ['MIROC6_historical_r2i1p1f1','MIROC6_historical_r3i1p1f1','MIROC6_historical_r4i1p1f1',\n",
    "       'MIROC6_historical_r5i1p1f1','MIROC6_historical_r6i1p1f1','MIROC6_historical_r7i1p1f1',\n",
    "       'MIROC6_historical_r8i1p1f1','MIROC6_historical_r9i1p1f1','MIROC6_historical_r10i1p1f1']\n",
    "\n",
    "mm3 = ['MPI-ESM1-2-HR_historical_r2i1p1f1','MPI-ESM1-2-HR_historical_r3i1p1f1','MPI-ESM1-2-HR_historical_r4i1p1f1',\n",
    "'MPI-ESM1-2-HR_historical_r5i1p1f1','MPI-ESM1-2-HR_historical_r6i1p1f1','MPI-ESM1-2-HR_historical_r7i1p1f1',\n",
    "'MPI-ESM1-2-HR_historical_r8i1p1f1','MPI-ESM1-2-HR_historical_r9i1p1f1','MPI-ESM1-2-HR_historical_r10i1p1f1']\n",
    "\n",
    "mm4 = ['CanESM5_historical_r2i1p1f1','CanESM5_historical_r3i1p1f1','CanESM5_historical_r4i1p1f1',\n",
    "      'CanESM5_historical_r5i1p1f1','CanESM5_historical_r6i1p1f1','CanESM5_historical_r7i1p1f1',\n",
    "      'CanESM5_historical_r8i1p1f1','CanESM5_historical_r9i1p1f1','CanESM5_historical_r10i1p1f1']\n",
    "\n",
    "mm5 = ['UKESM1-0-LL_historical_r1i1p1f2','UKESM1-0-LL_historical_r2i1p1f2','UKESM1-0-LL_historical_r3i1p1f2',\n",
    "'UKESM1-0-LL_historical_r4i1p1f2','UKESM1-0-LL_historical_r8i1p1f2','UKESM1-0-LL_historical_r9i1p1f2']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77dea1e0",
   "metadata": {},
   "source": [
    "### CMIP5 MME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcc688e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "cmip5dir = '/g/data/w97/azh561/cmip_heatwave/cmip5_ehf/'\n",
    "\n",
    "# Define the models\n",
    "models5 = ['ACCESS1-0','ACCESS1.3','bcc-csm1-1','bcc-csm1-1-m','BNU-ESM',\n",
    "          'CanESM2','CCSM4','CESM1-BGC','CESM1-CAM5','CMCC-CESM',\n",
    "'CMCC-CM','CMCC-CMS','CNRM-CM5','CSIRO-Mk3-6-0','FGOALS_g2','GFDL-CM3',\n",
    "'GFDL-ESM2G','GFDL-ESM2M',#'HadGEM2-AO','HadGEM2-CC','HadGEM2-ES',\n",
    "          'inmcm4',\n",
    "'IPSL-CM5A-LR','IPSL-CM5A-MR','IPSL-CM5B-LR','MIROC5','MIROC-ESM-CHEM',\n",
    "'MIROC-ESM','MPI-ESM-LR','MPI-ESM-MR','MRI-CGCM3','MRI-ESM1','NorESM1-M']\n",
    "\n",
    "nmod5 = len(models5)\n",
    "\n",
    "# Define the filename convention:\n",
    "mendfix5 = '_historical_r1i1p1_yearly_summer.nc'\n",
    "\n",
    "# The EHI files\n",
    "ehiefix5 = '_historical_r1i1p1_daily.nc'\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90ce9aa9",
   "metadata": {},
   "source": [
    "### Define some functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48c28a7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# If longitudes are 0 to 360 need to shift to -180 to 180\n",
    "# https://stackoverflow.com/questions/53345442/about-changing-longitude-array-from-0-360-to-180-to-180-with-python-xarray\n",
    "def fix_lon(dm,lon_name):\n",
    "    '''Function to adjust longitudes'''\n",
    "    dm['_longitude_adjusted'] = xr.where(dm[lon_name] > 180,dm[lon_name] - 360,dm[lon_name])\n",
    "    # reassign the new coords to as the main lon coords and sort DataArray using new coordinate values\n",
    "    dm = (\n",
    "        dm\n",
    "        .swap_dims({lon_name: '_longitude_adjusted'})\n",
    "        .sel(**{'_longitude_adjusted': sorted(dm._longitude_adjusted)})\n",
    "        .drop(lon_name))\n",
    "    dm = dm.rename({'_longitude_adjusted': lon_name})\n",
    "    \n",
    "    return dm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f8ef79b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def hwseason(data,spell,yr,variable):\n",
    "    \n",
    "    # Define the output array\n",
    "    dataout = np.empty((len(data['lat'].values),len(data['lon'].values)),dtype=np.float64)\n",
    "\n",
    "    # Select relevant data, mask for non-HW days and sum across the Northern Hemisphere MJJAS\n",
    "    nhtmp = data[variable].loc['%s-05-01' %(yr):'%s-09-30' %(yr),:,:]\n",
    "    nhspell = spell['event'].loc['%s-05-01' %(yr):'%s-09-30' %(yr),:,:]\n",
    "    nhma = np.ma.masked_array(nhtmp,nhspell<=0).filled(np.nan)\n",
    "\n",
    "    dataout[90:,:] = np.nansum(nhma[:,90:,:],axis=0)\n",
    "    del nhtmp,nhspell,nhma\n",
    "    \n",
    "    # Select relevant data, mask for non-HW days and sum across the Southern Hemisphere NDJFM\n",
    "    shtmp1 = data[variable].loc['%s-11-01' %(yr):'%s-12-31' %(yr),:,:]\n",
    "    shspell1 = spell['event'].loc['%s-11-01' %(yr):'%s-12-31' %(yr),:,:]\n",
    "    shma1 = np.ma.masked_array(shtmp1,shspell1<=0).filled(np.nan)\n",
    "\n",
    "    if yr == 2014:\n",
    "        dataout[:90,:] = np.nansum(shma1[:,:90,:],axis=0)\n",
    "    else:\n",
    "        shtmp2 = data[variable].loc['%s-01-01' %(yr+1):'%s-03-31' %(yr+1),:,:]\n",
    "        shspell2 = spell['event'].loc['%s-01-01' %(yr+1):'%s-03-31' %(yr+1),:,:]\n",
    "        shma2 = np.ma.masked_array(shtmp2,shspell2<=0).filled(np.nan)\n",
    "        dataout[:90,:] = np.nansum(shma1[:,:90,:],axis=0) + np.nansum(shma2[:,:90,:],axis=0)\n",
    "        del shtmp2,shspell2,shma2\n",
    "    del shtmp1,shspell1,shma1\n",
    "    \n",
    "    return dataout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "273ed182",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_model(hwind,datadir,prefix,modellist,endfix):\n",
    "    \n",
    "    nmod = len(modellist)\n",
    "    \n",
    "    moddata = np.empty((nmod,nyrs,nlat,nlon),dtype=np.float64)\n",
    "    \n",
    "    for mind in range(nmod):\n",
    "        dm = xr.open_dataset('%s%s%s%s' %(datadir,prefix,modellist[mind],endfix),decode_times=False)\n",
    "        \n",
    "        if np.nanmax(dm['lon']) > 200.:\n",
    "            dm = fix_lon(dm,'lon')\n",
    "\n",
    "        if dm[hwind].shape[0] > nyrs:\n",
    "            dm = dm.isel(time=slice(100,165))\n",
    "            \n",
    "        # regrid data\n",
    "        datavalue = dm[hwind]\n",
    "        regridder = xe.Regridder(dm, ds_out, 'bilinear')\n",
    "        moddata[mind,:,:,:] = regridder(datavalue)\n",
    "        del dm, datavalue,regridder\n",
    "            \n",
    "    return moddata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c03625aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_model_hwc(hwind,datadir,prefix1,prefix2,modellist,endfix):\n",
    "    \n",
    "    nmod = len(modellist)\n",
    "    \n",
    "    moddata = np.empty((nmod,nyrs,nlat,nlon),dtype=np.float64)\n",
    "\n",
    "    for mind in range(nmod):\n",
    "        \n",
    "        # Open data\n",
    "        dm = xr.open_dataset('%s%s%s%s' %(datadir,prefix1,modellist[mind],endfix))\n",
    "        dms = xr.open_dataset('%s%s%s%s' %(datadir,prefix2,modellist[mind],endfix))\n",
    "        \n",
    "        # Fix longitudes if required\n",
    "        if np.nanmax(dm['lon']) > 200.:                \n",
    "            dm = fix_lon(dm,'lon')\n",
    "        if np.nanmax(dms['lon']) > 200.:\n",
    "            dms = fix_lon(dms,'lon')\n",
    "\n",
    "        # calculate HWD\n",
    "        datahwseason = np.empty((nyrs,len(dm['lat'].values),len(dm['lon'].values)),dtype=np.float64)\n",
    "        yr=1950\n",
    "        for yy in range(nyrs):\n",
    "            datahwseason[yy,:,:] = hwseason(dm,dms,yr,'EHIsig')\n",
    "            yr += 1\n",
    "\n",
    "        dm_new = xr.Dataset({\n",
    "            'HWC': xr.DataArray(\n",
    "                    data   = datahwseason,\n",
    "                    dims   = ['year','lat','lon'],\n",
    "                    coords = {'year': np.arange(1950,2015),'lat': dm['lat'].values,'lon': dm['lon'].values},\n",
    "                    )\n",
    "                })\n",
    "        \n",
    "        dm_new.to_netcdf('/g/data/w97/azh561/cmip_heatwave/HWC/HWC_%s.nc' %(modellist[mind]))\n",
    "        \n",
    "        # Regrid data \n",
    "        datavalue = dm_new['HWC']\n",
    "        regridder = xe.Regridder(dm_new, ds_out, 'bilinear')\n",
    "        moddata[mind,:,:,:] = regridder(datavalue)\n",
    "        del dm,dms,dm_new,datavalue,regridder,datahwseason\n",
    "        \n",
    "    moddata[:,-1,:,:] = np.nan\n",
    "    \n",
    "    return moddata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19f791ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "process_model_hwc('HWC',mmdir,'EHI_heatwaves_','EHF_heatwaves_',mm1,'_daily.nc')\n",
    "process_model_hwc('HWC',mmdir,'EHI_heatwaves_','EHF_heatwaves_',mm2,'_daily.nc')\n",
    "process_model_hwc('HWC',mmdir,'EHI_heatwaves_','EHF_heatwaves_',mm3,'_daily.nc')\n",
    "process_model_hwc('HWC',mmdir,'EHI_heatwaves_','EHF_heatwaves_',mm4,'_daily.nc')\n",
    "process_model_hwc('HWC',mmdir,'EHI_heatwaves_','EHF_heatwaves_',mm5,'_daily.nc')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97b78f8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to plot data\n",
    "def plot_kde(bedata,cmip6,cmip5,mm1,mm2,mm3,mm4,mm5,vlabel,rlabel,xlim,ylim,figurename):\n",
    "    \"\"\"This function plots the kde \"\"\"\n",
    "          \n",
    "    # Create figure object and subplots\n",
    "    fig, ax = plt.subplots(figsize=(5,5))\n",
    "\n",
    "    plt.rcParams['savefig.dpi']=300\n",
    "    plt.rcParams[\"font.weight\"] = \"bold\"\n",
    "    plt.rcParams[\"axes.labelweight\"] = \"bold\"\n",
    "    plt.rcParams[\"axes.titleweight\"] = \"bold\"\n",
    "    plt.rcParams[\"axes.titlepad\"] = 5.0\n",
    "\n",
    "    # Plot CMIP6\n",
    "    sns.kdeplot(np.reshape(cmip6,-1), ax=ax, shade=False, color='red', linestyle='-',linewidth=3,label='CMIP6') # All years\n",
    "\n",
    "    # Plot CMIP5\n",
    "    sns.kdeplot(np.reshape(cmip5,-1), ax=ax, shade=False, color='blue', linestyle='-',linewidth=3,label='CMIP5') # All years\n",
    "\n",
    "    # Plot Individual Model Ensembles\n",
    "    \n",
    "    evenly_spaced_interval = np.linspace(0.2, 0.8, 5)\n",
    "    mycolors = [plt.cm.binary(x) for x in evenly_spaced_interval]\n",
    "   \n",
    "    sns.kdeplot(np.reshape(mm1,-1), ax=ax, shade=False, color=mycolors[0], linestyle='-',linewidth=1,label='IPSL') # All years\n",
    "    sns.kdeplot(np.reshape(mm2,-1), ax=ax, shade=False, color=mycolors[1], linestyle='-',linewidth=1,label='MIROC') # All years\n",
    "    sns.kdeplot(np.reshape(mm3,-1), ax=ax, shade=False, color=mycolors[2], linestyle='-',linewidth=1,label='MPI-ESM') # All years\n",
    "    sns.kdeplot(np.reshape(mm4,-1), ax=ax, shade=False, color=mycolors[3], linestyle='-',linewidth=1,label='CanESM') # All years\n",
    "    sns.kdeplot(np.reshape(mm5,-1), ax=ax, shade=False, color=mycolors[4], linestyle='-',linewidth=1,label='UKESM') # All years\n",
    "\n",
    "    # Plot Berkeley Earth\n",
    "    sns.kdeplot(np.reshape(bedata,-1), ax=ax, shade=False, color='black', linestyle='-',linewidth=3,label='BE') # All years\n",
    "\n",
    "    # Labelling of subplot panels\n",
    "    ax.set_xlabel('%s' %(vlabel), fontweight = 'bold')\n",
    "    ax.set_ylabel(\"Density Function\", fontweight = 'bold')\n",
    "\n",
    "    # Title\n",
    "    ax.set_title('%s' %(rlabel), fontweight = 'bold')\n",
    "\n",
    "    # Set common axis limits\n",
    "    ax.set_ylim(0,ylim)\n",
    "    ax.set_xlim(0,xlim)\n",
    "\n",
    "    # Add legend\n",
    "    plt.legend()\n",
    "    \n",
    "    fig.savefig(figurename, transparent=True)\n",
    "    plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abd7c64c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def kde_reg_wgt_avg(indata,st,fi):\n",
    "\n",
    "    # region mask\n",
    "    maskreg = ar6_land.mask(ds_out)\n",
    "    \n",
    "    if len(indata.shape) == 3:\n",
    "    \n",
    "        nyrs,nlat,nlon = indata.shape\n",
    "        mask = np.repeat(maskreg.values[np.newaxis,...],nyrs,axis=0)\n",
    "\n",
    "    else:\n",
    "\n",
    "        nmod,nyrs,nlat,nlon = indata.shape\n",
    "        masktmp = np.repeat(maskreg.values[np.newaxis,...],nyrs,axis=0)\n",
    "        mask = np.repeat(masktmp[np.newaxis,...],nmod,axis=0)\n",
    "\n",
    "    regdata = np.ma.masked_array(indata,(mask<st)|(mask>fi)).filled(np.nan)\n",
    "\n",
    "    # calculate the area-weighted average\n",
    "    latr = np.deg2rad(lat2d)\n",
    "    weights = np.cos(latr)\n",
    "\n",
    "    # Loop through the years\n",
    "    \n",
    "    if len(indata.shape) == 3:\n",
    "        datawgt = np.empty((nyrs),dtype=np.float64)\n",
    "        for yy in range(nyrs):\n",
    "            datawgt[yy] = np.ma.average(np.ma.MaskedArray(regdata[yy,:,:], mask=np.isnan(regdata[yy,:,:])),weights=weights) \n",
    "    else:\n",
    "        datawgt = np.empty((nmod,nyrs),dtype=np.float64)\n",
    "\n",
    "        for yy in range(nyrs):\n",
    "            for mm in range(nmod):\n",
    "                datawgt[mm,yy] = np.ma.average(np.ma.MaskedArray(regdata[mm,yy,:,:], mask=np.isnan(regdata[mm,yy,:,:])),weights=weights)\n",
    "\n",
    "    return datawgt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f28e104d",
   "metadata": {},
   "source": [
    "### Loop through the heatwave indices "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7590e62a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the index, label and axis limit\n",
    "var = ['HWF','HWD','HWM','HWC']\n",
    "nvar = len(var)\n",
    "modvar = ['HWF_EHF','HWD_EHF','HWM_EHF','HWC_EHF']\n",
    "vlabels = ['HWF [days]','HWD [days]','HWM [\\xb0 $C^{2}$]','HWC [\\xb0 C]']\n",
    "vlim = [50,20,20,50]\n",
    "vxlim = [0.20,0.5,0.5,0.20]\n",
    "\n",
    "# Define dimensions\n",
    "nyrs = 65\n",
    "nlat = 180\n",
    "nlon = 360\n",
    "\n",
    "regst = [1,9,16,20,28,31,39]\n",
    "regfi = [6,15,19,27,30,37,42]\n",
    "regnm = [\"North America\",\"South America\",\"Europe\",\"Africa\",\"Russia\",\"Asia\",\"Australia\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e400a8a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "for eind in range(nvar):\n",
    "    \n",
    "    print('Calculating and plotting %s' %(var[eind]))\n",
    "    \n",
    "    # Read in previously calculated HWC\n",
    "    if var[eind] in ['HWC']:\n",
    "        \n",
    "        # Read in the Berkeley Earth data \n",
    "        berk = xr.open_dataset('/g/data/w97/azh561/cmip_heatwave/HWC/HWC_Berkeley_Earth.nc')\n",
    "        berkma = np.ma.masked_array(berk[var[eind]].values,berk[var[eind]].values<=0).filled(np.nan)\n",
    "        berkma[-1,:,:] = np.nan # Make the last year nan as model data for last summer season in the Southern Hemisphere is incomplete\n",
    "        ds_out = xr.Dataset({'lat': (['lat'], berk['lat'].values),'lon': (['lon'], berk['lon'].values),})\n",
    "        \n",
    "        # CMIP6 models\n",
    "        moddata6 = process_model('HWC','/g/data/w97/azh561/cmip_heatwave/HWC/cmip6/','EHI_heatwaves_',models6,'_historical_r1i1p1f1_HWC.nc')\n",
    "        moddata6[:,-1,:,:] = np.nan\n",
    "        \n",
    "        # CMIP5 models\n",
    "        moddata5 = process_model('HWC','/g/data/w97/azh561/cmip_heatwave/HWC/cmip5/','EHI_heatwaves_',models5,'_historical_r1i1p1_HWC.nc')\n",
    "        moddata5[:,-1,:,:] = np.nan\n",
    "        \n",
    "        # Individual Model Ensembles\n",
    "        mm1data = process_model('HWC','/g/data/w97/azh561/cmip_heatwave/HWC/','HWC_',mm1,'.nc')\n",
    "        mm1data[:,-1,:,:] = np.nan\n",
    "        mm2data = process_model('HWC','/g/data/w97/azh561/cmip_heatwave/HWC/','HWC_',mm2,'.nc')\n",
    "        mm2data[:,-1,:,:] = np.nan\n",
    "        mm3data = process_model('HWC','/g/data/w97/azh561/cmip_heatwave/HWC/','HWC_',mm3,'.nc')\n",
    "        mm3data[:,-1,:,:] = np.nan\n",
    "        mm4data = process_model('HWC','/g/data/w97/azh561/cmip_heatwave/HWC/','HWC_',mm4,'.nc')\n",
    "        mm4data[:,-1,:,:] = np.nan\n",
    "        mm5data = process_model('HWC','/g/data/w97/azh561/cmip_heatwave/HWC/','HWC_',mm5,'.nc')\n",
    "        mm5data[:,-1,:,:] = np.nan\n",
    "\n",
    "    else:\n",
    "        \n",
    "        # Read in the Berkeley Earth data - decode times necessary as HWF and HWD in units of days\n",
    "        berk = xr.open_dataset('%s%s%s%s' %(berkdir,bprefix,var[eind],bendfix),decode_times=False)\n",
    "        berkma = np.ma.masked_array(berk[var[eind]].values,berk[var[eind]].values<=0).filled(np.nan)\n",
    "\n",
    "        # Define the common grid for the models to regrid too!\n",
    "        ds_out = xr.Dataset({'lat': (['lat'], berk['latitude'].values),'lon': (['lon'], berk['longitude'].values),})\n",
    "\n",
    "        # CMIP6 models\n",
    "        moddata6 = process_model(modvar[eind],cmip6dir,'EHF_heatwaves_',models6,mendfix6)\n",
    "        \n",
    "        # CMIP5 models\n",
    "        moddata5 = process_model(modvar[eind],cmip5dir,'EHF_heatwaves_',models5,mendfix5)\n",
    "\n",
    "        # Individual Model Ensembles\n",
    "        mm1data = process_model(modvar[eind],mmdir,'EHF_heatwaves_',mm1,mmendfix)\n",
    "        mm2data = process_model(modvar[eind],mmdir,'EHF_heatwaves_',mm2,mmendfix)\n",
    "        mm3data = process_model(modvar[eind],mmdir,'EHF_heatwaves_',mm3,mmendfix)\n",
    "        mm4data = process_model(modvar[eind],mmdir,'EHF_heatwaves_',mm4,mmendfix)\n",
    "        mm5data = process_model(modvar[eind],mmdir,'EHF_heatwaves_',mm5,mmendfix)\n",
    "        \n",
    "    _, lat2d = np.meshgrid(ds_out['lon'].values, ds_out['lat'].values)\n",
    "\n",
    "    # Define the land-sea mask (region = 0, elsewhere is nan)\n",
    "    land_110 = regionmask.defined_regions.natural_earth.land_110\n",
    "    land_mask = land_110.mask(ds_out)\n",
    "    mask2d = ar6_land.mask(ds_out) * land_mask.squeeze(drop=True)\n",
    "\n",
    "    # Define the land-sea mask (region = 0, elsewhere is nan)\n",
    "    maskreg = ar6_land.mask(ds_out)\n",
    "    maskland = maskreg * land_mask.squeeze(drop=True)\n",
    "    maskland3d = np.repeat(maskland.values[np.newaxis,...],nyrs,axis=0)\n",
    "        \n",
    "    # Mask the missing values as these have not been set to nan in the files\n",
    "    modma6 = np.ma.masked_array(moddata6,moddata6<=0).filled(np.nan)\n",
    "    modma5 = np.ma.masked_array(moddata5,moddata5<=0).filled(np.nan)\n",
    "    mm1ma5 = np.ma.masked_array(mm1data,mm1data<=0).filled(np.nan)\n",
    "    mm2ma5 = np.ma.masked_array(mm2data,mm2data<=0).filled(np.nan)\n",
    "    mm3ma5 = np.ma.masked_array(mm3data,mm3data<=0).filled(np.nan)\n",
    "    mm4ma5 = np.ma.masked_array(mm4data,mm4data<=0).filled(np.nan)\n",
    "    mm5ma5 = np.ma.masked_array(mm5data,mm5data<=0).filled(np.nan)\n",
    "           \n",
    "    # Mask ocean points\n",
    "    berklm = np.ma.masked_array(berkma,maskland3d!=0).filled(np.nan)\n",
    "    modlm6 = np.ma.masked_array(modma6,np.repeat(maskland3d[np.newaxis,...],nmod6,axis=0)!=0).filled(np.nan)\n",
    "    modlm5 = np.ma.masked_array(modma5,np.repeat(maskland3d[np.newaxis,...],nmod5,axis=0)!=0).filled(np.nan)\n",
    "    mm1lm5 = np.ma.masked_array(mm1ma5,np.repeat(maskland3d[np.newaxis,...],len(mm1),axis=0)!=0).filled(np.nan)\n",
    "    mm2lm5 = np.ma.masked_array(mm2ma5,np.repeat(maskland3d[np.newaxis,...],len(mm2),axis=0)!=0).filled(np.nan)\n",
    "    mm3lm5 = np.ma.masked_array(mm3ma5,np.repeat(maskland3d[np.newaxis,...],len(mm3),axis=0)!=0).filled(np.nan)\n",
    "    mm4lm5 = np.ma.masked_array(mm4ma5,np.repeat(maskland3d[np.newaxis,...],len(mm4),axis=0)!=0).filled(np.nan)\n",
    "    mm5lm5 = np.ma.masked_array(mm5ma5,np.repeat(maskland3d[np.newaxis,...],len(mm5),axis=0)!=0).filled(np.nan)\n",
    "    \n",
    "    # Plot the KDEs for each region\n",
    "    for rr in range(len(regnm)):\n",
    "        berkwgt = kde_reg_wgt_avg(berklm,regst[rr],regfi[rr])\n",
    "        mwgt6 = kde_reg_wgt_avg(modlm6,regst[rr],regfi[rr])\n",
    "        mwgt5 = kde_reg_wgt_avg(modlm5,regst[rr],regfi[rr])\n",
    "        m1wgt5 = kde_reg_wgt_avg(mm1lm5,regst[rr],regfi[rr])\n",
    "        m2wgt5 = kde_reg_wgt_avg(mm2lm5,regst[rr],regfi[rr])\n",
    "        m3wgt5 = kde_reg_wgt_avg(mm3lm5,regst[rr],regfi[rr])\n",
    "        m4wgt5 = kde_reg_wgt_avg(mm4lm5,regst[rr],regfi[rr])\n",
    "        m5wgt5 = kde_reg_wgt_avg(mm5lm5,regst[rr],regfi[rr])\n",
    "        fname = \"%s_%s.png\" %(var[eind],regnm[rr])\n",
    "        plot_kde(berkwgt,mwgt6,mwgt5,m1wgt5,m2wgt5,m3wgt5,m4wgt5,m5wgt5,vlabels[eind],regnm[rr],vlim[eind],vxlim[eind],fname)\n",
    "        del berkwgt,mwgt6,mwgt5,m1wgt5,m2wgt5,m3wgt5,m4wgt5,m5wgt5\n",
    " \n",
    "    # Clean up\n",
    "    del berkma,berklm,moddata6,modma6,modlm6,moddata5,modma5,modlm5\n",
    "    del mm1data,mm1ma5,mm1lm5\n",
    "    del mm2data,mm2ma5,mm2lm5\n",
    "    del mm3data,mm3ma5,mm3lm5\n",
    "    del mm4data,mm4ma5,mm4lm5\n",
    "    del mm5data,mm5ma5,mm5lm5\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f51947a2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:analysis3-21.01]",
   "language": "python",
   "name": "conda-env-analysis3-21.01-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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}