parameter_optimization.jl 7.53 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
11
12
13
14
    new_params = merge(default_p, NamedTuple{p_names}(ntuple(i -> new_p_list[i],length(p_names))))
    out = DebugRecorder(0,default_p.sim_length)
    model = abm(new_params,out)
    return out, model
end
Peter Jentsch's avatar
Peter Jentsch committed
15

16
function fit_all_parameters(p_tuple)
Peter Jentsch's avatar
Peter Jentsch committed
17
    p_names = (:ω,:β_y,:β_m,:β_o,:π_base_y,:π_base_m,:π_base_o,:α_y,:α_m,:α_o)
18
    priors = Factored(
Peter Jentsch's avatar
Peter Jentsch committed
19
20
        Uniform(0.0,0.01),
        Uniform(0.0,0.2),
21
        Uniform(0.0,0.1),
22
23
24
25
26
27
28
        Uniform(0.0,1.0), 
        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)   
29
30
31
32
33
34
35
36
    )
    #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,
Peter Jentsch's avatar
Peter Jentsch committed
37
            I_0_fraction = 0.005,
38
            immunization_begin_day = 60, 
39
40
41
42
43
44
45
46
            infection_introduction_day = 180,
            immunizing = true,
        )
    )
    function cost(p)
        output,model = solve_w_parameters(p_tuple_adjust, p_names,p)
        target_ymo_vac = ymo_vac .* sum(vaccination_data[1:end]) .* length.(model.index_vectors)
        ymo_vaccination_ts = output.daily_immunized_by_age
Peter Jentsch's avatar
Peter Jentsch committed
47
48

        normalize_by_pop(v) = v./length.(model.index_vectors)
49
50
        total_postinf_vaccination = sum.(eachrow(ymo_vaccination_ts[:,180:end]))

Peter Jentsch's avatar
Peter Jentsch committed
51
        final_size = sum.(eachrow(output.daily_unvac_cases_by_age))
52
53
54

        target_final_size = ymo_attack_rate .* length.(model.index_vectors)
        target_ymo_vac = ymo_vac .* sum(vaccination_data[1:4]) .* length.(model.index_vectors)
Peter Jentsch's avatar
Peter Jentsch committed
55

56
57
        ymo_vaccination_ts = output.daily_immunized_by_age
        total_preinfection_vaccination = sum.(eachrow(ymo_vaccination_ts))
58
59
60
61
62
63
64
65
66
67
68
69
70
71
    
        # return sum((normalize_by_pop(total_preinfection_vaccination .- target_ymo_vac)).^2)  
        # + sum(normalize_by_pop((total_postinf_vaccination .- target_ymo_vac)).^2)  
        # + sum(normalize_by_pop((final_size .- target_final_size)).^2)
        display(final_size)
        display( sum.(eachrow(ymo_vaccination_ts[:,1:end])))
        display( length.(model.index_vectors))
        return sum(normalize_by_pop((final_size .- target_final_size)).^2)
    end
    display(cost((0.0,0.0005,0.0005,1.0,-3.0,-3.0,-0.2,0.25,0.25,0.0)))
    # out = smc(priors,cost; verbose = true, nparticles = 1000, parallel = true)#this one has better NaN handling
    return NamedTuple{p_names}(ntuple(i -> out.P[i].particles,length(p_names)))
end
function fit_distribution_parameters(p_tuple)
72
    p_names = (:ω,:β_y,:β_m,:β_o,:π_base_y,:π_base_m,:π_base_o,:α_y,:α_m,:α_o,:O_distribution_shift)
73
74
75
    priors = Factored(
        Uniform(0.0,0.01),
        Uniform(0.0,0.5),  
76
77
        Uniform(0.0,0.5),  
        Uniform(0.0,1.0),  
78
79
80
81
82
83
        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),   
84
        Uniform(0.0,1.0),   
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
    )
    #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)
        p_mergedbeta = merge(p_tuple_adjust, (β_m = p[2],β_o = p[2], )) #assume one beta
        output,model = solve_w_parameters(p_mergedbeta, p_names,p)
        target_ymo_vac = ymo_vac .* sum(vaccination_data[1:end]) .* length.(model.index_vectors)
        ymo_vaccination_ts = output.daily_immunized_by_age

        normalize_by_pop(v) = v./length.(model.index_vectors)
        total_postinf_vaccination = sum.(eachrow(ymo_vaccination_ts[:,180:end]))

        final_size = sum.(eachrow(output.daily_unvac_cases_by_age))

        target_final_size = ymo_attack_rate .* length.(model.index_vectors)
        target_ymo_vac = ymo_vac .* sum(vaccination_data[1:4]) .* length.(model.index_vectors)

        ymo_vaccination_ts = output.daily_immunized_by_age
        total_preinfection_vaccination = sum.(eachrow(ymo_vaccination_ts))
    
Peter Jentsch's avatar
Peter Jentsch committed
116
117
        return sum((normalize_by_pop(total_preinfection_vaccination .- target_ymo_vac)).^2)  
        + sum(normalize_by_pop((total_postinf_vaccination .- target_ymo_vac)).^2)  
118
        + 5*sum(normalize_by_pop((final_size .- target_final_size)).^2)
119
120
121
122
        # display(final_size)
        # display( sum.(eachrow(ymo_vaccination_ts[:,1:end])))
        # display( length.(model.index_vectors))
        # return sum(normalize_by_pop((final_size .- target_final_size)).^2)
123
    end
124
125
    # display(cost((0.0,0.0005,-3.0,-3.0,-0.2,0.25,0.25,0.0,0.1)))
    #
126
    out = smc(priors,cost; verbose = true, nparticles = 800, parallel = true)
127
128
    return NamedTuple{p_names}(ntuple(i -> out.P[i].particles,length(p_names)))
end
129
function fit_parameters(default_parameters)
Peter Jentsch's avatar
Peter Jentsch committed
130
131
    # 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")
132
    fit_all_parameters_path = joinpath(PACKAGE_FOLDER,"abm_parameter_fits","fit_all_parameters.dat")
133
    fit_dist_parameters_path = joinpath(PACKAGE_FOLDER,"abm_parameter_fits","fit_dist_parameters.dat")
134

135

136
    # output = fit_distribution_parameters(default_parameters)
137
    # serialize(fit_all_parameters_path,output)
138
139

    
140
141
    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
142

143
144
145
146
    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)
147
    return fitted_sol
148
149
150
end 

function plot_fitting_posteriors(fname,particles_tuple,parameters)
Peter Jentsch's avatar
Peter Jentsch committed
151
    # p_adjust = merge(parameters,(β_y = p.β_y*0.1, β_m = p.β_m*0.1 ,β_o = p.β_o)) 
152
153
154
155
156
    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...))
157
    for p_set in sample(samples,20)
158
159
160
161
162
        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
163
    avg_populations ./= 20
164
    p = plot_model(nothing,[nothing],[stat_recorder],parameters.infection_introduction_day,parameters.immunization_begin_day)
165
166
167
168
    savefig(p, "$fname.pdf")
    
    plts = [plot() for i in 1:length(particles_tuple)]
    for (plt,(k,v)) in zip(plts,pairs(particles_tuple))
169
        hist = StatsBase.fit(Histogram,v; nbins = 30)
170
        plot!(plt,hist;legend = false,xlabel = k)
171
    end
172
    p = plot(plts...; size = (1400,800),bottommargin = 5Plots.mm)
173
    savefig(p,"$(fname)_posteriors.pdf")
174
    return stat_recorder,avg_populations
175
end