MAGE_GA_COVID19.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load dependencies\n",
    "#load in dependencies to read data\n",
    "using CSV, Plots, DataFrames, Dates\n",
    "using LinearAlgebra,Pkg\n",
    "#using DifferentialEquations # do not need the whole package necessarilly. Using just the ODE functions will save some memory\n",
    "using OrdinaryDiffEq\n",
    "#show versions\n",
    "Pkg.status()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "versioninfo()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Spatial Epidemic Model for GA for COVID-19\n",
    "\n",
    "Here, we extend the model to look at infection spread through a spatial metapopulation model.\n",
    "To do this, we consider that individuals spend time in the local community, but also may commute - for work, errands, social activities or other things. Following this we split the population $N_{tot}$ into two compartments $N_{local}$ and $N_{commute}$. $N_{local}$ is the population who live in a particular county, whilst $N_{commute} = N_{local} - N_{leave} + N_{arrive}$. Here $N_{leave}$ are those who leave the county, whilst $N_{arrive}$ are those who travel into the county. We do not explicitly account for those arriving or leaving in the local ODE's.\n",
    "\n",
    "We assume that those in the county see the local population for a proportion of time $w_{local}$ and the commuter population for a proportion of time $w_{commmute}$, where we assume $w_{local} + w_{commute} = 1$.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The system of nonlinear differential equations governing this patch and age-structured epidemiological for patch $i$ and age $a$ are as follows:\n",
    "\n",
    "\\begin{align}\n",
    "\\frac{dS(i,a)}{dt} & =  -\\overset{new\\ cases\\ from\\ symptomatic\\ contact}{\\overbrace{\\kappa\\beta_S S(i,a) \\left( w_{local}\\frac{ I_{S,local}(i)}{N_{local}(i)} + w_{commute}\\frac{ I_{S,commute}(i)}{N_{commute}(i)} \\right) }}  - \\overset{new\\ cases\\ from\\ asymptomatic\\ contact}{\\overbrace{\\kappa\\beta_A S(i,a) \\left( w_{local}\\frac{ I_{A,local}(i)}{N_{local}(i)} + w_{commute}\\frac{ I_{A,commute}(i)}{N_{commute}(i)} \\right) }}  \\\\ \\\\\n",
    "\\frac{dE(i,a)}{dt} & =  \\overset{new\\ cases\\ from\\ symptomatic\\ contact}{\\overbrace{\\kappa\\beta_S S(i,a) \\left( w_{local}\\frac{ I_{S,local}(i)}{N_{local}(i)} + w_{commute}\\frac{ I_{S,commute}(i)}{N_{commute}(i)} \\right) }}  + \\overset{new\\ cases\\ from\\ asymptomatic\\ contact}{\\overbrace{\\kappa\\beta_A S(i,a) \\left( w_{local}\\frac{ I_{A,local}(i)}{N_{local}(i))} + w_{commute}\\frac{ I_{A,commute}(i)}{N_{commute}(i)} \\right) }}  - \\overset{onset\\ of\\ infectiousness}{\\overbrace{\\gamma_E E(i,a)}} \\\\\n",
    "\\frac{dI_A(i,a)}{dt} & =  \\overset{asymptomatic\\ onset}{\\overbrace{p\\gamma_E E(i,a)}} - \\overset{recovery}{\\overbrace{\\gamma_A I_A(i,a)}} \\\\\n",
    "\\frac{dI_S(i,a)}{dt} & =  \\overset{symptomatic\\ onset}{\\overbrace{(1-p)\\gamma_E E(i,a)}} - \\overset{transfer\\ from\\ I_S}{\\overbrace{\\gamma_S I_S(i,a)}} \\\\\n",
    "\\frac{dI_{h,sub}(i,a)}{dt} & =  \\overset{subcritical\\ cases}{\\overbrace{h(a)(1-\\xi(a)) \\gamma_S I_S(i,a)}} - \\overset{transfer\\ from\\ I_{h,sub}}{\\overbrace{\\gamma_h I_{h,sub}(i,a)}} \\\\\n",
    "\\frac{dI_{h,crit}(i,a)}{dt} & =  \\overset{critical\\ (ICU)\\ cases}{\\overbrace{h(a)\\xi(a) \\gamma_S I_S(i,a)}} - \\overset{transfer\\ from\\ I_{h,crit}}{\\overbrace{\\gamma_h I_{h,crit}(i,a)}} \\\\\n",
    "\\frac{dR(a)}{dt} & =  \\overset{recovery\\ from\\ I_A}{\\overbrace{\\gamma_A I_A(i,a)}} + \\overset{recovery\\ from\\ I_S}{\\overbrace{(1-h(a))\\gamma_S I_S(i,a)}} + \\overset{recovery\\ from\\ I_{h,sub}}{\\overbrace{\\gamma_h I_{h,sub}(i,a)}} + \\overset{recovery\\ from\\ I_{h,crit}}{\\overbrace{(1-\\mu)\\gamma_h I_{h,crit}(i,a)}} \\\\\n",
    "\\frac{dD(i,a)}{dt} & = \\overset{mortality}{\\overbrace{\\mu \\gamma_h I_{h,crit}(i,a)}}\n",
    "\\end{align}\n",
    "\n",
    "where $I_{S,local}$ and $I_{S,commute}$ are the total number (across all age classes) of sympomatic infectious individuals; $I_{A,local}$ and $I_{S,commute}$ are the total number of asymptomatic infectious individuals; and $N_{local}$ and $N_{commute}$ are the total number of alive individuals during local and commute times."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Read in functions called in the scripts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#ODE model code\n",
    "#assumes age-dependent journeys between counties\n",
    "function SpatialAgeStructuredEpidemiologicalModelMixAge!(du,u,par,t)\n",
    "    SDI,β_a,β_s,p,γ_e,γ_a,γ_s,γ_h,h,ξ,μ,S_ind,E_ind,\n",
    "    Ia_ind,Is_ind,Ihsub_ind,Ihcrit_ind,R_ind,D_ind,\n",
    "    NumPatches,NumAgeClasses,patchstates,patchstatesALIVE,\n",
    "    patchstatesMOVEABLE,w,wy,wo,transport,ISMOVEFACTOR,w1,w2 = par\n",
    "    \n",
    "    #NumPatches is the number of spatial \"nodes\"\n",
    "    #w is a matrix of size NxN determining flows between nodes.\n",
    "    #transport is a matrix to store movements vectors (don't want to recreate this object throughout this function)\n",
    "    #patchstates - total number of states in a patch\n",
    "    #patchstatesALIVE - number of non-dead states\n",
    "    #patchstatesMOVEABLE - number of states that are allowed to move between patches\n",
    "    #S_ind,E_ind,Ia_ind,Is_ind,Ihsub_ind,Ihcrit_ind,R_ind,D_ind - indexes of the states for the first patch i.e. number 1:80\n",
    "\n",
    "    \n",
    "    #have matrix w, representing movement of commuters i.e. X1 stay in county, X2 go to county 2, X3 go to county3\n",
    "    #total travel is fraction travelled applied to people in moving states\n",
    "    #some amounts of these will be IA and IS - keep track!\n",
    "    #want to know N,Is,IA leaving each county; and arriving at each county\n",
    "    #w11 S1 + w12 S2  #net change for each state so long as diagonal is -ve.\n",
    "\n",
    "    #w is those aged 30-59\n",
    "    #wy is those 0-29\n",
    "    #wo is those 60+\n",
    "    \n",
    "    #Stage 1:\n",
    "    #calculate transport between patches\n",
    "    \n",
    "    #transport = transport .* 0.0;\n",
    "    #want this to be age structured\n",
    "    for ii = 1:patchstatesMOVEABLE #for each age class of alive states that can move.\n",
    "        INDS = ii:patchstates:(ii+(NumPatches-1)*patchstates ) #gives 1xpatch indices for transport multiplication for each age class and each state that can move\n",
    "        MOD = mod(ii,NumAgeClasses-1)\n",
    "        if MOD>0 & MOD<4\n",
    "        transport[:,ii] = wy*u[INDS] \n",
    "        elseif MOD==0  || MOD >=7\n",
    "        transport[:,ii] = wo*u[INDS]\n",
    "        else\n",
    "        transport[:,ii] =  w*u[INDS]  #add the 1*patch vector of movement fluxes to transport\n",
    "        end\n",
    "    end\n",
    "    \n",
    "    #modify tranport matrix to account for Is people deciding to move or not.\n",
    "    transport[:,Is_ind] = ISMOVEFACTOR .* transport[:,Is_ind]\n",
    "\n",
    "    #calculate net change for each patch for N,Ia and Is\n",
    "    C_TOT_NETCHANGE = sum(transport,dims=2)\n",
    "    C_I_A_NETCHANGE = sum(transport[:,Ia_ind],dims=2)\n",
    "    C_I_S_NETCHANGE = sum(transport[:,Is_ind],dims=2)\n",
    "    \n",
    "    #Stage 2:\n",
    "    #Go through the ODE system node by node. As have found fluxes between nodes already, this will be a synchronous update of local and spatial dynamics.\n",
    "\n",
    "    for aa=1:(Int64(NumPatches)) #for each patch node\n",
    "        \n",
    "        #Local patch dynamics depend on patch quantitites\n",
    "        NTOT_Loc = sum(u[((aa-1)*patchstates+1):((aa-1)*patchstates+patchstatesALIVE)]) #total alive\n",
    "        TOTI_a_Loc = sum(u[(aa-1)*patchstates .+ Ia_ind]) #total asymptomatic infected\n",
    "        TOTI_s_Loc = sum(u[(aa-1)*patchstates .+ Is_ind]) #total symptomatic infected\n",
    "        \n",
    "        NTOT_Com = NTOT_Loc + C_TOT_NETCHANGE[aa]#total alive\n",
    "        TOTI_a_Com = TOTI_a_Loc + C_I_A_NETCHANGE[aa] #total asymptomatic infected\n",
    "        TOTI_s_Com = TOTI_s_Loc + C_I_S_NETCHANGE[aa] #total symptomatic infected\n",
    "        \n",
    "        LOCAL = w1/(NTOT_Loc)\n",
    "        COMMUTE = w2/(NTOT_Com)\n",
    "        propI_A = ( TOTI_a_Loc*LOCAL + TOTI_a_Com*COMMUTE )\n",
    "        propI_S = ( TOTI_s_Loc*LOCAL + TOTI_s_Com*COMMUTE )\n",
    "        \n",
    "        \n",
    "        for bb = 1:(Int64(NumAgeClasses))  #for each age class in this patch\n",
    "        INDS = (aa-1)*patchstates .+ [S_ind[bb],E_ind[bb],Ia_ind[bb],Is_ind[bb],R_ind[bb],Ihsub_ind[bb],Ihcrit_ind[bb],D_ind[bb]]   \n",
    "            \n",
    "            \n",
    "            S,E,Ia,Is,R,Ihsub,Ihcrit,D = u[INDS]\n",
    "\n",
    "            du[INDS[1]] =  -β_a*SDI*S*propI_A - β_s*SDI*S*propI_S   #dS/dt\n",
    "            du[INDS[2]] =  +β_a*SDI*S*propI_A + β_s*SDI*S*propI_S - γ_e*E  #dE/dt\n",
    "            du[INDS[3]] = p[bb]*γ_e*E - γ_a*Ia   #dI_a /dt\n",
    "            du[INDS[4]] = (1-p[bb])*γ_e*E - γ_s*Is  #dI_s/dt\n",
    "            du[INDS[5]] = γ_a*Ia + (1-h[bb])*γ_s*Is + γ_h*Ihsub + (1-μ[bb])*γ_h*Ihcrit #dR/dt\n",
    "            du[INDS[6]] = h[bb]*(1-ξ[bb])*γ_s*Is - γ_h*Ihsub #dIhsub/dt\n",
    "            du[INDS[7]] = h[bb]*ξ[bb]*γ_s*Is -  γ_h*Ihcrit #dIhcrit/dt\n",
    "            du[INDS[8]] = μ[bb]*γ_h*Ihcrit #dD/dt\n",
    "        end\n",
    "    end\n",
    "    \n",
    "end\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#ODE model code - this is used to model at the level of the whole of Georgia! (helpful in setting initial conditions)\n",
    "function SingleNode!(du,u,par,t)\n",
    "    SDI,β_a,β_s,p,γ_e,γ_a,γ_s,γ_h,h,ξ,μ,S_ind,E_ind,Ia_ind,Is_ind,R_ind,Ihsub_ind,Ihcrit_ind,D_ind = par\n",
    "    \n",
    "    TOTI_a = sum(u[Ia_ind]) #Total asymptomatic infectious individuals\n",
    "    TOTI_s = sum(u[Is_ind]) #Total symptomatic infectious individuals\n",
    "    NTOT = sum(u[1:63]) #All individuals not in the dead (D) state\n",
    "\n",
    "    \n",
    "    for aa = 1:length(S_ind)#for each age class\n",
    "        INDS = [S_ind[aa],E_ind[aa],Ia_ind[aa],Is_ind[aa],R_ind[aa],Ihsub_ind[aa],Ihcrit_ind[aa],D_ind[aa]]\n",
    "        S,E,Ia,Is,R,Ihsub,Ihcrit,D = u[INDS]\n",
    "        du[INDS[1]] = -β_a*SDI*S*TOTI_a/(NTOT) - β_s*SDI*S*TOTI_s/(NTOT) #dS/dt\n",
    "        du[INDS[2]] = +β_a*SDI*S*TOTI_a/(NTOT) + β_s*SDI*S*TOTI_s/(NTOT) -  γ_e*E   #dE/dt\n",
    "        du[INDS[3]] = p[aa]*γ_e*E - γ_a*Ia  #dI_a /dt\n",
    "        du[INDS[4]] = (1-p[aa])*γ_e*E - γ_s*Is  #dI_s/dt\n",
    "        du[INDS[5]] = γ_a*Ia + (1-h[aa])*γ_s*Is + γ_h*Ihsub + (1-μ[aa])*γ_h*Ihcrit  #dR/dt\n",
    "        du[INDS[6]] = h[aa]*(1-ξ[aa])*γ_s*Is - γ_h*Ihsub #dIhsub/dt\n",
    "        du[INDS[7]] = h[aa]*ξ[aa]*γ_s*Is -  γ_h*Ihcrit #dIhcrit/dt\n",
    "        du[INDS[8]] = μ[aa]*γ_h*Ihcrit #dD/dt\n",
    "    end\n",
    " \n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#reads in ODE solution; and saves output summary figures and CSV files\n",
    "function WriteOutput!(PatchFIPS,PatchNames,Regions,GAPOP,TODAY,TIMERESOLTION,T_END,sol,pars,SaveName)\n",
    "    SDI,β_a,β_s,p,γ_e,γ_a,γ_s,γ_h,h,ξ,μ,S_ind,E_ind,\n",
    "        NumPatches,NumAgeClasses,patchstates,patchstatesALIVE,\n",
    "        patchstatesMOVEABLE,w,wy,wo,transport,ISMOVEFACTOR,w1,w2 = pars \n",
    "    \n",
    "    Gath = gatherStates!(sol,pars)\n",
    "    STATES = GAPOP * Gath[1] #organise the solution\n",
    "    Times = sol.t[Gath[2]]\n",
    "        \n",
    "    \n",
    "    #patches  x EpiStates x AgeClasses x Time\n",
    "    SIZES = size(STATES) \n",
    "    \n",
    "    #Make summary tables for each epi state, and total hospitalised\n",
    "    #S,E,Ia,Is,R,Ihsub,Ihcrit,D, Ihsub+Ihcrit\n",
    "    S = zeros(SIZES[1],SIZES[4]) # pacthes x time\n",
    "    E = zeros(SIZES[1],SIZES[4]) # pacthes x time\n",
    "    Ia = zeros(SIZES[1],SIZES[4]) # pacthes x time\n",
    "    Is = zeros(SIZES[1],SIZES[4]) # pacthes x time\n",
    "    R = zeros(SIZES[1],SIZES[4]) # pacthes x time\n",
    "    Ihsub = zeros(SIZES[1],SIZES[4]) # pacthes x time\n",
    "    Ihcrit = zeros(SIZES[1],SIZES[4]) # pacthes x time\n",
    "    D = zeros(SIZES[1],SIZES[4]) # pacthes x time\n",
    "    Htot = zeros(SIZES[1],SIZES[4]) # pacthes x time\n",
    "    Dperday = zeros(SIZES[1],SIZES[4]) # pacthes x time\n",
    "    NewCperday = zeros(SIZES[1],SIZES[4]) # pacthes x time\n",
    "    \n",
    "    for aa = 1:SIZES[1] #for each patch\n",
    "        S[aa,:] = dropdims(sum(STATES[aa,1,1:9,:],dims=1),dims=1)\n",
    "        E[aa,:] = dropdims(sum(STATES[aa,2,1:9,:],dims=1),dims=1)\n",
    "        Ia[aa,:] = dropdims(sum(STATES[aa,3,1:9,:],dims=1),dims=1)\n",
    "        Is[aa,:] = dropdims(sum(STATES[aa,4,1:9,:],dims=1),dims=1)\n",
    "        R[aa,:] = dropdims(sum(STATES[aa,5,1:9,:],dims=1),dims=1)\n",
    "        Ihsub[aa,:] = dropdims(sum(STATES[aa,6,1:9,:],dims=1),dims=1)\n",
    "        Ihcrit[aa,:] = dropdims(sum(STATES[aa,7,1:9,:],dims=1),dims=1)\n",
    "        D[aa,:] = dropdims(sum(STATES[aa,8,1:9,:],dims=1),dims=1)\n",
    "        Htot[aa,:] =  Ihsub[aa,:] .+ Ihcrit[aa,:]\n",
    "        for bb = 1:SIZES[3] #for age classes\n",
    "            #Dperday[aa,:] .+=  μ[bb].*γ_h .*STATES[aa,7,bb,:].*1.0;  #deaths per day\n",
    "            Dperday[aa,:] .+= [( STATES[aa,8,bb,2:end] .- STATES[aa,8,bb,1:(end-1)]) ./ (Times[2:end] - Times[1:(end-1)]); 0]\n",
    "            NewCperday[aa,:] .+= [( STATES[aa,1,bb,1:(end-1)] .- STATES[aa,1,bb,2:end]) ./ (Times[2:end] - Times[1:(end-1)]); 0]\n",
    "        end\n",
    "    end\n",
    "    \n",
    "    #build Time column names\n",
    "    colnames = Array{Date,1}(undef,1)\n",
    "    colnames[1] = TODAY\n",
    "    for aa = 2:(SIZES[4])#for each time\n",
    "        push!(colnames,colnames[aa-1]+Dates.Day(convert(Int,Times[aa]-Times[aa-1])))\n",
    "    end\n",
    "\n",
    "\n",
    "    #build dataframes\n",
    "    SSdata = DataFrame(S)\n",
    "    rename!(SSdata , Symbol.(colnames))\n",
    "    insertcols!(SSdata, 1, :FIP => PatchFIPS)\n",
    "    insertcols!(SSdata, 2, :County => PatchNames)\n",
    "    \n",
    "    EEdata = DataFrame(E)\n",
    "    rename!(EEdata , Symbol.(colnames))\n",
    "    insertcols!(EEdata, 1, :FIP => PatchFIPS)\n",
    "    insertcols!(EEdata, 2, :County => PatchNames)\n",
    "    \n",
    "    Iadata = DataFrame(Ia)\n",
    "    rename!(Iadata , Symbol.(colnames))\n",
    "    insertcols!(Iadata, 1, :FIP => PatchFIPS)\n",
    "    insertcols!(Iadata, 2, :County => PatchNames)\n",
    "    \n",
    "    Isdata = DataFrame(Is)\n",
    "    rename!(Isdata , Symbol.(colnames))\n",
    "    insertcols!(Isdata, 1, :FIP => PatchFIPS)\n",
    "    insertcols!(Isdata, 2, :County => PatchNames)\n",
    "   \n",
    "    RRdata = DataFrame(R)\n",
    "    rename!(RRdata , Symbol.(colnames))\n",
    "    insertcols!(RRdata, 1, :FIP => PatchFIPS)\n",
    "    insertcols!(RRdata, 2, :County => PatchNames)\n",
    "    \n",
    "    Ihsubdata = DataFrame(Ihsub)\n",
    "    rename!(Ihsubdata , Symbol.(colnames))\n",
    "    insertcols!(Ihsubdata, 1, :FIP => PatchFIPS)\n",
    "    insertcols!(Ihsubdata, 2, :County => PatchNames)\n",
    "    \n",
    "    Ihcritdata = DataFrame(Ihcrit)\n",
    "    rename!(Ihcritdata , Symbol.(colnames))\n",
    "    insertcols!(Ihcritdata, 1, :FIP => PatchFIPS)\n",
    "    insertcols!(Ihcritdata, 2, :County => PatchNames)\n",
    "    \n",
    "    DDdata = DataFrame(D)\n",
    "    rename!(DDdata , Symbol.(colnames))\n",
    "    insertcols!(DDdata, 1, :FIP => PatchFIPS)\n",
    "    insertcols!(DDdata, 2, :County => PatchNames)\n",
    "    \n",
    "    HHdata = DataFrame(Htot)\n",
    "    rename!(HHdata , Symbol.(colnames))\n",
    "    insertcols!(HHdata, 1, :FIP => PatchFIPS)\n",
    "    insertcols!(HHdata, 2, :County => PatchNames)\n",
    "    \n",
    "    DpDdata = DataFrame(Dperday)\n",
    "    rename!(DpDdata , Symbol.(colnames))\n",
    "    insertcols!(DpDdata, 1, :FIP => PatchFIPS)\n",
    "    insertcols!(DpDdata, 2, :County => PatchNames)\n",
    "    \n",
    "    CpDdata = DataFrame(NewCperday)\n",
    "    rename!(CpDdata , Symbol.(colnames))\n",
    "    insertcols!(CpDdata, 1, :FIP => PatchFIPS)\n",
    "    insertcols!(CpDdata, 2, :County => PatchNames)\n",
    "    \n",
    "    #make a directory to store files\n",
    "    mkdir(SaveName)\n",
    "    \n",
    "    #write data to file\n",
    "    CSV.write(string(SaveName,\"/GA_county_Susceptibles.csv\"),SSdata)\n",
    "    CSV.write(string(SaveName,\"/GA_county_Exposed.csv\"),EEdata)\n",
    "    CSV.write(string(SaveName,\"/GA_county_AsymptInfected.csv\"),Iadata)\n",
    "    CSV.write(string(SaveName,\"/GA_county_SymptInfected.csv\"),Isdata)\n",
    "    CSV.write(string(SaveName,\"/GA_county_Recovered.csv\"),RRdata)\n",
    "    CSV.write(string(SaveName,\"/GA_county_subacuteHospitalised.csv\"),Ihsubdata)\n",
    "    CSV.write(string(SaveName,\"/GA_county_criticalHospitalised.csv\"),Ihcritdata)\n",
    "    CSV.write(string(SaveName,\"/GA_county_Killed.csv\"),DDdata)\n",
    "    CSV.write(string(SaveName,\"/GA_county_totalHospitalised.csv\"),HHdata)\n",
    "    CSV.write(string(SaveName,\"/GA_county_deathrateperday.csv\"),DpDdata)\n",
    "    CSV.write(string(SaveName,\"/GA_county_newcaserateperday.csv\"),CpDdata)\n",
    "    \n",
    "    ##save summary data at RegionalLevel\n",
    "    \n",
    "    RegionList = sort(unique(Regions))\n",
    "    SSdataR = zeros(length(RegionList),SIZES[4])\n",
    "    EEdataR = zeros(length(RegionList),SIZES[4])\n",
    "    IadataR = zeros(length(RegionList),SIZES[4])\n",
    "    IsdataR = zeros(length(RegionList),SIZES[4])\n",
    "    IhsubdataR = zeros(length(RegionList),SIZES[4])\n",
    "    IhcritdataR = zeros(length(RegionList),SIZES[4])\n",
    "    RRdataR = zeros(length(RegionList),SIZES[4])\n",
    "    DDdataR = zeros(length(RegionList),SIZES[4])\n",
    "    HHdataR = zeros(length(RegionList),SIZES[4])\n",
    "    DpDdataR = zeros(length(RegionList),SIZES[4])\n",
    "    CpDdataR = zeros(length(RegionList),SIZES[4])\n",
    "     \n",
    "    for aa = 1:length(RegionList) #for each region\n",
    "        RegionInd = findall(Regions.== RegionList[aa])\n",
    "        SSdataR[aa,:] = sum(S[RegionInd,:],dims=1)\n",
    "        EEdataR[aa,:] = sum(E[RegionInd,:],dims=1)\n",
    "        IadataR[aa,:] = sum(Ia[RegionInd,:],dims=1)\n",
    "        IsdataR[aa,:] = sum(Is[RegionInd,:],dims=1)\n",
    "        RRdataR[aa,:] = sum(R[RegionInd,:],dims=1)\n",
    "        DDdataR[aa,:] = sum(D[RegionInd,:],dims=1)\n",
    "        IhsubdataR[aa,:] = sum(Ihsub[RegionInd,:],dims=1)\n",
    "        IhcritdataR[aa,:] = sum(Ihcrit[RegionInd,:],dims=1)\n",
    "        HHdataR[aa,:] = sum(Htot[RegionInd,:],dims=1)\n",
    "        DpDdataR[aa,:] = sum(Dperday[RegionInd,:],dims=1)\n",
    "        CpDdataR[aa,:] = sum(NewCperday[RegionInd,:],dims=1)\n",
    "    end\n",
    "    \n",
    "    SSdataRR = DataFrame(SSdataR)\n",
    "    rename!(SSdataRR , Symbol.(colnames))\n",
    "    insertcols!(SSdataRR, 1, :Region => RegionList)\n",
    "    \n",
    "    EEdataRR = DataFrame(EEdataR)\n",
    "    rename!(EEdataRR , Symbol.(colnames))\n",
    "    insertcols!(EEdataRR, 1, :Region => RegionList)\n",
    "    \n",
    "    IadataRR = DataFrame(IadataR)\n",
    "    rename!(IadataRR , Symbol.(colnames))\n",
    "    insertcols!(IadataRR, 1, :Region => RegionList)\n",
    "    \n",
    "    IsdataRR = DataFrame(IsdataR)\n",
    "    rename!(IsdataRR , Symbol.(colnames))\n",
    "    insertcols!(IsdataRR, 1, :Region => RegionList)\n",
    "   \n",
    "    RRdataRR = DataFrame(RRdataR)\n",
    "    rename!(RRdataRR , Symbol.(colnames))\n",
    "    insertcols!(RRdataRR, 1, :Region => RegionList)\n",
    "    \n",
    "    IhsubdataRR = DataFrame(IhsubdataR)\n",
    "    rename!(IhsubdataRR , Symbol.(colnames))\n",
    "    insertcols!(IhsubdataRR, 1, :Region => RegionList)\n",
    "    \n",
    "    IhcritdataRR = DataFrame(IhcritdataR)\n",
    "    rename!(IhcritdataRR , Symbol.(colnames))\n",
    "    insertcols!(IhcritdataRR, 1, :Region => RegionList)\n",
    "    \n",
    "    DDdataRR = DataFrame(DDdataR)\n",
    "    rename!(DDdataRR , Symbol.(colnames))\n",
    "    insertcols!(DDdataRR, 1, :Region => RegionList)\n",
    "    \n",
    "    HHdataRR = DataFrame(HHdataR)\n",
    "    rename!(HHdataRR , Symbol.(colnames))\n",
    "    insertcols!(HHdataRR, 1, :Region => RegionList)\n",
    "    \n",
    "    DpDdataRR = DataFrame(DpDdataR)\n",
    "    rename!(DpDdataRR , Symbol.(colnames))\n",
    "    insertcols!(DpDdataRR, 1, :Region => RegionList)\n",
    "    \n",
    "    CpDdataRR = DataFrame(CpDdataR)\n",
    "    rename!(CpDdataRR , Symbol.(colnames))\n",
    "    insertcols!(CpDdataRR, 1, :Region => RegionList)\n",
    "        \n",
    "    \n",
    "    CSV.write(string(SaveName,\"/GA_region_Susceptibles.csv\"),SSdataRR)\n",
    "    CSV.write(string(SaveName,\"/GA_region_Exposed.csv\"),EEdataRR)\n",
    "    CSV.write(string(SaveName,\"/GA_region_AsymptInfected.csv\"),IadataRR)\n",
    "    CSV.write(string(SaveName,\"/GA_region_SymptInfected.csv\"),IsdataRR)\n",
    "    CSV.write(string(SaveName,\"/GA_region_Recovered.csv\"),RRdataRR)\n",
    "    CSV.write(string(SaveName,\"/GA_region_subacuteHospitalised.csv\"),IhsubdataRR)\n",
    "    CSV.write(string(SaveName,\"/GA_region_criticalHospitalised.csv\"),IhcritdataRR)\n",
    "    CSV.write(string(SaveName,\"/GA_region_Killed.csv\"),DDdataRR)\n",
    "    CSV.write(string(SaveName,\"/GA_region_totalHospitalised.csv\"),HHdataRR)\n",
    "    CSV.write(string(SaveName,\"/GA_region_deathrateperday.csv\"),DpDdataRR)\n",
    "    CSV.write(string(SaveName,\"/GA_region_newcaserateperday.csv\"),CpDdataRR) \n",
    "    \n",
    "    \n",
    "    \n",
    "    #Save summary data at County Level\n",
    "    #sum across counties\n",
    "    GASummary = zeros(11,SIZES[4])\n",
    "    GASummary[1,:] = sum(S,dims=1)\n",
    "    GASummary[2,:] = sum(E,dims=1)\n",
    "    GASummary[3,:] = sum(Ia,dims=1)\n",
    "    GASummary[4,:] = sum(Is,dims=1)\n",
    "    GASummary[5,:] = sum(R,dims=1)\n",
    "    GASummary[6,:] = sum(Ihsub,dims=1)\n",
    "    GASummary[7,:] = sum(Ihcrit,dims=1)\n",
    "    GASummary[8,:] = sum(D,dims=1)\n",
    "    GASummary[9,:] = sum(Htot,dims=1)\n",
    "    GASummary[10,:] = sum(Dperday,dims=1)\n",
    "    GASummary[11,:] = sum(NewCperday,dims=1)\n",
    "    \n",
    "    GASum = DataFrame(GASummary)\n",
    "    rename!(GASum , Symbol.(colnames))\n",
    "    insertcols!(GASum, 1,:EpiState => [\"S\",\"E\",\"Ia\",\"Is\",\"R\",\"Ihsub\",\"Ihcrit\",\"D\",\"Htot\",\"Dperday\",\"NCperday\"])\n",
    "    \n",
    "    CSV.write(string(SaveName,\"/GA_Summary.csv\"),GASum)\n",
    "    \n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Used to turn the solution variable 1x(159*9*8) into an array that can be more easily parsed\n",
    "function gatherStates!(sol,par)\n",
    "    SDI,β_a,β_s,p,γ_e,γ_a,γ_s,γ_h,h,ξ,μ,S_ind,E_ind,\n",
    "    Ia_ind,Is_ind,Ihsub_ind,Ihcrit_ind,R_ind,D_ind,\n",
    "    NumPatches,NumAgeClasses,patchstates,patchstatesALIVE,\n",
    "    patchstatesMOVEABLE,w,wy,wo,transport,ISMOVEFACTOR,w1,w2 = par\n",
    "    \n",
    "    #find time indices (in case of duplicates)\n",
    "    UU = unique(sol.t)\n",
    "    TS = Array{Int64,1}(undef,0)\n",
    "    for aa = 1:length(UU)\n",
    "       push!(TS,findall(sol.t .== UU[aa])[1])\n",
    "    end\n",
    "    \n",
    "    Sols =  hcat(sol.u...)\n",
    "    \n",
    "    ARRAYSTORE = zeros(NumPatches,8,NumAgeClasses,length(TS)) #array of patches x epi states x ages x time\n",
    "    \n",
    "    for aa = 1:NumPatches\n",
    "        for bb = 1:NumAgeClasses\n",
    "        INDS = (aa-1)*patchstates .+ [S_ind[bb],E_ind[bb],Ia_ind[bb],Is_ind[bb],R_ind[bb],Ihsub_ind[bb],Ihcrit_ind[bb],D_ind[bb]]       \n",
    "            for cc=1:8\n",
    "                for dd = 1:length(TS)\n",
    "                    ARRAYSTORE[aa,cc,bb,dd] = Sols[INDS[cc],TS[dd]]\n",
    "                end\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "    return ARRAYSTORE,TS\n",
    "end          "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#summary plot - plots the summed epi. states at the Georgia level i.e. dropping spatial dimension\n",
    "function SummaryPlot!(JAM,T)\n",
    "SZ = size(JAM)\n",
    "#AA=dropdims(sum(JAM[aa,1,1:9,:],dims=1),dims=1)[end]\n",
    "S = zeros(SZ[4])\n",
    "E = zeros(SZ[4])\n",
    "Ia = zeros(SZ[4])\n",
    "Is = zeros(SZ[4])\n",
    "R = zeros(SZ[4])\n",
    "Ihs = zeros(SZ[4])\n",
    "Ihc = zeros(SZ[4])\n",
    "D = zeros(SZ[4])\n",
    "\n",
    "for aa = 1:SZ[1]\n",
    "    S .+= dropdims(sum(JAM[aa,1,1:9,:],dims=1),dims=1)\n",
    "    E .+= dropdims(sum(JAM[aa,2,1:9,:],dims=1),dims=1)\n",
    "    Ia .+= dropdims(sum(JAM[aa,3,1:9,:],dims=1),dims=1)\n",
    "    Is .+= dropdims(sum(JAM[aa,4,1:9,:],dims=1),dims=1)\n",
    "    R .+= dropdims(sum(JAM[aa,5,1:9,:],dims=1),dims=1)\n",
    "    Ihs .+= dropdims(sum(JAM[aa,6,1:9,:],dims=1),dims=1)\n",
    "    Ihc .+= dropdims(sum(JAM[aa,7,1:9,:],dims=1),dims=1)\n",
    "    D .+= dropdims(sum(JAM[aa,8,1:9,:],dims=1),dims=1)\n",
    "end\n",
    "\n",
    "pl = plot(T,S,label=\"S\",yguidefontsize=10,xlabel=\"Time (days)\",ylabel=\"People\",legend=:outertopright)\n",
    "plot!(T,E,label=\"E\")\n",
    "plot!(T,Ia,label=\"Ia\")\n",
    "plot!(T,Is,label=\"Is\")\n",
    "plot!(T,R,label=\"R\")\n",
    "plot!(T,Ihs,label=\"Ihsub\")\n",
    "plot!(T,Ihc,label=\"Ihcrit\")\n",
    "plot!(T,D,label=\"D\")\n",
    "\n",
    "return pl\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set up the model scripts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#######################################STUFFTO EDIT ###################### \n",
    "#Time stepping stuff\n",
    "TIMERESOLUTION = 1 #MAKE THIS AN INTEGER IN DAYS!!!!!! REQUIRED FOR SAVING LATER ON\n",
    "T_END = 150\n",
    "tspanrun = (0.0,T_END*1.0) #run the simulations between 0 days and 400 days\n",
    "TODAY = Date(2020,3,28) #Date format of time 0\n",
    "\n",
    "#Load in Data used to set time 0\n",
    "CasesbyCounty = CSV.read(\"./data/28March2020.csv\");\n",
    "HospitalRegions = CSV.read(\"./data/ga_hosp_regions.csv\")\n",
    "HospRegions = HospitalRegions.HospRegion\n",
    "\n",
    "#String denoting where to save output\n",
    "\n",
    "#SAVEFOLD = \"SimulationOutput/March28_sim\"\n",
    "SAVEFOLD = \"SimulationOutput/March28_eig\"\n",
    "\n",
    "#initial condition setting to use (1 - uses a simulation method county by county; 2 - uses the state level model to partition to counties)\n",
    "#other choices are archives right now from previous attempts\n",
    "\n",
    "#CHOICES = 1\n",
    "CHOICES = 2\n",
    "\n",
    "#######################################STUFFTO EDIT ###################### "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Model states\n",
    "NumPatches = 159 ; #159 counties\n",
    "NumAgeClasses = 9; # 0-9,10-19,20-29,30-39,40-49,50-59,60-69,70-79,80+\n",
    "NumEpiStates = 8; # S,E,Ia,Is,R,Ihsub,Ihcrit,D\n",
    "\n",
    "patchstates = NumAgeClasses*NumEpiStates\n",
    "patchstatesMOVEABLE = (NumEpiStates-3)*NumAgeClasses  #D cannot move. Ihsub, Ihcrit cannot move.\n",
    "#Say that I_s may be able to move - can modify this assumption later.\n",
    "patchstatesALIVE = (NumEpiStates-1)*NumAgeClasses #all but D are alive\n",
    "transport = zeros(NumPatches,patchstatesMOVEABLE);\n",
    "\n",
    "#Epidemiological state indexing - just for within one patch.\n",
    "S_ind = 1:NumAgeClasses\n",
    "E_ind = (1*NumAgeClasses+1):2*NumAgeClasses\n",
    "Ia_ind = (2*NumAgeClasses+1):3*NumAgeClasses\n",
    "Is_ind = (3*NumAgeClasses+1):4*NumAgeClasses\n",
    "R_ind = (4*NumAgeClasses+1):5*NumAgeClasses\n",
    "Ihsub_ind = (5*NumAgeClasses+1):6*NumAgeClasses\n",
    "Ihcrit_ind = (6*NumAgeClasses+1):7*NumAgeClasses\n",
    "D_ind = (7*NumAgeClasses+1):8*NumAgeClasses\n",
    "\n",
    "#Age class information\n",
    "MinAge = [0,10,20,30,40,50,60,70,80]\n",
    "MiddleAge = [5,15,25,35,45,55,65,75,90];"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#read in the GA level data\n",
    "\n",
    "ODL29 = CSV.read(\"./data/ga_countyOD_29_younger.csv\") # transport matrix for 0-29\n",
    "OD30_55 = CSV.read(\"./data/ga_countyOD_30_54.csv\") #transport matrix for 30-55 (assume till 59)\n",
    "OD55_up = CSV.read(\"./data/ga_countyOD_55_older.csv\") #transport matrix for 55+ (assume for 60+)\n",
    "GA_county_size = CSV.read(\"./data/GA_county_decade_age.csv\") #age distributions by county\n",
    "\n",
    "#FIP is a way of designating counties (https://en.wikipedia.org/wiki/List_of_FIPS_codes_for_Georgia_(U.S._state))\n",
    "PatchFIPS = GA_county_size.FIP\n",
    "PatchCounty = GA_county_size.County\n",
    "#NOTE - will assign initial states by GA_county_size\n",
    "#USE THE FIP indices here to keep everything ordered.\n",
    "\n",
    "\n",
    "\n",
    "#Here we load in origin-destination data (159x159) showing trip numbers between places (including self-loops) which is sampled (not normalised)\n",
    "#we need to translate into a form useful to the model\n",
    "#to do this we want to:\n",
    "#a) normalise such that we know the proportion of trips from an origin to a destination\n",
    "#b) replace the diagonal (self trips) with negative elements that sum to the proportion of those leaving\n",
    "#c) transpose the data such that we can use it in a matrix multiplication to calculate net pop. change (- leavers + incomers) for each type of state e.g. multiply transport matrix by all S individuals in age class 20-29 (159x1) [ then repeat for all epi states and age classes]\n",
    "#we are going to get transport matrices w for each of the age structured commuting data\n",
    "\n",
    "#load and create commuter matrix\n",
    "## Commuter data for those aged 30-55\n",
    "COMMRULES = OD30_55[:,2:end] #matrix of data showing origin (row) and destination (col) patterns by county.\n",
    "#data is total number of trips from using a field sampling approach. Want to turn into a proportion.\n",
    "#Note 1: this is a sample. Needs normalising. Note 2: first column was FIP\n",
    "JJ = 1.0*convert(Array,COMMRULES)#convert to matrix of Floats\n",
    "sJJ = sum(JJ,dims=1) #sum over columns - this is the total arrivals (including self)\n",
    "for aa =1:159\n",
    "    JJ[:,aa] = JJ[:,aa] ./ sJJ[aa]\n",
    "end\n",
    "#find proportion of those staying in county\n",
    "PROPSTAY = JJ[1:size(JJ, 1)+1:end]\n",
    "#reassign diagonal elements as those who will be leaving county\n",
    "JJ[1:size(JJ, 1)+1:end] = -(1.0 .- PROPSTAY)\n",
    "#transpose matrix as we require destination to be the rows and origin to be the columns\n",
    "w = JJ';\n",
    "\n",
    "## Commuter data for those aged less than 30\n",
    "COMMRULES = ODL29[:,2:end] #matrix of data showing origin (row) and destination (col) patterns by county.\n",
    "#Note 1: this is a sample. Needs normalising. Note 2: first column was FIP\n",
    "JJ = 1.0*convert(Array,COMMRULES)#convert to matrix of Floats\n",
    "sJJ = sum(JJ,dims=1) #sum over columns - this is the total arrivals (including self)\n",
    "for aa =1:159\n",
    "    JJ[:,aa] = JJ[:,aa] ./ sJJ[aa]\n",
    "end\n",
    "#find proportion of those staying in county\n",
    "PROPSTAY = JJ[1:size(JJ, 1)+1:end]\n",
    "#reassign diagonal elements as those who will be leaving county\n",
    "JJ[1:size(JJ, 1)+1:end] = -(1.0 .- PROPSTAY)\n",
    "#transpose matrix as we require destination to be the rows and origin to be the columns\n",
    "wy = JJ';\n",
    "\n",
    "## Commuter data for those 55 and older\n",
    "COMMRULES = OD55_up[:,2:end] #matrix of data showing origin (row) and destination (col) patterns by county.\n",
    "#Note 1: this is a sample. Needs normalising. Note 2: first column was FIP\n",
    "JJ = 1.0*convert(Array,COMMRULES)#convert to matrix of Floats\n",
    "sJJ = sum(JJ,dims=1) #sum over columns - this is the total arrivals (including self)\n",
    "for aa =1:159\n",
    "    JJ[:,aa] = JJ[:,aa] ./ sJJ[aa]\n",
    "end\n",
    "#find proportion of those staying in county\n",
    "PROPSTAY = JJ[1:size(JJ, 1)+1:end]\n",
    "#reassign diagonal elements as those who will be leaving county\n",
    "JJ[1:size(JJ, 1)+1:end] = -(1.0 .- PROPSTAY)\n",
    "#transpose matrix as we require destination to be the rows and origin to be the columns\n",
    "wo = JJ';\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Make plot of the age-distribution across counties\n",
    "makeplot=1\n",
    "if makeplot==1\n",
    "gr()\n",
    "plot(MiddleAge,100 .* convert(Array,(GA_county_size[1,3:11])),xlim=(0,100),title=\"Age distribution by county\",\n",
    "    ylabel=\"Percent of county population (%)\",xlabel=\"Age\",legend=false,yguidefontsize=8,markershape = :auto,markersize=2)\n",
    "for aa=2:159\n",
    "    plot!(MiddleAge,100 .* convert(Array,(GA_county_size[aa,3:11])),markershape = :auto,markersize=2)\n",
    "end\n",
    "savefig(\"./figures/GAagestructure.pdf\")    \n",
    "current()\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Basic epidemiology to apply across all patches\n",
    "SocialDistancingInterventions=[1.0,0.75,0.5,0.25] #proportionalities to apply to both β_A and β_S (1.0 is no distancing)\n",
    "β_A = 0.3 #asymptomatic transmission rate (per day)\n",
    "β_S = 0.6 #symptomatic transmission rate (per day)\n",
    "γ_e = (1.0/4.0) #mean exposed period (days)\n",
    "γ_a = (1.0/6.0) #mean asymptomatic period (days)\n",
    "γ_s = (1.0/6.0) #mean symptomatic period (days)\n",
    "γ_h = (1.0/10.0) # mean hospitalised period (days)\n",
    "\n",
    "p = [0.95, 0.95,0.9,0.8,0.7,0.6,0.4,0.2,0.2]  #fraction of cases that are asymptomatic\n",
    "h = [0.1, 0.3, 1.2, 3.2, 4.9, 10.2, 16.6, 24.3, 27.3]./100; #fraction of cases that require a hospital visit (by age class)\n",
    "ξ = [5, 5, 5, 5, 6.3, 12.2, 27.4, 43.2, 70.9]./100; #fraction of hospitalised cases that become critical and require critical care (by age class)\n",
    "μ = 0.5.*ones(NumAgeClasses) #mortality fraction\n",
    "\n",
    "ISMOVEFACTOR = 1.0 # how much to let Is people commute (must be 0-1)\n",
    "w1 = 16/24 ; #fraction of time mixing with local pop.\n",
    "w2 = 1.0 - w1 ;#fraction of time mixing with commuter pop.\n",
    "\n",
    "\n",
    "\n",
    "#set parameter sets to use to initialise/drive the model\n",
    "param = [1.0,β_A,β_S,p,γ_e,γ_a,γ_s,γ_h,h,ξ,μ,S_ind,E_ind,\n",
    "    Ia_ind,Is_ind,Ihsub_ind,Ihcrit_ind,R_ind,D_ind,\n",
    "    NumPatches,NumAgeClasses,patchstates,patchstatesALIVE,\n",
    "    patchstatesMOVEABLE,w,wy,wo,transport,ISMOVEFACTOR,w1,w2]\n",
    "\n",
    "singlenodepar =[1.0,β_A,β_S,p,γ_e,γ_a,γ_s,γ_h,h,ξ,μ,S_ind,E_ind,Ia_ind,Is_ind,R_ind,Ihsub_ind,Ihcrit_ind,D_ind]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ### Set the initial conditions\n",
    " \n",
    "We now have the ingredients. Now we need to mix in the ingredients.\n",
    "Want to set initial conditions, being cognisant of the GA data.\n",
    "Two main ways right now are: a simulation approach (CHOICE=5); and an \"eigenvector\" approach (CHOICE=7)\n",
    "\n",
    "Simulation starts an epidemic in Fulton & Dougherty, then runs the model forward until deaths match up those recorded in GA data for each county one-by-one. The states in each county at these times are then patched together into an initial condition.\n",
    "\n",
    "Eigenvector approach runs the model at the Georgia level till total deaths are matched. Then assigns the final vector of values proportionally across counties as an intial condition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Setting initial conditions\n",
    "#Assign county age structured data to initial state \n",
    "GAPOP = sum(GA_county_size.POPSIZE) #total population of GA\n",
    "\n",
    "MAT = convert(Array, (GA_county_size[:,3:(end-1)] .* GA_county_size.POPSIZE))\n",
    "MAT2 = MAT ./ GAPOP\n",
    "#sum(MAT2)#should = 1. Yes!\n",
    "#MAT2 is the fraction of GA in eaah age class in each county.\n",
    "\n",
    "#want to now assign MAT2 to initial vector u00\n",
    "u00 = zeros(NumPatches*NumAgeClasses*NumEpiStates)\n",
    "\n",
    "for aa = 1:NumPatches\n",
    "        u00[(aa-1)*patchstates .+ S_ind] = MAT2[aa,:]\n",
    "end\n",
    "\n",
    "# load case data and figure out initial conditions\n",
    "tspan = (0.0,300.0)\n",
    "\n",
    "#start epidemic with 1 exposed individual ages 20-29 in Fulton county ...\n",
    "INDFult = findall(GA_county_size.FIP .== 13121)\n",
    "INDFult=INDFult[1]#patch index\n",
    "u00[S_ind[3] + patchstates*(INDFult-1)] = u00[S_ind[3] + patchstates*(INDFult-1)] - 1.0/GAPOP #minus from S\n",
    "u00[E_ind[3] + patchstates*(INDFult-1)] = u00[E_ind[3] + patchstates*(INDFult-1)] + 1.0/GAPOP #add to E\n",
    "#...and with 1 exposed individual ages 20-29 in Dougherty county\n",
    "INDDougherty = findall(GA_county_size.FIP .== 13095) #dougherty county (Albany)\n",
    "INDDougherty=INDDougherty[1]#patch index\n",
    "u00[S_ind[3] + patchstates*(INDDougherty-1)] = u00[S_ind[3] + patchstates*(INDDougherty-1)] - 1.0/GAPOP #minus from S\n",
    "u00[E_ind[3] + patchstates*(INDDougherty-1)] = u00[E_ind[3] + patchstates*(INDDougherty-1)] + 1.0/GAPOP #add to E\n",
    "initu = u00\n",
    "\n",
    "#set the ODE problem using this initial condition (1 infection in Dougherty and 1 infection in Fulton from people aged 20-29)\n",
    "prob_init = ODEProblem(SpatialAgeStructuredEpidemiologicalModelMixAge!,initu,tspan,param)\n",
    "\n",
    "#set an terminate integration affect.\n",
    "affect!(integrator) = terminate!(integrator)\n",
    "\n",
    "#state wide deaths and cases\n",
    "TOT_deaths = sum(CasesbyCounty.deaths)/GAPOP\n",
    "TOT_cases = sum(CasesbyCounty.cases)/GAPOP\n",
    "\n",
    "\n",
    "#now use this information to initialise the model.\n",
    "#Currently done with two options:\n",
    "###CHOICES = 1 :  County simulation method (CSM)\n",
    "#uses the initial outbreak defined above to find epi. distributions \n",
    "#in each county that hit death data in the GA data input county-by-county.\n",
    "#These conditions are then merged together to make the final initial condition.\n",
    "###CHOICES = 2 :  State level model (SLM)\n",
    "#uses a non-spatial model for the whole of GA to find epi. distributions\n",
    "# that hit deaths for the whole of GA. Then apportions these out to counties\n",
    "# given the proportions of deaths recorded in each county.\n",
    "\n",
    "\n",
    "if  CHOICES == 1\n",
    "    ## County Simulation Method (CSM) ##\n",
    "    \n",
    "    #in general should use hospitalisations rather than cases, unless death available\n",
    "    #case numbers highly likely to be underestimates\n",
    "    Hprop = 0.2 #20.48% proportion of cases that are hospitalised - appears consistently around 20% (April 15)\n",
    "    #- this is taken from Georgia Department of Public Health COVID-19 Daily Status Report\n",
    "    #https://dph.georgia.gov/covid-19-daily-status-report\n",
    "    \n",
    "    #cases county by county\n",
    "    for aa = 1:length(GA_county_size.FIP) #for each county in GA\n",
    "    print(aa) #print index to help show progress\n",
    "        \n",
    "        IND = findall(CasesbyCounty.FIP .== GA_county_size.FIP[aa]) #find patch index\n",
    "        DEATH = 0.0\n",
    "        if length(IND) == 1 #ignore unknown category in case data\n",
    "            DEATH = CasesbyCounty.deaths[IND[1]] #deaths recorded in this patch\n",
    "            HOSP = Hprop * CasesbyCounty.cases[IND[1]] #hsopitalised assumed in this patch\n",
    "        else\n",
    "            HOSP = 0.5 # if no cases recorded. Assign some proportion to be hospitalised. Most likely non-zero.\n",
    "        end\n",
    "        \n",
    "        #use this info. to simulate to epi. distributions in this patch\n",
    "        \n",
    "        if DEATH >0.0\n",
    "            #If there are deaths simulate until the number of deaths is reached\n",
    "             #set up callback to stop integration when number of cases reaches a threshold\n",
    "            CaseThreshold = DEATH/GAPOP\n",
    "            #sum of S in patch - (all in patch - thresh)\n",
    "            condition(u,t,integrator)  = sum(u[D_ind .+ patchstates*(aa-1)]) - CaseThreshold  #this must be equal to zero at stopping condition\n",
    "            affect!(integrator) = terminate!(integrator) \n",
    "            cb = ContinuousCallback(condition,affect!)\n",
    "            #solve the initial spin up, checking for positivity and using callback definition \n",
    "            sol_init = solve(prob_init,Tsit5(), abstol=1e-8,reltol=1e-8,saveat=1,isoutofdomain=(y,p,t)->any(x->x<0,y), callback = cb)\n",
    "            #assign end spinup states per county into new initial condition.\n",
    "            END = sol_init.u[end]\n",
    "            u00[(1+patchstates*(aa-1)):(patchstates*aa)] = END[(1+patchstates*(aa-1)):(patchstates*aa)]\n",
    "        else\n",
    "            #If no deaths simulate until the number hospitalised (both Ihsub and Ihcrit) is reached\n",
    "            #set up callback to stop integration when number of cases reaches a threshold\n",
    "            CaseThreshold = HOSP/GAPOP\n",
    "            #sum of S in patch - (all in patch - thresh)\n",
    "            condition2(u,t,integrator)  = sum(u[[Ihsub_ind;Ihcrit_ind] .+ patchstates*(aa-1)]) - CaseThreshold  #this must be equal to zero at stopping condition\n",
    "            affect2!(integrator) = terminate!(integrator) \n",
    "            cb2 = ContinuousCallback(condition2,affect2!)\n",
    "            #solve the initial spin up, checking for positivity and using callback definition \n",
    "            sol_init = solve(prob_init,Tsit5(), abstol=1e-8,reltol=1e-8,saveat=1,isoutofdomain=(y,p,t)->any(x->x<0,y), callback = cb2)\n",
    "            #assign end spinup states per county into new initial condition.\n",
    "            END = sol_init.u[end]\n",
    "            u00[(1+patchstates*(aa-1)):(patchstates*aa)] = END[(1+patchstates*(aa-1)):(patchstates*aa)]\n",
    "        end\n",
    "    end\n",
    "   \n",
    "elseif CHOICES == 2\n",
    "    ## State Level Model (SLM) ##\n",
    "    \n",
    "    #1.set up initial conditions for a GA level model\n",
    "    #Simple model to total deaths at state level. Then reapportion in the county model.\n",
    "    Singleu0 = zeros(NumAgeClasses*NumEpiStates)\n",
    "    Singleu0[S_ind] = sum(MAT2,dims=1)\n",
    "    Singleu0[S_ind[3]] = Singleu0[S_ind[3]] - 1.0/GAPOP\n",
    "    Singleu0[E_ind[3]] = 1.0/GAPOP  #seed the epidemic (1 person 20-29).\n",
    "\n",
    "    prob_s_init = ODEProblem(SingleNode!,Singleu0,tspan,singlenodepar)\n",
    "\n",
    "    CaseThreshold = TOT_deaths #as a fraction of GA pop.\n",
    "\n",
    "    condition3(u,t,integrator)  = CaseThreshold - sum(u[D_ind])   #this must be equal to zero at stopping condition\n",
    "    affect3!(integrator) = terminate!(integrator) \n",
    "    cb3 = ContinuousCallback(condition3,affect3!)\n",
    "\n",
    "    #solve the initial spin up, checking for positivity and using callback definition \n",
    "    sol_s_init = solve(prob_s_init,Tsit5(), abstol=1e-8,reltol=1e-8,saveat=0.1,isoutofdomain=(y,p,t)->any(x->x<0,y), callback = cb3)\n",
    "\n",
    "    #2. Now use this epi. distribution and assign it to each county\n",
    "    START = sol_s_init.u[end] #initial distributions at state scale (age*epi classes)\n",
    "    SumD = sum(START[D_ind])*GAPOP #sum of all deaths in model\n",
    "    TOTD =  sum(CasesbyCounty.deaths) #total deaths in data\n",
    "    TOTU =  sum(CasesbyCounty.deaths[CasesbyCounty.FIP.==-999]) #total deaths with unknown assignment\n",
    "    L1 = length(GA_county_size.FIP) #all counties\n",
    "    NUMUNKNOWN = sum(CasesbyCounty.FIP.==-999)\n",
    "    L2 = length(CasesbyCounty.FIP)- NUMUNKNOWN - sum(CasesbyCounty.deaths .== 0) #num counties with cases\n",
    "   \n",
    "    for aa = 1:159\n",
    "        countypop = GA_county_size.POPSIZE[aa]\n",
    "        \n",
    "        IND = findall(CasesbyCounty.FIP .== GA_county_size.FIP[aa]) #find patch index\n",
    "        if length(IND) == 1 #ignore unknown category in case data\n",
    "            D_county = 1.0 * CasesbyCounty.deaths[IND[1]]\n",
    "            if D_county == 0.0\n",
    "                D_county = TOTU/(L1-L2) #if this county has no recorded cases assign from the unknown pool.\n",
    "            end\n",
    "        else\n",
    "            D_county = TOTU/(L1-L2) #if this county has no recorded cases assign from the unknown pool.\n",
    "        end\n",
    "        \n",
    "        RATIO = 1.0*(D_county/TOTD) #county ratio - ratio of deaths in county to total\n",
    "       \n",
    "        NEW = START .* RATIO  .* GAPOP\n",
    "        JAM = zeros(NumAgeClasses)\n",
    "        Ages = MAT[aa,:] # age distribution in this county\n",
    "        \n",
    "        for bb=1:NumAgeClasses\n",
    "           #calculate sum of all but S this age class\n",
    "            JAM[bb] = sum(NEW[[E_ind[bb];Ia_ind[bb];Is_ind[bb];R_ind[bb];Ihsub_ind[bb];Ihcrit_ind[bb];D_ind[bb]]])\n",
    "           #Now find what S for this age class should be\n",
    "            NEW[S_ind[bb]] = Ages[bb] - JAM[bb]\n",
    "        end\n",
    "    \n",
    "        NEW = NEW ./ GAPOP\n",
    "        \n",
    "        #check if any element in NEW is less than 0 - would indicate problems in preserving county demographics.\n",
    "        if sum(isless.(NEW,0)) < 0\n",
    "            print(\"WARNING. u LESS THAN 0\")\n",
    "            NEW[isless.(NEW,0)] = 0.0\n",
    "        end\n",
    "        \n",
    "        #assign to intiial condition\n",
    "        u00[(1+patchstates*(aa-1)):(patchstates*aa)] = NEW\n",
    "\n",
    "    end\n",
    "    \n",
    "end\n",
    "\n",
    "#assign end spinup states as new initial condition.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#clean up memory by removing tables and objects no longer of interest. Julia has no explicit removal of variables (that I know of).\n",
    "#assigning large matrices to integers should save a fair amount of memory.\n",
    "#varinfo()\n",
    "COMMRULES = 0\n",
    "ODL29 = 0\n",
    "OD30_55 = 0\n",
    "OD55_up = 0\n",
    "#GA_county_size = 0\n",
    "#CasesbyCounty = 0\n",
    "MAT = 0\n",
    "MAT2 = 0\n",
    "JJ =0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run the simulations\n",
    "\n",
    "Now that the initial conditions are set, we are ready to bake the model.\n",
    "Note that in the above conditions are set without any social distancing interventions. Currently, these are applied at time =0. That is here.\n",
    "Here, we have a loop that goes through different social distancing levels (that reduce βS and βA); runs the model and saves the output to file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#NOW DO INTERVENTIONS IN A LOOP\n",
    "\n",
    "mkdir(SAVEFOLD)\n",
    "for bb = 1:length(SocialDistancingInterventions)\n",
    "    \n",
    "    SAVEFOLD2 = string(SAVEFOLD,\"/Social_\",bb)\n",
    "    SAVEFOLD3 = string(SAVEFOLD,\"/Social_\",bb,\"relaxed\")\n",
    "    \n",
    "    #set parameters for this intervention\n",
    "    pars = [SocialDistancingInterventions[bb],β_A,β_S,p,γ_e,γ_a,γ_s,γ_h,h,ξ,μ,S_ind,E_ind,\n",
    "    Ia_ind,Is_ind,Ihsub_ind,Ihcrit_ind,R_ind,D_ind,\n",
    "    NumPatches,NumAgeClasses,patchstates,patchstatesALIVE,\n",
    "    patchstatesMOVEABLE,w,wy,wo,transport,ISMOVEFACTOR,w1,w2]\n",
    "   \n",
    "    ## NO RELAXATION\n",
    "    #define ODE problems to solve\n",
    "    prob = ODEProblem(SpatialAgeStructuredEpidemiologicalModelMixAge!,u00,tspanrun,pars)\n",
    "    #set New Problem    \n",
    "    #solve the ODEs\n",
    "    SOL = solve(prob,Tsit5(),abstol=1e-8,reltol=1e-8,saveat=TIMERESOLUTION,isoutofdomain=(y,p,t)->any(x->x<0,y))\n",
    " \n",
    "    \n",
    "    ##RELAXATION ON MAY 15th\n",
    "    prob_a = ODEProblem(SpatialAgeStructuredEpidemiologicalModelMixAge!,u00,tspanrun,pars)\n",
    "    condition(u,t,integrator) = t== 48  #48 days to May 15th\n",
    "    affect!(integrator) = integrator.p[1] = 1.0\n",
    "    cb = DiscreteCallback(condition,affect!)\n",
    "    SOL_a = solve(prob_a,Tsit5(),abstol=1e-8,reltol=1e-8,saveat=TIMERESOLUTION,isoutofdomain=(y,p,t)->any(x->x<0,y),callback = cb, tstops=[48.0],save_positions=(false,false))\n",
    "    \n",
    "    #save output\n",
    "\n",
    "    WriteOutput!(PatchFIPS,PatchCounty,HospRegions,GAPOP,TODAY,TIMERESOLUTION,T_END,SOL,pars,SAVEFOLD2)\n",
    "    WriteOutput!(PatchFIPS,PatchCounty,HospRegions,GAPOP,TODAY,TIMERESOLUTION,T_END,SOL_a,pars,SAVEFOLD3)\n",
    "    \n",
    "    JAM = GAPOP * gatherStates!(SOL,param)[1]\n",
    "    Gath = gatherStates!(SOL_a,param)\n",
    "    JAM2 = GAPOP  * Gath[1]\n",
    "    time2 = SOL_a.t[Gath[2]]\n",
    "    \n",
    "    \n",
    "    #county-level projected deaths\n",
    "    gr()\n",
    "    plot(SOL.t,dropdims(sum(JAM[1,8,1:9,:],dims=1),dims=1),legend=false,ylabel=\"D per county\",xlabel=\"Time (days)\")\n",
    "\n",
    "    for aa = 2:159\n",
    "        plot!(SOL.t,dropdims(sum(JAM[aa,8,1:9,:],dims=1),dims=1))\n",
    "    end\n",
    "    savefig(string(SAVEFOLD2,\"/deaths.png\"))\n",
    "    #current()\n",
    "\n",
    "    #county-level projected ICU\n",
    "    plot(SOL.t,dropdims(sum(JAM[1,7,1:9,:],dims=1),dims=1),legend=false,ylabel=\"ICU per county\",xlabel=\"Time (days)\")\n",
    "\n",
    "    for aa = 2:159\n",
    "        plot!(SOL.t,dropdims(sum(JAM[aa,7,1:9,:],dims=1),dims=1))\n",
    "    end\n",
    "    savefig(string(SAVEFOLD2,\"/ICU.png\"))\n",
    "    #current()\n",
    "\n",
    "    #county-level projected hosp.s\n",
    "    plot(SOL.t,dropdims(sum(JAM[1,6,1:9,:],dims=1),dims=1).+dropdims(sum(JAM[1,7,1:9,:],dims=1),dims=1),legend=false,ylabel=\"Hosp. per county\",xlabel=\"Time (days)\")\n",
    "\n",
    "    for aa = 2:159\n",
    "        plot!(SOL.t,dropdims(sum(JAM[aa,6,1:9,:],dims=1),dims=1).+dropdims(sum(JAM[1,7,1:9,:],dims=1),dims=1))\n",
    "    end\n",
    "    savefig(string(SAVEFOLD2,\"/totalHosp.png\"))\n",
    "    #current()\n",
    "\n",
    "    \n",
    "    p1 = SummaryPlot!(JAM,SOL.t)\n",
    "    savefig(string(SAVEFOLD2,\"/GAscenarios.pdf\"))\n",
    "    #current()\n",
    "    p2 = SummaryPlot!(JAM2,time2)\n",
    "    savefig(string(SAVEFOLD3,\"/GAscenarios.pdf\"))\n",
    "    #current()\n",
    "    \n",
    "    \n",
    "    \n",
    "end    "
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Julia 1.3.1",
   "language": "julia",
   "name": "julia-1.3"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.3.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}