runPoliciesOnEnsemble.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run policies on a base ensemble"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Packages\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from scipy import integrate, stats\n",
    "from scipy.special import expit, binom\n",
    "\n",
    "import pandas as pd\n",
    "import xlrd\n",
    "\n",
    "import copy\n",
    "import warnings\n",
    "\n",
    "from datetime import datetime\n",
    "import random\n",
    "import string\n",
    "import os\n",
    "import shutil\n",
    "import sys\n",
    "import cloudpickle\n",
    "import dask\n",
    "import distributed\n",
    "from dask.distributed import Client\n",
    "import itertools\n",
    "\n",
    "import pymc3\n",
    "\n",
    "import inspect\n",
    "from collections import OrderedDict\n",
    "\n",
    "import dask.dataframe as dd\n",
    "\n",
    "from datetime import datetime\n",
    "import random\n",
    "import string\n",
    "import os\n",
    "import shutil\n",
    "import sys\n",
    "import cloudpickle\n",
    "import dask\n",
    "import distributed\n",
    "from dask.distributed import Client\n",
    "from dask.distributed import as_completed\n",
    "import itertools\n",
    "\n",
    "\n",
    "\n",
    "# # import helper funcs\n",
    "# modulePath = \"/mnt/efs/modules\"\n",
    "# shutil.copyfile(\"modelHelperFuncs.py\", modulePath + \"modelHelperFuncs.py\")\n",
    "# if not modulePath in sys.path:\n",
    "#     sys.path.insert(0, modulePath)\n",
    "\n",
    "# from modelHelperFuncs import regroup_by_age, build_paramDict, paramDict_toTable, paramTable_toDict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Helper functions\n",
    "\n",
    "Distributed dask workers seem very fiddly at importing non-standard modules, so it's better to just copy-paste the functions into this notebook sadly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To get relatve age-related risks, we have to first re-group into our basic age groups, \n",
    "# then devide by total population (here's a non-well-defined subset of UK, so absolute values wont be used, only relative)\n",
    "import numpy as np\n",
    "import inspect\n",
    "from collections import OrderedDict\n",
    "\n",
    "def regroup_by_age(inp, fromAgeSplits, toAgeSplits, maxAge=100., maxAgeWeight = 5.):\n",
    "    fromAgeSplits = np.concatenate([np.array([0]), fromAgeSplits, np.array([maxAge])]) # Add a zero at beginning for calculations\n",
    "    toAgeSplits = np.concatenate([np.array([0]), toAgeSplits, np.array([maxAge])]) # Add inf at end for calculations\n",
    "    def getOverlap(a, b):\n",
    "        return max(0, min(a[1], b[1]) - max(a[0], b[0]))\n",
    "    out = np.zeros((len(toAgeSplits)-1,)+inp.shape[1:])\n",
    "    for from_ind in range(1, len(fromAgeSplits)):\n",
    "        # Redistribute to the new bins by calculating how many years in from_ind-1:from_ind falls into each output bin\n",
    "        cur_out_distribution = (\n",
    "        [getOverlap(toAgeSplits[cur_to_ind-1:cur_to_ind+1],fromAgeSplits[from_ind-1:from_ind+1])  for cur_to_ind in range(1, len(toAgeSplits))]\n",
    "        )\n",
    "        \n",
    "        if cur_out_distribution[-1] > 0:\n",
    "            cur_out_distribution[-1] = maxAgeWeight # Define the relative number of ages if we have to distribute between second to last and last age groups\n",
    "\n",
    "        cur_out_distribution = cur_out_distribution/np.sum(cur_out_distribution)\n",
    "        \n",
    "        for to_ind in range(len(out)):\n",
    "            out[to_ind] += cur_out_distribution[to_ind] * inp[from_ind-1]\n",
    "            \n",
    "    return out\n",
    "\n",
    "\n",
    "# PARAMETER DICTIONARIES AND TABLES\n",
    "# -----------------------------------------------------------------------------------------\n",
    "\n",
    "\n",
    "def build_paramDict(cur_func):\n",
    "    \"\"\"\n",
    "    This function iterates through all inputs of a function, \n",
    "    and saves the default argument names and values into a dictionary.\n",
    "    \n",
    "    If any of the default arguments are functions themselves, then recursively (depth-first) adds an extra field to\n",
    "    the dictionary, named <funcName + \"_params\">, that contains its inputs and arguments.\n",
    "    \n",
    "    The output of this function can then be passed as a \"kwargs\" object to the highest level function, \n",
    "    which will then pass the parameter values to the lower dictionary levels appropriately\n",
    "    \"\"\"\n",
    "    \n",
    "    paramDict = OrderedDict()\n",
    "    \n",
    "    allArgs = inspect.getfullargspec(cur_func)\n",
    "    \n",
    "    # Check if there are any default parameters, if no, just return empty dict\n",
    "    if allArgs.defaults is None:\n",
    "        return paramDict\n",
    "    \n",
    "    \n",
    "    for argname, argval in zip(allArgs.args[-len(allArgs.defaults):], allArgs.defaults):\n",
    "        # Save the default argument\n",
    "        paramDict[argname] = argval\n",
    "        # If the default argument is a function, inspect it for further \n",
    "        \n",
    "        if callable(argval):\n",
    "            # print(argname)\n",
    "            paramDict[argname+\"_params\"] = build_paramDict(argval)\n",
    "\n",
    "    return paramDict\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# Do a mapping between dictionary and parameter table row (for convenient use)\n",
    "\n",
    "# Flatten the dictionary into a table with a single row (but many headers):\n",
    "def paramDict_toTable(paramDict):\n",
    "    paramTable = pd.DataFrame()\n",
    "    def paramDictRecurseIter(cur_table, cur_dict, preString):\n",
    "        # Iterate through the dictionary to find all keys not ending in \"_params\", \n",
    "        # and add them to the table with name <preString + key>\n",
    "        # \n",
    "        # If the key doesn end in \"_params\", then append the key to preString, in call this function on the value (that is a dict)\n",
    "        for key, value in cur_dict.items():\n",
    "            if key.endswith(\"_params\"):\n",
    "                paramDictRecurseIter(cur_table, value, preString+key+\"_\")\n",
    "            else:\n",
    "                paramTable[preString+key] = [value]\n",
    "                \n",
    "        # For the rare case where we want to keep an empty dictionary, the above for cycle doesn't keep it\n",
    "        if len(cur_dict)==0:\n",
    "            paramTable[preString] = [OrderedDict()]\n",
    "                \n",
    "        return cur_table\n",
    "    \n",
    "    return paramDictRecurseIter(paramTable, paramDict, preString=\"\")\n",
    "\n",
    "# Example dict -> table\n",
    "# paramTable_default = paramDict_toTable(paramDict_default)\n",
    "    \n",
    "\n",
    "def paramTable_toDict(paramTable, defaultDict=None, to_flat=False):\n",
    "    # enable to pass a default dict (if paramTable is incomplete), in which we'll just add / overwrite the values\n",
    "    paramDict = defaultDict if defaultDict is not None else OrderedDict() \n",
    "    def placeArgInDictRecurse(argName, argVal, cur_dict):\n",
    "        # Find all \"_params_\" in the argName, and for each step more and more down in the dictionary\n",
    "        strloc = argName.find(\"_params_\")\n",
    "        if strloc == -1 or to_flat:\n",
    "            # We're at the correct level of dictionary\n",
    "            cur_dict[argName] = argVal\n",
    "            return cur_dict\n",
    "        else:\n",
    "            # step to the next level of dictionary\n",
    "            nextKey = argName[:strloc+len(\"_params_\")-1]\n",
    "            nextArgName = argName[strloc+len(\"_params_\"):]\n",
    "            if not nextKey in cur_dict:\n",
    "                cur_dict[nextKey] = OrderedDict()\n",
    "            placeArgInDictRecurse(nextArgName, argVal, cur_dict[nextKey])\n",
    "            \n",
    "        return cur_dict\n",
    "            \n",
    "    for key in paramTable.columns:\n",
    "        paramDict = placeArgInDictRecurse(key, paramTable.at[0,key], paramDict)\n",
    "            \n",
    "    return paramDict\n",
    "\n",
    "# Example table -> dict \n",
    "# paramDict_new = paramTable_toDict(paramTable_default)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## # Load a full ensemble and select the base ensemble with queries\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load a full ensemble and select the base ensemble with queries\n",
    "base_ensemble_params = OrderedDict()\n",
    "\n",
    "base_ensemble_params[\"ensembleDir\"] = \"/mnt/efs/results/run_20200421T002117/\"\n",
    "base_ensemble_params[\"ensembleQuery\"] = \"likelihood_0 > -260 & likelihood_2 > -280\"\n",
    "base_ensemble_params[\"ensembleSortby\"] = \"likelihood_total\"\n",
    "base_ensemble_params[\"ensembleMaxnumber\"] = 200\n",
    "\n",
    "paramTable_ensemble_dd = dd.from_pandas(\n",
    "    pd.read_hdf(base_ensemble_params[\"ensembleDir\"] + 'paramTable_part0', key=\"paramTable\"),chunksize=1000)\n",
    "\n",
    "i = 1\n",
    "while True:\n",
    "    print(\"Loading {}\".format(i))\n",
    "    try:\n",
    "        paramTable_ensemble_dd = paramTable_ensemble_dd.append(\n",
    "            pd.read_hdf(base_ensemble_params[\"ensembleDir\"] + 'paramTable_part{}'.format(i), key=\"paramTable\"))\n",
    "        i += 1\n",
    "    except:\n",
    "        break\n",
    "        \n",
    "        \n",
    "paramTable_ensemble_dd_queriedSorted = (\n",
    "    paramTable_ensemble_dd\n",
    "    .query(base_ensemble_params[\"ensembleQuery\"])\n",
    "    .compute()\n",
    "    .sort_values(base_ensemble_params[\"ensembleSortby\"])\n",
    "    .tail(base_ensemble_params[\"ensembleMaxnumber\"])\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "paramTable_ensemble_dd_queriedSorted"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load all the functions from the base ensemble (default_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(base_ensemble_params[\"ensembleDir\"] + \"paramDict_default.cpkl\", 'rb') as fh:\n",
    "    paramDict_default = cloudpickle.load(fh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up new policy parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "paramTable_default = paramDict_toTable(paramDict_default)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Case isolation policy parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "newPolicyParams_caseIsolation = OrderedDict()\n",
    "\n",
    "newPolicyParams_caseIsolation[\"tStartQuarantineCaseIsolation\"] = [\n",
    "    #pd.to_datetime(\"2020-04-20\", format=\"%Y-%m-%d\"),\n",
    "    pd.to_datetime(\"2020-05-01\", format=\"%Y-%m-%d\")\n",
    "]\n",
    "newPolicyParams_caseIsolation['tStopSocialDistancing'] = [\n",
    "    #pd.to_datetime(\"2020-04-30\", format=\"%Y-%m-%d\"),\n",
    "    pd.to_datetime(\"2020-05-30\", format=\"%Y-%m-%d\"),\n",
    "    pd.to_datetime(\"2020-06-30\", format=\"%Y-%m-%d\"),\n",
    "    pd.to_datetime(\"2020-09-30\", format=\"%Y-%m-%d\")\n",
    "]\n",
    "\n",
    "newPolicyParams_caseIsolation['trFunc_testing_params_policyFunc_params_basic_policyFunc_params_antibody_testing_policy'] = [\n",
    "    \"virus_positive_only_hospworker_first\"\n",
    "]\n",
    "\n",
    "newPolicyParams_caseIsolation['trFunc_testing_params_trFunc_testCapacity_params_testCapacity_antigenratio_country'] = [1.]\n",
    "\n",
    "\n",
    "newPolicyParams_caseIsolation['trFunc_testing_params_trFunc_testCapacity_params_testCapacity_antibody_country_total'] = [\n",
    "    1e5, 5e5, 1e6, 2e6, 3e6, 5e6, 10e6]\n",
    "\n",
    "newPolicyParams_caseIsolation['trFunc_testing_params_trFunc_testCapacity_params_testCapacity_antibody_country_inflexday'] = [\n",
    "     #pd.to_datetime(\"2020-04-30\", format=\"%Y-%m-%d\"),\n",
    "     pd.to_datetime(\"2020-05-30\", format=\"%Y-%m-%d\")\n",
    "]\n",
    "\n",
    "newPolicyParams_caseIsolation['trFunc_testing_params_trFunc_testCapacity_params_testCapacity_antibody_country_inflexslope'] = [\n",
    "    10.#, 20\n",
    "]\n",
    "\n",
    "newPolicyParams_caseIsolation['trFunc_quarantine_params_nDaysInHomeIsolation'] = [\n",
    "    14.#, 24.\n",
    "]\n",
    "\n",
    "\n",
    "\n",
    "newPolicyParams_caseIsolation[\"trFunc_testing_params_inpFunc_testSpecifications_params_antigen_FNR_I1_to_R2\"]= np.array(\n",
    "    [\n",
    "        # Good test\n",
    "        [0.8,   0.3,   0.1, 0.15, 0.3, 0.7],\n",
    "        # Medium test\n",
    "        [0.9,   0.5,   0.25,  0.35, 0.5, 0.9],\n",
    "        #Bad test\n",
    "        [0.99,   0.8,   0.4,  0.5, 0.6, 0.99]\n",
    "    ]\n",
    ")\n",
    "\n",
    "newPolicyParams_caseIsolation[\"trFunc_testing_params_inpFunc_testSpecifications_params_antigen_FPR\"] = [0.05, 0.1]\n",
    "\n",
    "\n",
    "\n",
    "# Make the joint dataframe\n",
    "df_newPolicyParams_caseIsolation = pd.DataFrame(\n",
    "    data=list(itertools.product(*newPolicyParams_caseIsolation.values())), \n",
    "    columns=newPolicyParams_caseIsolation.keys()\n",
    "    )\n",
    "\n",
    "\n",
    "# Filter rows to retain only zipped combinations (TODO - there's probably a way to do it during the itertools.product, but this is ok for now)\n",
    "zippedColumnSets = [\n",
    "#    [\"trFunc_testing_params_inpFunc_testSpecifications_params_antigen_FNR_I1_to_R2\", \"trFunc_testing_params_inpFunc_testSpecifications_params_antigen_FPR\"]\n",
    "]\n",
    "\n",
    "all_validInds = df_newPolicyParams_caseIsolation.index\n",
    "for zcs in zippedColumnSets:\n",
    "    zippedVals = list(zip(*[newPolicyParams_caseIsolation[key] for key in zcs]))\n",
    "    \n",
    "    curZCS_validInds = []\n",
    "    for zv in zippedVals:\n",
    "        validInds = df_newPolicyParams_caseIsolation.loc[all_validInds].index\n",
    "        for ind, key in enumerate(zcs):\n",
    "            tmp = df_newPolicyParams_caseIsolation.loc[validInds]\n",
    "            validInds = tmp.index[\n",
    "                            tmp[key].apply(lambda x: np.all(x==zv[ind]))\n",
    "                        ]\n",
    "            \n",
    "        curZCS_validInds = curZCS_validInds + list(validInds)\n",
    "        \n",
    "    #curZCS_validInds = [b for a in curZCS_validInds for b in a]\n",
    "    \n",
    "    all_validInds = list(set(all_validInds) & set(curZCS_validInds))\n",
    "            \n",
    "df_newPolicyParams_caseIsolation = df_newPolicyParams_caseIsolation.loc[all_validInds] \n",
    "df_newPolicyParams_caseIsolation = df_newPolicyParams_caseIsolation.reset_index(drop=True)\n",
    "    \n",
    "\n",
    "len(df_newPolicyParams_caseIsolation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Merge the new policy params with the selected ensemble params\n",
    "for i in range(len(paramTable_ensemble_dd_queriedSorted)):\n",
    "    tmp = copy.deepcopy(df_newPolicyParams_caseIsolation)\n",
    "    for colname in (paramTable_ensemble_dd_queriedSorted\n",
    "                    .reset_index()\n",
    "                    .loc[i:i]\n",
    "                    #.rename(columns={\"Index\":\"EnsembleIndex\", \"out_fname\":\"Ensemble_out_fname\"})\n",
    "                    .columns\n",
    "                   ):\n",
    "        tmp[colname] = [paramTable_ensemble_dd_queriedSorted.reset_index().at[i, colname]]*len(df_newPolicyParams_caseIsolation)\n",
    "        tmp = tmp.rename(columns={\"index\":\"EnsembleIndex\", \"out_fname\":\"Ensemble_out_fname\"})\n",
    "        \n",
    "    if i==0:\n",
    "        paramTable_merged = tmp\n",
    "    else:\n",
    "        paramTable_merged = paramTable_merged.append(tmp)\n",
    "        \n",
    "paramTable_merged = paramTable_merged.reset_index(drop=True)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "paramTable_merged"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Also make an alternative paramTable to evaluate the \"keep social distancing\" policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a single row on top that runs with base settings (ie policy completely off)\n",
    "\n",
    "paramTable_baseParams_keepSocialDistancing = copy.deepcopy(paramTable_ensemble_dd_queriedSorted)\n",
    "\n",
    "paramTable_baseParams_keepSocialDistancing = (\n",
    "    paramTable_baseParams_keepSocialDistancing\n",
    "    .reset_index()\n",
    "    .rename(columns={\"index\":\"EnsembleIndex\", \"out_fname\":\"Ensemble_out_fname\"})\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "paramTable_baseParams_keepSocialDistancing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Also make an alternative paramTable to evaluate the \"stop social distancing\" policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a single row on top that runs with base settings (ie policy completely off)\n",
    "# Make a single table with base params but stopping social distancing at some points\n",
    "\n",
    "newPolicyParams_stopSocialDistancing = OrderedDict()\n",
    "\n",
    "newPolicyParams_stopSocialDistancing['tStopSocialDistancing'] = [\n",
    "    pd.to_datetime(\"2020-04-30\", format=\"%Y-%m-%d\"),\n",
    "    pd.to_datetime(\"2020-05-30\", format=\"%Y-%m-%d\"),\n",
    "    pd.to_datetime(\"2020-06-30\", format=\"%Y-%m-%d\"),\n",
    "    pd.to_datetime(\"2020-09-30\", format=\"%Y-%m-%d\")\n",
    "]\n",
    "\n",
    "df_newPolicyParams_stopSocialDistancing = pd.DataFrame(\n",
    "    data=list(itertools.product(*newPolicyParams_stopSocialDistancing.values())), \n",
    "    columns=newPolicyParams_stopSocialDistancing.keys()\n",
    "    )\n",
    "\n",
    "\n",
    "# Merge the new policy params with the selected ensemble params\n",
    "for i in range(len(paramTable_ensemble_dd_queriedSorted)):\n",
    "    tmp = copy.deepcopy(df_newPolicyParams_stopSocialDistancing)\n",
    "    for colname in (paramTable_ensemble_dd_queriedSorted\n",
    "                    .reset_index()\n",
    "                    .loc[i:i]\n",
    "                    #.rename(columns={\"Index\":\"EnsembleIndex\", \"out_fname\":\"Ensemble_out_fname\"})\n",
    "                    .columns\n",
    "                   ):\n",
    "        tmp[colname] = [paramTable_ensemble_dd_queriedSorted.reset_index().at[i, colname]]*len(df_newPolicyParams_stopSocialDistancing)\n",
    "        tmp = tmp.rename(columns={\"index\":\"EnsembleIndex\", \"out_fname\":\"Ensemble_out_fname\"})\n",
    "        \n",
    "    if i==0:\n",
    "        paramTable_baseParams_stopSocialDistancing = tmp\n",
    "    else:\n",
    "        paramTable_baseParams_stopSocialDistancing = paramTable_baseParams_stopSocialDistancing.append(tmp)\n",
    "        \n",
    "paramTable_baseParams_stopSocialDistancing = paramTable_baseParams_stopSocialDistancing.reset_index(drop=True)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "paramTable_baseParams_stopSocialDistancing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "# Set up and save policy control dicts for dashboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('paramTypes.cpkl', 'rb') as fh:\n",
    "    paramTypes = cloudpickle.load(fh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "policy_control_parameters = {\n",
    "        # Timings\n",
    "        'tStopSocialDistancing': [\"stop_social_distancing\", \"Social Distancing End Date\", \"datetime64\", \"dropdown\", None, True],\n",
    "        'tStartImmunityPassports': [\"start_immunity_passports\", \"Immunity Passports Start Date\", \"datetime64\", \"dropdown\", None, False],\n",
    "        'tStopImmunityPassports': [\"stop_immunity_passports\", \"Immunity Passports End Date\", \"datetime64\", \"dropdown\", None, False],\n",
    "        'tStartQuarantineCaseIsolation': [\"start_case_isolation\", \"Case Isolation Start Date\", \"datetime64\", \"dropdown\", None, False],\n",
    "        'tStopQuarantineCaseIsolation': [\"stop_case_isolation\", \"Case Isolation End Date\", \"datetime64\", \"dropdown\", None, True],\n",
    "        # Quarantine\n",
    "        'trFunc_quarantine_params_nDaysInHomeIsolation': [\n",
    "            \"caseiso_ndayshome\", \"Length of strict home quarantine (days)\", \"int\", \"slider\", None, False\n",
    "        ],\n",
    "        # 'trFunc_newInfections_params_ageSocialMixingIsolation': [\"STATIC\"],\n",
    "        'trFunc_quarantine_params_timeToIsolation': [\n",
    "            \"caseiso_timetoiso\", \"Time between test and quarantine start (days)\", \"float\", \"slider\", None, False\n",
    "        ],\n",
    "        # 'trFunc_quarantine_params_symptomHospitalisedRate_ageAdjusted' : [\"STATIC\"],\n",
    "        # 'trFunc_quarantine_params_symptomaticHealthStates' : [\"STATIC\"],\n",
    "        # Testing\n",
    "        'trFunc_testing_params_trFunc_testCapacity_params_testCapacity_pcr_phe_total': [\n",
    "            \"testcapacity_pcr_phe_total\", \"PHE lab PCR tests - maximum capacity per day\", \"float\", \"slider\", None, False\n",
    "        ],\n",
    "        'trFunc_testing_params_trFunc_testCapacity_params_testCapacity_pcr_phe_inflexday': [\n",
    "            \"testcapacity_pcr_phe_inflexday\", \"PHE lab PCR tests - date of reaching half maximum capacity\", \"datetime64\", \"dropdown\", None, False\n",
    "        ],\n",
    "        'trFunc_testing_params_trFunc_testCapacity_params_testCapacity_pcr_phe_inflexslope': [\n",
    "            \"testcapacity_pcr_phe_inflexslope\", \"PHE lab PCR tests - days to reach maximum capacity\", \"float\", \"slider\", None, False\n",
    "        ],\n",
    "        'trFunc_testing_params_trFunc_testCapacity_params_testCapacity_pcr_country_total': [\n",
    "            \"testcapacity_pcr_country_total\", \"UK non-PHE PCR tests - maximum capacity per day\", \"float\", \"slider\", None, False\n",
    "        ],\n",
    "        'trFunc_testing_params_trFunc_testCapacity_params_testCapacity_pcr_country_inflexday': [\n",
    "            \"testcapacity_pcr_country_inflexday\", \"UK non-PHE PCR tests - date of reaching half maximum capacity\", \"datetime64\", \"dropdown\", None, False\n",
    "        ],\n",
    "        'trFunc_testing_params_trFunc_testCapacity_params_testCapacity_pcr_country_inflexslope': [\n",
    "            \"testcapacity_pcr_country_inflexslope\", \"UK non-PHE PCR tests - days to reach maximum capacity\", \"float\", \"slider\", None, False\n",
    "        ],\n",
    "        'trFunc_testing_params_trFunc_testCapacity_params_testCapacity_antibody_country_firstday': [\n",
    "            \"testcapacity_lfa_country_firstday\", \"UK home tests (LFA) - first date deployed\", \"datetime64\", \"dropdown\", None, False\n",
    "        ],\n",
    "        'trFunc_testing_params_trFunc_testCapacity_params_testCapacity_antibody_country_total': [\n",
    "            \"testcapacity_lfa_country_total\", \"UK home tests (LFA) - maximum capacity per day\", \"float\", \"slider\", None, False\n",
    "        ],\n",
    "        'trFunc_testing_params_trFunc_testCapacity_params_testCapacity_antibody_country_inflexday': [\n",
    "            \"testcapacity_lfa_country_inflexday\", \"UK home tests (LFA) - date of reaching half maximum capacity\", \"datetime64\", \"dropdown\", None, False\n",
    "        ],\n",
    "        'trFunc_testing_params_trFunc_testCapacity_params_testCapacity_antibody_country_inflexslope': [\n",
    "            \"testcapacity_lfa_country_inflexslope\", \"UK home tests (LFA) - days to reach maximum capacity\", \"float\", \"slider\", None, False\n",
    "        ],\n",
    "        'trFunc_testing_params_trFunc_testCapacity_params_testCapacity_antigenratio_country': [\n",
    "            \"testcapacity_lfa_country_antigenratio\", \"UK home tests (LFA) - ratio of virus and immunity tests\", \"float\", \"slider\", None, True\n",
    "        ],\n",
    "        'trFunc_testing_params_policyFunc_params_basic_policyFunc_params_antibody_testing_policy': [\n",
    "            \"antibody_testing_policy\", \"Antibody test distribution policy\", \"string\", \"dropdown\", None, False\n",
    "        ],\n",
    "        # 'trFunc_testing_params_policyFunc_params_f_symptoms_nonCOVID_params_' : [\"STATIC\"],\n",
    "        # 'trFunc_testing_params_policyFunc_params_distributeRemainingToRandom' : [\"STATIC\"],\n",
    "        # Test specs\n",
    "        'trFunc_testing_params_inpFunc_testSpecifications_params_antigen_FNR_I1_to_R2': [\n",
    "            \"testspecs_antigen_FNR\", \"LFA antigen test specifications (FNR)\", \"spec_list\", \"dropdown\", \n",
    "            [[np.array([0.8, 0.3, 0.1, 0.15, 0.3, 0.7]), \n",
    "              np.array([0.9, 0.5, 0.25, 0.35, 0.5, 0.9]), \n",
    "              np.array([0.99, 0.8, 0.4, 0.5, 0.6, 0.99])], [\"Good\", \"Medium\", \"Bad\"]], True\n",
    "        ],\n",
    "        'trFunc_testing_params_inpFunc_testSpecifications_params_antigen_FPR' : [\n",
    "            \"testspecs_antigen_FPR\",\"LFA antigen test specifications (FPR)\", \"float\", \"slider\", None, True\n",
    "        ],\n",
    "        'trFunc_testing_params_inpFunc_testSpecifications_params_antibody_FNR_I1_to_R2': [\n",
    "            \"testspecs_antibody_FNR\", \"LFA antibody test specifications (FNR)\", \"spec_list\", \"dropdown\", [[np.array([0.99, 0.85, 0.8, 0.65, 0.3, 0.05])], [\"Baseline\"]], True\n",
    "        ],\n",
    "        'trFunc_testing_params_inpFunc_testSpecifications_params_antibody_FPR_S_to_I4': [\n",
    "            \"testspecs_antibody_FPR\",\"LFA antibody test specifications (FPR)\", \"float\", \"slider\", None, True\n",
    "        ],\n",
    "    \n",
    "        'trFunc_testing_params_policyFunc_params_retesting_antigen_viruspos_ratio': [\n",
    "            \"retesting_antigen_viruspos_ratio\", \"Ratio of leftover home virus tests used for re-testing virus pos\", \"float\", \"slider\", None, False            \n",
    "        ],\n",
    "    \n",
    "        'trFunc_testing_params_policyFunc_params_retesting_antigen_immunepos_ratio': [\n",
    "                \"retesting_antigen_viruspos_ratio\", \"Ratio of leftover home virus tests used for re-testing immune pos\", \"float\", \"slider\", None, False            \n",
    "            ],\n",
    "    \n",
    "        'trFunc_testing_params_policyFunc_params_retesting_antibody_immunepos_ratio': [\n",
    "            \"retesting_antibody_immunepos_ratio\", \"Ratio of leftover home immunity tests used for re-testing\", \"float\", \"slider\", None, False            \n",
    "        ]\n",
    "    \n",
    "    }\n",
    "\n",
    "\n",
    "[print(i) for i in (set(policy_control_parameters.keys()) - set(paramTypes[\"policy\"]))]\n",
    "print(\"-------------\")\n",
    "[print(i) for i in (set(paramTypes[\"policy\"]) - set(policy_control_parameters.keys()))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the config dicts for each table\n",
    "# Policy params\n",
    "controlDict_caseIsolation = OrderedDict()\n",
    "for key in paramTable_merged.columns:\n",
    "    if key in policy_control_parameters:\n",
    "        if len(newPolicyParams_caseIsolation[key])>1: # keep only the ones that actually vary!\n",
    "            controlDict_caseIsolation[key] = policy_control_parameters[key]\n",
    "df_controlDict_caseIsolation = paramDict_toTable(\n",
    "    controlDict_caseIsolation\n",
    ")\n",
    "\n",
    "df_controlDict_caseIsolation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the config dicts for each table\n",
    "# Policy params\n",
    "controlDict_keepSocialDistancing = OrderedDict()\n",
    "df_controlDict_keepSocialDistancing = paramDict_toTable(\n",
    "    controlDict_keepSocialDistancing\n",
    ")\n",
    "\n",
    "df_controlDict_keepSocialDistancing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the config dicts for each table\n",
    "# Policy params\n",
    "controlDict_stopSocialDistancing = OrderedDict()\n",
    "for key in paramTable_baseParams_stopSocialDistancing.columns:\n",
    "    if key in policy_control_parameters:\n",
    "        if len(newPolicyParams_stopSocialDistancing[key])>1: # keep only the ones that actually vary!\n",
    "            controlDict_stopSocialDistancing[key] = policy_control_parameters[key]\n",
    "df_controlDict_stopSocialDistancing = paramDict_toTable(\n",
    "    controlDict_stopSocialDistancing\n",
    ")\n",
    "\n",
    "df_controlDict_stopSocialDistancing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate the full policy grid on all ensemble members"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Load the saved initialisation\n",
    "stateTensor_init = paramDict_default[\"INIT_stateTensor_init\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solveSystem(stateTensor_init, total_days = 200, samplesPerDay=np.inf, **kwargs):\n",
    "    # Run the simulation\n",
    "    \n",
    "    if kwargs[\"debugReturnNewPerDay\"]: # Keep the second copy as well\n",
    "        cur_stateTensor = np.reshape(\n",
    "            np.stack([copy.deepcopy(stateTensor_init), copy.deepcopy(stateTensor_init)], axis=0),-1)\n",
    "    else:\n",
    "        cur_stateTensor = np.reshape(copy.deepcopy(stateTensor_init),-1)\n",
    "    \n",
    "    if np.isinf(samplesPerDay):\n",
    "        # Run precise integrator takes forever\n",
    "        out = integrate.solve_ivp(\n",
    "            fun = lambda t,y: kwargs[\"dydt_Complete\"](t,y, **kwargs),\n",
    "            t_span=(0.,total_days),\n",
    "            y0 = cur_stateTensor,\n",
    "            method='RK23',\n",
    "            t_eval=range(total_days),\n",
    "            #max_step = 1.,\n",
    "            #first_step = 1e-1,\n",
    "            rtol = 1e-3, #default 1e-3\n",
    "            atol = 1e-3, # default 1e-6\n",
    "        )\n",
    "        \n",
    "        out = out.y\n",
    "        \n",
    "    else:\n",
    "        # Run simple Euler method with given step size (1/samplesPerDay)\n",
    "        deltaT = 1./samplesPerDay\n",
    "        out = np.zeros((np.prod(stateTensor_init.shape),total_days))\n",
    "                       \n",
    "        for tt in range(total_days*samplesPerDay):\n",
    "            if tt % samplesPerDay==0:\n",
    "                out[:, int(tt/samplesPerDay)] = cur_stateTensor\n",
    "                       \n",
    "            cur_stateTensor += deltaT * kwargs[\"dydt_Complete\"]((tt*1.)/(1.*samplesPerDay),cur_stateTensor, **kwargs)\n",
    "            \n",
    "    \n",
    "    # Reshape to reasonable format\n",
    "    if kwargs[\"debugReturnNewPerDay\"]:\n",
    "        out = np.reshape(out, (2,) + stateTensor_init.shape+(-1,))\n",
    "    else:\n",
    "        out = np.reshape(out, stateTensor_init.shape+(-1,))\n",
    "    \n",
    "    \n",
    "    return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up dask parallel run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client(\"127.0.0.1:8786\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up where to save and save default parameters\n",
    "\n",
    "timeOfRunning = datetime.now().strftime(\"%Y%m%dT%H%M%S\")\n",
    "\n",
    "saveDir = \"/mnt/efs/results/run_\" + timeOfRunning + \"/\"\n",
    "os.makedirs(saveDir, exist_ok=True)\n",
    "os.chmod(saveDir, 0o777) # enable workers to write the files\n",
    "\n",
    "# Save the default parameter dictionary that we'll merge with new inputs\n",
    "with open(saveDir+'paramDict_default.cpkl', 'wb') as fh:\n",
    "    cloudpickle.dump(paramDict_default, fh)\n",
    "\n",
    "\n",
    "# save also the ensemble parameters dictionary\n",
    "with open(saveDir+'ensembleSelectionSettings.cpkl', 'wb') as fh:\n",
    "    cloudpickle.dump(base_ensemble_params, fh)\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create file names for the policy grid\n",
    "out_fnames = []\n",
    "for ind in range(len(paramTable_merged)):\n",
    "    out_fnames.append(\"outTensor_\" + timeOfRunning + \"_\" + ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(20))+\".npy\")\n",
    "    \n",
    "paramTable_merged[\"out_fname\"] = out_fnames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create file names for the \"keep social distancing\" run\n",
    "out_fnames = []\n",
    "for ind in range(len(paramTable_baseParams_keepSocialDistancing)):\n",
    "    out_fnames.append(\"outTensor_\" + timeOfRunning + \"_\" + ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(20))+\".npy\")\n",
    "    \n",
    "paramTable_baseParams_keepSocialDistancing[\"out_fname\"] = out_fnames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create file names for the \"keep social distancing\" run\n",
    "out_fnames = []\n",
    "for ind in range(len(paramTable_baseParams_stopSocialDistancing)):\n",
    "    out_fnames.append(\"outTensor_\" + timeOfRunning + \"_\" + ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(20))+\".npy\")\n",
    "    \n",
    "paramTable_baseParams_stopSocialDistancing[\"out_fname\"] = out_fnames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save also the paramTable, with the appropriate keys to determine \"policy type\"\n",
    "paramTable_merged.to_hdf(saveDir + \"paramTable_part{}\".format(0), key=\"Case Isolation Varying Test Numbers\")\n",
    "df_controlDict_caseIsolation.to_hdf(saveDir + \"paramTable_part{}\".format(0), key=\"DashboardConfig-Case Isolation Varying Test Numbers\")\n",
    "\n",
    "paramTable_baseParams_keepSocialDistancing.to_hdf(saveDir + \"paramTable_part{}\".format(0), key=\"Keep Social Distancing\")\n",
    "df_controlDict_keepSocialDistancing.to_hdf(saveDir + \"paramTable_part{}\".format(0), key=\"DashboardConfig-Keep Social Distancing\")\n",
    "\n",
    "paramTable_baseParams_stopSocialDistancing.to_hdf(saveDir + \"paramTable_part{}\".format(0), key=\"Stop Social Distancing\")\n",
    "df_controlDict_stopSocialDistancing.to_hdf(saveDir + \"paramTable_part{}\".format(0), key=\"DashboardConfig-Stop Social Distancing\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run parallel for each parameter setting and save to out_fname\n",
    "def runAll(newParams_row, stateTensor_init=stateTensor_init, defaultDict=paramDict_default, timeOfRunning=timeOfRunning):\n",
    "    # Run model \n",
    "    # Make sure the newOnly stuff is saved as well\n",
    "    curDict = copy.deepcopy(defaultDict)\n",
    "    curDict[\"debugReturnNewPerDay\"] = True\n",
    "    \n",
    "    out = solveSystem(stateTensor_init, \n",
    "                total_days = 365, \n",
    "                **paramTable_toDict(\n",
    "                            # sub-select allowed columns\n",
    "                           newParams_row[list(set(newParams_row.columns) & set(paramDict_toTable(defaultDict).columns))].reset_index(drop=True),\n",
    "                           defaultDict=copy.deepcopy(curDict)\n",
    "                    )\n",
    "               )\n",
    "    # The out is now both the states and the cumsum\n",
    "    out_newOnly = np.diff(np.concatenate([np.expand_dims(copy.deepcopy(out[0][:,:,:,:,0]),axis=4), copy.deepcopy(out[1])], axis=-1), axis=-1)\n",
    "    out = out[0]    \n",
    "    \n",
    "    return out, out_newOnly, newParams_row"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### First run the short ones"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "futures = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Submit all futures that were not yet completed (check if files exist)\n",
    "for index in range(len(paramTable_baseParams_keepSocialDistancing)):\n",
    "    \n",
    "    tmp_params_row = copy.deepcopy(paramTable_baseParams_keepSocialDistancing.loc[index:index])\n",
    "    fut = client.submit(runAll, tmp_params_row)\n",
    "    futures.append(fut)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Submit all futures that were not yet completed (check if files exist)\n",
    "for index in range(len(paramTable_baseParams_stopSocialDistancing)):\n",
    "    \n",
    "    tmp_params_row = copy.deepcopy(paramTable_baseParams_stopSocialDistancing.loc[index:index])\n",
    "    fut = client.submit(runAll, tmp_params_row)\n",
    "    futures.append(fut)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Then the variations on the policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Submit all futures that were not yet completed (check if files exist)\n",
    "for index in range(len(paramTable_merged)):\n",
    "    \n",
    "    tmp_params_row = copy.deepcopy(paramTable_merged.loc[index:index])\n",
    "    fut = client.submit(runAll, tmp_params_row)\n",
    "    futures.append(fut)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Monitor and save results as they arrive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Monitor and save\n",
    "seq = as_completed(futures)\n",
    "\n",
    "curIndex = 0\n",
    "\n",
    "for future in seq:\n",
    "    if future.status == \"finished\":\n",
    "        out, out_newOnly, newParams_row = future.result()\n",
    "        \n",
    "#         if not \"out_fname\" in newParams_row.columns:\n",
    "#             orig_row_ind = findRow(newParams_row.iloc[0], df = paramTable_merged, policyColumns=paramTypes[\"policy\"]+paramTypes[\"ensemble\"])\n",
    "#             newParams_row[\"out_fname\"] = paramTable_merged.loc[orig_row_ind].out_fname\n",
    "\n",
    "        # Save all the files\n",
    "        np.save(file = saveDir + newParams_row[\"out_fname\"].values[0],\n",
    "                        arr= out\n",
    "                    )\n",
    "\n",
    "        np.save(file = saveDir + newParams_row[\"out_fname\"].values[0][:-4]+\"_newOnly.npy\",\n",
    "                        arr= out_newOnly\n",
    "                    )\n",
    "            \n",
    "        curIndex += 1\n",
    "    \n",
    "        client.cancel(future) # remove from memory after saved"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.cancel(futures)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:.conda-py-corona-2-ensemble]",
   "language": "python",
   "name": "conda-env-.conda-py-corona-2-ensemble-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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}