bptest_1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "abdba210",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<p>To authorize access needed by Earth Engine, open the following\n",
       "        URL in a web browser and follow the instructions:</p>\n",
       "        <p><a href=https://code.earthengine.google.com/client-auth?scopes=https%3A//www.googleapis.com/auth/earthengine%20https%3A//www.googleapis.com/auth/devstorage.full_control&request_id=B5DdHdFFDDPUvfW3yIkBJg2eZn0b-9Q3CBZAr4rIljQ&tc=aFi3rxApM1WRCXPf0JAlhodaEl6C1ZPkf436Px8YJ1Q&cc=YiAOwXWDzuxVPSSFv8gpR64c4EGS8jrRGxMenz2Yap4>https://code.earthengine.google.com/client-auth?scopes=https%3A//www.googleapis.com/auth/earthengine%20https%3A//www.googleapis.com/auth/devstorage.full_control&request_id=B5DdHdFFDDPUvfW3yIkBJg2eZn0b-9Q3CBZAr4rIljQ&tc=aFi3rxApM1WRCXPf0JAlhodaEl6C1ZPkf436Px8YJ1Q&cc=YiAOwXWDzuxVPSSFv8gpR64c4EGS8jrRGxMenz2Yap4</a></p>\n",
       "        <p>The authorization workflow will generate a code, which you\n",
       "        should paste in the box below</p>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Enter verification code: 4/1Adeu5BWfWmJRR8ZXwwzjPonfJRMcQ9iRlYJTTKNQIHGLFrg1K_laKbAhONE\n",
      "\n",
      "Successfully saved authorization token.\n"
     ]
    }
   ],
   "source": [
    "import ee\n",
    "%matplotlib inline\n",
    "import math\n",
    "import warnings\n",
    "import intake\n",
    "warnings.filterwarnings('ignore')\n",
    "from collections import defaultdict\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "from scipy import stats\n",
    "\n",
    "import xarray as xr\n",
    "xr.set_options(display_style='html')\n",
    "\n",
    "import datetime, calendar\n",
    "\n",
    "plt.style.use(\"seaborn-darkgrid\")\n",
    "\n",
    "#ee.Authenticate()\n",
    "ee.Initialize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "28f44297",
   "metadata": {},
   "outputs": [],
   "source": [
    "HIST_START = 1980\n",
    "HIST_END = 2014\n",
    "FUTURE_START = 2080\n",
    "FUTURE_END = 2099\n",
    "\n",
    "PERCENTILE_STARTYEAR = 1980\n",
    "PERCENTILE_ENDYEAR = 2019\n",
    "\n",
    "NUM_BEST_MODELS = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "9951f379",
   "metadata": {},
   "outputs": [],
   "source": [
    "FUTURE_YEARS = [(2080, 2099)]\n",
    "FUTURE_SCENARIOS = ('ssp119', 'ssp126')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5d04f9a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODELS = {\n",
    "    'tasmax': ('CanESM5', 'MRI-ESM2-0', 'CAMS-CSM1-0'),\n",
    "    'tasmin': ('CanESM5', 'MRI-ESM2-0'),\n",
    "    'pr': ('CanESM5', 'MRI-ESM2-0', 'CAMS-CSM1-0' ),\n",
    "    'hurs': ('CanESM5', 'MRI-ESM2-0', 'UKESM1-0-LL')\n",
    "}\n",
    "RUNS = {\n",
    "    'CanESM5': 'r1i1p1f1',\n",
    "    'MRI-ESM2-0': 'r1i1p1f1',\n",
    "    'CAMS-CSM1-0': 'r2i1p1f1',\n",
    "    #'UKESM1-0-LL': 'r1i1p1f2'  # Uses 360 day\n",
    "}\n",
    "\n",
    "YEARLENGTH = {\n",
    "    'ERA5': 366,\n",
    "    'CanESM5': 365,\n",
    "    'MRI-ESM2-0': 366,\n",
    "    'CAMS-CSM1-0': 366\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "137f5d97",
   "metadata": {},
   "outputs": [],
   "source": [
    "cat_url = \"https://storage.googleapis.com/cmip6/pangeo-cmip6.json\"\n",
    "col = intake.open_esm_datastore(cat_url)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "66e61532",
   "metadata": {},
   "outputs": [],
   "source": [
    "VARIABLES = {\n",
    "    'tasmax': {\n",
    "        'era_varname': 'maximum_2m_air_temperature',\n",
    "        'nex_transform': lambda x: x - 273.5,\n",
    "        'era_transform': lambda x: x - 273.5\n",
    "    },\n",
    "    'tasmin': {\n",
    "        'era_varname': 'minimum_2m_air_temperature',\n",
    "        'nex_transform': lambda x: x - 273.5,\n",
    "        'era_transform': lambda x: x - 273.5\n",
    "    },\n",
    "    'pr': {\n",
    "        'era_varname': 'total_precipitation',\n",
    "        'nex_transform': lambda x: x * 86400,\n",
    "        'era_transform': lambda x: x * 1000\n",
    "    }   \n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b9af88e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calendardate_percentiles(nex_varname, q, latlon, sh_hem=False):\n",
    "    era_varname = VARIABLES[nex_varname]['era_varname']\n",
    "    hist_start = PERCENTILE_STARTYEAR\n",
    "    hist_end = PERCENTILE_ENDYEAR\n",
    "    allyears = []\n",
    "    for year in range(hist_start, hist_end):\n",
    "        allyears.append(VARIABLES[nex_varname]['era_transform'](get_var(era_varname, 'ERA5', latlon, start_year=year, end_year=year, southern_hem=False)))\n",
    "    if not sh_hem:\n",
    "        return np.percentile(np.vstack(allyears), q, axis=0)\n",
    "    else:\n",
    "        res = np.percentile(np.vstack(allyears), q, axis=0)\n",
    "        return np.concatenate([res[152:], res[:152]])\n",
    "\n",
    "def wholeyear_percentile(nex_varname, q, latlon):\n",
    "    era_varname = VARIABLES[nex_varname]['era_varname']\n",
    "    hist_start = PERCENTILE_STARTYEAR\n",
    "    hist_end = PERCENTILE_ENDYEAR\n",
    "    allyears = []\n",
    "    for year in range(hist_start, hist_end):\n",
    "        allyears.append(VARIABLES[nex_varname]['era_transform'](get_var(era_varname, 'ERA5', latlon, start_year=year, end_year=year, southern_hem=False)))\n",
    "    return np.percentile(np.concatenate(allyears).flatten(), q)\n",
    "\n",
    "def yearextreme_percentile(nex_varname, q, latlon, wantmax):\n",
    "    era_varname = VARIABLES[nex_varname]['era_varname']\n",
    "    hist_start = PERCENTILE_STARTYEAR\n",
    "    hist_end = PERCENTILE_ENDYEAR\n",
    "    allyears = []\n",
    "    for year in range(hist_start, hist_end):\n",
    "        allyears.append([np.min, np.max][int(wantmax)](VARIABLES[nex_varname]['era_transform'](get_var(era_varname, 'ERA5', latlon, start_year=year, end_year=year, southern_hem=False))))\n",
    "    return np.percentile(np.array(allyears), q)\n",
    "\n",
    "def d2j(datestring):\n",
    "    d = datetime.date.fromisoformat(datestring)\n",
    "    jday = d.timetuple().tm_yday\n",
    "    if calendar.isleap(d.year) and jday > 59:\n",
    "        jday -= 1\n",
    "    return jday\n",
    "\n",
    "def get_rmsd(d1, d2):\n",
    "    c1 = seasonal_means(d1)\n",
    "    c2 = seasonal_means(d2)\n",
    "    return np.sqrt(np.mean(np.sum((c1 - c2)**2)))\n",
    "\n",
    "def count_runs(tf_array, min_runsize):\n",
    "    falses = np.zeros(tf_array.shape[0]).reshape((tf_array.shape[0],1))\n",
    "    extended_a = np.concatenate([[0], tf_array, [0]])\n",
    "    df = np.diff(extended_a)\n",
    "    starts = np.nonzero(df == 1)[0]\n",
    "    ends = np.nonzero(df == -1)[0]\n",
    "    count = 0\n",
    "    for idx in range(starts.size):\n",
    "        if ends[idx] - starts[idx] >= min_runsize:\n",
    "            count += 1\n",
    "    return count\n",
    "\n",
    "\n",
    "\n",
    "def get_var(varname, model, latlon, start_year, end_year, southern_hem=False, extralong=False, scenario='ssp585'):\n",
    "    def removeLeapDays(arr, extralong_=False, southern_hem_=False):\n",
    "        if extralong_:\n",
    "            indices = list(range(184))\n",
    "            jan1_idx = 184\n",
    "            \n",
    "            for year in range(start_year, end_year+1):\n",
    "                indices += [jan1_idx + i for i in range(365)]\n",
    "                jan1_idx += 365\n",
    "                if calendar.isleap(year):\n",
    "                    jan1_idx += 1\n",
    "            return arr[indices]\n",
    "        elif not southern_hem_:\n",
    "            indices = []\n",
    "            jan1_idx = 0\n",
    "            for year in range(start_year, end_year+1):\n",
    "                indices += [jan1_idx + i for i in range(365)]\n",
    "                jan1_idx += 365\n",
    "                if calendar.isleap(year):\n",
    "                    jan1_idx += 1\n",
    "            return arr[indices]\n",
    "        else:\n",
    "            indices = []\n",
    "            jul1_idx = 0\n",
    "            for year in range(start_year-1, end_year):\n",
    "                indices += [jul1_idx + i for i in range(365)]\n",
    "                jul1_idx += 365\n",
    "                if calendar.isleap(year):\n",
    "                    jul1_idx += 1\n",
    "            return arr[indices]\n",
    "    \n",
    "    if model != 'ERA5' and start_year < 2015 and end_year >= 2015:\n",
    "        raise Exception(\"Requesting hist and non-hist variables in one query\")\n",
    "    if model == 'ERA5':\n",
    "        dataset = ee.ImageCollection(\"ECMWF/ERA5/DAILY\")\n",
    "        gee_geom = ee.Geometry.Point((latlon[1], latlon[0]))\n",
    "        if extralong:\n",
    "            data_vars = dataset.select(varname).filter(ee.Filter.date('{0}-07-01'.format(start_year-1), '{0}-01-01'.format(end_year+1)))\n",
    "            result = [i[4] for i in data_vars.getRegion(gee_geom, 2500, 'epsg:4326').getInfo()[1:]]\n",
    "            return removeLeapDays(np.array(result), extralong_=True, southern_hem_=False)\n",
    "        elif not southern_hem:\n",
    "            data_vars = dataset.select(varname).filter(ee.Filter.date('{0}-01-01'.format(start_year), '{0}-01-01'.format(end_year+1)))\n",
    "            result = [i[4] for i in data_vars.getRegion(gee_geom, 2500, 'epsg:4326').getInfo()[1:]]\n",
    "            return removeLeapDays(np.array(result), extralong_=False, southern_hem_=False)\n",
    "        else:\n",
    "            data_vars = dataset.select(varname).filter(ee.Filter.date('{0}-07-01'.format(start_year-1), '{0}-07-01'.format(end_year)))\n",
    "            result = [i[4] for i in data_vars.getRegion(gee_geom, 2500, 'epsg:4326').getInfo()[1:]]\n",
    "            return removeLeapDays(np.array(result), extralong_=False, southern_hem_=True)\n",
    "    else:  # Retrieve NetCDF from pangeo archive\n",
    "        \n",
    "    # NOTE: IT WOULD BE BETTER TO GET DATA FROM AWS, LIKE THIS:\n",
    "        \n",
    "    #    def s3open(path):\n",
    "    #        fs = s3fs.S3FileSystem(anon=True, default_fill_cache=False, \n",
    "    #                               config_kwargs = {'max_pool_connections': 50})\n",
    "    #        return s3fs.S3Map(path, s3=fs)\n",
    "    #    testfile = s3open('s3://cmip6-pds/CMIP6/ScenarioMIP/NOAA-GFDL/GFDL-ESM4/ssp119/r1i1p1f1/day/tasmax/gr1/v20180701/')\n",
    "    #    # (Get path from data catalog csv)\n",
    "    #    ds = xr.open_mfdataset([testfile], engine='zarr', parallel=True)\n",
    "    #    a = ds['tasmax'].sel(time=slice('2030-01-01', '2039-12-31')).sel(lat=lat, lon=lon, method='nearest').to_numpy()\n",
    "        \n",
    "    # BUT THIS LAST STEP TAKES A VERY LONG TIME\n",
    "        \n",
    "        scenario = [scenario, 'historical'][int(start_year < 2015)]\n",
    "        cat = col.search(experiment_id=scenario, table_id='day', variable_id=varname, source_id=model, member_id=RUNS[model], grid_label='gn')\n",
    "        ds_dict = cat.to_dataset_dict(zarr_kwargs={'consolidated': True})\n",
    "        if extralong:\n",
    "            ds = list(ds_dict.values())[0][varname][0][0].sel(time=slice('{0}-07-01'.format(start_year-1), '{0}-12-31'.format(end_year))).sel(lat=latlon[0], lon=latlon[1], method='nearest').to_numpy()\n",
    "            if YEARLENGTH[model] != 365:\n",
    "                return removeLeapDays(np.array(ds), extralong_=True)\n",
    "            else:\n",
    "                return np.array(ds)\n",
    "        elif not southern_hem:\n",
    "            ds = list(ds_dict.values())[0][varname][0][0].sel(time=slice('{0}-01-01'.format(start_year), '{0}-12-31'.format(end_year))).sel(lat=latlon[0], lon=latlon[1], method='nearest').to_numpy()\n",
    "            if YEARLENGTH[model] != 365:\n",
    "                return removeLeapDays(np.array(ds), extralong_=False, southern_hem_=False)\n",
    "            else:\n",
    "                return np.array(ds)\n",
    "        else:\n",
    "            ds = list(ds_dict.values())[0][varname][0][0].sel(time=slice('{0}-07-01'.format(start_year-1), '{0}-06-30'.format(end_year))).sel(lat=latlon[0], lon=latlon[1], method='nearest').to_numpy()\n",
    "            if YEARLENGTH[model] != 365:\n",
    "                return removeLeapDays(np.array(ds), southern_hem_=True)\n",
    "            else:\n",
    "                return np.array(ds)\n",
    "\n",
    "    \n",
    "def quarters(d, start_year, end_year, southern_hem=False):\n",
    "    q2 = []  # 60-151\n",
    "    q3 = []  # 152-243\n",
    "    q4 = []  # 244-334\n",
    "    q1 = []  # 335-59\n",
    "    if not southern_hem:\n",
    "        jan1_idx = 365\n",
    "        for year in range(start_year, end_year):\n",
    "            tmp = np.concatenate((d[jan1_idx - 365 : jan1_idx - 365 + 60], d[jan1_idx + 335 : jan1_idx + 365]), axis=0)\n",
    "            q1.append(tmp)\n",
    "            q2.append(d[jan1_idx + 60 : jan1_idx + 152])\n",
    "            q3.append(d[jan1_idx + 152 : jan1_idx + 244])\n",
    "            q4.append(d[jan1_idx + 244 : jan1_idx + 335])\n",
    "\n",
    "            jan1_idx += 365 + [0, 0][int(False and calendar.isleap(year))]\n",
    "        mam_res = np.vstack(q2)\n",
    "        jja_res = np.vstack(q3)\n",
    "        son_res = np.vstack(q4)\n",
    "        djf_res = np.vstack(q1)\n",
    "    else:\n",
    "        jul1_idx = 365\n",
    "        for year in range(start_year, end_year):\n",
    "            tmp = np.concatenate((d[jul1_idx - 365 : jul1_idx - 365 + 60], d[jul1_idx + 335 : jul1_idx + 365]), axis=0)\n",
    "            q3.append(tmp)\n",
    "            q4.append(d[jul1_idx + 60 : jul1_idx + 152])\n",
    "            q1.append(d[jul1_idx + 152 : jul1_idx + 244])\n",
    "            q2.append(d[jul1_idx + 244 : jul1_idx + 335])\n",
    "\n",
    "            jul1_idx += 365 + [0, 0][int(False and calendar.isleap(year))]\n",
    "        mam_res = np.vstack(q4)\n",
    "        jja_res = np.vstack(q1)\n",
    "        son_res = np.vstack(q2)\n",
    "        djf_res = np.vstack(q3)\n",
    "    return mam_res, jja_res, son_res, djf_res\n",
    "    \n",
    "def seasonal_means(d):\n",
    "    q = quarters(d, HIST_START, HIST_END)\n",
    "    return np.array([np.mean(q[0], axis=1), np.mean(q[1], axis=1), np.mean(q[2], axis=1), np.mean(q[3], axis=1)])\n",
    "\n",
    "def calibration_function(hist_obs, hist_mod):\n",
    "# Calibration functions are P-P plots of historical and modeled values\n",
    "\n",
    "    source = np.sort(hist_obs.flatten())\n",
    "    target= np.sort(hist_mod.flatten())\n",
    "   \n",
    "    if (np.max(source) == 0 and np.min(source) == 0):\n",
    "        return np.arange(0, target.size) / target.size\n",
    "    if (np.max(target) == 0 and np.min(target) == 0):\n",
    "        return np.arange(0, source.size) / source.size\n",
    "    new_indices = []\n",
    "\n",
    "    for target_idx, target_value in enumerate(target):\n",
    "        if target_idx < len(source):\n",
    "            source_value = source[target_idx]\n",
    "            if source_value > target[-1]:\n",
    "                new_indices.append(target.size - 1)\n",
    "            else:\n",
    "                new_indices.append(np.argmax(target >= source_value))\n",
    "    return np.array(new_indices) / source.size\n",
    "\n",
    "def calibrate_component(uncalibrated_data, calibration_fxn):\n",
    "    N = len(uncalibrated_data)\n",
    "    unsorted_uncalib = [(i, idx) for idx, i in enumerate(uncalibrated_data)]\n",
    "    sorted_uncalib = sorted(unsorted_uncalib)\n",
    "    result = [0] * N\n",
    "    for j in range(N):\n",
    "        X_j = j / (N + 1)\n",
    "        Y_jprime = calibration_fxn[math.floor(X_j * len(calibration_fxn))]\n",
    "        jprime = math.floor(Y_jprime * (N + 1))\n",
    "        result[sorted_uncalib[j][1]] = sorted_uncalib[min(len(sorted_uncalib)-1, jprime)][0]\n",
    "    \n",
    "    return result\n",
    "\n",
    "def calibrate(uncalibrated_data, calibration_fxn):\n",
    "    mam = []\n",
    "    jja = []\n",
    "    son = []\n",
    "    djf = []\n",
    "    mam_idx = []\n",
    "    jja_idx = []\n",
    "    son_idx = []\n",
    "    djf_idx = []\n",
    "    for idx, i in enumerate(uncalibrated_data):\n",
    "        if idx % 365 >= 60 and idx % 365 < 152:\n",
    "            mam.append(uncalibrated_data[idx])\n",
    "            mam_idx.append(idx)\n",
    "        elif idx % 365 >= 152 and idx % 365 < 244:\n",
    "            jja.append(uncalibrated_data[idx])\n",
    "            jja_idx.append(idx)\n",
    "        elif idx % 365 >= 244 and idx % 365 < 335:\n",
    "            son.append(uncalibrated_data[idx])\n",
    "            son_idx.append(idx)\n",
    "        else:\n",
    "            djf.append(uncalibrated_data[idx])\n",
    "            djf_idx.append(idx)\n",
    "    \n",
    "    mam_calib = calibrate_component(np.array(mam), calibration_fxn[0])\n",
    "    jja_calib = calibrate_component(np.array(jja), calibration_fxn[1])\n",
    "    son_calib = calibrate_component(np.array(son), calibration_fxn[2])\n",
    "    djf_calib = calibrate_component(np.array(djf), calibration_fxn[3])\n",
    "    \n",
    "    result = [0] * len(uncalibrated_data)\n",
    "    for i in range(len(mam_idx)):\n",
    "        result[mam_idx[i]] = mam_calib[i]\n",
    "    for i in range(len(jja_idx)):\n",
    "        result[jja_idx[i]] = jja_calib[i]\n",
    "    for i in range(len(son_idx)):\n",
    "        result[son_idx[i]] = son_calib[i]\n",
    "    for i in range(len(djf_idx)):\n",
    "        result[djf_idx[i]] = djf_calib[i]\n",
    "\n",
    "    return np.array(result)\n",
    "\n",
    "def get_gamma(count, size):\n",
    "    return np.random.gamma(shape = count + 0.5, size=size)\n",
    "def get_beta(count, num, size):\n",
    "    return np.random.beta(a = count + 0.5, b = num - count + 0.5, size=size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "05aa86a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Hazard:\n",
    "    pass\n",
    "\n",
    "class Tempwave_simple(Hazard):\n",
    "    def __init__(self, varname, min_duration, threshold, want_gte=True):\n",
    "        if type(threshold) == np.ndarray and threshold.size % 365 != 0:\n",
    "            raise Exception('Comparison array length is not an integer multiple of 365')\n",
    "        self.varname = varname\n",
    "        self.want_gte = want_gte\n",
    "        self.min_duration = min_duration\n",
    "        self.threshold = threshold  # May be scalar or 365-long array\n",
    "        self.probmodel = 'Poisson'\n",
    "    def count(self, datalist):\n",
    "        data = datalist[0]\n",
    "        if type(self.threshold) in (float, int, np.float64, np.int32):\n",
    "            threshold = self.threshold\n",
    "        else:   # type is np array\n",
    "            threshold = np.array([])\n",
    "            while threshold.size < data.size:\n",
    "                threshold = np.concatenate([threshold, self.threshold])\n",
    "        if self.want_gte:\n",
    "            tf_array = data >= threshold\n",
    "        else:\n",
    "            tf_array = data <= threshold\n",
    "        return count_runs(tf_array, self.min_duration)\n",
    "    \n",
    "class Heatwave_highlow(Hazard):\n",
    "    def __init__(self, hightemp, lowtemp, min_duration):\n",
    "        self.varname = 'tasmax+tasmin'\n",
    "        self.min_duration = min_duration\n",
    "        self.hightemp = hightemp\n",
    "        self.lowtemp = lowtemp\n",
    "        self.probmodel = 'Poisson'\n",
    "    def count(self, datalist):\n",
    "        data_tx = datalist[0]\n",
    "        data_tn = datalist[1]\n",
    "        if type(self.hightemp) in (float, int, np.float64, np.int32):\n",
    "            high_threshold = self.hightemp\n",
    "        else:   # type is np array\n",
    "            high_threshold = np.array([])\n",
    "            while high_threshold.size < data_tx.size:\n",
    "                high_threshold = np.concatenate([high_threshold, self.hightemp])\n",
    "        if type(self.lowtemp) in (float, int, np.float64, np.int32):\n",
    "            low_threshold = self.lowtemp\n",
    "        else:   # type is np array\n",
    "            low_threshold = np.array([])\n",
    "            while low_threshold.size < data_tn.size:\n",
    "                low_threshold = np.concatenate([low_threshold, self.lowtemp])\n",
    "        tf_array_tx = data_tx >= high_threshold\n",
    "        tf_array_tn = data_tn >= low_threshold\n",
    "        return count_runs(tf_array_tx * tf_array_tn, self.min_duration)\n",
    "\n",
    "class Threshold_simple(Hazard):\n",
    "    def __init__(self, varname, threshold, want_gte):\n",
    "        self.varname = varname\n",
    "        self.threshold = threshold\n",
    "        self.want_gte = want_gte\n",
    "        self.probmodel = 'binomial'\n",
    "    def count(self, datalist):\n",
    "        data = datalist[0]\n",
    "        if data.size % 365 != 0:\n",
    "            raise Exception('Data array length is not an integer multiple of 365')\n",
    "        byyear = data.reshape(data.size//365, 365)\n",
    "        if self.want_gte:\n",
    "            return np.sum((np.max(byyear, axis=1) >= self.threshold) * 1)\n",
    "        else:\n",
    "            return np.sum((np.max(byyear, axis=1) <= self.threshold) * 1)\n",
    "\n",
    "class Hotdays_inrange(Hazard):\n",
    "    def __init__(self, hightemp, lowtemp):\n",
    "        self.varname = 'tasmax'\n",
    "        self.hightemp = hightemp\n",
    "        self.lowtemp = lowtemp\n",
    "        self.probmodel = 'binomial'\n",
    "    def count(self, datalist):\n",
    "        data = datalist[0]\n",
    "        if data.size % 365 != 0:\n",
    "            raise Exception('Data array length is not an integer multiple of 365')\n",
    "        tf_array_high = data <= self.hightemp\n",
    "        tf_array_low = data >= self.lowtemp\n",
    "        return runs(tf_array_high * tf_array_low, self.min_duration, 'count')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "ae03622d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Location:\n",
    "    def __init__(self, name, latlon):\n",
    "        self.name = name\n",
    "        self.latlon = latlon\n",
    "        self.hist_observed = {}\n",
    "        self.hist_modeled = {scenario: {} for scenario in FUTURE_SCENARIOS}\n",
    "        self.best_models = {scenario: {} for scenario in FUTURE_SCENARIOS}\n",
    "        self.fut_modeled = {scenario: {} for scenario in FUTURE_SCENARIOS}\n",
    "        self.calib_fxns = {scenario: {} for scenario in FUTURE_SCENARIOS}\n",
    "        \n",
    "    def get_data(self, varname, future_start, future_end):\n",
    "    # Gets historical observations from ERA5 Daily Aggregate (from GEE)\n",
    "    # Goes through all NEX-GDDP-CMIP6 models in GEE and gets historical model outputs\n",
    "    # Chooses three best models based on quarterly RMSD\n",
    "        hist_obs = VARIABLES[varname]['era_transform'](get_var(VARIABLES[varname]['era_varname'], 'ERA5', self.latlon, HIST_START, HIST_END, southern_hem=False))\n",
    "        self.hist_observed[varname] = hist_obs\n",
    "        for scenario in FUTURE_SCENARIOS:\n",
    "            hist_mods = {}\n",
    "            rmsds = []\n",
    "            for model in MODELS[varname]:\n",
    "                hist_mod = VARIABLES[varname]['nex_transform'](get_var(varname, model, self.latlon, HIST_START, HIST_END, southern_hem=False, scenario=scenario))\n",
    "                hist_mods[model] = hist_mod\n",
    "                rmsds.append((get_rmsd(hist_obs, hist_mod), model))\n",
    "            rmsds.sort()\n",
    "\n",
    "            '''best_models = []\n",
    "            families = []\n",
    "            idx = 0\n",
    "            while len(best_models) < 3:\n",
    "                if not MODEL_INFO[rmsds[idx][1]] in families:\n",
    "                    best_models.append(rmsds[idx][1])\n",
    "                    families.append(MODEL_INFO[rmsds[idx][1]])\n",
    "                idx += 1\n",
    "\n",
    "            for m in best_models:\n",
    "                print(m, [i[0] for i in rmsds if i[1]==m][0])'''\n",
    "            best_models = []\n",
    "            for idx in range(min(NUM_BEST_MODELS, len(MODELS[varname]))):\n",
    "                best_models.append(rmsds[idx][1])\n",
    "\n",
    "            self.hist_modeled[scenario][varname] = hist_mods\n",
    "            self.best_models[scenario][varname] = best_models\n",
    "\n",
    "        # Get calibration functions\n",
    "            self.calib_fxns[varname] = {}\n",
    "            hist_obs = self.hist_observed[varname]\n",
    "            hist_mod = self.hist_modeled[scenario][varname]\n",
    "            for model in self.best_models[scenario][varname]:\n",
    "                o_quarters = quarters(hist_obs, HIST_START, HIST_END)\n",
    "                m_quarters = quarters(hist_mod[model], HIST_START, HIST_END)\n",
    "                self.calib_fxns[scenario][varname][model] = [calibration_function(o_quarters[i].flatten(), m_quarters[i].flatten()) for i in range(4)]\n",
    "\n",
    "        # Get future model outputs\n",
    "            self.fut_modeled[scenario][varname] = {\n",
    "                model: VARIABLES[varname]['nex_transform'](get_var(varname, model, self.latlon, future_start, FUTURE_END, extralong=True, scenario=scenario)) for model in best_models\n",
    "            }\n",
    "        \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "41ad76b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Estimate:\n",
    "    def __init__(self, location, hazard, future_start, future_end, sh_year):\n",
    "        self.location = location\n",
    "        self.hazard = hazard\n",
    "        self.future_start = future_start\n",
    "        self.future_end = future_end\n",
    "        self.sh_year = sh_year \n",
    "        self.estimate = {scenario: {} for scenario in FUTURE_SCENARIOS}\n",
    "        \n",
    "        for varname in self.hazard.varname.split('+'):\n",
    "            if not varname in list(self.location.hist_observed.keys()):\n",
    "                self.location.get_data(varname)\n",
    "        \n",
    "    def future_count(self):\n",
    "    # Gets future model outputs for three best models\n",
    "    # Calibrates based on stored model-specific quarterly calibration functions\n",
    "    # Calculates event count within future time series\n",
    "    # Draws 10,000 posterior rate parameters from appropriate Jeffrys prior distribution (parameterized with count)\n",
    "    # For each rate parameter, draw one event count\n",
    "        for scenario in FUTURE_SCENARIOS:\n",
    "            fut_mod = {}\n",
    "            varnames = self.hazard.varname.split('+')\n",
    "            for varname in varnames:\n",
    "                fut_mod[varname] = {}\n",
    "                for model in self.location.best_models[varname]:\n",
    "                    if self.sh_year:\n",
    "                        fut_mod[varname][model] = self.location.fut_modeled[varname][model][:-184]\n",
    "                    else:\n",
    "                        fut_mod[varname][model] = self.location.fut_modeled[varname][model][184:]\n",
    "            best_models = []\n",
    "            for idx in range(min(len(MODELS[varname]), NUM_BEST_MODELS)):\n",
    "                best_models.append('+'.join([self.location.best_models[scenario][varname][idx] for varname in varnames]))\n",
    "            posterior_rateparams = {}\n",
    "            posterior_draws = {}\n",
    "            estimate = {}\n",
    "            for modelplus in best_models:\n",
    "                calib_data = []\n",
    "                for idx, varname in enumerate(varnames):\n",
    "                    model = modelplus.split('+')[idx]\n",
    "                    calib_data.append(np.array(calibrate(fut_mod[varname][model], self.location.calib_fxns[scenario][varname][model])))\n",
    "                if self.sh_year:\n",
    "                    count = self.hazard.count(calib_data)\n",
    "                else:\n",
    "                    count = self.hazard.count([cd[152:-213] for cd in calib_data])\n",
    "                if self.hazard.probmodel == 'Poisson':\n",
    "                    posterior_rateparams[modelplus] = get_gamma(count, 10000)\n",
    "                    posterior_draws[modelplus] = np.random.poisson(posterior_rateparams[modelplus], 10000)\n",
    "                else:  # self.hazard.probmodel == 'binomial'\n",
    "                    posterior_rateparams[modelplus] = get_beta(count, self.future_end - self.future_start + 1, 10000)\n",
    "                    posterior_draws[modelplus] = np.random.binomial(self.future_end - self.future_start + 1, posterior_rateparams[modelplus], 10000)\n",
    "                self.estimate[scenario][modelplus] = np.sum(posterior_draws[modelplus]) / 10000\n",
    "\n",
    "            print()\n",
    "            print(scenario)\n",
    "            for modelplus in best_models:\n",
    "                if self.hazard.probmodel == 'Poisson':\n",
    "                    print('{0}: {1}'.format(modelplus, self.estimate[scenario][modelplus]))\n",
    "                else:\n",
    "                    print('{0}: {1}'.format(modelplus, '{0:.1f}%'.format(self.estimate[scenario][modelplus] / (self.future_end - self.future_start + 1) * 100)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44bf2b40",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dubai AE\n",
      "  95th pctl maxtemp = 42.9\n",
      "  5th pctl mintemp = 14.5\n",
      "Heat wave 95th percentile\n",
      "2080 2099\n",
      "\n",
      "--> The keys in the returned dictionary of datasets are constructed as follows:\n",
      "\t'activity_id.institution_id.source_id.experiment_id.table_id.grid_label'\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "<style>\n",
       "    /* Turns off some styling */\n",
       "    progress {\n",
       "        /* gets rid of default border in Firefox and Opera. */\n",
       "        border: none;\n",
       "        /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "        background-size: auto;\n",
       "    }\n",
       "    progress:not([value]), progress:not([value])::-webkit-progress-bar {\n",
       "        background: repeating-linear-gradient(45deg, #7e7e7e, #7e7e7e 10px, #5c5c5c 10px, #5c5c5c 20px);\n",
       "    }\n",
       "    .progress-bar-interrupted, .progress-bar-interrupted::-webkit-progress-bar {\n",
       "        background: #F44336;\n",
       "    }\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      <progress value='1' class='' max='1' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      100.00% [1/1 00:00&lt;00:00]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "with open('citiesonly.csv', 'r', encoding='utf-8') as ifile:\n",
    "    locations = {}\n",
    "    for line in ifile.readlines()[1:]:\n",
    "        items = line.split(',')\n",
    "        locations[items[0]] = (float(items[1]), float(items[2]))\n",
    "for location_name in locations:\n",
    "    print(location_name)\n",
    "    latlon = locations[location_name]\n",
    "    \n",
    "    high_95c = calendardate_percentiles('tasmax', 95, latlon, sh_hem=latlon[0] < 0)\n",
    "    cold_5n = wholeyear_percentile('tasmin', 5, latlon)\n",
    "    test_hazards = [\n",
    "        {'name': 'Heat wave 95th percentile', 'obj': Tempwave_simple('tasmax', 5, high_95c, True), 'sh_year': latlon[0] < 0},\n",
    "        {'name': 'Days warmer than 35C', 'obj': Threshold_simple('tasmax', 35, want_gte=True), 'sh_year': latlon[0] < 0},\n",
    "        {'name': 'Days colder than than 5th pctle yearlong', 'obj': Threshold_simple('tasmin', cold_5n, want_gte=False), 'sh_year': False},\n",
    "        {'name': 'Days rainier than than 500mm', 'obj': Threshold_simple('pr', 500, want_gte=True), 'sh_year': False},\n",
    "    ]\n",
    "    print('  95th pctl maxtemp = {0:.1f}'.format(np.max(high_95c)))\n",
    "    print('  5th pctl mintemp = {0:.1f}'.format(cold_5n))\n",
    "    loc = Location(location_name, latlon)\n",
    "    with open('test_outputs_1.csv', 'w') as ofile:\n",
    "        ofile.write('City,Hazard,')\n",
    "        for year_range in FUTURE_YEARS:\n",
    "            for scenario in FUTURE_SCENARIOS:\n",
    "                for prefix in ('low', 'mid', 'high'):\n",
    "                    ofile.write('{0}_{1}-{2}_{3},'.format(prefix, year_range[0], year_range[1], scenario))\n",
    "        ofile.write('\\n')\n",
    "        for haz in test_hazards:\n",
    "            ofile.write('{0},{1}'.format(loc.name, haz['name']))\n",
    "            for fut_start, fut_end in FUTURE_YEARS:\n",
    "                print(haz['name'])\n",
    "                print(fut_start, fut_end)\n",
    "                est = Estimate(loc, haz['obj'], fut_start, fut_end, haz['sh_year'])\n",
    "                est.future_count()\n",
    "                res = {scenario: list(est.estimate[scenario].values()) for scenario in FUTURE_SCENARIOS}\n",
    "                for scenario in res:\n",
    "                    res[scenario].sort()\n",
    "                for scenario in FUTURE_SCENARIOS:\n",
    "                    if haz['obj'].probmodel == 'binomial':\n",
    "                        res[scenario] = ['{0:.1f}%'.format(i / (fut_end - fut_start + 1) * 100) for i in res[scenario]]\n",
    "                    else:\n",
    "                        res[scenario] = ['{0:.1f}'.format(i / (fut_end - fut_start + 1)) for i in res[scenario]]\n",
    "                    ofile.write(',{0}'.format(res[scenario][0]))\n",
    "                    if len(res) == 3:\n",
    "                        ofile.write(',{0}'.format(res[scenario][1]))\n",
    "                    else:\n",
    "                        ofile.write(',')\n",
    "                    ofile.write(',{0}'.format(res[scenario][-1]))\n",
    "            ofile.write('\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b910b13e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54618e6f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6dc5f4e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59045705",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b62d743",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "c7dd43ed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GFDL-CM4: min modeled value does not exceed observed 10th percentile  True\n",
      "GFDL-CM4: max modeled value does not exceed observed 90th percentile  True\n",
      "CanESM5: min modeled value does not exceed observed 10th percentile  True\n",
      "CanESM5: max modeled value does not exceed observed 90th percentile  True\n",
      "ACCESS-CM2: min modeled value does not exceed observed 10th percentile  True\n",
      "ACCESS-CM2: max modeled value does not exceed observed 90th percentile  True\n",
      "GFDL-CM4: min modeled value does not exceed observed 10th percentile  True\n",
      "GFDL-CM4: max modeled value does not exceed observed 90th percentile  True\n",
      "CanESM5: min modeled value does not exceed observed 10th percentile  True\n",
      "CanESM5: max modeled value does not exceed observed 90th percentile  True\n",
      "ACCESS-CM2: min modeled value does not exceed observed 10th percentile  True\n",
      "ACCESS-CM2: max modeled value does not exceed observed 90th percentile  True\n",
      "GFDL-CM4: min modeled value does not exceed observed 10th percentile  True\n",
      "GFDL-CM4: max modeled value does not exceed observed 90th percentile  True\n",
      "CanESM5: min modeled value does not exceed observed 10th percentile  True\n",
      "CanESM5: max modeled value does not exceed observed 90th percentile  True\n",
      "ACCESS-CM2: min modeled value does not exceed observed 10th percentile  True\n",
      "ACCESS-CM2: max modeled value does not exceed observed 90th percentile  True\n",
      "GFDL-CM4: min modeled value does not exceed observed 10th percentile  True\n",
      "GFDL-CM4: max modeled value does not exceed observed 90th percentile  True\n",
      "CanESM5: min modeled value does not exceed observed 10th percentile  False\n",
      "CanESM5: max modeled value does not exceed observed 90th percentile  True\n",
      "ACCESS-CM2: min modeled value does not exceed observed 10th percentile  False\n",
      "ACCESS-CM2: max modeled value does not exceed observed 90th percentile  True\n"
     ]
    }
   ],
   "source": [
    "for quarter in range(4):\n",
    "    obs_10 = np.percentile(quarters(hist_obs_tx, HIST_START, HIST_END)[quarter], 10)\n",
    "    obs_90 = np.percentile(quarters(hist_obs_tx, HIST_START, HIST_END)[quarter], 90)\n",
    "    for model in best_models_tx:\n",
    "        mod = quarters(hist_mods_tx[model] - 273.15, HIST_START, HIST_END)[quarter].flatten()\n",
    "        print('{0}: min modeled value does not exceed observed 10th percentile  {1}'.format(model, min(mod) <= obs_10))\n",
    "        print('{0}: max modeled value does not exceed observed 90th percentile  {1}'.format(model, max(mod) >= obs_90))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2847140a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}