ProbabilityRatio_calcu_NorESM_part3.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "08c59d05",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import scipy.io as scio\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "import cartopy.feature as cfeature\n",
    "import cartopy.crs as ccrs\n",
    "import netCDF4 as nc\n",
    "import numpy as np\n",
    "from Load_data import Data_from_nc\n",
    "import xarray as xr\n",
    "import math\n",
    "import gc\n",
    "\n",
    "data_land_mask = Data_from_nc('/dodrio/scratch/projects/2022_200/project_output/cesm/yi_yao_IRRMIP/surfdata_0.9x1.25_hist_78pfts_CMIP6_simyr1850_c201126.nc')\n",
    "data_mask = data_land_mask.load_variable('PFTDATA_MASK')\n",
    "data_mask = data_mask[29:, :]\n",
    "\n",
    "#create func to compute PR\n",
    "def comp_PR(ctl,scen,q):\n",
    "    #first get quantile from ctl\n",
    "    qctl=ctl.quantile(q,dim='time')\n",
    "    \n",
    "    del ctl\n",
    "    gc.collect()\n",
    "    \n",
    "\n",
    "    #get subsample scenarios\n",
    "    qscen=xr.where(scen>qctl,1,0)\n",
    "    pscen=qscen.sum('time')/len(qscen.time)\n",
    "    \n",
    "    del qscen\n",
    "    gc.collect()\n",
    "\n",
    "    #compute PR\n",
    "    PR_scen=pscen/(1-q)\n",
    "    PR_scen = np.where(data_mask==0, np.nan, PR_scen)\n",
    "    PR_scen = np.where(PR_scen==0, np.nan, PR_scen)\n",
    "    \n",
    "    return PR_scen\n",
    "\n",
    "def get_data_from_mat(file, variable):\n",
    "    var_dict = scio.loadmat(file)\n",
    "    var = var_dict[variable]\n",
    "    var = var[:, 29:, :]\n",
    "    return var\n",
    "\n",
    "data_surface = Data_from_nc(\\\n",
    "                            '/dodrio/scratch/projects/2022_200/project_output/cesm/yi_yao_IRRMIP/' + \\\n",
    "                            'surfdata_irrigation_method.nc'\\\n",
    "                           )   #load the data\n",
    "\n",
    "data_irrigation_method = data_surface.load_variable('irrigation_method')\n",
    "\n",
    "data_lon = data_surface.load_variable('LONGXY')\n",
    "data_lon = data_lon[0,:]\n",
    "data_lat = data_surface.load_variable('LATIXY')\n",
    "data_lat = data_lat[29:,0]\n",
    "data_lat = data_lat[:]\n",
    "data_time = np.array(range(262800))\n",
    "\n",
    "def get_data_for_this_study_1(metric, period, exp, ens):\n",
    "    data_out = get_data_from_mat('/dodrio/scratch/projects/2022_200/project_output/cesm/yi_yao_IRRMIP/after_process/NorESM/post_process/3_hourly/' +\\\n",
    "                                             metric + '_' + period + '_3_hourly_' + exp + ens + '_1.mat', metric)\n",
    "    data_out = np.array(data_out)\n",
    "    return data_out\n",
    "\n",
    "def get_data_for_this_study_2(metric, period, exp, ens):\n",
    "    data_out = get_data_from_mat('/dodrio/scratch/projects/2022_200/project_output/cesm/yi_yao_IRRMIP/after_process/NorESM/post_process/3_hourly/' +\\\n",
    "                                             metric + '_' + period + '_3_hourly_' + exp + ens + '_2.mat', metric)\n",
    "    data_out = np.array(data_out)\n",
    "    return data_out\n",
    "\n",
    "def get_data_for_this_study_3(metric, period, exp, ens):\n",
    "    data_out = get_data_from_mat('/dodrio/scratch/projects/2022_200/project_output/cesm/yi_yao_IRRMIP/after_process/NorESM/post_process/3_hourly/' +\\\n",
    "                                             metric + '_' + period + '_3_hourly_' + exp + ens + '_3.mat', metric)\n",
    "    data_out = np.array(data_out)\n",
    "    return data_out\n",
    "\n",
    "def get_data_for_this_study_4(metric, period, exp, ens):\n",
    "    data_out = get_data_from_mat('/dodrio/scratch/projects/2022_200/project_output/cesm/yi_yao_IRRMIP/after_process/NorESM/post_process/3_hourly/' +\\\n",
    "                                             metric + '_' + period + '_3_hourly_' + exp + ens + '_4.mat', metric)\n",
    "    data_out = np.array(data_out)\n",
    "    return data_out\n",
    "\n",
    "def nparray_to_xarray_for_this_study(np_array):\n",
    "    data_xrray = xr.DataArray(np_array, coords={'time': data_time, 'y': data_lat, 'x': data_lon}, dims=[\"time\", \"y\", \"x\"])\n",
    "    return data_xrray\n",
    "    \n",
    "def calcu_PR_between_two_array(data_ref, data_new, perc):\n",
    "    data_xr_ref = xr.DataArray(data_ref, coords={'time': data_time, 'y': data_lat, 'x': data_lon}, dims=[\"time\", \"y\", \"x\"])\n",
    "    data_xr_new = xr.DataArray(data_new, coords={'time': data_time, 'y': data_lat, 'x': data_lon}, dims=[\"time\", \"y\", \"x\"])\n",
    "    PR_result = comp_PR(data_xr_ref,data_xr_new, perc)\n",
    "    return PR_result\n",
    "\n",
    "\n",
    "def calcu_PR_6_ensembles(metric, perc):\n",
    "    data_NOI01_1985_1994_1 = get_data_for_this_study_1(metric, '1985_1994', 'NOI', '01')\n",
    "    data_NOI01_1985_1994_2 = get_data_for_this_study_2(metric, '1985_1994', 'NOI', '01')\n",
    "    data_NOI01_1985_1994_3 = get_data_for_this_study_3(metric, '1985_1994', 'NOI', '01')\n",
    "    data_NOI01_1985_1994_4 = get_data_for_this_study_4(metric, '1985_1994', 'NOI', '01')\n",
    "    \n",
    "    data_NOI01_1995_2004_1 = get_data_for_this_study_1(metric, '1995_2004', 'NOI', '01')\n",
    "    data_NOI01_1995_2004_2 = get_data_for_this_study_2(metric, '1995_2004', 'NOI', '01')\n",
    "    data_NOI01_1995_2004_3 = get_data_for_this_study_3(metric, '1995_2004', 'NOI', '01')\n",
    "    data_NOI01_1995_2004_4 = get_data_for_this_study_4(metric, '1995_2004', 'NOI', '01')\n",
    "    \n",
    "    data_NOI01_2005_2014_1 = get_data_for_this_study_1(metric, '2005_2014', 'NOI', '01')\n",
    "    data_NOI01_2005_2014_2 = get_data_for_this_study_2(metric, '2005_2014', 'NOI', '01')\n",
    "    data_NOI01_2005_2014_3 = get_data_for_this_study_3(metric, '2005_2014', 'NOI', '01')\n",
    "    data_NOI01_2005_2014_4 = get_data_for_this_study_4(metric, '2005_2014', 'NOI', '01')\n",
    "\n",
    "    data_NOI02_1985_1994_1 = get_data_for_this_study_1(metric, '1985_1994', 'NOI', '02')\n",
    "    data_NOI02_1985_1994_2 = get_data_for_this_study_2(metric, '1985_1994', 'NOI', '02')\n",
    "    data_NOI02_1985_1994_3 = get_data_for_this_study_3(metric, '1985_1994', 'NOI', '02')\n",
    "    data_NOI02_1985_1994_4 = get_data_for_this_study_4(metric, '1985_1994', 'NOI', '02')\n",
    "    \n",
    "    data_NOI02_1995_2004_1 = get_data_for_this_study_1(metric, '1995_2004', 'NOI', '02')\n",
    "    data_NOI02_1995_2004_2 = get_data_for_this_study_2(metric, '1995_2004', 'NOI', '02')\n",
    "    data_NOI02_1995_2004_3 = get_data_for_this_study_3(metric, '1995_2004', 'NOI', '02')\n",
    "    data_NOI02_1995_2004_4 = get_data_for_this_study_4(metric, '1995_2004', 'NOI', '02')\n",
    "    \n",
    "    data_NOI02_2005_2014_1 = get_data_for_this_study_1(metric, '2005_2014', 'NOI', '02')\n",
    "    data_NOI02_2005_2014_2 = get_data_for_this_study_2(metric, '2005_2014', 'NOI', '02')\n",
    "    data_NOI02_2005_2014_3 = get_data_for_this_study_3(metric, '2005_2014', 'NOI', '02')\n",
    "    data_NOI02_2005_2014_4 = get_data_for_this_study_4(metric, '2005_2014', 'NOI', '02')\n",
    "\n",
    "    data_NOI03_1985_1994_1 = get_data_for_this_study_1(metric, '1985_1994', 'NOI', '03')\n",
    "    data_NOI03_1985_1994_2 = get_data_for_this_study_2(metric, '1985_1994', 'NOI', '03')\n",
    "    data_NOI03_1985_1994_3 = get_data_for_this_study_3(metric, '1985_1994', 'NOI', '03')\n",
    "    data_NOI03_1985_1994_4 = get_data_for_this_study_4(metric, '1985_1994', 'NOI', '03')\n",
    "    \n",
    "    data_NOI03_1995_2004_1 = get_data_for_this_study_1(metric, '1995_2004', 'NOI', '03')\n",
    "    data_NOI03_1995_2004_2 = get_data_for_this_study_2(metric, '1995_2004', 'NOI', '03')\n",
    "    data_NOI03_1995_2004_3 = get_data_for_this_study_3(metric, '1995_2004', 'NOI', '03')\n",
    "    data_NOI03_1995_2004_4 = get_data_for_this_study_4(metric, '1995_2004', 'NOI', '03')\n",
    "    \n",
    "    data_NOI03_2005_2014_1 = get_data_for_this_study_1(metric, '2005_2014', 'NOI', '03')\n",
    "    data_NOI03_2005_2014_2 = get_data_for_this_study_2(metric, '2005_2014', 'NOI', '03')\n",
    "    data_NOI03_2005_2014_3 = get_data_for_this_study_3(metric, '2005_2014', 'NOI', '03')\n",
    "    data_NOI03_2005_2014_4 = get_data_for_this_study_4(metric, '2005_2014', 'NOI', '03')\n",
    "\n",
    "    data_NOI_1985_2014 = np.concatenate((data_NOI01_1985_1994_1,\n",
    "                                         data_NOI01_1985_1994_2,\n",
    "                                         data_NOI01_1995_2004_1,\n",
    "                                         data_NOI01_1995_2004_2,\n",
    "                                         data_NOI01_2005_2014_1,\n",
    "                                         data_NOI01_2005_2014_2,\n",
    "                                         \n",
    "                                         data_NOI01_1985_1994_3,\n",
    "                                         data_NOI01_1985_1994_4,\n",
    "                                         data_NOI01_1995_2004_3,\n",
    "                                         data_NOI01_1995_2004_4,\n",
    "                                         data_NOI01_2005_2014_3,\n",
    "                                         data_NOI01_2005_2014_4,\n",
    "                                         \n",
    "                                         data_NOI02_1985_1994_1,\n",
    "                                         data_NOI02_1985_1994_2,\n",
    "                                         data_NOI02_1995_2004_1,\n",
    "                                         data_NOI02_1995_2004_2,\n",
    "                                         data_NOI02_2005_2014_1,\n",
    "                                         data_NOI02_2005_2014_2,\n",
    "                                         \n",
    "                                         data_NOI02_1985_1994_3,\n",
    "                                         data_NOI02_1985_1994_4,\n",
    "                                         data_NOI02_1995_2004_3,\n",
    "                                         data_NOI02_1995_2004_4,\n",
    "                                         data_NOI02_2005_2014_3,\n",
    "                                         data_NOI02_2005_2014_4,\n",
    "                                         \n",
    "                                         data_NOI03_1985_1994_1,\n",
    "                                         data_NOI03_1985_1994_2,\n",
    "                                         data_NOI03_1995_2004_1,\n",
    "                                         data_NOI03_1995_2004_2,\n",
    "                                         data_NOI03_2005_2014_1,\n",
    "                                         data_NOI03_2005_2014_2,\n",
    "                                        \n",
    "                                         data_NOI03_1985_1994_3,\n",
    "                                         data_NOI03_1985_1994_4,\n",
    "                                         data_NOI03_1995_2004_3,\n",
    "                                         data_NOI03_1995_2004_4,\n",
    "                                         data_NOI03_2005_2014_3,\n",
    "                                         data_NOI03_2005_2014_4,), axis = 0)\n",
    "    \n",
    "    del data_NOI01_1985_1994_1, data_NOI01_1985_1994_2, data_NOI02_1985_1994_1, data_NOI02_1985_1994_2, data_NOI03_1985_1994_1, data_NOI03_1985_1994_2\n",
    "    gc.collect()\n",
    "    del data_NOI01_1985_1994_3, data_NOI01_1985_1994_4, data_NOI02_1985_1994_4, data_NOI02_1985_1994_3, data_NOI03_1985_1994_3, data_NOI03_1985_1994_4\n",
    "    gc.collect()\n",
    "    \n",
    "    del data_NOI01_1995_2004_1, data_NOI01_1995_2004_2, data_NOI02_1995_2004_1, data_NOI02_1995_2004_2, data_NOI03_1995_2004_1, data_NOI03_1995_2004_2\n",
    "    gc.collect()\n",
    "    del data_NOI01_1995_2004_3, data_NOI01_1995_2004_4, data_NOI02_1995_2004_3, data_NOI02_1995_2004_4, data_NOI03_1995_2004_3, data_NOI03_1995_2004_4\n",
    "\n",
    "    \n",
    "    del data_NOI01_2005_2014_1, data_NOI01_2005_2014_2, data_NOI02_2005_2014_1, data_NOI02_2005_2014_2, data_NOI03_2005_2014_1, data_NOI03_2005_2014_2\n",
    "    gc.collect()\n",
    "    del data_NOI01_2005_2014_3, data_NOI01_2005_2014_4, data_NOI02_2005_2014_3, data_NOI02_2005_2014_4, data_NOI03_2005_2014_3, data_NOI03_2005_2014_4\n",
    "    gc.collect()\n",
    "    \n",
    "    data_time = np.array(range(data_NOI_1985_2014.shape[0]))\n",
    "    data_xr_ref = xr.DataArray(data_NOI_1985_2014, coords={'time': data_time, 'y': data_lat, 'x': data_lon}, dims=[\"time\", \"y\", \"x\"])\n",
    "\n",
    "    \n",
    "    del data_NOI_1985_2014\n",
    "    gc.collect()\n",
    "    \n",
    "    qctl=data_xr_ref.quantile(perc,dim='time')\n",
    "    del data_xr_ref\n",
    "    gc.collect()\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    data_IRR01_1985_1994_1 = get_data_for_this_study_1(metric, '1985_1994', 'IRR', '01')\n",
    "    data_IRR01_1985_1994_2 = get_data_for_this_study_2(metric, '1985_1994', 'IRR', '01')\n",
    "    data_IRR01_1985_1994_3 = get_data_for_this_study_3(metric, '1985_1994', 'IRR', '01')\n",
    "    data_IRR01_1985_1994_4 = get_data_for_this_study_4(metric, '1985_1994', 'IRR', '01')\n",
    "    \n",
    "    data_IRR01_1995_2004_1 = get_data_for_this_study_1(metric, '1995_2004', 'IRR', '01')\n",
    "    data_IRR01_1995_2004_2 = get_data_for_this_study_2(metric, '1995_2004', 'IRR', '01')\n",
    "    data_IRR01_1995_2004_3 = get_data_for_this_study_3(metric, '1995_2004', 'IRR', '01')\n",
    "    data_IRR01_1995_2004_4 = get_data_for_this_study_4(metric, '1995_2004', 'IRR', '01')\n",
    "    \n",
    "    data_IRR01_2005_2014_1 = get_data_for_this_study_1(metric, '2005_2014', 'IRR', '01')\n",
    "    data_IRR01_2005_2014_2 = get_data_for_this_study_2(metric, '2005_2014', 'IRR', '01')\n",
    "    data_IRR01_2005_2014_3 = get_data_for_this_study_3(metric, '2005_2014', 'IRR', '01')\n",
    "    data_IRR01_2005_2014_4 = get_data_for_this_study_4(metric, '2005_2014', 'IRR', '01')\n",
    "    \n",
    "    data_IRR02_1985_1994_1 = get_data_for_this_study_1(metric, '1985_1994', 'IRR', '02')\n",
    "    data_IRR02_1985_1994_2 = get_data_for_this_study_2(metric, '1985_1994', 'IRR', '02')\n",
    "    data_IRR02_1985_1994_3 = get_data_for_this_study_3(metric, '1985_1994', 'IRR', '02')\n",
    "    data_IRR02_1985_1994_4 = get_data_for_this_study_4(metric, '1985_1994', 'IRR', '02')\n",
    "    \n",
    "    data_IRR02_1995_2004_1 = get_data_for_this_study_1(metric, '1995_2004', 'IRR', '02')\n",
    "    data_IRR02_1995_2004_2 = get_data_for_this_study_2(metric, '1995_2004', 'IRR', '02')\n",
    "    data_IRR02_1995_2004_3 = get_data_for_this_study_3(metric, '1995_2004', 'IRR', '02')\n",
    "    data_IRR02_1995_2004_4 = get_data_for_this_study_4(metric, '1995_2004', 'IRR', '02')\n",
    "    \n",
    "    data_IRR02_2005_2014_1 = get_data_for_this_study_1(metric, '2005_2014', 'IRR', '02')\n",
    "    data_IRR02_2005_2014_2 = get_data_for_this_study_2(metric, '2005_2014', 'IRR', '02')\n",
    "    data_IRR02_2005_2014_3 = get_data_for_this_study_3(metric, '2005_2014', 'IRR', '02')\n",
    "    data_IRR02_2005_2014_4 = get_data_for_this_study_4(metric, '2005_2014', 'IRR', '02')\n",
    "    \n",
    "    data_IRR03_1985_1994_1 = get_data_for_this_study_1(metric, '1985_1994', 'IRR', '03')\n",
    "    data_IRR03_1985_1994_2 = get_data_for_this_study_2(metric, '1985_1994', 'IRR', '03')\n",
    "    data_IRR03_1985_1994_3 = get_data_for_this_study_3(metric, '1985_1994', 'IRR', '03')\n",
    "    data_IRR03_1985_1994_4 = get_data_for_this_study_4(metric, '1985_1994', 'IRR', '03')\n",
    "    \n",
    "    data_IRR03_1995_2004_1 = get_data_for_this_study_1(metric, '1995_2004', 'IRR', '03')\n",
    "    data_IRR03_1995_2004_2 = get_data_for_this_study_2(metric, '1995_2004', 'IRR', '03')\n",
    "    data_IRR03_1995_2004_3 = get_data_for_this_study_3(metric, '1995_2004', 'IRR', '03')\n",
    "    data_IRR03_1995_2004_4 = get_data_for_this_study_4(metric, '1995_2004', 'IRR', '03')\n",
    "    \n",
    "    data_IRR03_2005_2014_1 = get_data_for_this_study_1(metric, '2005_2014', 'IRR', '03')\n",
    "    data_IRR03_2005_2014_2 = get_data_for_this_study_2(metric, '2005_2014', 'IRR', '03')\n",
    "    data_IRR03_2005_2014_3 = get_data_for_this_study_3(metric, '2005_2014', 'IRR', '03')\n",
    "    data_IRR03_2005_2014_4 = get_data_for_this_study_4(metric, '2005_2014', 'IRR', '03')\n",
    "    \n",
    "\n",
    "    \n",
    "    \n",
    "    data_IRR_1985_2014 = np.concatenate((data_IRR01_1985_1994_1,\n",
    "                                         data_IRR01_1985_1994_2,\n",
    "                                         data_IRR01_1995_2004_1,\n",
    "                                         data_IRR01_1995_2004_2,\n",
    "                                         data_IRR01_2005_2014_1,\n",
    "                                         data_IRR01_2005_2014_2,\n",
    "                                         \n",
    "                                         data_IRR01_1985_1994_3,\n",
    "                                         data_IRR01_1985_1994_4,\n",
    "                                         data_IRR01_1995_2004_3,\n",
    "                                         data_IRR01_1995_2004_4,\n",
    "                                         data_IRR01_2005_2014_3,\n",
    "                                         data_IRR01_2005_2014_4,\n",
    "                                         \n",
    "                                         data_IRR02_1985_1994_1,\n",
    "                                         data_IRR02_1985_1994_2,\n",
    "                                         data_IRR02_1995_2004_1,\n",
    "                                         data_IRR02_1995_2004_2,\n",
    "                                         data_IRR02_2005_2014_1,\n",
    "                                         data_IRR02_2005_2014_2,\n",
    "                                         \n",
    "                                         data_IRR02_1985_1994_3,\n",
    "                                         data_IRR02_1985_1994_4,\n",
    "                                         data_IRR02_1995_2004_3,\n",
    "                                         data_IRR02_1995_2004_4,\n",
    "                                         data_IRR02_2005_2014_3,\n",
    "                                         data_IRR02_2005_2014_4,\n",
    "                                         \n",
    "                                         data_IRR03_1985_1994_1,\n",
    "                                         data_IRR03_1985_1994_2,\n",
    "                                         data_IRR03_1995_2004_1,\n",
    "                                         data_IRR03_1995_2004_2,\n",
    "                                         data_IRR03_2005_2014_1,\n",
    "                                         data_IRR03_2005_2014_2,\n",
    "                                        \n",
    "                                         data_IRR03_1985_1994_3,\n",
    "                                         data_IRR03_1985_1994_4,\n",
    "                                         data_IRR03_1995_2004_3,\n",
    "                                         data_IRR03_1995_2004_4,\n",
    "                                         data_IRR03_2005_2014_3,\n",
    "                                         data_IRR03_2005_2014_4,), axis = 0)\n",
    "    \n",
    "    del data_IRR01_1985_1994_1, data_IRR01_1985_1994_2, data_IRR02_1985_1994_1, data_IRR02_1985_1994_2, data_IRR03_1985_1994_1, data_IRR03_1985_1994_2\n",
    "    gc.collect()\n",
    "    del data_IRR01_1985_1994_3, data_IRR01_1985_1994_4, data_IRR02_1985_1994_3, data_IRR02_1985_1994_4, data_IRR03_1985_1994_3, data_IRR03_1985_1994_4\n",
    "    gc.collect()\n",
    "    \n",
    "    del data_IRR01_1995_2004_1, data_IRR01_1995_2004_2, data_IRR02_1995_2004_1, data_IRR02_1995_2004_2, data_IRR03_1995_2004_1, data_IRR03_1995_2004_2\n",
    "    gc.collect()\n",
    "    del data_IRR01_1995_2004_3, data_IRR01_1995_2004_4, data_IRR02_1995_2004_3, data_IRR02_1995_2004_4, data_IRR03_1995_2004_3, data_IRR03_1995_2004_4\n",
    "    gc.collect()\n",
    "\n",
    "    \n",
    "    del data_IRR01_2005_2014_1, data_IRR01_2005_2014_2, data_IRR02_2005_2014_1, data_IRR02_2005_2014_2, data_IRR03_2005_2014_1, data_IRR03_2005_2014_2\n",
    "    gc.collect()\n",
    "    del data_IRR01_2005_2014_3, data_IRR01_2005_2014_4, data_IRR02_2005_2014_3, data_IRR02_2005_2014_4, data_IRR03_2005_2014_3, data_IRR03_2005_2014_4\n",
    "    gc.collect()\n",
    "    \n",
    "    data_time = np.array(range(data_IRR_1985_2014.shape[0]))\n",
    "    data_xr_new = xr.DataArray(data_IRR_1985_2014, coords={'time': data_time, 'y': data_lat, 'x': data_lon}, dims=[\"time\", \"y\", \"x\"])\n",
    "    \n",
    "    \n",
    "    del data_IRR_1985_2014\n",
    "    gc.collect()\n",
    "    \n",
    "    \n",
    "    #get subsample scenarios\n",
    "    qscen=xr.where(data_xr_new>qctl,1,0)\n",
    "    \n",
    "    del data_xr_new\n",
    "    gc.collect()\n",
    "    \n",
    "    pscen=qscen.sum('time')/len(qscen.time)\n",
    "    \n",
    "    del qscen\n",
    "    gc.collect()\n",
    "\n",
    "    #compute PR\n",
    "    PR_scen=pscen/(1-perc)\n",
    "    PR_scen = np.where(data_mask==0, np.nan, PR_scen)\n",
    "    PR_scen = np.where(PR_scen==0, np.nan, PR_scen)\n",
    "    \n",
    "    PR_3 = PR_scen\n",
    "    \n",
    "    del PR_scen\n",
    "    gc.collect()\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "        \n",
    "    \n",
    "    return PR_3\n",
    "\n",
    "def plot_PR_single(ax, data_xarray, title, metric, perc):\n",
    "    cmap='bwr'\n",
    "    vmin=1/4\n",
    "    vmax=4\n",
    "    alpha=0.05\n",
    "    levels = [1/4, 1/3, 1/2, 2/3, 1.5, 2, 3, 4]\n",
    "    im = data_xarray.plot(ax=ax, cmap=cmap,vmin=vmin,vmax=vmax,levels=levels, extend='both', add_colorbar=False, add_labels=False)\n",
    "    cb = plt.colorbar(im,fraction= 0.3, pad= 0.04, extend='both', orientation='horizontal')\n",
    "    cb.ax.tick_params(labelsize=12)\n",
    "    cb.set_ticks([0.25, 0.33333333, 0.5, 0.66666667, 1.5, 2, 3, 4])\n",
    "    cb.set_ticklabels([\"1/4\", \"1/3\", \"1/2\", \"2/3\", \"3/2\", \"2\", \"3\", \"4\"])\n",
    "    ax.coastlines(color='dimgray', linewidth=0.5)\n",
    "    ax.add_feature(cfeature.OCEAN, color='lightgrey')\n",
    "    ax.set_title(title, loc='right', fontsize=14)\n",
    "    ax.set_title(metric + '_' + perc, loc='left', fontsize=14)\n",
    "    ax.add_feature(cfeature.BORDERS, linestyle='-', linewidth=0.1)\n",
    "    \n",
    "def plot_PR_single_no_cb(ax, data_xarray, title, metric, perc):\n",
    "    cmap='bwr'\n",
    "    vmin=1/4\n",
    "    vmax=4\n",
    "    alpha=0.05\n",
    "    levels = [1/4, 1/3, 1/2, 2/3, 1, 1.5, 2, 3, 4]\n",
    "    im = data_xarray.plot(ax=ax, cmap=cmap,vmin=vmin,vmax=vmax,levels=levels, extend='both', add_colorbar=False, add_labels=False)\n",
    "    #cb = plt.colorbar(im,fraction= 0.02, pad= 0.04, extend='both')\n",
    "    #cb.set_ticks([0.25, 0.33333344, 0.5, 0.66666667, 1, 1.5, 2, 3, 4])\n",
    "    #cb.set_ticklabels([\"1/4\", \"1/3\", \"1/2\", \"2/3\", \"1\", \"3/2\", \"2\", \"3\", \"4\"])\n",
    "    ax.coastlines(color='dimgray', linewidth=0.5)\n",
    "    ax.add_feature(cfeature.OCEAN, color='lightgrey')\n",
    "    ax.set_title(title, loc='right', fontsize=10)\n",
    "    ax.set_title(metric + '_' + perc, loc='left', fontsize=10)\n",
    "    ax.add_feature(cfeature.BORDERS, linestyle='-', linewidth=0.1)\n",
    "    \n",
    "def plot_all(PR_1, PR_2, PR_3, metric, perc):\n",
    "    fig = plt.figure(figsize=(15,6),dpi=300)\n",
    "    proj=ccrs.PlateCarree()\n",
    "    data_xr_1 = xr.DataArray(PR_1, coords={'y': data_lat, 'x': data_lon}, dims=[\"y\", \"x\"])\n",
    "    ax = plt.subplot(311, projection=proj, frameon=False)\n",
    "    plot_PR_single(ax, data_xr_1, 'all forcings', metric, perc)\n",
    "    \n",
    "\n",
    "    data_xr_2 = xr.DataArray(PR_2, coords={'y': data_lat, 'x': data_lon}, dims=[\"y\", \"x\"])\n",
    "    ax = plt.subplot(312, projection=proj, frameon=False)\n",
    "    plot_PR_single(ax, data_xr_2, 'other forcings', metric, perc)\n",
    "\n",
    "    data_xr_3 = xr.DataArray(PR_3, coords={'y': data_lat, 'x': data_lon}, dims=[\"y\", \"x\"])\n",
    "    ax = plt.subplot(313, projection=proj, frameon=False)\n",
    "    plot_PR_single(ax, data_xr_3, 'irrigation expansion', metric, perc)\n",
    "\n",
    "\n",
    "    \n",
    "def just_do_all(metric, perc, perc_str):\n",
    "    PR_1, PR_2, PR_3 = calcu_PR_6_ensembles(metric, perc)\n",
    "    plot_all(PR_1, PR_2, PR_3, metric, perc_str)\n",
    "    \n",
    "def just_do_all_even_more(metric, varname):\n",
    "    PR_3 = calcu_PR_6_ensembles(metric, 0.99)\n",
    "    \n",
    "    scio.savemat(\"/dodrio/scratch/projects/2022_200/project_output/cesm/yi_yao_IRRMIP/after_process/NorESM/PR/PR_irr_\" + \\\n",
    "                     metric + \"_99.0%.mat\", {'PR_3': PR_3})\n",
    "    \n",
    "    PR_6 = calcu_PR_6_ensembles(metric, 0.995)\n",
    "    \n",
    "    scio.savemat(\"/dodrio/scratch/projects/2022_200/project_output/cesm/yi_yao_IRRMIP/after_process/NorESM/PR/PR_irr_\" + \\\n",
    "                     metric + \"_99.5%.mat\", {'PR_6': PR_6})\n",
    "    \n",
    "    PR_9 = calcu_PR_6_ensembles(metric, 0.999)\n",
    "    \n",
    "    scio.savemat(\"/dodrio/scratch/projects/2022_200/project_output/cesm/yi_yao_IRRMIP/after_process/NorESM/PR/PR_irr_\" + \\\n",
    "                     metric + \"_99.9%.mat\", {'PR_9': PR_9})\n",
    "    \n",
    "    fig = plt.figure(figsize=(18,12),dpi=300)\n",
    "    proj=ccrs.PlateCarree()\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    data_xr_3 = xr.DataArray(PR_3, coords={'y': data_lat, 'x': data_lon}, dims=[\"y\", \"x\"])\n",
    "    ax = plt.subplot(337, projection=proj, frameon=False)\n",
    "    plot_PR_single(ax, data_xr_3, 'irrigation expansion', varname, '99.0%')\n",
    "    ax.text(0.01, 0.92, 'g', color='dimgrey', fontsize=10, transform=ax.transAxes, weight='bold')\n",
    "    \n",
    "    \n",
    "    \n",
    "\n",
    "    data_xr_6 = xr.DataArray(PR_6, coords={'y': data_lat, 'x': data_lon}, dims=[\"y\", \"x\"])\n",
    "    ax = plt.subplot(338, projection=proj, frameon=False)\n",
    "    plot_PR_single(ax, data_xr_6, 'irrigation expansion', varname, '99.5%')\n",
    "    ax.text(0.01, 0.92, 'h', color='dimgrey', fontsize=10, transform=ax.transAxes, weight='bold')\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "\n",
    "    data_xr_9 = xr.DataArray(PR_9, coords={'y': data_lat, 'x': data_lon}, dims=[\"y\", \"x\"])\n",
    "    ax = plt.subplot(339, projection=proj, frameon=False)\n",
    "    plot_PR_single(ax, data_xr_9, 'irrigation expansion', varname, '99.9%')\n",
    "    ax.text(0.01, 0.92, 'i', color='dimgrey', fontsize=10, transform=ax.transAxes, weight='bold')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "249e6f7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "just_do_all_even_more('TSA', 'NorESM T_2m')\n",
    "just_do_all_even_more('AT', 'NorESM AT')\n",
    "just_do_all_even_more('DI', 'NorESM DI')\n",
    "just_do_all_even_more('HUMIDEX', 'NorESM HU')\n",
    "just_do_all_even_more('THIC', 'NorESM THIC')\n",
    "just_do_all_even_more('THIP', 'NorESM THIP')\n",
    "just_do_all_even_more('T_s_65', 'NorESM T_s_65')\n",
    "just_do_all_even_more('T_s_80', 'NorESM T_s_80')\n",
    "just_do_all_even_more('T_w', 'NorESM T_w')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}