added POLYMOD priors to IntervalsModel

CovidAlertVaccinationModel/src/mixing_distributions.jl

@@ -51,6 +51,24 @@ function from_mean(::Type{ZWDist{DistType,T}},α,μ) where {DistType <: Distribu
     return ZWDist(α,from_mean(DistType,μ/(1-α)))
 end
 StatsBase.mean(d::ZWDist{Dist,T}) where {Dist,T} = (1 - d.α)*StatsBase.mean(d.base_dist)
+function from_mean_workschool(_,μ)
+    return Poisson(μ)
+end
+
+function alpha_matrix(alphas)
+    M = zeros(length(alphas),length(alphas))
+    for i in 1:length(alphas), j in 1:length(alphas)
+        M[i,j] = alphas[i] + alphas[j] - alphas[j]*alphas[i]
+    end
+    return M
+end
+
+const alphas = alpha_matrix((
+    Y = 0.11,
+    M = 0.14,
+    O = 0.87
+))
+
 
 const initial_workschool_mixing_matrix = map(t->from_mean(t...),[
     (ZWDist{Geometric{Float64},Discrete}, 0.433835,4.104848) (ZWDist{Geometric{Float64},Discrete},0.406326,2.568782) (ZWDist{Geometric{Float64},Discrete},0.888015,0.017729)

IntervalsModel/src/IntervalsModel.jl

@@ -24,20 +24,17 @@ const YOUNG, MIDDLE,OLD = 1,2,3
 const durmax = 144
 const color_palette = palette(:seaborn_bright) #color theme for the plots
 const sparam = [60, 12]
-const priors = get_priors()
     # Swap age brackets for numbers
 const swap_dict = OrderedDict("Y" => YOUNG, "M" => MIDDLE, "O" => OLD)
-include("interval_overlap_sampling.jl")
 include("utils.jl")
+include("interval_overlap_sampling.jl")
+
 include("hh_durations_model.jl")
 include("ws_rest_durations_model.jl")
 include("plotting_functions.jl")
 
 
-
-const μ_bounds = (1,144) 
-const σ_bounds = (1,144)
-const α_bounds = (0.0,0.8) 
+const priors = get_priors()
 const α_priors_mean_ws = ( 
     Y = 0.364286,
     M = 0.385714,
@@ -50,22 +47,22 @@ const α_priors_mean_rest = (
     O = 0.092857
 )
 
+
+
 pgfplotsx()
 function main()
 
-    # do_hh(400)
-    do_ws(400)
-    do_rest(400)
+    do_hh(400)
+    # do_ws(400)
+    do_rest(1000)
 
 end
 using BenchmarkTools
-function bayesian_estimation(fname, err_func, priors_list, dists, particles; alpha = 0.995)
+function bayesian_estimation(fname, err_func, priors_list, dists, particles; alpha = 0.95)
     data_pairs = map(zip(dists,priors_list)) do (dist,priors)
-        # Random.seed!(rng,1234)
-        # init = rand(rng,priors)
-        # @btime $err_ws($init,$dist)
-        out = smc(priors,p -> err_ws(p, dist), verbose=true, nparticles=particles, M=2, alpha=alpha,  parallel = true)
-        
+        init = rand(priors)
+        @btime $err_func($init,$dist)
+        out = smc(priors,p -> err_func(p, dist), verbose=true, nparticles=particles, alpha=alpha,  parallel = true)
         return dist => out
     end |> OrderedDict
 
@@ -74,12 +71,4 @@ function bayesian_estimation(fname, err_func, priors_list, dists, particles; alp
 end
 
 
-function alpha_matrix(alphas)
-    M = zeros(length(alphas),length(alphas))
-    for i in 1:length(alphas), j in 1:length(alphas)
-        M[i,j] = alphas[i] + alphas[j] - alphas[j]*alphas[i]
-    end
-    return [M[1,1], M[1,2],M[1,3],M[2,2],M[2,3],M[3,3]] #lol
-end
-
 end # module
\ No newline at end of file

IntervalsModel/src/hh_durations_model.jl

 
 const HHYMO = DataFrame(CSV.File("$PACKAGE_FOLDER/network-data/Timeuse/HH/HHYMO.csv"))
 
-# This function applies pre-processing to the HHYMO data file, and splits it into a namedtuple, which should be faster to index.
-# In particular, we avoid having to modify any strings in the error function.
-
 const cnst_hh = (
     # Set parameters for intervals sample and subpopulation size
-    numsamples = 5_000,
+    numsamples = 100,
     subsize = size(HHYMO)[1],
     # Total weight in survey
     Wghttotal = sum(HHYMO[:,"WGHT_PER"]),
@@ -15,15 +12,18 @@ const cnst_hh = (
 
 function do_hh(particles)
     dists = [
-        Normal,
         Poisson,
     ]
-    priors_list = map(l -> Factored(vcat(l...)...),[
-        [[Uniform(μ_bounds...) for i=1:6], [Uniform(σ_bounds...)  for i = 1:6]],
-        [[Uniform(μ_bounds...) for i=1:6]],
-    ])
+
+
+    poisson_priors = filter_priors("rest")
+
+    # # Set parameter priors for fitting
+    priors_list = [
+        Factored(poisson_priors...)
+    ]
     fname = "hh"
-    bayesian_estimation(fname,err_hh,priors_list,dists,particles)
+    bayesian_estimation(fname,err_hh,priors_list,dists,particles; alpha = 0.95)
 end
 function make_dat_array()
     durs = hcat(
@@ -51,12 +51,14 @@ const dat = make_dat_array() #assign a constant data array
 
 #error function
 function err_hh(p,dist)
-
-    params = get_params(p)
-    age_dists = [dist(params[i,j]...) for i in YOUNG:OLD, j in YOUNG:OLD]
+    p_matrix = zeros(3,3)
+    for i in 1:9
+        p_matrix[i] = p[i]
+    end
+    age_dists = [dist(p_matrix[i,j]...) for i in YOUNG:OLD, j in YOUNG:OLD]
     duration_subarray =  dat.durs
     num_contacts_subarray = dat.nums
-
+    # display(p_matrix)
     AGERESP =  dat.AGERESP 
     errsum = 0
     @inbounds for i = 1:cnst_hh.subsize
@@ -67,5 +69,7 @@ function err_hh(p,dist)
             errsum += (expdur/cnst_hh.numsamples - duration_subarray[i,age_j])^2 #compute total 
         end
     end
+    # display("2")
+
     return errsum
 end
\ No newline at end of file

IntervalsModel/src/plotting_functions.jl

@@ -16,7 +16,15 @@ function plot_estimate(fname)
     end
 end
 function plot_dists(fname,dist_constructors,data)
-    p_estimate_as_arrays = map(d -> get_params(d.P),data)
+    function get_non_symmetric_params(p) 
+
+            p_matrix = [[zero(eltype(p))] for i in 1:3, j in 1:3]
+            for i in 1:9
+                p_matrix[i][1] = p[i]
+            end
+        return p_matrix
+    end
+    p_estimate_as_arrays = map(d -> get_non_symmetric_params(d.P),data)
     p_matrix = map(x -> plot(),p_estimate_as_arrays[1])
     ymo = ["Y","M","O"]
     x_range = 0.0:144.0
@@ -45,6 +53,7 @@ compute_x_pos(p) = xlims(p)[1] + 0.02*((xlims(p)[2]  - xlims(p)[1]))
 compute_y_pos(p) = ylims(p)[2] - 0.11*((ylims(p)[2]  - ylims(p)[1]))
 
 function plot_posteriors(fname,data)
+    display(data.P)
     p_list = map(x -> plot(),1:length(data.P))
     for i in 1:length(data.P)
         # display(data.P[i].particles)
@@ -56,7 +65,7 @@ function plot_posteriors(fname,data)
         # vline!(p_list[i],[argmax(kernel_data)]; seriescolor = color_palette[2])
 
     end   
-    p = plot(p_list...;layout=(length(p_list) ÷ 6,6), size = (1000,(length(p_list) ÷ 6)*300), seriescolor = color_palette[1])
+    p = plot(p_list...;layout=(length(p_list) ÷ 9,9), size = (1500,(length(p_list) ÷ 9)*300), seriescolor = color_palette[1])
     savefig(p,joinpath(PACKAGE_FOLDER,"plots","$fname.pdf"))
 end
 function add_subplot_letters!(plot_list; pos = :top)

IntervalsModel/src/utils.jl

@@ -11,5 +11,25 @@ import Pandas: read_csv
 using DataFrames
 function get_priors()
     df = DataFrame(read_csv("IntervalsModel/network-data/POLYMOD/AALPoisPriors.csv"))
-    return df
+    df_w_indicies = transform(
+        df,
+        :Age_in => e->map(x-> swap_dict[x],e),
+        :Age_out => e -> map(x -> swap_dict[x],e),
+        :Age_in,:Age_out
+    )
+    rename!(df_w_indicies, Dict(:Age_out_function => "index_out",:Age_in_function => "index_in"))
+    return df_w_indicies
+end
+
+function filter_priors(key)
+    priors_probs = filter(row-> row["location"] == key, eachrow(priors)) |>
+                    df -> map(r -> (r[:index_out],r[:index_in]) => convert(Vector,r[string.(0:143)]),df) |>
+                    Dict |>
+                    df -> Dict(map(v -> v => Categorical(df[v]), collect(keys(df))))
+
+    output_array = Array{eltype(values(priors_probs)),2}(undef,(3,3))
+    for (key,value) in priors_probs
+        output_array[key...] = value
+    end
+    return output_array
 end
\ No newline at end of file

IntervalsModel/src/ws_rest_durations_model.jl

@@ -33,42 +33,47 @@ rest_distributions = CovidAlertVaccinationModel.initial_rest_mixing_matrix
 
 function do_ws(particles)
     dists = [
-        ZWDist{Normal},
-        ZWDist{Poisson},
+        Poisson,
     ]
 
     # Set parameter priors for fitting
-    alpha_priors = [TriangularDist(α*(0.8),min(1,α*(1.2),α)) for α in alpha_matrix(α_priors_mean_ws)]
-    priors_list = map(l -> Factored(vcat(l...)...),[
-        [alpha_priors,[Uniform(μ_bounds...) for i=1:6], [Uniform(σ_bounds...)  for i = 1:6]],
-        [alpha_priors,[Uniform(μ_bounds...) for i=1:6]],
-        [alpha_priors,[Uniform(μ_bounds...) for i=1:6], [Uniform(σ_bounds...)  for i = 1:6]],
-    ])
+    # alpha_priors = [TriangularDist(α*(0.8),min(1,α*(1.2),α)) for α in alpha_matrix(α_priors_mean_ws)]
+    # priors_list = map(l -> Factored(vcat(l...)...),[
+    #     [alpha_priors,[Uniform(μ_bounds...) for i=1:6], [Uniform(σ_bounds...)  for i = 1:6]],
+    #     [alpha_priors,[Uniform(μ_bounds...) for i=1:6]],
+    #     [alpha_priors,[Uniform(μ_bounds...) for i=1:6], [Uniform(σ_bounds...)  for i = 1:6]],
+    # ])
+    # priors_list = map(l -> Factored(vcat(l...)...),[
+    #    filter(priors,
+    # ])
     bayesian_estimation("ws",err_ws,priors_list,dists,particles; alpha = 0.99)
 end
 
 function do_rest(particles)
     dists = [
-        ZWDist{Normal},
-        ZWDist{Poisson},
+        Poisson,
     ]
-    # Set parameter priors for fitting
-    alpha_priors = [TriangularDist(α*(0.8),min(1,α*(1.2),α)) for α in alpha_matrix(α_priors_mean_ws)]
-    priors_list = map(l -> Factored(vcat(l...)...),[
-        [alpha_priors,[Uniform(μ_bounds...) for i=1:6], [Uniform(σ_bounds...)  for i = 1:6]],
-        [alpha_priors,[Uniform(μ_bounds...) for i=1:6]],
-        [alpha_priors,[Uniform(μ_bounds...) for i=1:6], [Uniform(σ_bounds...)  for i = 1:6]],
-    ])
-
-    bayesian_estimation("rest",err_rest, priors_list, dists, particles, alpha = 0.99)
+
+
+    poisson_priors = filter_priors("rest")
+
+    # # Set parameter priors for fitting
+    priors_list = [
+        Factored(poisson_priors...)
+    ]
+    display(poisson_priors)
+    bayesian_estimation("rest",err_rest, priors_list, dists, particles, alpha = 0.98)
 end
 
 
 #error function for ws
 function err_ws(p,dist)
-    params = get_params(p)
+    p_matrix = zeros(3,3)
+    for i in 1:9
+        p_matrix[i] = p[i]
+    end
     neighourhoods = rand.(rng,ws_distributions)
-    age_dists = [dist(params[i,j]...) for i in YOUNG:OLD, j in YOUNG:OLD]
+    age_dists = [dist(p_matrix[i,j]...) for i in YOUNG:OLD, j in YOUNG:OLD]
     sample_list = zeros(comparison_samples)
     errsum = 0
     for age_sample in YOUNG:OLD
@@ -89,11 +94,16 @@ end
 
 #error function for rest
 function err_rest(p,dist)
-    params = get_params(p)
+    p_matrix = zeros(3,3)
+    for i in 1:9
+        p_matrix[i] = p[i]
+    end
+
     neighourhoods = rand.(rng,rest_distributions)
-    age_dists = [dist(params[i,j]...) for i in YOUNG:OLD, j in YOUNG:OLD]
+    age_dists = [dist(p_matrix[i,j]...) for i in YOUNG:OLD, j in YOUNG:OLD]
     sample_list = zeros(comparison_samples)
     errsum = 0
+
     for age_sample in YOUNG:OLD
         for age_j in YOUNG:OLD #for a given age_sample loop over possible contact ages
             if neighourhoods[age_sample,age_j] > 0