Clean_noNPI_Cluster_Results.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "43a863f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "import re\n",
    "from plotfx import *\n",
    "import pickle\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.gridspec as gridspec\n",
    "from matplotlib.ticker import FuncFormatter\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eacc4aaa",
   "metadata": {},
   "source": [
    "### Restructure Results: No NPI, All Betas, High Community Transmission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "db54af4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sims = ['beta003_noNPI_vax_hicom', 'beta006_noNPI_vax_hicom', 'beta009_noNPI_vax_hicom']\n",
    "partitions = np.arange(5)\n",
    "sim_names = [sim + '_detailed_' + str(part) + '.csv' for sim in sims for part in partitions]\n",
    "outfile = 'allbeta_noNPI_vax_hicom'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "370caa89",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['beta003_noNPI_vax_hicom_detailed_0.csv',\n",
       " 'beta003_noNPI_vax_hicom_detailed_1.csv',\n",
       " 'beta003_noNPI_vax_hicom_detailed_2.csv',\n",
       " 'beta003_noNPI_vax_hicom_detailed_3.csv',\n",
       " 'beta003_noNPI_vax_hicom_detailed_4.csv',\n",
       " 'beta006_noNPI_vax_hicom_detailed_0.csv',\n",
       " 'beta006_noNPI_vax_hicom_detailed_1.csv',\n",
       " 'beta006_noNPI_vax_hicom_detailed_2.csv',\n",
       " 'beta006_noNPI_vax_hicom_detailed_3.csv',\n",
       " 'beta006_noNPI_vax_hicom_detailed_4.csv',\n",
       " 'beta009_noNPI_vax_hicom_detailed_0.csv',\n",
       " 'beta009_noNPI_vax_hicom_detailed_1.csv',\n",
       " 'beta009_noNPI_vax_hicom_detailed_2.csv',\n",
       " 'beta009_noNPI_vax_hicom_detailed_3.csv',\n",
       " 'beta009_noNPI_vax_hicom_detailed_4.csv']"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sim_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c6178458",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame()\n",
    "for file in sim_names:\n",
    "    tmpdf = pd.read_csv(file, header = 0, dtype = {'rep': int, 'tstep': int, 'p_vax': float, \n",
    "                                                                  'test_freq': int})\n",
    "    part = int(file[-5])\n",
    "    beta = int(file[6])\n",
    "    tmpdf['partition'] = np.repeat(part, repeats = tmpdf.shape[0])\n",
    "    tmpdf['beta'] = np.repeat(beta, repeats = tmpdf.shape[0])\n",
    "    df = df.append(tmpdf, ignore_index = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3e9506fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Linspace adds trailing decimal places in some cases\n",
    "df['p_vax'] = df['p_vax'].round(decimals = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1708ef64",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(847500, 13)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "77e97809",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>rep</th>\n",
       "      <th>p_vax</th>\n",
       "      <th>tstep</th>\n",
       "      <th>ext_inf_ct</th>\n",
       "      <th>test_freq</th>\n",
       "      <th>ia_nodes</th>\n",
       "      <th>is_nodes</th>\n",
       "      <th>test_pos_schd</th>\n",
       "      <th>test_pos_symp</th>\n",
       "      <th>q_schd</th>\n",
       "      <th>q_symp</th>\n",
       "      <th>partition</th>\n",
       "      <th>beta</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>128</td>\n",
       "      <td>0</td>\n",
       "      <td>[]</td>\n",
       "      <td>[684]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1</td>\n",
       "      <td>128</td>\n",
       "      <td>0</td>\n",
       "      <td>[]</td>\n",
       "      <td>[684]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>2</td>\n",
       "      <td>128</td>\n",
       "      <td>0</td>\n",
       "      <td>[690, 159]</td>\n",
       "      <td>[684]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3</td>\n",
       "      <td>128</td>\n",
       "      <td>0</td>\n",
       "      <td>[690, 159]</td>\n",
       "      <td>[684, 661]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4</td>\n",
       "      <td>128</td>\n",
       "      <td>0</td>\n",
       "      <td>[159]</td>\n",
       "      <td>[684, 661]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[684]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   rep  p_vax  tstep  ext_inf_ct  test_freq    ia_nodes    is_nodes  \\\n",
       "0    1    0.0      0         128          0          []       [684]   \n",
       "1    1    0.0      1         128          0          []       [684]   \n",
       "2    1    0.0      2         128          0  [690, 159]       [684]   \n",
       "3    1    0.0      3         128          0  [690, 159]  [684, 661]   \n",
       "4    1    0.0      4         128          0       [159]  [684, 661]   \n",
       "\n",
       "  test_pos_schd test_pos_symp q_schd q_symp  partition  beta  \n",
       "0            []            []     []     []          0     3  \n",
       "1            []            []     []     []          0     3  \n",
       "2            []            []     []     []          0     3  \n",
       "3            []            []     []     []          0     3  \n",
       "4            []         [684]     []     []          0     3  "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a649900d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>rep</th>\n",
       "      <th>p_vax</th>\n",
       "      <th>tstep</th>\n",
       "      <th>ext_inf_ct</th>\n",
       "      <th>test_freq</th>\n",
       "      <th>ia_nodes</th>\n",
       "      <th>is_nodes</th>\n",
       "      <th>test_pos_schd</th>\n",
       "      <th>test_pos_symp</th>\n",
       "      <th>q_schd</th>\n",
       "      <th>q_symp</th>\n",
       "      <th>partition</th>\n",
       "      <th>beta</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>847495</th>\n",
       "      <td>100</td>\n",
       "      <td>0.8</td>\n",
       "      <td>108</td>\n",
       "      <td>47</td>\n",
       "      <td>28</td>\n",
       "      <td>[346, 303]</td>\n",
       "      <td>[305, 425, 17, 651, 294, 607, 344, 462, 474, 4...</td>\n",
       "      <td>[179, 111]</td>\n",
       "      <td>[444]</td>\n",
       "      <td>[521, 609, 131, 177, 294, 300, 336, 346, 440, ...</td>\n",
       "      <td>[521, 336, 474, 607, 230]</td>\n",
       "      <td>4</td>\n",
       "      <td>9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>847496</th>\n",
       "      <td>100</td>\n",
       "      <td>0.8</td>\n",
       "      <td>109</td>\n",
       "      <td>47</td>\n",
       "      <td>28</td>\n",
       "      <td>[346, 303]</td>\n",
       "      <td>[305, 425, 17, 294, 607, 344, 462, 444, 656, 4...</td>\n",
       "      <td>[411]</td>\n",
       "      <td>[656]</td>\n",
       "      <td>[521, 609, 111, 131, 177, 179, 300, 336, 346, ...</td>\n",
       "      <td>[521, 444, 336, 474]</td>\n",
       "      <td>4</td>\n",
       "      <td>9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>847497</th>\n",
       "      <td>100</td>\n",
       "      <td>0.8</td>\n",
       "      <td>110</td>\n",
       "      <td>47</td>\n",
       "      <td>28</td>\n",
       "      <td>[346, 303]</td>\n",
       "      <td>[305, 425, 17, 607, 344, 462, 444, 656, 488, 526]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[521, 609, 111, 131, 177, 179, 300, 336, 346, ...</td>\n",
       "      <td>[521, 656, 444, 474, 230]</td>\n",
       "      <td>4</td>\n",
       "      <td>9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>847498</th>\n",
       "      <td>100</td>\n",
       "      <td>0.8</td>\n",
       "      <td>111</td>\n",
       "      <td>47</td>\n",
       "      <td>28</td>\n",
       "      <td>[346, 303, 128]</td>\n",
       "      <td>[305, 425, 17, 607, 344, 462, 444, 488, 526, 1...</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[111, 131, 177, 179, 300, 346, 411, 440, 441]</td>\n",
       "      <td>[521, 656, 444, 474]</td>\n",
       "      <td>4</td>\n",
       "      <td>9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>847499</th>\n",
       "      <td>100</td>\n",
       "      <td>0.8</td>\n",
       "      <td>112</td>\n",
       "      <td>47</td>\n",
       "      <td>28</td>\n",
       "      <td>[303, 128]</td>\n",
       "      <td>[305, 425, 17, 607, 344, 462, 444, 488, 526, 1...</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[111, 179, 300, 346, 411, 440, 441]</td>\n",
       "      <td>[521, 656, 444, 474, 230]</td>\n",
       "      <td>4</td>\n",
       "      <td>9</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        rep  p_vax  tstep  ext_inf_ct  test_freq         ia_nodes  \\\n",
       "847495  100    0.8    108          47         28       [346, 303]   \n",
       "847496  100    0.8    109          47         28       [346, 303]   \n",
       "847497  100    0.8    110          47         28       [346, 303]   \n",
       "847498  100    0.8    111          47         28  [346, 303, 128]   \n",
       "847499  100    0.8    112          47         28       [303, 128]   \n",
       "\n",
       "                                                 is_nodes test_pos_schd  \\\n",
       "847495  [305, 425, 17, 651, 294, 607, 344, 462, 474, 4...    [179, 111]   \n",
       "847496  [305, 425, 17, 294, 607, 344, 462, 444, 656, 4...         [411]   \n",
       "847497  [305, 425, 17, 607, 344, 462, 444, 656, 488, 526]            []   \n",
       "847498  [305, 425, 17, 607, 344, 462, 444, 488, 526, 1...            []   \n",
       "847499  [305, 425, 17, 607, 344, 462, 444, 488, 526, 1...            []   \n",
       "\n",
       "       test_pos_symp                                             q_schd  \\\n",
       "847495         [444]  [521, 609, 131, 177, 294, 300, 336, 346, 440, ...   \n",
       "847496         [656]  [521, 609, 111, 131, 177, 179, 300, 336, 346, ...   \n",
       "847497            []  [521, 609, 111, 131, 177, 179, 300, 336, 346, ...   \n",
       "847498            []      [111, 131, 177, 179, 300, 346, 411, 440, 441]   \n",
       "847499            []                [111, 179, 300, 346, 411, 440, 441]   \n",
       "\n",
       "                           q_symp  partition  beta  \n",
       "847495  [521, 336, 474, 607, 230]          4     9  \n",
       "847496       [521, 444, 336, 474]          4     9  \n",
       "847497  [521, 656, 444, 474, 230]          4     9  \n",
       "847498       [521, 656, 444, 474]          4     9  \n",
       "847499  [521, 656, 444, 474, 230]          4     9  "
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Renumber reps to remove duplicates\n",
    "df['rep'] = df['rep'] + df['partition']*20\n",
    "df.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "50ba43b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"RESTRUCTURE DATA\"\"\"\n",
    "\n",
    "df['ia_nodes'] = df['ia_nodes'].apply(eval)\n",
    "df['is_nodes'] = df['is_nodes'].apply(eval)\n",
    "df['all_inf'] = df['ia_nodes'] + df['is_nodes']\n",
    "df['n_inf'] = df['all_inf'].apply(set).apply(len)\n",
    "df['cum_list_inf'] = df.groupby(['beta','rep','p_vax','test_freq'])['all_inf'].apply(lambda x: x.cumsum()).apply(set)\n",
    "df['n_cum_inf'] = df['cum_list_inf'].apply(len)\n",
    "\n",
    "df['test_pos_schd'] = df['test_pos_schd'].apply(eval)\n",
    "df['test_pos_symp'] = df['test_pos_symp'].apply(eval)\n",
    "df['all_test'] = df['test_pos_schd'] + df['test_pos_symp']\n",
    "df['n_test_pos'] = df['all_test'].apply(set).apply(len)\n",
    "df['cum_list_test'] = df.groupby(['beta','rep','p_vax','test_freq'])['all_test'].apply(lambda x: x.cumsum()).apply(set)\n",
    "df['n_cum_pos'] = df['cum_list_test'].apply(len)\n",
    "        \n",
    "df['q_schd'] = df['q_schd'].apply(eval)\n",
    "df['q_symp'] = df['q_symp'].apply(eval)\n",
    "df['all_quar'] = df['q_schd'] + df['q_symp']\n",
    "df['n_quar'] = df['all_quar'].apply(set).apply(len)\n",
    "df['cum_list_quar'] = df.groupby(['beta','rep','p_vax','test_freq'])['all_quar'].apply(lambda x: x.cumsum()).apply(set)\n",
    "df['n_cum_quar'] = df['cum_list_quar'].apply(len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6baedf99",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save out reformatted dataframe\n",
    "filehandler = open(outfile + \".pkl\",\"wb\")\n",
    "pickle.dump(df,filehandler)\n",
    "filehandler.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04484cf8",
   "metadata": {},
   "source": [
    "### Restructure Results: No NPI, All Betas, Low Community Transmission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "21fe83eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "sims = ['beta003_noNPI_vax_locom', 'beta006_noNPI_vax_locom', 'beta009_noNPI_vax_locom']\n",
    "partitions = np.arange(5)\n",
    "sim_names = [sim + '_detailed_' + str(part) + '.csv' for sim in sims for part in partitions]\n",
    "outfile = 'allbeta_noNPI_vax_locom'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "76c1218f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['beta003_noNPI_vax_locom_detailed_0.csv',\n",
       " 'beta003_noNPI_vax_locom_detailed_1.csv',\n",
       " 'beta003_noNPI_vax_locom_detailed_2.csv',\n",
       " 'beta003_noNPI_vax_locom_detailed_3.csv',\n",
       " 'beta003_noNPI_vax_locom_detailed_4.csv',\n",
       " 'beta006_noNPI_vax_locom_detailed_0.csv',\n",
       " 'beta006_noNPI_vax_locom_detailed_1.csv',\n",
       " 'beta006_noNPI_vax_locom_detailed_2.csv',\n",
       " 'beta006_noNPI_vax_locom_detailed_3.csv',\n",
       " 'beta006_noNPI_vax_locom_detailed_4.csv',\n",
       " 'beta009_noNPI_vax_locom_detailed_0.csv',\n",
       " 'beta009_noNPI_vax_locom_detailed_1.csv',\n",
       " 'beta009_noNPI_vax_locom_detailed_2.csv',\n",
       " 'beta009_noNPI_vax_locom_detailed_3.csv',\n",
       " 'beta009_noNPI_vax_locom_detailed_4.csv']"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sim_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "5a2d4cd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame()\n",
    "for file in sim_names:\n",
    "    tmpdf = pd.read_csv(file, header = 0, dtype = {'rep': int, 'tstep': int, 'p_vax': float, \n",
    "                                                                  'test_freq': int})\n",
    "    part = int(file[-5])\n",
    "    beta = int(file[6])\n",
    "    tmpdf['partition'] = np.repeat(part, repeats = tmpdf.shape[0])\n",
    "    tmpdf['beta'] = np.repeat(beta, repeats = tmpdf.shape[0])\n",
    "    df = df.append(tmpdf, ignore_index = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "33aa388a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Linspace adds trailing decimal places in some cases\n",
    "df['p_vax'] = df['p_vax'].round(decimals = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "60ef0ec1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(847500, 13)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "7cad1311",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>rep</th>\n",
       "      <th>p_vax</th>\n",
       "      <th>tstep</th>\n",
       "      <th>ext_inf_ct</th>\n",
       "      <th>test_freq</th>\n",
       "      <th>ia_nodes</th>\n",
       "      <th>is_nodes</th>\n",
       "      <th>test_pos_schd</th>\n",
       "      <th>test_pos_symp</th>\n",
       "      <th>q_schd</th>\n",
       "      <th>q_symp</th>\n",
       "      <th>partition</th>\n",
       "      <th>beta</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>44</td>\n",
       "      <td>0</td>\n",
       "      <td>[]</td>\n",
       "      <td>[684]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1</td>\n",
       "      <td>44</td>\n",
       "      <td>0</td>\n",
       "      <td>[]</td>\n",
       "      <td>[684]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>2</td>\n",
       "      <td>44</td>\n",
       "      <td>0</td>\n",
       "      <td>[]</td>\n",
       "      <td>[684]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3</td>\n",
       "      <td>44</td>\n",
       "      <td>0</td>\n",
       "      <td>[403]</td>\n",
       "      <td>[684, 661]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4</td>\n",
       "      <td>44</td>\n",
       "      <td>0</td>\n",
       "      <td>[403]</td>\n",
       "      <td>[684, 661]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[684]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   rep  p_vax  tstep  ext_inf_ct  test_freq ia_nodes    is_nodes  \\\n",
       "0    1    0.0      0          44          0       []       [684]   \n",
       "1    1    0.0      1          44          0       []       [684]   \n",
       "2    1    0.0      2          44          0       []       [684]   \n",
       "3    1    0.0      3          44          0    [403]  [684, 661]   \n",
       "4    1    0.0      4          44          0    [403]  [684, 661]   \n",
       "\n",
       "  test_pos_schd test_pos_symp q_schd q_symp  partition  beta  \n",
       "0            []            []     []     []          0     3  \n",
       "1            []            []     []     []          0     3  \n",
       "2            []            []     []     []          0     3  \n",
       "3            []            []     []     []          0     3  \n",
       "4            []         [684]     []     []          0     3  "
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8a2197a4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>rep</th>\n",
       "      <th>p_vax</th>\n",
       "      <th>tstep</th>\n",
       "      <th>ext_inf_ct</th>\n",
       "      <th>test_freq</th>\n",
       "      <th>ia_nodes</th>\n",
       "      <th>is_nodes</th>\n",
       "      <th>test_pos_schd</th>\n",
       "      <th>test_pos_symp</th>\n",
       "      <th>q_schd</th>\n",
       "      <th>q_symp</th>\n",
       "      <th>partition</th>\n",
       "      <th>beta</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>847495</th>\n",
       "      <td>100</td>\n",
       "      <td>0.8</td>\n",
       "      <td>108</td>\n",
       "      <td>10</td>\n",
       "      <td>28</td>\n",
       "      <td>[]</td>\n",
       "      <td>[266]</td>\n",
       "      <td>[74]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[218, 359, 489, 248]</td>\n",
       "      <td>[]</td>\n",
       "      <td>4</td>\n",
       "      <td>9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>847496</th>\n",
       "      <td>100</td>\n",
       "      <td>0.8</td>\n",
       "      <td>109</td>\n",
       "      <td>10</td>\n",
       "      <td>28</td>\n",
       "      <td>[]</td>\n",
       "      <td>[266]</td>\n",
       "      <td>[500, 49]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[74, 218, 359, 489, 248]</td>\n",
       "      <td>[]</td>\n",
       "      <td>4</td>\n",
       "      <td>9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>847497</th>\n",
       "      <td>100</td>\n",
       "      <td>0.8</td>\n",
       "      <td>110</td>\n",
       "      <td>10</td>\n",
       "      <td>28</td>\n",
       "      <td>[]</td>\n",
       "      <td>[266]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[266]</td>\n",
       "      <td>[49, 74, 218, 359, 489, 248]</td>\n",
       "      <td>[]</td>\n",
       "      <td>4</td>\n",
       "      <td>9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>847498</th>\n",
       "      <td>100</td>\n",
       "      <td>0.8</td>\n",
       "      <td>111</td>\n",
       "      <td>10</td>\n",
       "      <td>28</td>\n",
       "      <td>[]</td>\n",
       "      <td>[266]</td>\n",
       "      <td>[221]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[49, 74, 218, 359, 500, 248]</td>\n",
       "      <td>[266]</td>\n",
       "      <td>4</td>\n",
       "      <td>9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>847499</th>\n",
       "      <td>100</td>\n",
       "      <td>0.8</td>\n",
       "      <td>112</td>\n",
       "      <td>10</td>\n",
       "      <td>28</td>\n",
       "      <td>[]</td>\n",
       "      <td>[266]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[]</td>\n",
       "      <td>[49, 74, 218, 221, 359, 489, 500, 248]</td>\n",
       "      <td>[266]</td>\n",
       "      <td>4</td>\n",
       "      <td>9</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        rep  p_vax  tstep  ext_inf_ct  test_freq ia_nodes is_nodes  \\\n",
       "847495  100    0.8    108          10         28       []    [266]   \n",
       "847496  100    0.8    109          10         28       []    [266]   \n",
       "847497  100    0.8    110          10         28       []    [266]   \n",
       "847498  100    0.8    111          10         28       []    [266]   \n",
       "847499  100    0.8    112          10         28       []    [266]   \n",
       "\n",
       "       test_pos_schd test_pos_symp                                  q_schd  \\\n",
       "847495          [74]            []                    [218, 359, 489, 248]   \n",
       "847496     [500, 49]            []                [74, 218, 359, 489, 248]   \n",
       "847497            []         [266]            [49, 74, 218, 359, 489, 248]   \n",
       "847498         [221]            []            [49, 74, 218, 359, 500, 248]   \n",
       "847499            []            []  [49, 74, 218, 221, 359, 489, 500, 248]   \n",
       "\n",
       "       q_symp  partition  beta  \n",
       "847495     []          4     9  \n",
       "847496     []          4     9  \n",
       "847497     []          4     9  \n",
       "847498  [266]          4     9  \n",
       "847499  [266]          4     9  "
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Renumber reps to remove duplicates\n",
    "df['rep'] = df['rep'] + df['partition']*20\n",
    "df.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "7a35c32d",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"RESTRUCTURE DATA\"\"\"\n",
    "\n",
    "df['ia_nodes'] = df['ia_nodes'].apply(eval)\n",
    "df['is_nodes'] = df['is_nodes'].apply(eval)\n",
    "df['all_inf'] = df['ia_nodes'] + df['is_nodes']\n",
    "df['n_inf'] = df['all_inf'].apply(set).apply(len)\n",
    "df['cum_list_inf'] = df.groupby(['beta','rep','p_vax','test_freq'])['all_inf'].apply(lambda x: x.cumsum()).apply(set)\n",
    "df['n_cum_inf'] = df['cum_list_inf'].apply(len)\n",
    "\n",
    "df['test_pos_schd'] = df['test_pos_schd'].apply(eval)\n",
    "df['test_pos_symp'] = df['test_pos_symp'].apply(eval)\n",
    "df['all_test'] = df['test_pos_schd'] + df['test_pos_symp']\n",
    "df['n_test_pos'] = df['all_test'].apply(set).apply(len)\n",
    "df['cum_list_test'] = df.groupby(['beta','rep','p_vax','test_freq'])['all_test'].apply(lambda x: x.cumsum()).apply(set)\n",
    "df['n_cum_pos'] = df['cum_list_test'].apply(len)\n",
    "        \n",
    "df['q_schd'] = df['q_schd'].apply(eval)\n",
    "df['q_symp'] = df['q_symp'].apply(eval)\n",
    "df['all_quar'] = df['q_schd'] + df['q_symp']\n",
    "df['n_quar'] = df['all_quar'].apply(set).apply(len)\n",
    "df['cum_list_quar'] = df.groupby(['beta','rep','p_vax','test_freq'])['all_quar'].apply(lambda x: x.cumsum()).apply(set)\n",
    "df['n_cum_quar'] = df['cum_list_quar'].apply(len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f267958",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save out reformatted dataframe\n",
    "filehandler = open(outfile + \".pkl\",\"wb\")\n",
    "pickle.dump(df,filehandler)\n",
    "filehandler.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}