# Load dependencies

```julia
#load dependencies
#load in dependencies to read data
using CSV, Plots, DataFrames, Dates
using LinearAlgebra,Pkg
#using DifferentialEquations # do not need the whole package necessarilly. Using just the ODE functions will save some memory
using OrdinaryDiffEq
#show versions
Pkg.status()
```

        Status `C:\Users\Steve\.julia\environments\v1.3\Project.toml`
      [336ed68f] CSV v0.5.26
      [a93c6f00] DataFrames v0.20.2
      [aae7a2af] DiffEqFlux v1.3.2
      [0c46a032] DifferentialEquations v6.10.1
      [587475ba] Flux v0.10.1
      [7073ff75] IJulia v1.21.1
      [c601a237] Interact v0.10.3
      [b964fa9f] LaTeXStrings v1.1.0
      [47be7bcc] ORCA v0.3.1
      [429524aa] Optim v0.20.1
      [1dea7af3] OrdinaryDiffEq v5.34.1
      [f0f68f2c] PlotlyJS v0.13.1
      [91a5bcdd] Plots v0.28.4
      [d330b81b] PyPlot v2.8.2
      [6f49c342] RCall v0.13.4
      [90137ffa] StaticArrays v0.12.1
      [f3b207a7] StatsPlots v0.14.2
      [fd094767] Suppressor v0.1.1
      [fce5fe82] Turing v0.8.3
      [30578b45] URIParser v0.4.0
      [0ae4a718] VegaDatasets v2.0.0
      [112f6efa] VegaLite v2.0.1
      [0f1e0344] WebIO v0.8.13
      [fdbf4ff8] XLSX v0.6.0

```julia
versioninfo()
```

    Julia Version 1.3.1
    Commit 2d5741174c (2019-12-30 21:36 UTC)
    Platform Info:
      OS: Windows (x86_64-w64-mingw32)
      CPU: Intel(R) Core(TM) i7-8550U CPU @ 1.80GHz
      WORD_SIZE: 64
      LIBM: libopenlibm
      LLVM: libLLVM-6.0.1 (ORCJIT, skylake)

# Spatial Epidemic Model for GA for COVID-19

Here, we extend the model to look at infection spread through a spatial metapopulation model.
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.

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$.

The system of nonlinear differential equations governing this patch and age-structured epidemiological for patch $i$ and age $a$ are as follows:

\begin{align}
\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) }}  \\ \\
\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)}} \\
\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)}} \\
\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)}} \\
\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)}} \\
\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)}} \\
\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)}} \\
\frac{dD(i,a)}{dt} & = \overset{mortality}{\overbrace{\mu \gamma_h I_{h,crit}(i,a)}}
\end{align}

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.

# Read in functions called in the scripts

```julia
#ODE model code
#assumes age-dependent journeys between counties
function SpatialAgeStructuredEpidemiologicalModelMixAge!(du,u,par,t)
    SDI,β_a,β_s,p,γ_e,γ_a,γ_s,γ_h,h,ξ,μ,S_ind,E_ind,
    Ia_ind,Is_ind,Ihsub_ind,Ihcrit_ind,R_ind,D_ind,T_Hosp_ind,
    NumPatches,NumAgeClasses,patchstates,patchstatesALIVE,
    patchstatesMOVEABLE,w,wy,wo,transport,ISMOVEFACTOR,w1,w2 = par

    #NumPatches is the number of spatial "nodes"
    #w is a matrix of size NxN determining flows between nodes.
    #transport is a matrix to store movements vectors (don't want to recreate this object throughout this function)
    #patchstates - total number of states in a patch
    #patchstatesALIVE - number of non-dead states
    #patchstatesMOVEABLE - number of states that are allowed to move between patches
    #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

    #have matrix w, representing movement of commuters i.e. X1 stay in county, X2 go to county 2, X3 go to county3
    #total travel is fraction travelled applied to people in moving states
    #some amounts of these will be IA and IS - keep track!
    #want to know N,Is,IA leaving each county; and arriving at each county
    #w11 S1 + w12 S2  #net change for each state so long as diagonal is -ve.

    #w is those aged 30-59
    #wy is those 0-29
    #wo is those 60+

    #Stage 1:
    #calculate transport between patches

    #transport = transport .* 0.0;
    #want this to be age structured
    for ii = 1:patchstatesMOVEABLE #for each age class of alive states that can move.
        INDS = ii:patchstates:(ii+(NumPatches-1)*patchstates ) #gives 1xpatch indices for transport multiplication for each age class and each state that can move
        MOD = mod(ii,NumAgeClasses-1)
        if MOD>0 & MOD<4
        transport[:,ii] = wy*u[INDS]
        elseif MOD==0  || MOD >=7
        transport[:,ii] = wo*u[INDS]
        else
        transport[:,ii] =  w*u[INDS]  #add the 1*patch vector of movement fluxes to transport
        end
    end

    #modify tranport matrix to account for Is people deciding to move or not.
    transport[:,Is_ind] = ISMOVEFACTOR .* transport[:,Is_ind]

    #calculate net change for each patch for N,Ia and Is
    C_TOT_NETCHANGE = sum(transport,dims=2)
    C_I_A_NETCHANGE = sum(transport[:,Ia_ind],dims=2)
    C_I_S_NETCHANGE = sum(transport[:,Is_ind],dims=2)

    #Stage 2:
    #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.

    for aa=1:(Int64(NumPatches)) #for each patch node

        #Local patch dynamics depend on patch quantitites
        NTOT_Loc = sum(u[((aa-1)*patchstates+1):((aa-1)*patchstates+patchstatesALIVE)]) #total alive
        TOTI_a_Loc = sum(u[(aa-1)*patchstates .+ Ia_ind]) #total asymptomatic infected
        TOTI_s_Loc = sum(u[(aa-1)*patchstates .+ Is_ind]) #total symptomatic infected

        NTOT_Com = NTOT_Loc + C_TOT_NETCHANGE[aa]#total alive
        TOTI_a_Com = TOTI_a_Loc + C_I_A_NETCHANGE[aa] #total asymptomatic infected
        TOTI_s_Com = TOTI_s_Loc + C_I_S_NETCHANGE[aa] #total symptomatic infected

        LOCAL = w1/(NTOT_Loc)
        COMMUTE = w2/(NTOT_Com)
        propI_A = ( TOTI_a_Loc*LOCAL + TOTI_a_Com*COMMUTE )
        propI_S = ( TOTI_s_Loc*LOCAL + TOTI_s_Com*COMMUTE )

        for bb = 1:(Int64(NumAgeClasses))  #for each age class in this patch
        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],T_Hosp_ind[bb]]

            S,E,Ia,Is,R,Ihsub,Ihcrit,D = u[INDS]

            du[INDS[1]] =  -β_a*SDI*S*propI_A - β_s*SDI*S*propI_S   #dS/dt
            du[INDS[2]] =  +β_a*SDI*S*propI_A + β_s*SDI*S*propI_S - γ_e*E  #dE/dt
            du[INDS[3]] = p[bb]*γ_e*E - γ_a*Ia   #dI_a /dt
            du[INDS[4]] = (1-p[bb])*γ_e*E - γ_s*Is  #dI_s/dt
            du[INDS[5]] = γ_a*Ia + (1-h[bb])*γ_s*Is + γ_h*Ihsub + (1-μ[bb])*γ_h*Ihcrit #dR/dt
            du[INDS[6]] = h[bb]*(1-ξ[bb])*γ_s*Is - γ_h*Ihsub #dIhsub/dt
            du[INDS[7]] = h[bb]*ξ[bb]*γ_s*Is -  γ_h*Ihcrit #dIhcrit/dt
            du[INDS[8]] = μ[bb]*γ_h*Ihcrit #dD/dt
            du[INDS[9]] = h[bb]*γ_s*Is  #tot Hosp
        end
    end

end

```

    SpatialAgeStructuredEpidemiologicalModelMixAge! (generic function with 1 method)

```julia
#ODE model code - this is used to model at the level of the whole of Georgia! (helpful in setting initial conditions)
function SingleNode!(du,u,par,t)
    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,T_Hosp_ind = par

    TOTI_a = sum(u[Ia_ind]) #Total asymptomatic infectious individuals
    TOTI_s = sum(u[Is_ind]) #Total symptomatic infectious individuals
    NTOT = sum(u[1:63]) #All individuals not in the dead (D) state

    for aa = 1:length(S_ind)#for each age class
        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],T_Hosp_ind[aa]]
        S,E,Ia,Is,R,Ihsub,Ihcrit,D = u[INDS]
        du[INDS[1]] = -β_a*SDI*S*TOTI_a/(NTOT) - β_s*SDI*S*TOTI_s/(NTOT) #dS/dt
        du[INDS[2]] = +β_a*SDI*S*TOTI_a/(NTOT) + β_s*SDI*S*TOTI_s/(NTOT) -  γ_e*E   #dE/dt
        du[INDS[3]] = p[aa]*γ_e*E - γ_a*Ia  #dI_a /dt
        du[INDS[4]] = (1-p[aa])*γ_e*E - γ_s*Is  #dI_s/dt
        du[INDS[5]] = γ_a*Ia + (1-h[aa])*γ_s*Is + γ_h*Ihsub + (1-μ[aa])*γ_h*Ihcrit  #dR/dt
        du[INDS[6]] = h[aa]*(1-ξ[aa])*γ_s*Is - γ_h*Ihsub #dIhsub/dt
        du[INDS[7]] = h[aa]*ξ[aa]*γ_s*Is -  γ_h*Ihcrit #dIhcrit/dt
        du[INDS[8]] = μ[aa]*γ_h*Ihcrit #dD/dt
        du[INDS[9]] = h[aa]*γ_s*Is #Cumulative hospitalised
    end

end
```

    SingleNode! (generic function with 1 method)

```julia
#reads in ODE solution; and saves output summary figures and CSV files
function WriteOutput!(PatchFIPS,PatchNames,Regions,GAPOP,TODAY,TIMERESOLTION,T_END,sol,pars,SaveName)

    Gath = gatherStates!(sol,pars)
    STATES = GAPOP * Gath[1] #organise the solution
    Times = sol.t[Gath[2]]

    #patches  x EpiStates x AgeClasses x Time
    SIZES = size(STATES)

    #Make summary tables for each epi state, and total hospitalised
    #S,E,Ia,Is,R,Ihsub,Ihcrit,D, Ihsub+Ihcrit
    S = zeros(SIZES[1],SIZES[4]) # patches x time
    E = zeros(SIZES[1],SIZES[4]) # patches x time
    Ia = zeros(SIZES[1],SIZES[4]) # patches x time
    Is = zeros(SIZES[1],SIZES[4]) # patches x time
    R = zeros(SIZES[1],SIZES[4]) # patches x time
    Ihsub = zeros(SIZES[1],SIZES[4]) # patches x time
    Ihcrit = zeros(SIZES[1],SIZES[4]) # patches x time
    D = zeros(SIZES[1],SIZES[4]) # patches x time
    Htot = zeros(SIZES[1],SIZES[4]) # patches x time
    Dperday = zeros(SIZES[1],SIZES[4]) # patches x time
    NewCperday = zeros(SIZES[1],SIZES[4]) # patches x time
    Hcumulative = zeros(SIZES[1],SIZES[4]) # patches x time

    for aa = 1:SIZES[1] #for each patch
        S[aa,:] = dropdims(sum(STATES[aa,1,1:9,:],dims=1),dims=1)
        E[aa,:] = dropdims(sum(STATES[aa,2,1:9,:],dims=1),dims=1)
        Ia[aa,:] = dropdims(sum(STATES[aa,3,1:9,:],dims=1),dims=1)
        Is[aa,:] = dropdims(sum(STATES[aa,4,1:9,:],dims=1),dims=1)
        R[aa,:] = dropdims(sum(STATES[aa,5,1:9,:],dims=1),dims=1)
        Ihsub[aa,:] = dropdims(sum(STATES[aa,6,1:9,:],dims=1),dims=1)
        Ihcrit[aa,:] = dropdims(sum(STATES[aa,7,1:9,:],dims=1),dims=1)
        D[aa,:] = dropdims(sum(STATES[aa,8,1:9,:],dims=1),dims=1)
        Htot[aa,:] =  Ihsub[aa,:] .+ Ihcrit[aa,:]
        for bb = 1:SIZES[3] #for age classes
            #Dperday[aa,:] .+=  μ[bb].*γ_h .*STATES[aa,7,bb,:].*1.0;  #deaths per day
            Dperday[aa,:] .+= [( STATES[aa,8,bb,2:end] .- STATES[aa,8,bb,1:(end-1)]) ./ (Times[2:end] - Times[1:(end-1)]); 0]
            NewCperday[aa,:] .+= [( STATES[aa,1,bb,1:(end-1)] .- STATES[aa,1,bb,2:end]) ./ (Times[2:end] - Times[1:(end-1)]); 0]
        end
        Hcumulative[aa,:] = dropdims(sum(STATES[aa,9,1:9,:],dims=1),dims=1)
    end

    #build Time column names
    colnames = Array{Date,1}(undef,1)
    colnames[1] = TODAY
    for aa = 2:(SIZES[4])#for each time
        push!(colnames,colnames[aa-1]+Dates.Day(convert(Int,Times[aa]-Times[aa-1])))
    end

    #build dataframes
    SSdata = DataFrame(S)
    rename!(SSdata , Symbol.(colnames))
    insertcols!(SSdata, 1, :FIP => PatchFIPS)
    insertcols!(SSdata, 2, :County => PatchNames)

    EEdata = DataFrame(E)
    rename!(EEdata , Symbol.(colnames))
    insertcols!(EEdata, 1, :FIP => PatchFIPS)
    insertcols!(EEdata, 2, :County => PatchNames)

    Iadata = DataFrame(Ia)
    rename!(Iadata , Symbol.(colnames))
    insertcols!(Iadata, 1, :FIP => PatchFIPS)
    insertcols!(Iadata, 2, :County => PatchNames)

    Isdata = DataFrame(Is)
    rename!(Isdata , Symbol.(colnames))
    insertcols!(Isdata, 1, :FIP => PatchFIPS)
    insertcols!(Isdata, 2, :County => PatchNames)

    RRdata = DataFrame(R)
    rename!(RRdata , Symbol.(colnames))
    insertcols!(RRdata, 1, :FIP => PatchFIPS)
    insertcols!(RRdata, 2, :County => PatchNames)

    Ihsubdata = DataFrame(Ihsub)
    rename!(Ihsubdata , Symbol.(colnames))
    insertcols!(Ihsubdata, 1, :FIP => PatchFIPS)
    insertcols!(Ihsubdata, 2, :County => PatchNames)

    Ihcritdata = DataFrame(Ihcrit)
    rename!(Ihcritdata , Symbol.(colnames))
    insertcols!(Ihcritdata, 1, :FIP => PatchFIPS)
    insertcols!(Ihcritdata, 2, :County => PatchNames)

    DDdata = DataFrame(D)
    rename!(DDdata , Symbol.(colnames))
    insertcols!(DDdata, 1, :FIP => PatchFIPS)
    insertcols!(DDdata, 2, :County => PatchNames)

    HHdata = DataFrame(Htot)
    rename!(HHdata , Symbol.(colnames))
    insertcols!(HHdata, 1, :FIP => PatchFIPS)
    insertcols!(HHdata, 2, :County => PatchNames)

    DpDdata = DataFrame(Dperday)
    rename!(DpDdata , Symbol.(colnames))
    insertcols!(DpDdata, 1, :FIP => PatchFIPS)
    insertcols!(DpDdata, 2, :County => PatchNames)

    CpDdata = DataFrame(NewCperday)
    rename!(CpDdata , Symbol.(colnames))
    insertcols!(CpDdata, 1, :FIP => PatchFIPS)
    insertcols!(CpDdata, 2, :County => PatchNames)

    Hcumdata = DataFrame(Hcumulative)
    rename!(Hcumdata , Symbol.(colnames))
    insertcols!(Hcumdata, 1, :FIP => PatchFIPS)
    insertcols!(Hcumdata, 2, :County => PatchNames)

    #make a directory to store files
    mkdir(SaveName)

    #write data to file
    CSV.write(string(SaveName,"/GA_county_Susceptibles.csv"),SSdata)
    CSV.write(string(SaveName,"/GA_county_Exposed.csv"),EEdata)
    CSV.write(string(SaveName,"/GA_county_AsymptInfected.csv"),Iadata)
    CSV.write(string(SaveName,"/GA_county_SymptInfected.csv"),Isdata)
    CSV.write(string(SaveName,"/GA_county_Recovered.csv"),RRdata)
    CSV.write(string(SaveName,"/GA_county_subacuteHospitalised.csv"),Ihsubdata)
    CSV.write(string(SaveName,"/GA_county_criticalHospitalised.csv"),Ihcritdata)
    CSV.write(string(SaveName,"/GA_county_Killed.csv"),DDdata)
    CSV.write(string(SaveName,"/GA_county_totalHospitalised.csv"),HHdata)
    CSV.write(string(SaveName,"/GA_county_deathrateperday.csv"),DpDdata)
    CSV.write(string(SaveName,"/GA_county_newcaserateperday.csv"),CpDdata)
    CSV.write(string(SaveName,"/GA_county_cumulativeHospitalised.csv"),Hcumdata)

    ##save summary data at RegionalLevel

    RegionList = sort(unique(Regions))
    SSdataR = zeros(length(RegionList),SIZES[4])
    EEdataR = zeros(length(RegionList),SIZES[4])
    IadataR = zeros(length(RegionList),SIZES[4])
    IsdataR = zeros(length(RegionList),SIZES[4])
    IhsubdataR = zeros(length(RegionList),SIZES[4])
    IhcritdataR = zeros(length(RegionList),SIZES[4])
    RRdataR = zeros(length(RegionList),SIZES[4])
    DDdataR = zeros(length(RegionList),SIZES[4])
    HHdataR = zeros(length(RegionList),SIZES[4])
    DpDdataR = zeros(length(RegionList),SIZES[4])
    CpDdataR = zeros(length(RegionList),SIZES[4])
    HcumdataR = zeros(length(RegionList),SIZES[4])

    for aa = 1:length(RegionList) #for each region
        RegionInd = findall(Regions.== RegionList[aa])
        SSdataR[aa,:] = sum(S[RegionInd,:],dims=1)
        EEdataR[aa,:] = sum(E[RegionInd,:],dims=1)
        IadataR[aa,:] = sum(Ia[RegionInd,:],dims=1)
        IsdataR[aa,:] = sum(Is[RegionInd,:],dims=1)
        RRdataR[aa,:] = sum(R[RegionInd,:],dims=1)
        DDdataR[aa,:] = sum(D[RegionInd,:],dims=1)
        IhsubdataR[aa,:] = sum(Ihsub[RegionInd,:],dims=1)
        IhcritdataR[aa,:] = sum(Ihcrit[RegionInd,:],dims=1)
        HHdataR[aa,:] = sum(Htot[RegionInd,:],dims=1)
        DpDdataR[aa,:] = sum(Dperday[RegionInd,:],dims=1)
        CpDdataR[aa,:] = sum(NewCperday[RegionInd,:],dims=1)
        HcumdataR[aa,:] = sum(Hcumulative[RegionInd,:],dims=1)
    end

    SSdataRR = DataFrame(SSdataR)
    rename!(SSdataRR , Symbol.(colnames))
    insertcols!(SSdataRR, 1, :Region => RegionList)

    EEdataRR = DataFrame(EEdataR)
    rename!(EEdataRR , Symbol.(colnames))
    insertcols!(EEdataRR, 1, :Region => RegionList)

    IadataRR = DataFrame(IadataR)
    rename!(IadataRR , Symbol.(colnames))
    insertcols!(IadataRR, 1, :Region => RegionList)

    IsdataRR = DataFrame(IsdataR)
    rename!(IsdataRR , Symbol.(colnames))
    insertcols!(IsdataRR, 1, :Region => RegionList)

    RRdataRR = DataFrame(RRdataR)
    rename!(RRdataRR , Symbol.(colnames))
    insertcols!(RRdataRR, 1, :Region => RegionList)

    IhsubdataRR = DataFrame(IhsubdataR)
    rename!(IhsubdataRR , Symbol.(colnames))
    insertcols!(IhsubdataRR, 1, :Region => RegionList)

    IhcritdataRR = DataFrame(IhcritdataR)
    rename!(IhcritdataRR , Symbol.(colnames))
    insertcols!(IhcritdataRR, 1, :Region => RegionList)

    DDdataRR = DataFrame(DDdataR)
    rename!(DDdataRR , Symbol.(colnames))
    insertcols!(DDdataRR, 1, :Region => RegionList)

    HHdataRR = DataFrame(HHdataR)
    rename!(HHdataRR , Symbol.(colnames))
    insertcols!(HHdataRR, 1, :Region => RegionList)

    DpDdataRR = DataFrame(DpDdataR)
    rename!(DpDdataRR , Symbol.(colnames))
    insertcols!(DpDdataRR, 1, :Region => RegionList)

    CpDdataRR = DataFrame(CpDdataR)
    rename!(CpDdataRR , Symbol.(colnames))
    insertcols!(CpDdataRR, 1, :Region => RegionList)

    HHcumdataRR = DataFrame(HcumdataR)
    rename!(HHcumdataRR , Symbol.(colnames))
    insertcols!(HHcumdataRR, 1, :Region => RegionList)

    CSV.write(string(SaveName,"/GA_region_Susceptibles.csv"),SSdataRR)
    CSV.write(string(SaveName,"/GA_region_Exposed.csv"),EEdataRR)
    CSV.write(string(SaveName,"/GA_region_AsymptInfected.csv"),IadataRR)
    CSV.write(string(SaveName,"/GA_region_SymptInfected.csv"),IsdataRR)
    CSV.write(string(SaveName,"/GA_region_Recovered.csv"),RRdataRR)
    CSV.write(string(SaveName,"/GA_region_subacuteHospitalised.csv"),IhsubdataRR)
    CSV.write(string(SaveName,"/GA_region_criticalHospitalised.csv"),IhcritdataRR)
    CSV.write(string(SaveName,"/GA_region_Killed.csv"),DDdataRR)
    CSV.write(string(SaveName,"/GA_region_totalHospitalised.csv"),HHdataRR)
    CSV.write(string(SaveName,"/GA_region_deathrateperday.csv"),DpDdataRR)
    CSV.write(string(SaveName,"/GA_region_newcaserateperday.csv"),CpDdataRR)
    CSV.write(string(SaveName,"/GA_region_cumulativeHospitalised.csv"),HHcumdataRR)

    #Save summary data at County Level
    #sum across counties
    GASummary = zeros(12,SIZES[4])
    GASummary[1,:] = sum(S,dims=1)
    GASummary[2,:] = sum(E,dims=1)
    GASummary[3,:] = sum(Ia,dims=1)
    GASummary[4,:] = sum(Is,dims=1)
    GASummary[5,:] = sum(R,dims=1)
    GASummary[6,:] = sum(Ihsub,dims=1)
    GASummary[7,:] = sum(Ihcrit,dims=1)
    GASummary[8,:] = sum(D,dims=1)
    GASummary[9,:] = sum(Htot,dims=1)
    GASummary[10,:] = sum(Dperday,dims=1)
    GASummary[11,:] = sum(NewCperday,dims=1)
    GASummary[12,:] = sum(Hcumulative,dims=1)

    GASum = DataFrame(GASummary)
    rename!(GASum , Symbol.(colnames))
    insertcols!(GASum, 1,:EpiState => ["S","E","Ia","Is","R","Ihsub","Ihcrit","D","Htot","Dperday","NCperday","Hcumulative"])

    CSV.write(string(SaveName,"/GA_Summary.csv"),GASum)

end
```

    WriteOutput! (generic function with 1 method)

```julia
#Used to turn the solution variable 1x(159*9*8) into an array that can be more easily parsed
function gatherStates!(sol,par)
    SDI,β_a,β_s,p,γ_e,γ_a,γ_s,γ_h,h,ξ,μ,S_ind,E_ind,
    Ia_ind,Is_ind,Ihsub_ind,Ihcrit_ind,R_ind,D_ind,T_Hosp_ind,
    NumPatches,NumAgeClasses,patchstates,patchstatesALIVE,
    patchstatesMOVEABLE,w,wy,wo,transport,ISMOVEFACTOR,w1,w2 = par

    #find time indices (in case of duplicates)
    UU = unique(sol.t)
    TS = Array{Int64,1}(undef,0)
    for aa = 1:length(UU)
       push!(TS,findall(sol.t .== UU[aa])[1])
    end

    Sols =  hcat(sol.u...)

    ARRAYSTORE = zeros(NumPatches,9,NumAgeClasses,length(TS)) #array of patches x epi states x ages x time

    for aa = 1:NumPatches
        for bb = 1:NumAgeClasses
        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],T_Hosp_ind[bb]]
            for cc=1:9
                for dd = 1:length(TS)
                    ARRAYSTORE[aa,cc,bb,dd] = Sols[INDS[cc],TS[dd]]
                end
            end
        end
    end
    return ARRAYSTORE,TS
end
```

    gatherStates! (generic function with 1 method)

```julia
#summary plot - plots the summed epi. states at the Georgia level i.e. dropping spatial dimension
function SummaryPlot!(JAM,T)
SZ = size(JAM)
#AA=dropdims(sum(JAM[aa,1,1:9,:],dims=1),dims=1)[end]
S = zeros(SZ[4])
E = zeros(SZ[4])
Ia = zeros(SZ[4])
Is = zeros(SZ[4])
R = zeros(SZ[4])
Ihs = zeros(SZ[4])
Ihc = zeros(SZ[4])
D = zeros(SZ[4])

for aa = 1:SZ[1]
    S .+= dropdims(sum(JAM[aa,1,1:9,:],dims=1),dims=1)
    E .+= dropdims(sum(JAM[aa,2,1:9,:],dims=1),dims=1)
    Ia .+= dropdims(sum(JAM[aa,3,1:9,:],dims=1),dims=1)
    Is .+= dropdims(sum(JAM[aa,4,1:9,:],dims=1),dims=1)
    R .+= dropdims(sum(JAM[aa,5,1:9,:],dims=1),dims=1)
    Ihs .+= dropdims(sum(JAM[aa,6,1:9,:],dims=1),dims=1)
    Ihc .+= dropdims(sum(JAM[aa,7,1:9,:],dims=1),dims=1)
    D .+= dropdims(sum(JAM[aa,8,1:9,:],dims=1),dims=1)
end

pl = plot(T,S,label="S",yguidefontsize=10,xlabel="Time (days)",ylabel="People",legend=:outertopright)
plot!(T,E,label="E")
plot!(T,Ia,label="Ia")
plot!(T,Is,label="Is")
plot!(T,R,label="R")
plot!(T,Ihs,label="Ihsub")
plot!(T,Ihc,label="Ihcrit")
plot!(T,D,label="D")

return pl
end
```

    SummaryPlot! (generic function with 1 method)

# Set up the model scripts

```julia
#######################################STUFFTO EDIT ######################
#Time stepping stuff
TIMERESOLUTION = 1 #MAKE THIS AN INTEGER IN DAYS!!!!!! REQUIRED FOR SAVING LATER ON
T_END = 150
tspanrun = (0.0,T_END*1.0) #run the simulations between 0 days and 400 days
TODAY = Date(2020,3,28) #Date format of time 0

#Load in Data used to set time 0
CasesbyCounty = CSV.read("./data/28March2020.csv");
HospitalRegions = CSV.read("./data/ga_hosp_regions.csv")
HospRegions = HospitalRegions.HospRegion

#String denoting where to save output

#SAVEFOLD = "SimulationOutput/March28_sim"
SAVEFOLD = "SimulationOutput/March28_eig_RunApr23"

#initial condition setting to use (1 - uses a simulation method county by county; 2 - uses the state level model to partition to counties)
#other choices are archives right now from previous attempts

#CHOICES = 1
CHOICES = 2

#######################################STUFFTO EDIT ######################
```

    2

```julia
#Model states
NumPatches = 159 ; #159 counties
NumAgeClasses = 9; # 0-9,10-19,20-29,30-39,40-49,50-59,60-69,70-79,80+
NumEpiStates = 9; # S,E,Ia,Is,R,Ihsub,Ihcrit,D, cumulative hosp

patchstates = NumAgeClasses*NumEpiStates
patchstatesMOVEABLE = (NumEpiStates-4)*NumAgeClasses  #D cannot move. Ihsub, Ihcrit cannot move.
#Say that I_s may be able to move - can modify this assumption later.
patchstatesALIVE = (NumEpiStates-2)*NumAgeClasses #all but D and cumulative hosp are alive
transport = zeros(NumPatches,patchstatesMOVEABLE);

#Epidemiological state indexing - just for within one patch.
S_ind = 1:NumAgeClasses
E_ind = (1*NumAgeClasses+1):2*NumAgeClasses
Ia_ind = (2*NumAgeClasses+1):3*NumAgeClasses
Is_ind = (3*NumAgeClasses+1):4*NumAgeClasses
R_ind = (4*NumAgeClasses+1):5*NumAgeClasses
Ihsub_ind = (5*NumAgeClasses+1):6*NumAgeClasses
Ihcrit_ind = (6*NumAgeClasses+1):7*NumAgeClasses
D_ind = (7*NumAgeClasses+1):8*NumAgeClasses
T_Hosp_ind = (8*NumAgeClasses+1):9*NumAgeClasses

#Age class information
MinAge = [0,10,20,30,40,50,60,70,80]
MiddleAge = [5,15,25,35,45,55,65,75,90];
```

```julia
#read in the GA level data

ODL29 = CSV.read("./data/ga_countyOD_29_younger.csv") # transport matrix for 0-29
OD30_55 = CSV.read("./data/ga_countyOD_30_54.csv") #transport matrix for 30-55 (assume till 59)
OD55_up = CSV.read("./data/ga_countyOD_55_older.csv") #transport matrix for 55+ (assume for 60+)
GA_county_size = CSV.read("./data/GA_county_decade_age.csv") #age distributions by county

#FIP is a way of designating counties (https://en.wikipedia.org/wiki/List_of_FIPS_codes_for_Georgia_(U.S._state))
PatchFIPS = GA_county_size.FIP
PatchCounty = GA_county_size.County
#NOTE - will assign initial states by GA_county_size
#USE THE FIP indices here to keep everything ordered.

#Here we load in origin-destination data (159x159) showing trip numbers between places (including self-loops) which is sampled (not normalised)
#we need to translate into a form useful to the model
#to do this we want to:
#a) normalise such that we know the proportion of trips from an origin to a destination
#b) replace the diagonal (self trips) with negative elements that sum to the proportion of those leaving
#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]
#we are going to get transport matrices w for each of the age structured commuting data

#load and create commuter matrix
## Commuter data for those aged 30-55
COMMRULES = OD30_55[:,2:end] #matrix of data showing origin (row) and destination (col) patterns by county.
#data is total number of trips from using a field sampling approach. Want to turn into a proportion.
#Note 1: this is a sample. Needs normalising. Note 2: first column was FIP
JJ = 1.0*convert(Array,COMMRULES)#convert to matrix of Floats
sJJ = sum(JJ,dims=1) #sum over columns - this is the total arrivals (including self)
for aa =1:159
    JJ[:,aa] = JJ[:,aa] ./ sJJ[aa]
end
#find proportion of those staying in county
PROPSTAY = JJ[1:size(JJ, 1)+1:end]
#reassign diagonal elements as those who will be leaving county
JJ[1:size(JJ, 1)+1:end] = -(1.0 .- PROPSTAY)
#transpose matrix as we require destination to be the rows and origin to be the columns
w = JJ';

## Commuter data for those aged less than 30
COMMRULES = ODL29[:,2:end] #matrix of data showing origin (row) and destination (col) patterns by county.
#Note 1: this is a sample. Needs normalising. Note 2: first column was FIP
JJ = 1.0*convert(Array,COMMRULES)#convert to matrix of Floats
sJJ = sum(JJ,dims=1) #sum over columns - this is the total arrivals (including self)
for aa =1:159
    JJ[:,aa] = JJ[:,aa] ./ sJJ[aa]
end
#find proportion of those staying in county
PROPSTAY = JJ[1:size(JJ, 1)+1:end]
#reassign diagonal elements as those who will be leaving county
JJ[1:size(JJ, 1)+1:end] = -(1.0 .- PROPSTAY)
#transpose matrix as we require destination to be the rows and origin to be the columns
wy = JJ';

## Commuter data for those 55 and older
COMMRULES = OD55_up[:,2:end] #matrix of data showing origin (row) and destination (col) patterns by county.
#Note 1: this is a sample. Needs normalising. Note 2: first column was FIP
JJ = 1.0*convert(Array,COMMRULES)#convert to matrix of Floats
sJJ = sum(JJ,dims=1) #sum over columns - this is the total arrivals (including self)
for aa =1:159
    JJ[:,aa] = JJ[:,aa] ./ sJJ[aa]
end
#find proportion of those staying in county
PROPSTAY = JJ[1:size(JJ, 1)+1:end]
#reassign diagonal elements as those who will be leaving county
JJ[1:size(JJ, 1)+1:end] = -(1.0 .- PROPSTAY)
#transpose matrix as we require destination to be the rows and origin to be the columns
wo = JJ';

```

```julia
#Make plot of the age-distribution across counties
makeplot=1
if makeplot==1
gr()
plot(MiddleAge,100 .* convert(Array,(GA_county_size[1,3:11])),xlim=(0,100),title="Age distribution by county",
    ylabel="Percent of county population (%)",xlabel="Age",legend=false,yguidefontsize=8,markershape = :auto,markersize=2)
for aa=2:159
    plot!(MiddleAge,100 .* convert(Array,(GA_county_size[aa,3:11])),markershape = :auto,markersize=2)
end
savefig("./figures/GAagestructure.pdf")
current()
end
```

```julia
#Basic epidemiology to apply across all patches
SocialDistancingInterventions=[1.0,0.75,0.5,0.25] #proportionalities to apply to both β_A and β_S (1.0 is no distancing)
β_A = 0.3 #asymptomatic transmission rate (per day)
β_S = 0.6 #symptomatic transmission rate (per day)
γ_e = (1.0/4.0) #mean exposed period (days)
γ_a = (1.0/6.0) #mean asymptomatic period (days)
γ_s = (1.0/6.0) #mean symptomatic period (days)
γ_h = (1.0/10.0) # mean hospitalised period (days)

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
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)
ξ = [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)
μ = 0.5.*ones(NumAgeClasses) #mortality fraction

ISMOVEFACTOR = 1.0 # how much to let Is people commute (must be 0-1)
w1 = 16/24 ; #fraction of time mixing with local pop.
w2 = 1.0 - w1 ;#fraction of time mixing with commuter pop.

#set parameter sets to use to initialise/drive the model
param = [1.0,β_A,β_S,p,γ_e,γ_a,γ_s,γ_h,h,ξ,μ,S_ind,E_ind,
    Ia_ind,Is_ind,Ihsub_ind,Ihcrit_ind,R_ind,D_ind,T_Hosp_ind,
    NumPatches,NumAgeClasses,patchstates,patchstatesALIVE,
    patchstatesMOVEABLE,w,wy,wo,transport,ISMOVEFACTOR,w1,w2]

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,T_Hosp_ind]

```

    20-element Array{Any,1}:
     1.0
     0.3
     0.6
      [0.95, 0.95, 0.9, 0.8, 0.7, 0.6, 0.4, 0.2, 0.2]
     0.25
     0.16666666666666666
     0.16666666666666666
     0.1
      [0.001, 0.003, 0.012, 0.032, 0.049, 0.102, 0.166, 0.243, 0.273]
      [0.05, 0.05, 0.05, 0.05, 0.063, 0.122, 0.27399999999999997, 0.43200000000000005, 0.7090000000000001]
      [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5]
      1:9
      10:18
      19:27
      28:36
      37:45
      46:54
      55:63
      64:72
      73:81

 ### Set the initial conditions

We now have the ingredients. Now we need to mix in the ingredients.
Want to set initial conditions, being cognisant of the GA data.
Two main ways right now are: a simulation approach (CHOICE=5); and an "eigenvector" approach (CHOICE=7)

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.

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.

```julia
### Setting initial conditions
#Assign county age structured data to initial state
GAPOP = sum(GA_county_size.POPSIZE) #total population of GA

MAT = convert(Array, (GA_county_size[:,3:(end-1)] .* GA_county_size.POPSIZE))
MAT2 = MAT ./ GAPOP
#sum(MAT2)#should = 1. Yes!
#MAT2 is the fraction of GA in eaah age class in each county.

#want to now assign MAT2 to initial vector u00
u00 = zeros(NumPatches*NumAgeClasses*NumEpiStates)

for aa = 1:NumPatches
        u00[(aa-1)*patchstates .+ S_ind] = MAT2[aa,:]
end

# load case data and figure out initial conditions
tspan = (0.0,3000.0)

#start epidemic with 1 exposed individual ages 20-29 in Fulton county ...
INDFult = findall(GA_county_size.FIP .== 13121)
INDFult=INDFult[1]#patch index
u00[S_ind[3] + patchstates*(INDFult-1)] = u00[S_ind[3] + patchstates*(INDFult-1)] - 1.0/GAPOP #minus from S
u00[E_ind[3] + patchstates*(INDFult-1)] = u00[E_ind[3] + patchstates*(INDFult-1)] + 1.0/GAPOP #add to E
#...and with 1 exposed individual ages 20-29 in Dougherty county
INDDougherty = findall(GA_county_size.FIP .== 13095) #dougherty county (Albany)
INDDougherty=INDDougherty[1]#patch index
u00[S_ind[3] + patchstates*(INDDougherty-1)] = u00[S_ind[3] + patchstates*(INDDougherty-1)] - 1.0/GAPOP #minus from S
u00[E_ind[3] + patchstates*(INDDougherty-1)] = u00[E_ind[3] + patchstates*(INDDougherty-1)] + 1.0/GAPOP #add to E
initu = u00

#set the ODE problem using this initial condition (1 infection in Dougherty and 1 infection in Fulton from people aged 20-29)
prob_init = ODEProblem(SpatialAgeStructuredEpidemiologicalModelMixAge!,initu,tspan,param)

#set an terminate integration affect.
affect!(integrator) = terminate!(integrator)

#state wide deaths and cases
TOT_deaths = sum(CasesbyCounty.deaths)/GAPOP
TOT_cases = sum(CasesbyCounty.cases)/GAPOP

#now use this information to initialise the model.
#Currently done with two options:
###CHOICES = 1 :  County simulation method (CSM)
#uses the initial outbreak defined above to find epi. distributions
#in each county that hit death data in the GA data input county-by-county.
#These conditions are then merged together to make the final initial condition.
###CHOICES = 2 :  State level model (SLM)
#uses a non-spatial model for the whole of GA to find epi. distributions
# that hit deaths for the whole of GA. Then apportions these out to counties
# given the proportions of deaths recorded in each county.

if  CHOICES == 1
    ## County Simulation Method (CSM) ##

    #in general should use hospitalisations rather than cases, unless death available
    #case numbers highly likely to be underestimates
    Hprop = 0.2 #20.48% proportion of cases that are hospitalised - appears consistently around 20% (April 15)
    #- this is taken from Georgia Department of Public Health COVID-19 Daily Status Report
    #https://dph.georgia.gov/covid-19-daily-status-report

    #cases county by county
    for aa = 1:length(GA_county_size.FIP) #for each county in GA
    print(aa) #print index to help show progress

        IND = findall(CasesbyCounty.FIP .== GA_county_size.FIP[aa]) #find patch index
        DEATH = 0.0
        if length(IND) == 1 #ignore unknown category in case data
            DEATH = CasesbyCounty.deaths[IND[1]] #deaths recorded in this patch
            HOSP = Hprop * CasesbyCounty.cases[IND[1]] #hsopitalised assumed in this patch
        else
            HOSP = 0.5 # if no cases recorded. Assign some proportion to be hospitalised. Most likely non-zero.
        end

        #use this info. to simulate to epi. distributions in this patch

        if DEATH >0.0
            #If there are deaths simulate until the number of deaths is reached
             #set up callback to stop integration when number of cases reaches a threshold
            CaseThreshold = DEATH/GAPOP
            #sum of S in patch - (all in patch - thresh)
            condition(u,t,integrator)  = sum(u[D_ind .+ patchstates*(aa-1)]) - CaseThreshold  #this must be equal to zero at stopping condition
            affect!(integrator) = terminate!(integrator)
            cb = ContinuousCallback(condition,affect!)
            #solve the initial spin up, checking for positivity and using callback definition
            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)
            #assign end spinup states per county into new initial condition.
            END = sol_init.u[end]
            u00[(1+patchstates*(aa-1)):(patchstates*aa)] = END[(1+patchstates*(aa-1)):(patchstates*aa)]
        else
            #If no deaths simulate until the number hospitalised (both Ihsub and Ihcrit) is reached
            #set up callback to stop integration when number of cases reaches a threshold
            CaseThreshold = HOSP/GAPOP
            #sum of S in patch - (all in patch - thresh)
            condition2(u,t,integrator)  = sum(u[T_Hosp_ind .+ patchstates*(aa-1)]) - CaseThreshold  #this must be equal to zero at stopping condition
            #condition2(u,t,integrator)  = sum(u[[Ihsub_ind;Ihcrit_ind] .+ patchstates*(aa-1)]) - CaseThreshold  #this must be equal to zero at stopping condition
            affect2!(integrator) = terminate!(integrator)
            cb2 = ContinuousCallback(condition2,affect2!)
            #solve the initial spin up, checking for positivity and using callback definition
            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)
            #assign end spinup states per county into new initial condition.
            END = sol_init.u[end]
            u00[(1+patchstates*(aa-1)):(patchstates*aa)] = END[(1+patchstates*(aa-1)):(patchstates*aa)]
        end
    end

elseif CHOICES == 2
    ## State Level Model (SLM) ##

    #1.set up initial conditions for a GA level model
    #Simple model to total deaths at state level. Then reapportion in the county model.
    Singleu0 = zeros(NumAgeClasses*NumEpiStates)
    Singleu0[S_ind] = sum(MAT2,dims=1)
    Singleu0[S_ind[3]] = Singleu0[S_ind[3]] - 1.0/GAPOP
    Singleu0[E_ind[3]] = 1.0/GAPOP  #seed the epidemic (1 person 20-29).

    prob_s_init = ODEProblem(SingleNode!,Singleu0,tspan,singlenodepar)

    CaseThreshold = TOT_deaths #as a fraction of GA pop.

    condition3(u,t,integrator)  = CaseThreshold - sum(u[D_ind])   #this must be equal to zero at stopping condition
    affect3!(integrator) = terminate!(integrator)
    cb3 = ContinuousCallback(condition3,affect3!)

    #solve the initial spin up, checking for positivity and using callback definition
    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)

    #2. Now use this epi. distribution and assign it to each county
    START = sol_s_init.u[end] #initial distributions at state scale (age*epi classes)
    SumD = sum(START[D_ind])*GAPOP #sum of all deaths in model
    TOTD =  sum(CasesbyCounty.deaths) #total deaths in data
    TOTU =  sum(CasesbyCounty.deaths[CasesbyCounty.FIP.==-999]) #total deaths with unknown assignment
    L1 = length(GA_county_size.FIP) #all counties
    NUMUNKNOWN = sum(CasesbyCounty.FIP.==-999)
    L2 = length(CasesbyCounty.FIP)- NUMUNKNOWN - sum(CasesbyCounty.deaths .== 0) #num counties with cases

    for aa = 1:159
        countypop = GA_county_size.POPSIZE[aa]

        IND = findall(CasesbyCounty.FIP .== GA_county_size.FIP[aa]) #find patch index
        if length(IND) == 1 #ignore unknown category in case data
            D_county = 1.0 * CasesbyCounty.deaths[IND[1]]
            if D_county == 0.0
                D_county = TOTU/(L1-L2) #if this county has no recorded cases assign from the unknown pool.
            end
        else
            D_county = TOTU/(L1-L2) #if this county has no recorded cases assign from the unknown pool.
        end

        RATIO = 1.0*(D_county/TOTD) #county ratio - ratio of deaths in county to total

        NEW = START .* RATIO  .* GAPOP
        JAM = zeros(NumAgeClasses)
        Ages = MAT[aa,:] # age distribution in this county

        for bb=1:NumAgeClasses
           #calculate sum of all but S and cumulative Hosp. this age class
            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]]])
           #Now find what S for this age class should be
            NEW[S_ind[bb]] = Ages[bb] - JAM[bb]
        end

        NEW = NEW ./ GAPOP

        #check if any element in NEW is less than 0 - would indicate problems in preserving county demographics.
        if sum(isless.(NEW,0)) < 0
            print("WARNING. u LESS THAN 0")
            NEW[isless.(NEW,0)] = 0.0
        end

        #assign to intiial condition
        u00[(1+patchstates*(aa-1)):(patchstates*aa)] = NEW

    end

end

#assign end spinup states as new initial condition.

```

### Run the simulations

Now that the initial conditions are set, we are ready to bake the model.
Note that in the above conditions are set without any social distancing interventions. Currently, these are applied at time =0. That is here.
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.

## Forecasting phase

To run the forecasting phase, we train the data with no social distancing to March 28th using the non-spatial model; then we train the model to April 28th using a scenario with a 50% reduction in transmission rates. We then distribute to the spatial model, MAGE, and use this to forecast forward.

```julia
#Model Calibration phase for forecasting
SAVEFOLD = "SimulationOutput/28thApr_Run_int_1stMay"
CasesbyCounty2 = CSV.read("./data/7pm_28Apr2020.csv");
TODAY = Date(2020,4,28) #Date format of time 0

#1. Run state level model to March 28 with no SD
    TOT_deaths = sum(CasesbyCounty.deaths)/GAPOP
    TOT_deaths2 = sum(CasesbyCounty2.deaths)/GAPOP

    tspan = (0.0,5000.0)
    Singleu0 = zeros(NumAgeClasses*NumEpiStates)
    Singleu0[S_ind] = sum(MAT2,dims=1)
    Singleu0[S_ind[3]] = Singleu0[S_ind[3]] - 1.0/GAPOP
    Singleu0[E_ind[3]] = 1.0/GAPOP  #seed the epidemic (1 person 20-29).
    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,T_Hosp_ind]

    prob_s_init = ODEProblem(SingleNode!,Singleu0,tspan,singlenodepar)

    CaseThreshold = TOT_deaths #as a fraction of GA pop.

    condition3(u,t,integrator)  = CaseThreshold - sum(u[D_ind])   #this must be equal to zero at stopping condition
    affect3!(integrator) = terminate!(integrator)
    cb3 = ContinuousCallback(condition3,affect3!)

    #solve the initial spin up, checking for positivity and using callback definition
    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)

#2. Run state level model from this point foward to Apr 28 with 50% SD
    STAR = sol_s_init.u[end]
    singlenodepar =[0.5,β_A,β_S,p,γ_e,γ_a,γ_s,γ_h,h,ξ,μ,S_ind,E_ind,Ia_ind,Is_ind,R_ind,Ihsub_ind,Ihcrit_ind,D_ind,T_Hosp_ind]

    prob_t_init = ODEProblem(SingleNode!,STAR,tspan,singlenodepar)
    CaseThreshold = TOT_deaths2 #as a fraction of GA pop.

    condition3(u,t,integrator)  = CaseThreshold - sum(u[D_ind])   #this must be equal to zero at stopping condition
    affect3!(integrator) = terminate!(integrator)
    cb3 = ContinuousCallback(condition3,affect3!)

    #solve the initial spin up, checking for positivity and using callback definition
    sol_t_init = solve(prob_t_init,Tsit5(), abstol=1e-8,reltol=1e-8,saveat=0.1,isoutofdomain=(y,p,t)->any(x->x<0,y), callback = cb3)

#3. Distribute to counties
    u22 = zeros(NumPatches*NumAgeClasses*NumEpiStates)
    # Now use this epi. distribution and assign it to each county
    START = sol_t_init.u[end] #initial distributions at state scale (age*epi classes)
    SumD = sum(START[D_ind])*GAPOP #sum of all deaths in model
    TOTD =  sum(CasesbyCounty2.deaths) #total deaths in data
    TOTU =  sum(CasesbyCounty2.deaths[CasesbyCounty2.FIP.==-999]) #total deaths with unknown assignment
    L1 = length(GA_county_size.FIP) #all counties
    NUMUNKNOWN = sum(CasesbyCounty2.FIP.==-999)
    L2 = length(CasesbyCounty2.FIP)- NUMUNKNOWN - sum(CasesbyCounty2.deaths .== 0) #num counties with cases

    for aa = 1:159
        countypop = GA_county_size.POPSIZE[aa]

        IND = findall(CasesbyCounty2.FIP .== GA_county_size.FIP[aa]) #find patch index
        if length(IND) == 1 #ignore unknown category in case data
            D_county = 1.0 * CasesbyCounty2.deaths[IND[1]]
            if D_county == 0.0
                D_county = TOTU/(L1-L2) #if this county has no recorded cases assign from the unknown pool.
            end
        else
            D_county = TOTU/(L1-L2) #if this county has no recorded cases assign from the unknown pool.
        end

        RATIO = 1.0*(D_county/TOTD) #county ratio - ratio of deaths in county to total

        NEW = START .* RATIO  .* GAPOP
        JAM = zeros(NumAgeClasses)
        Ages = MAT[aa,:] # age distribution in this county

        for bb=1:NumAgeClasses
           #calculate sum of all but S and cumulative Hosp. this age class
            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]]])
           #Now find what S for this age class should be
            NEW[S_ind[bb]] = Ages[bb] - JAM[bb]
        end

        NEW = NEW ./ GAPOP

        #check if any element in NEW is less than 0 - would indicate problems in preserving county demographics.
        if sum(isless.(NEW,0)) < 0
            print("WARNING. u LESS THAN 0")
            NEW[isless.(NEW,0)] = 0.0
        end

        #assign to intiial condition
        u22[(1+patchstates*(aa-1)):(patchstates*aa)] = NEW

    end

#4. forecast forward from Apr 28 to May 1 at 50% SD, then introduce interventions

mkdir(SAVEFOLD)
for bb = 1:length(SocialDistancingInterventions)

    SAVEFOLD2 = string(SAVEFOLD,"/Social_",bb)

    #set parameters for this intervention
    pars = [0.5,β_A,β_S,p,γ_e,γ_a,γ_s,γ_h,h,ξ,μ,S_ind,E_ind,
    Ia_ind,Is_ind,Ihsub_ind,Ihcrit_ind,R_ind,D_ind,T_Hosp_ind,
    NumPatches,NumAgeClasses,patchstates,patchstatesALIVE,
    patchstatesMOVEABLE,w,wy,wo,transport,ISMOVEFACTOR,w1,w2]

    ##RELAXATION ON MAY 1st
    prob_a = ODEProblem(SpatialAgeStructuredEpidemiologicalModelMixAge!,u22,tspanrun,pars)
    condition(u,t,integrator) = t == 3.0  #3 days to May 1st
    affect!(integrator) = integrator.p[1] = SocialDistancingInterventions[bb]
    cb = DiscreteCallback(condition,affect!)
    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=[3.0],save_positions=(false,false))

    #save output

    WriteOutput!(PatchFIPS,PatchCounty,HospRegions,GAPOP,TODAY,TIMERESOLUTION,T_END,SOL_a,pars,SAVEFOLD2)

end

```

```julia

```


Sascha Spors,
Professorship Signal Theory and Digital Signal Processing,
Institute of Communications Engineering (INT),
Faculty of Computer Science and Electrical Engineering (IEF),
University of Rostock,
Germany

# Tutorial Digital Signal Processing

**Uniform Quantization, Dithering, Noiseshaping**,
Winter Semester 2021/22 (Master Course #24505)

- lecture: https://github.com/spatialaudio/digital-signal-processing-lecture
- tutorial: https://github.com/spatialaudio/digital-signal-processing-exercises

Feel free to contact lecturer frank.schultz@uni-rostock.de

# Fundamentals of Quantization

## Packages / Functions

We import the required packages first and put some functions here that we will frequently use.

```python
# most common used packages for DSP, have a look into other scipy submodules
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from scipy import signal

# audio write and play stuff
import soundfile as sf  # requires 'pip install soundfile'
# last tested with soundfile-0.10.3
```

```python
def my_xcorr2(x, y, scaleopt='none'):
    r""" Cross Correlation function phixy[kappa] -> x[k+kappa] y

    input:
    x   input signal shifted by +kappa
    y   input signal
    scaleopt   scaling of CCF estimator
    output:
    kappa   sample index
    ccf   correlation result
    """
    N = len(x)
    M = len(y)
    kappa = np.arange(0, N+M-1) - (M-1)
    ccf = signal.correlate(x, y, mode='full', method='auto')
    if N == M:
        if scaleopt == 'none' or scaleopt == 'raw':
            ccf /= 1
        elif scaleopt == 'biased' or scaleopt == 'bias':
            ccf /= N
        elif scaleopt == 'unbiased' or scaleopt == 'unbias':
            ccf /= (N - np.abs(kappa))
        elif scaleopt == 'coeff' or scaleopt == 'normalized':
            ccf /= np.sqrt(np.sum(x**2) * np.sum(y**2))
        else:
            print('scaleopt unknown: we leave output unnormalized')
    return kappa, ccf
```

```python
def uniform_midtread_quantizer(x, deltaQ):
    r"""uniform_midtread_quantizer from the lecture:
    https://github.com/spatialaudio/digital-signal-processing-lecture/blob/master/quantization/linear_uniform_quantization_error.ipynb
    commit: b00e23e
    note: we renamed the second input to deltaQ, since this is what the variable
    actually represents, i.e. the quantization step size

    input:
    x   input signal to be quantized
    deltaQ    quantization step size
    output:
    xq   quantized signal
    """
    # [-1...1) amplitude limiter
    x = np.copy(x)
    idx = np.where(x <= -1)
    x[idx] = -1
    idx = np.where(x > 1 - deltaQ)
    x[idx] = 1 - deltaQ
    # linear uniform quantization
    xq = deltaQ * np.floor(x/deltaQ + 1/2)
    return xq
```

```python
def my_quant(x, Q):
    r"""Saturated uniform midtread quantizer

    input:
    x  input signal
    Q  number of quantization steps
    output:
    xq quantized signal

    Note: for even Q in order to retain midtread characteristics,
    we must omit one quantization step, either that for lowest or the highest
    amplitudes. Typically the highest signal amplitudes are saturated to
    the 'last' quantization step. Then, in the special case of log2(N)
    being an integer the quantization can be represented with bits.
    """
    tmp = Q//2  # integer div
    quant_steps = (np.arange(Q) - tmp) / tmp  # we don't use this

    # forward quantization, round() and inverse quantization
    xq = np.round(x*tmp) / tmp
    # always saturate to -1
    xq[xq < -1.] = -1.
    # saturate to ((Q-1) - (Q\2)) / (Q\2), note that \ is integer div
    tmp2 = ((Q-1) - tmp) / tmp  # for odd N this always yields 1
    xq[xq > tmp2] = tmp2
    return xq
```

## Quantization Process and Error

Quantization generates signals that have discrete values $x_q[k]$, $x_q(t)$ from signals with continuous values $x[k]$, $x(t)$.

For quantization, the signals can be both, discrete and continuous in time.
However, a signal that is discrete in time **and** discrete in value is termed a **digital** signal.
Only digital signals can be processed by computers.
Here the quantization of discrete-time signals is treated due to practical importance.

To describe quantization analytically, the model in the figure below is used.

The input and output signal differ by the so called quantization error (quantization noise) $e[k]$, that is defined as

\begin{equation}
e[k] = x_q[k] - x[k],
\end{equation}

so that the error constitutes an additive superposition

\begin{equation}
x[k] + e[k] = x_q[k]
\end{equation}

To use this error model, some assumption have to be made.
The quantization noise shall be uniformly distributed, which then can be modeled with the probability density function (PDF) $p_e(\theta) = \frac{1}{\Delta Q} \mathrm{rect}(\frac{\theta_e}{\Delta Q})$, where $\Delta Q$ denotes the quantization step size and $\theta_e$ the amplitudes of the quantization error signal.
This PDF is shown below.

```python
plt.figure(figsize=(4, 2))
plt.plot((-1, -1/2, -1/2, +1/2, +1/2, +1), (0, 0, 1, 1, 0, 0), lw=3)
plt.xlim(-1, 1)
plt.ylim(-0.1, 1.1)
plt.xticks((-0.5, +0.5), [r'-$\frac{\Delta Q}{2}$', r'+$\frac{\Delta Q}{2}$'])
plt.yticks((0, 1), [r'0', r'+$\frac{1}{\Delta Q}$'])
plt.xlabel(r'$\theta_e$')
plt.ylabel(r'$p_e(\theta)$')
plt.title(
    r'$p_e(\theta) = \frac{1}{\Delta Q} \mathrm{rect}(\frac{\theta_e}{\Delta Q})$')
plt.grid(True)
```

Furthermore, it is assumed that $e[k]$ is not correlated with $x[k]$.
That this is not necessarily the case can be demonstrated with the help of some straightforward examples.

```python
Q = 9  # odd, number of quantization steps
N = 100
k = np.arange(2*N)
x = np.sin(2*np.pi/N*k)
xq = my_quant(x, Q)
e = xq-x
# actually stem plots would be correct, for convenience we plot as line style
plt.plot(k, x, 'C2', lw=3, label=r'$x$')
plt.plot(k, xq, 'C0o-', label=r'$x_q$')
plt.plot(k, e, 'C3', label=r'$e=x_q-x$')
plt.plot(k, k*0+1/(Q-1), 'k:', label=r'$\frac{\Delta Q}{2}$')
plt.xlabel(r'$k$')
plt.legend()
plt.grid(True)
```

A sine signal is quantized with $Q=9$ quantization steps.
A periodicity of the quantization noise can be easily identified.
For odd $Q$, the maximum amplitude of the quantization error can be estimated to

$$\frac{\Delta Q}{2}=\frac{\frac{2}{Q-1}}{2}=\frac{1}{Q-1}=\frac{1}{8}=0.125.$$

The auto-correlation function of the error signal $e[k]$ is presented next.

```python
kappa, acf = my_xcorr2(e, e, 'unbiased')
plt.plot(kappa, acf)
plt.xlim(-175, +175)
plt.xlabel(r'$\kappa$')
plt.ylabel(r'$\phi_{ee}[\kappa]$')
plt.title('ACF of quantization error')
plt.grid(True)
```

If $e[k]$ would be exactly following the probability density function $p_e(\theta) = \frac{1}{\Delta Q} \mathrm{rect}(\frac{\theta_e}{\Delta Q})$, the auto-correlation function $\phi_{ee}[\kappa]=\delta[\kappa]$ results.

However, this is not observable in this example!

Instead, from the above plot, we can deduce that $e[k]$ is correlated to itself, i.e. it exhibits periodicity each 100 samples in phase, and each 50 sample out of phase.
The sine period is precisely 100 samples, thus the input signal and the quantization error are somewhat linked and not independent.
Thus, the error model assumption is violated. That is bad, since the sine signal allows for otherwise comparable simple analytical calculus.

The links between the signals can be furthermore confirmed with the help of the cross-correlation functions.
Their oscillating characteristics reveal that quantization error is highly correlated.

```python
plt.figure(figsize=(9, 3))

plt.subplot(1, 2, 1)
kappa, acf = my_xcorr2(e, x, 'unbiased')
plt.plot(kappa, acf)
plt.xlim(-170, +170)
plt.xlabel(r'$\kappa$')
plt.ylabel(r'$\phi_{e,x}[\kappa]$')
plt.title('CCF quantization error and input signal')
plt.grid(True)

plt.subplot(1, 2, 2)
kappa, acf = my_xcorr2(e, xq, 'unbiased')
plt.plot(kappa, acf)
plt.xlim(-170, +170)
plt.xlabel(r'$\kappa$')
plt.ylabel(r'$\phi_{e,xq}[\kappa]$')
plt.title('CCF quantization error and quantized signal')
plt.grid(True)
```

Therefore, the special case of sine signals is in fact not suited for the quantization model above.
Because of the simplicity of the involved calculation it is common practice to conduct this analysis for sine signals nevertheless, and signal-to-noise ratios in the data sheets of A/D converters are mostly stated for excitation with sine signals.
For random signals, the quantization model is only valid for high levels in the quantizer. For more information see
[Udo Zölzer, Digital Audio Signal Processing, Wiley](https://onlinelibrary.wiley.com/doi/book/10.1002/9780470680018)
(might be available as free access in your uni network)

-  Task:
Increase the (odd) number of quantization steps $Q$ and check what happens with the shape and amplitudes of the correlations functions. Hint: closer look to the amplitudes of the correlation signals.

## Quantization Modeling / Mapping

The mapping of the infinitely large continuous set of values to a discrete number of amplitude steps is realized with a transfer characteristic.
The height of the amplitude steps is $\Delta Q$.

From the lecture, we know that the following mapping is used in order to quantize the continuous amplitude signal $x[k]$

towards

\begin{equation}
x_Q[k] = g( \; \lfloor \, f(x[k]) \, \rfloor \; ),
\end{equation}

where $g(\cdot)$ and $f(\cdot)$ denote real-valued mapping functions, and $\lfloor \cdot \rfloor$ a rounding operation (**not necessarily the plain floor operation**).

### Uniform Saturated Midtread Quantization Characteristic Curve

With the introduced mapping, the uniform saturated midtread quantizer can be discussed.
This is probably the most important curve for uniform quantization due to its practical relevance for coding quantized amplitude values as bits. In general, the uniform midtread quantizer can be given as the mapping

\begin{equation}
x_Q[k] =  \frac{1}{Q \backslash 2} \cdot \lfloor (Q \backslash 2) \cdot x[k]\rfloor,
\end{equation}

where for $\lfloor \cdot \rfloor$ a rounding operation might used and $\backslash$ denotes integer division.
So the mapping functions $g$ and $f$ are simple multiplications.
At the beginning of this notebook, the function `my_quant` is implemented that realizes quantization based on this mapping.
The approach uses `numpy`'s `round` operation.
When asking for rounding, care has to be taken, which [approach](https://en.wikipedia.org/wiki/Rounding) shall be used.
Numpy rounds to the nearest **even** integer in contrast to e.g. Matlab's rounding to nearest integer.

Detailed analysis for `my_quant`:

- the quantization should be properly performed only for $-1 \leq x < 1$
- thus, it always saturates $x<-1$ towards $x_q = -1$
- in the case of an **odd** number of quantization steps $Q$, it saturates $x>+1$ towards $x_q = +1$. The quantization step size is $\Delta Q = \frac{2}{Q-1}$.
- In the case of an **even** number of quantization steps $Q$, it saturates $x>\frac{Q - 1 - \frac{Q}{2}}{\frac{Q}{2}} = 1-\frac{2}{Q}$ towards $x_q = \frac{Q - 1 - \frac{Q}{2}}{\frac{Q}{2}}=1-\frac{2}{Q}$. The quantization step size is $\Delta Q = \frac{2}{Q}$.

### AD / DA Converter Convention

The case of **even** $Q$ is practically used for virtually all analog/digital (AD) and digital/analog (DA) converters.

When additionally to the above statements

\begin{equation}
\log_2(Q)\in\mathbb{N}
\end{equation}

holds, it is meaningful to code the even and power of two $Q$ possible quantization steps with bits.

With $B\in\mathbb{N}$ denoting the number of bits, the number range convention for AD and DA converters is

\begin{equation}
\begin{split}
&-1\leq x \leq 1-2^{-(B-1)}\\
&-1\leq x \leq 1-\frac{2}{Q}
\end{split}
\end{equation}

using

\begin{equation}
Q=2^B
\end{equation}

quantization steps.
Values of $x$ outside this range will be saturated to the minimum $-1$ and maximum $1-\frac{2}{Q}$ quantization values in the quantization process.

For example, $B = 16$ bits are used to code [PCM data for CD](https://en.wikipedia.org/wiki/Compact_disc) audio quality.
Then we get the following quantities.

```python
B = 16  # number of bits
Q = 2**B  # number of quantization steps

# for even Q only:
deltaQ = 2/Q
# maximum quantize value:
xqmax = 1-2**(-(B-1))
# or more general for even Q:
xqmax = 1-deltaQ

print(' B = %d bits\n quantization steps Q = %d\n quantization step size %e' %
      (B, Q, deltaQ))
print(' smallest quantization value xqmin = -1')
print(' largest quantization value xqmax = %16.15f' % xqmax)
# B = 16 bits
# quantization steps Q = 65536
# quantization step size 3.051758e-05
# smallest quantization value xqmin = -1
# largest quantization value xqmax = 0.999969482421875
```

So called high definition audio uses 24 Bit. Video and photo typically uses 8-12 Bit quantization per color channel.

### Plotting the Midtread Curve

We now can visualize the characteristic curve for a simple, made up input signal, i.e. a monotonic increasing signal between $x_{max} = -x_{min}$ using an equidistant increment $\Delta x$ over sample index $k$.

Here, we use $x_{max} = 1.25$ and $\Delta x=0.001$ and assume that we start with $x_{min} = -1.25$ at $k=0$.
If $\Delta x$ is sufficiently small, the signal's amplitude can be interpreted as continuous straight line.
This straight line is degraded in a quantization process.
Plotting the quantization result over the input, results in the characteristic curve, in our example in the curve of the uniform saturated midtread quantizer.

Let us plot this.

**Please note:**
The quantizer `uniform_midtread_quantizer` known from lecture and `my_quant` yield the same results besides a slight detail: `uniform_midtread_quantizer` always exhibits an **even** number of quantization steps $Q$.
So, only for even $Q$ results are exactly identical.

We might verify this in the next plots as well.

```python
x = np.arange(-1.25, +1.25, 1e-3)
plt.figure(figsize=(4, 2))
plt.plot(x)  # actually a stem plot is correct
plt.ylim(-1.25, +1.25)
plt.xlabel(r'$k$')
plt.ylabel(r'$x[k]$')
plt.grid(True)
```

```python
Q = 9  # number of quantization steps, odd or even
deltaQ = 1/(Q//2)  # quantization step size, even/odd Q

xq = my_quant(x, Q)  # used in exercise
xumq = uniform_midtread_quantizer(x, deltaQ)  # as used in lecture

plt.figure(figsize=(6, 6))
plt.plot(x, xumq, 'C0', lw=2, label='uniform_midtread_quantizer()')
plt.plot(x, xq, 'C3', label='my_quant()')
plt.xticks(np.arange(-1, 1.25, 0.25))
plt.yticks(np.arange(-1, 1.25, 0.25))
plt.xlabel(r'input amplitude of $x$')
plt.ylabel(r'output ampliude of $x_q$')
plt.title(
    r'uniform saturated midtread quantization, Q={0:d}, $\Delta Q$={1:3.2f}'.format(Q, deltaQ))
plt.axis('equal')
plt.legend()
plt.grid(True)
```

The following exercises used to be a homework assignment as exam's prerequisite.

# Exercise 1: Uniform Saturated Midtread Characteristic Curve of Quantization

## Task

Check this quantizer curve for $Q=7$ and $Q=8$.
Make sure that you get the idea of the midtread concept (the zero is always quantized to zero) and saturation (for even $Q$) largest quantization step is saturated).

```python
def check_my_quant(Q):
    N = 5e2
    x = 2*np.arange(N)/N - 1
    xq = my_quant(x, Q)
    e = xq - x

    plt.plot(x, x, color='C2', lw=3, label=r'$x[k]$')
    plt.plot(x, xq, color='C3', label=r'$x_q[k]$')
    plt.plot(x, e, color='C0', label=r'$e[k] = x_q[k] - x[k]$')
    plt.xticks(np.arange(-1, 1.25, 0.25))
    plt.yticks(np.arange(-1, 1.25, 0.25))
    plt.xlabel('input amplitude')
    plt.ylabel('output amplitude')
    if np.mod(Q, 2) == 0:
        s = ' saturated '
    else:
        s = ' '
    plt.title(
        'uniform'+s+'midtread quantization with Q=%d steps, $\Delta Q$=%4.3e' % (Q, 1/(Q//2)))
    plt.axis('equal')
    plt.legend(loc='upper left')
    plt.grid(True)
```

```python
Q = 7  # number of quantization steps
deltaQ = 1 / (Q//2)  # general rule
deltaQ = 2 / (Q-1)  # for odd Q only
plt.figure(figsize=(5, 5))
check_my_quant(Q)
```

```python
Q = 8  # number of quantization steps
deltaQ = 1 / (Q//2)  # general rule
deltaQ = 2 / Q  # for even Q only
plt.figure(figsize=(5, 5))
check_my_quant(Q)
```

# Exercise 2: Quantization and Signal-to-Noise Ratio

From theory the **6dB / Bit rule of thumb** is well known for uniform quantization. It states that the signal-to-noise ratio increases by 6 dB for every additional bit that is spent to quantize the input data.
Hence,

\begin{equation}
\text{SNR in dB} = 6 \cdot B + \gamma,
\end{equation}

where $\gamma$ is a offset value in dB that depends on the PDF of the signal to be quantized.
Note that this rule of thumb assumes that the quantization error exhibits uniform PDF and is not correlated with the quantized signal.
We can see that this rule of thumb is inaccurate when quantizing a sine signal with small number of bits or an amplitude in the range of the quantization step. Then, the mentioned assumptions are not fulfilled. We will observe this in exercise 3.

We plot the function SNR of bits below for uniform, normal and Laplace PDF noises and a sine signal.
We should observe the slope of always 6 dB per Bit.
We should note the different absolute values of the SNR depending on the varying $\gamma$.

The `dBoffset` values are discussed in the lecture and in the textbook [Udo Zölzer, Digital Audio Signal Processing, Wiley](https://onlinelibrary.wiley.com/doi/book/10.1002/9780470680018).

```python
def check_quant_SNR(x, dBoffset, title):
    print('std: {0:f}, var: {1:f}, mean: {2:f} of x'.format(np.std(x), np.var(x), np.mean(x)))
    Bmax = 24
    SNR = np.zeros(Bmax+1)
    SNR_ideal = np.zeros(Bmax+1)

    for B in range(1, Bmax+1):  # start at 1, since zero Q is not meaningful
        xq = my_quant(x, 2**B)
        SNR[B] = 10*np.log10(np.var(x) / np.var(xq-x))
        SNR_ideal[B] = B*20*np.log10(2) + dBoffset  # 6dB/bit + offset rule

    plt.figure(figsize=(5, 5))
    plt.plot(SNR_ideal, 'o-', label='theoretical', lw=3)
    plt.plot(SNR, 'x-', label='simulation')
    plt.xticks(np.arange(0, 26, 2))
    plt.yticks(np.arange(0, 156, 12))
    plt.xlim(2, 24)
    plt.ylim(6, 148)
    plt.xlabel('number of bits')
    plt.ylabel('SNR / dB')
    plt.title(title)
    plt.legend()
    plt.grid(True)
    print('maximum achievable SNR = {0:4.1f} dB at 24 Bit (i.e. HD audio)'.format(SNR[-1]))
```

```python
N = 10000
k = np.arange(N)
```

```python
np.random.seed(4)
x = np.random.rand(N)
x -= np.mean(x)
x *= np.sqrt(1/3) / np.std(x)
dBoffset = 0
check_quant_SNR(x, dBoffset, 'Uniform PDF')
```

```python
Omega = 2*np.pi * 997/44100  # use a rather odd ratio: e.g. in audio 997 Hz / 44100 Hz
sigma2 = 1/2
dBoffset = -10*np.log10(2 / 3)
x = np.sqrt(2*sigma2) * np.sin(Omega*k)
check_quant_SNR(x, dBoffset, 'Sine')
```

```python
np.random.seed(4)
x = np.random.randn(N)
x -= np.mean(x)
x *= np.sqrt(0.0471) / np.std(x)
dBoffset = -8.5  # from clipping propability 1e-5
check_quant_SNR(x, dBoffset, 'Normal PDF')
```

```python
np.random.seed(4)
x = np.random.laplace(size=N)
pClip = 1e-5  # clipping propability
sigma = -np.sqrt(2) / np.log(pClip)
x -= np.mean(x)
x *= sigma / np.std(x)
dBoffset = -13.5  # empircially found for pClip = 1e-5
check_quant_SNR(x, dBoffset, 'Laplace PDF')
```

# Exercise 3: Dithering

The discrete-time sine signal

- $x[k]=A \cdot\sin(\frac{2\pi f_\text{sin}}{f_s}k)$ for
- $0\leq k<96000$ with
- sampling frequency $f_s=48\,\text{kHz}$ and
- $f_\text{sin}=997\,\text{Hz}$

shall be quantized with the saturated uniform midtread quantizer for $-1\leq x_q \leq 1-\Delta Q$ using $B$ bits, i.e. $Q=2^B$ number of quantization steps and quantization step size of $\Delta Q = \frac{1}{Q\backslash 2}$.

We should discuss different parametrizations for signal amplitude $A$ and number of bits $B$.

Before quantizing $x[k]$, a dither noise signal $d[k]$ shall be added to $x[k]$ according figure below.

This dither signal with small amplitudes aims at de-correlating the quantization error $e[k]$ from the quantized signal $x_q[k]$, which is especially important for small amplitudes of $x[k]$.
This technique is called **dithering**.
For $d[k]=0$ no dithering is applied.

Since the quantization error may be in the range $-\frac{\Delta Q}{2}\leq e[k]\leq \frac{\Delta Q}{2}$ (assuming uniform distribution), it appears reasonable to use a dither noise with a probability density function (PDF) of

\begin{equation}
p_\text{RECT}(d)=\frac{1}{\Delta Q}\,\text{rect}\left(\frac{d}{\Delta Q}\right),
\end{equation}

i.e. a **zero-mean, uniformly distributed noise** with maximum amplitude $|d[k]|=\frac{\Delta Q}{2}$.
It can be shown that this dither noise improves the quality of the quantized signal.
However, there is still a noise modulation (i.e. a too high correlation between $x_q[k]$ and $e[k]$) that depends on the amplitude of the input signal.

The noise modulation can be almost completely eliminated with a **zero-mean noise** signal exhibiting a **symmetric triangular PDF**

\begin{equation}
p_\text{TRI}(d)=\frac{1}{\Delta Q}\,\text{tri}\left(\frac{d}{\Delta Q}\right)
\end{equation}

with maximum amplitude $|d[k]|=Q$.
By doing so, an almost ideal decorrelation between $x_q[k]$ and $e[k]$ is realized.
In audio, this technique is called TPDF-Dithering (Triangular Probability Density Function Dithering) and can be applied in the mastering process of audio material that is to be distributed e.g. on a CD or via streaming.

## Task

To get an impression on how dithering may be implemented and what quantized signals sound like, the following exercises shall be performed.

- Generate the sine signal $x[k]$ defined above.

- Generate the dither noise $d_\text{RECT}[k]$ according to the PDF $p_\text{RECT}(d)$. Check the resulting amplitude and distribution carefully. The length of $d_\text{RECT}[k]$ and $x[k]$ must be equal.

- Generate the dither noise $d_\text{TRI}[k]$ according to the PDF $p_\text{TRI}(d)$. Check the resulting amplitude and distribution carefully. The length of $d_\text{TRI}[k]$ and $x[k]$ must be equal.

- Add each dither noise $d_\text{RECT}[k]$ and $d_\text{TRI}[k]$ individually to $x[k]$. Together with the case of no dithering we now have three signals to be quantized.

- Quantize these signals individually using `my_quant(x,Q)` with $Q$ quantization steps.

- Plot the midtread characteristic curve.

- Plot the histogram of the dither noises as estimate of its PDF.

- Plot the histogram of the error noises as estimate of its PDF.

- Plot the sine signal, the dithered signal, the quantized signal and the quantization error signal in one diagram for all three cases.

- Calculate and plot the CCF of the signals $x_q[k]$ and $e[k]$ for all three cases.

- Interpret the obtained graphics.

- For each case, render WAV files from $x[k]$, $x[k]+d[k]$, $x_q[k]$ und $e[k]$ and listen to them. **Be careful! Do not harm your ears!** Pay special attention to the sound of the quantization error, how it is correlated with the quantized signal and how loud it appears.

- Consider the 5 cases

    1. $B=16$ Bit, $A=1-\Delta Q$
    2. $B=16$ Bit, $A=\Delta Q$
    3. $B=3$ Bit, $A=1-\Delta Q$
    4. $B=3$ Bit, $A=\Delta Q$
    5. $B=3$ Bit, $A=\frac{\Delta Q}{2}$

In the last case the signal has amplitude even below the quantization step size $\Delta Q$. You might verify by listening that the sine is still perceivable if dithering is applied, but not if no dithering is applied.

**Again: Be careful! Do not harm your ears!**
The signal amplitude $A$ and chosen $B$ is directly related to the playback level!

**Warning again: start with very very low playback level, find the loudest signal first and then increase volume to your convenience**

## Solution

The task asks for repeated steps.
This is perfectly handled by a little function that solves the repeating subtasks.

```python
fs = 48000
N = 2*fs
k = np.arange(0, N)
fsin = 997
```

```python
def check_dithering(x, dither, Q, case):

    deltaQ = 1 / (Q//2)  # general rule

    # dither noise
    pdf_dither, edges_dither = np.histogram(dither, bins='auto', density=True)
    xd = x + dither

    # quantization
    xq = my_quant(xd, Q)
    e = xq-x
    pdf_error, edges_error = np.histogram(e, bins='auto', density=True)

    # write wavs
    sf.write(file='x_'+case+'.wav', data=x,
             samplerate=48000, subtype='PCM_24')
    sf.write(file='xd_'+case+'.wav', data=xd,
             samplerate=48000, subtype='PCM_24')
    sf.write(file='xq_'+case+'.wav', data=xq,
             samplerate=48000, subtype='PCM_24')
    sf.write(file='e_'+case+'.wav', data=e,
             samplerate=48000, subtype='PCM_24')

    # CCF
    kappa, ccf = my_xcorr2(xq, e, scaleopt='biased')

    plt.figure(figsize=(12, 3))
    if case == 'nodither':
        plt.subplot(1, 2, 1)
        # nothing to plot for the zero signal
        # the PDF would be a weighted Dirac at amplitude zero
    else:
        # plot dither noise PDF estimate as histogram
        plt.subplot(1, 2, 1)
        plt.plot(edges_dither[:-1], pdf_dither, 'o-', ms=5)
        plt.ylim(-0.1, np.max(pdf_dither)*1.1)
        plt.grid(True)
        plt.xlabel(r'$\theta$')
        plt.ylabel(r'$\hat{p}(\theta)$')
        plt.title('PDF estimate of dither noise')

    # plot error noise PDF estimate as histogram
    plt.subplot(1, 2, 2)
    plt.plot(edges_error[:-1], pdf_error, 'o-', ms=5)
    plt.ylim(-0.1, np.max(pdf_error)*1.1)
    plt.grid(True)
    plt.xlabel(r'$\theta$')
    plt.ylabel(r'$\hat{p}(\theta)$')
    plt.title('PDF estimate of error noise')

    # plot signals
    plt.figure(figsize=(12, 3))
    plt.subplot(1, 2, 1)
    plt.plot(k, x, color='C2', label=r'$x[k]$')
    plt.plot(k, xd,  color='C1', label=r'$x_d[k] = x[k] + dither[k]$')
    plt.plot(k, xq,  color='C3', label=r'$x_q[k]$')
    plt.plot(k, e,  color='C0', label=r'$e[k] = x_q[k] - x[k]$')
    plt.plot(k, k*0+deltaQ, ':k', label=r'$\Delta Q$')
    plt.xlabel('k')
    plt.title('signals')
    plt.xticks(np.arange(0, 175, 25))
    plt.xlim(0, 150)
    plt.legend(loc='lower left')
    plt.grid(True)

    # plot CCF
    plt.subplot(1, 2, 2)
    plt.plot(kappa, ccf)
    plt.xlabel(r'$\kappa$')
    plt.ylabel(r'$\varphi_{xq,e}[\kappa]$')
    plt.title('CCF betwen xq and e=xq-x')
    plt.xticks(np.arange(-100, 125, 25))
    plt.xlim(-100, 100)
    plt.grid(True)
```

Chose one of the 5 cases and evaluate no dither, uniform PDF dither and triangular PDF dither noises below.

```python
# case 1
B = 16  # Bit
Q = 2**B  # number of quantization steps
deltaQ = 1 / (Q//2)  # quantization step size
x = (1-deltaQ) * np.sin(2*np.pi*fsin/fs*k)  # largest positive amplitude
```

```python
# case 2
B = 16
Q = 2**B
deltaQ = 1 / (Q//2)
x = deltaQ * np.sin(2*np.pi*fsin/fs*k)  # smallest amplitude
```

```python
# case 3
B = 3
Q = 2**B
deltaQ = 1 / (Q//2)
x = (1-deltaQ) * np.sin(2*np.pi*fsin/fs*k)
```

```python
# case 4  this is the default case when running the whole notebook
B = 3
Q = 2**B
deltaQ = 1 / (Q//2)
x = deltaQ * np.sin(2*np.pi*fsin/fs*k)
```

```python
# case 5
if False:
    B = 3
    Q = 2**B
    deltaQ = 1 / (Q//2)
    # amplitude below quantization step!
    x = deltaQ/2 * np.sin(2*np.pi*fsin/fs*k)
```

```python
plt.figure(figsize=(4, 4))
check_my_quant(Q)
```

### No Dither Noise

```python
# no dither
check_dithering(x=x, dither=x*0, Q=Q, case='nodither')
```

**Be very careful! Do not harm your ears!**

| Signal            | Audio Player  |
| ----------------- | :------------ |
| $x[k]$            | <audio type="audio/wave" src="x_nodither.wav" controls></audio> |
| $x_q[k]$          | <audio type="audio/wave" src="xq_nodither.wav" controls></audio> |
| $e[k]$            | <audio type="audio/wave" src="e_nodither.wav" controls></audio> |

### Uniform PDF Dither Noise

```python
# uniform dither with max amplitude of deltaQ/2
np.random.seed(1)
dither_uni = (np.random.rand(N) - 0.5) * 2 * deltaQ/2

check_dithering(x=x, dither=dither_uni, Q=Q, case='unidither')
```

**Be very careful! Do not harm your ears!**

| Signal            | Audio Player  |
| ----------------- | :------------ |
| $x[k]$            | <audio type="audio/wave" src="x_unidither.wav" controls></audio> |
| $x_d[k]$          | <audio type="audio/wave" src="xd_unidither.wav" controls></audio> |
| $x_q[k]$          | <audio type="audio/wave" src="xq_unidither.wav" controls></audio> |
| $e[k]$            | <audio type="audio/wave" src="e_unidither.wav" controls></audio> |

### Triangular PDF Dither Noise

```python
np.random.seed(1)
# uniform PDF for amplitudes -1...+1:
dither_uni1 = (np.random.rand(N) - 0.5) * 2
dither_uni2 = (np.random.rand(N) - 0.5) * 2
# triangular PDF with max amplitude of deltaQ
dither_tri = (dither_uni1 + dither_uni2) * deltaQ/2

check_dithering(x=x, dither=dither_tri, Q=Q, case='tridither')
```

**Be very careful! Do not harm your ears!**

| Signal            | Audio Player  |
| ----------------- | :------------ |
| $x[k]$            | <audio type="audio/wave" src="x_tridither.wav" controls></audio> |
| $x_d[k]$          | <audio type="audio/wave" src="xd_tridither.wav" controls></audio> |
| $x_q[k]$          | <audio type="audio/wave" src="xq_tridither.wav" controls></audio> |
| $e[k]$            | <audio type="audio/wave" src="e_tridither.wav" controls></audio> |

# **Copyright**

The notebooks are provided as [Open Educational Resources](https://en.wikipedia.org/wiki/Open_educational_resources). Feel free to use the notebooks for your own purposes. The text is licensed under [Creative Commons Attribution 4.0](https://creativecommons.org/licenses/by/4.0/), the code of the IPython examples under the [MIT license](https://opensource.org/licenses/MIT). Please attribute the work as follows: *Frank Schultz, Digital Signal Processing - A Tutorial Featuring Computational Examples* with the URL https://github.com/spatialaudio/digital-signal-processing-exercises