parameter_optimization.jl 4.79 KB
Newer Older
Peter Jentsch's avatar
Peter Jentsch committed
1
2
using KissABC

3
4
const vaccination_data = [0.0,0.043,0.385,0.424,0.115,0.03,0.005] #by month starting in august
const ymo_vac = [0.255,0.278,0.602]
5
const ymo_attack_rate = [10.376,5.636,7.2]./100
6

Peter Jentsch's avatar
Peter Jentsch committed
7

8
function solve_w_parameters(default_p, p_names, new_p_list)
9

10
    new_params = merge(default_p, NamedTuple{p_names}(ntuple(i -> new_p_list[i],length(p_names))))
11
    
12
13
14
15
    out = DebugRecorder(0,default_p.sim_length)
    model = abm(new_params,out)
    return out, model
end
Peter Jentsch's avatar
Peter Jentsch committed
16

17
function fit_distribution_parameters(p_tuple)
18
    p_names = (:ω,:β_y,:β_m,:β_o,:π_base_y,:π_base_m,:π_base_o,:α_y,:α_m,:α_o,:O_distribution_shift)
19
20
21
    priors = Factored(
        Uniform(0.0,0.01),
        Uniform(0.0,0.5),  
22
23
        Uniform(0.0,0.5),  
        Uniform(0.0,1.0),  
24
25
26
27
28
29
        Uniform(-5.0,3.0),
        Uniform(-5.0,3.0),
        Uniform(-5.0,3.0),
        TriangularDist(23.0, 44.0, 34.0),
        TriangularDist(22.0, 54.0, 40.0),
        TriangularDist(9.0, 59.0, 39.0),   
30
        Uniform(0.0,1.0),   
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
    )
    #simulation begins in july
    #60 days for opinion dynamics to stabilize, then immunization begins in september,
    #infection introduced at beginning of december
    sim_length = 300
    p_tuple_adjust = merge(p_tuple,
        (
            sim_length = sim_length,
            I_0_fraction = 0.005,
            immunization_begin_day = 60, 
            infection_introduction_day = 180,
            immunizing = true,
        )
    )
    function cost(p)
Peter Jentsch's avatar
Peter Jentsch committed
46
        output,model = solve_w_parameters(p_tuple_adjust, p_names,p)
47
48
        ymo_vaccination_ts = output.daily_immunized_by_age

Peter Jentsch's avatar
Peter Jentsch committed
49
        total_postinf_vaccination = sum.(eachrow(ymo_vaccination_ts[:,181:end]))
50
        final_size = sum.(eachrow(output.daily_unvac_cases_by_age))
Peter Jentsch's avatar
Peter Jentsch committed
51
        total_preinf_vaccination = sum.(eachrow(ymo_vaccination_ts[:,1:180]))
52
53

        target_final_size = ymo_attack_rate .* length.(model.index_vectors)
Peter Jentsch's avatar
Peter Jentsch committed
54
55
        target_preinf_vac = ymo_vac .* sum(vaccination_data[1:4]) .* length.(model.index_vectors)
        target_postinf_vac = ymo_vac .* sum(vaccination_data[5:end]) .* length.(model.index_vectors)
56
    
Peter Jentsch's avatar
Peter Jentsch committed
57
58
59
60
61
62
63
64
65
        display((final_size,target_final_size))
        display((total_preinf_vaccination,target_preinf_vac))
        display((total_postinf_vaccination,target_postinf_vac))
        display(sum.(eachrow(ymo_vaccination_ts)) ./length.(model.index_vectors))

        return sum((total_preinf_vaccination .- target_preinf_vac).^2)  
        + sum((total_postinf_vaccination .- total_postinf_vaccination).^2)  
        + 5*sum((final_size .- target_final_size).^2)
    end 
66
       display(cost((0.0000,0.00051,0.00046,0.18,-1.30,-1.24,-0.8,0.35,0.35,0.35,0.2)))
67
    #
Peter Jentsch's avatar
Peter Jentsch committed
68
    # out = smc(priors,cost; verbose = true, nparticles = 600, parallel = true)
69
70
    return NamedTuple{p_names}(ntuple(i -> out.P[i].particles,length(p_names)))
end
71
function fit_parameters(default_parameters)
Peter Jentsch's avatar
Peter Jentsch committed
72
73
    # pre_inf_behaviour_parameters_path =joinpath(PACKAGE_FOLDER,"abm_parameter_fits","pre_inf_behaviour_parameters.dat")
    # post_inf_behaviour_parameters_path = joinpath(PACKAGE_FOLDER,"abm_parameter_fits","post_inf_behaviour_parameters.dat")
74
    fit_all_parameters_path = joinpath(PACKAGE_FOLDER,"abm_parameter_fits","fit_all_parameters.dat")
Peter Jentsch's avatar
Peter Jentsch committed
75
76
    output = fit_distribution_parameters(default_parameters)
    serialize(fit_all_parameters_path,output)
77
78

    
79
80
    fitted_parameter_tuple = deserialize(fit_all_parameters_path)
    fitted_sol,avg_populations = plot_fitting_posteriors("post_inf_fitting",fitted_parameter_tuple,default_parameters)
Peter Jentsch's avatar
Peter Jentsch committed
81

82
83
84
85
    final_vac = sum.(eachrow(mean.(fitted_sol.daily_immunized_by_age)))
    final_size = sum.(eachrow(mean.(fitted_sol.daily_unvac_cases_by_age)))
    display(final_vac./avg_populations)
    display(final_size./avg_populations)
86
    return fitted_sol
87
88
89
end 

function plot_fitting_posteriors(fname,particles_tuple,parameters)
90
91
92
93
94
    stat_recorder = DebugRecorder(Variance(), parameters.sim_length)
    output_recorder = DebugRecorder(parameters.sim_length)
    avg_populations = [0.0,0.0,0.0]
    names = keys(particles_tuple)
    samples = collect(zip(particles_tuple...))
95
    for p_set in sample(samples,20)
96
97
98
99
100
        p_set_named = NamedTuple{names}(p_set)
        sol = abm(merge(parameters,p_set_named),output_recorder)
        avg_populations .+= length.(sol.index_vectors)
        fit!(stat_recorder,output_recorder)
    end
101
    avg_populations ./= 20
102
    p = plot_model(nothing,[nothing],[stat_recorder],parameters.infection_introduction_day,parameters.immunization_begin_day)
103
104
105
106
    savefig(p, "$fname.pdf")
    
    plts = [plot() for i in 1:length(particles_tuple)]
    for (plt,(k,v)) in zip(plts,pairs(particles_tuple))
107
        hist = StatsBase.fit(Histogram,v; nbins = 30)
108
        plot!(plt,hist;legend = false,xlabel = k)
109
    end
110
    p = plot(plts...; size = (1400,800),bottommargin = 5Plots.mm)
111
    savefig(p,"$(fname)_posteriors.pdf")
112
    return stat_recorder,avg_populations
113
end