parameter_optimization.jl 9.5 KB
Newer Older
Peter Jentsch's avatar
Peter Jentsch committed
1
2
3
4
using KissABC
using CurveFit
using Loess

5
6
7
const target_cumulative_vac_proportion = 0.33
const vaccination_data = @SVector [0.0,0.043,0.385,0.424,0.115,0.03,0.005] #by month starting in august
const ymo_vac = @SVector [0.255,0.278,0.602]
Peter Jentsch's avatar
Peter Jentsch committed
8

9
10
11
12
13
14
function solve_w_parameters(default_p, p_names, new_p_list)
    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
17
18
19
20
21
22
23
24
25
26
@views function exp_growth_rate_estimation(incident_cases)
    diff(l) = [l[i] - l[i-1] for i=2:length(l)]
    l = length(incident_cases)
    loess_interp = loess(1:l,incident_cases)
    interp_incident_cases = Loess.predict(loess_interp, collect(1.0:l))
    second_deriative =diff(diff(interp_incident_cases))
    inflection_pt = any(second_deriative.<=-1) ? findfirst(<=(-1),second_deriative) : l  #completely heuristic
    ic,exp_growth_rate = exp_fit(1:inflection_pt, incident_cases[1:inflection_pt]) #fit y = a exp(b x), b is exp growth rate
    return ic,exp_growth_rate
end

function fit_epi_parameters(p_tuple, target_growth_rate)
Peter Jentsch's avatar
Peter Jentsch committed
27
28
29
30
31
32
33
34
    p_tuple_without_vac = merge(p_tuple,
        (
            sim_length = 30,
            immunization_begin_day = 0,
            infection_introduction_day = 1,
            immunizing = false,
        )
    )
35
    p_names = (:base_transmission_probability,)
Peter Jentsch's avatar
Peter Jentsch committed
36
37
    priors = Factored(Uniform(0.0001,0.005))    
    function cost(p)
38
39
        output,model = solve_w_parameters(p_tuple_without_vac,p_names ,p)
        incident_cases = output.daily_cases
Peter Jentsch's avatar
Peter Jentsch committed
40
41
42
        _,exp_growth_rate = exp_growth_rate_estimation(incident_cases)
        return (target_growth_rate - exp_growth_rate)^2
    end
Peter Jentsch's avatar
Peter Jentsch committed
43

44
45
    out = ABCDE(priors,cost,0.01; verbose=true, nparticles= 200,generations=20,  parallel = true) #this one has better NaN handling
    return NamedTuple{p_names}((out.P.particles,))
Peter Jentsch's avatar
Peter Jentsch committed
46
47
end

48
function fit_pre_inf_behavioural_parameters(p_tuple)
Peter Jentsch's avatar
Peter Jentsch committed
49
    samples = 1
50
51
52
53
54
55
56
57
58
59
60
61
    p_names = (:π_base_y,:π_base_m,:π_base_o)
    priors = Factored(
        Uniform(-10.0,2.0),
        Uniform(-10.0,2.0),
        Uniform(-10.0,2.0)
    )
    
    #simulation begins in july
    #60 days for opinion dynamics to stabilize, then immunization begins in september,
    #infection is not considered
    
    sim_length = 180
Peter Jentsch's avatar
Peter Jentsch committed
62
63
64
65
    p_tuple_adjust = merge(p_tuple,
        (
            sim_length = sim_length,
            I_0_fraction = 0.000,
66
            immunization_begin_day =60, 
67
            infection_introduction_day = 180,
Peter Jentsch's avatar
Peter Jentsch committed
68
69
70
            immunizing = true,
        )
    )
71

72
73
    function cost(p)
        output,model = solve_w_parameters(p_tuple_adjust, p_names,p)
Peter Jentsch's avatar
Peter Jentsch committed
74
        target_cumulative_vaccinations = target_cumulative_vac_proportion*model.nodes
75
        target_ymo_vac = ymo_vac .* sum(vaccination_data[1:4]) .* target_cumulative_vaccinations
76
        ymo_vaccination_ts = output.new_ymo_immunization
77
78
79
80
        total_preinfection_vaccination = sum.(eachrow(ymo_vaccination_ts))
        # display(target_ymo_vac)
        # display(total_preinfection_vaccination)
        return sum((total_preinfection_vaccination .- target_ymo_vac).^2)    
Peter Jentsch's avatar
Peter Jentsch committed
81
    end
82
83
84
    #smc(priors,cost; verbose = true, nparticles = 100, parallel = true)# smc(priors,cost; verbose = true, nparticles =300, parallel = true)#
    out = smc(priors,cost; verbose = true, nparticles = 800, parallel = true)#
    return NamedTuple{p_names}(ntuple(i -> out.P[i].particles,length(p_names)))
Peter Jentsch's avatar
Peter Jentsch committed
85
86
end

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
116
117
function fit_post_inf_behavioural_parameters(p_tuple)
    samples = 1
    p_names = (:ω,)
    priors = Factored(Uniform(0.0,0.1))
    #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.002,
            immunization_begin_day =60, 
            infection_introduction_day = 180,
            immunizing = true,
        )
    )
    target_cumulative_vac_proportion = 0.33
    vaccination_data = @SVector [0.0,0.043,0.385,0.424,0.115,0.03,0.005] #by month starting in august
    ymo_vac = @SVector [0.255,0.278,0.602]
    function cost(p)
        output,model = solve_w_parameters(p_tuple_adjust, p_names,p)
        target_cumulative_vaccinations = target_cumulative_vac_proportion*model.nodes
        target_ymo_vac = ymo_vac .* sum(vaccination_data[5:end]) .* target_cumulative_vaccinations
        ymo_vaccination_ts = output.new_ymo_immunization
        total_vaccination = sum.(eachrow(ymo_vaccination_ts[:,180:end]))
        return sum((total_vaccination .- target_ymo_vac).^2)    
    end
    out =smc(priors,cost; verbose = true, nparticles = 200, parallel = true)# ABCDE(priors,cost,1e6; verbose=true, nparticles=300,generations=1000,  parallel = true) #this one has better NaN handling
    return NamedTuple{p_names}((out.P.particles,))
end
Peter Jentsch's avatar
Peter Jentsch committed
118
119

function plot_behavioural_fit(particles,p_tuple)
120
    p_names = (:π_base_y,:π_base_m,:π_base_o)
Peter Jentsch's avatar
Peter Jentsch committed
121
122
123
    sim_length = 210
    samples = 1
    p_tuple_adjust = merge(p_tuple,
124
        (
Peter Jentsch's avatar
Peter Jentsch committed
125
            sim_length = sim_length,
126
127
            I_0_fraction = 0.000,
            immunization_begin_day =60, 
Peter Jentsch's avatar
Peter Jentsch committed
128
129
130
131
132
133
134
            immunizing = true,
        )
    )
    p = map(e -> mode(e.particles),particles.P)
    display(p)
    new_params = merge(p_tuple_adjust, NamedTuple{p_names}(ntuple(i -> p[i],length(p_names))))
    out = mean_solve(samples, new_params ,DebugRecorder)
135
136
137
138
139
140
    target_cumulative_vac_proportion = 0.33
    vaccination_data = @SVector [0.0,0.043,0.385,0.424,0.115,0.03,0.005] #by month starting in august
    ymo_vac = @SVector [0.255,0.278,0.602]
    ymo_vaccination_ts = mean.(out.new_ymo_immunization)
    total_preinfection_vaccination = sum.(eachrow(ymo_vaccination_ts))
    display(total_preinfection_vaccination)
141
142
    p = [plot(),plot(),plot()]

Peter Jentsch's avatar
Peter Jentsch committed
143
144
    p = plot_model(nothing,[nothing],[out],new_params.infection_introduction_day,new_params.immunization_begin_day)
    savefig(p,"behaviour_fit.pdf")
Peter Jentsch's avatar
Peter Jentsch committed
145
    return out
146
end
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
# outbreak_transmission_dist = CovidAlertVaccinationModel.fit_epi_parameters(default_parameters,0.241) ##outbreak
# serialize(joinpath(PACKAGE_FOLDER,"abm_parameter_fits","outbreak_inf_parameters.dat"),outbreak_transmission_dist)
# plot_max_posterior("outbreak", outbreak_transmission_dist,default_parameters)

function fit_parameters(default_parameters)
    seasonal_parameters_path = joinpath(PACKAGE_FOLDER,"abm_parameter_fits","outbreak_inf_parameters.dat")
    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")

    seasonal_transmission_dist = CovidAlertVaccinationModel.fit_epi_parameters(default_parameters,0.073) ##seasonal
    plot_max_posterior("seasonal", seasonal_transmission_dist,default_parameters)
    pre_inf_behaviour_parameters = CovidAlertVaccinationModel.fit_pre_inf_behavioural_parameters(default_parameters)
    fitted_parameters = map(mode,((;seasonal_transmission_dist...,pre_inf_behaviour_parameters...)))
    fitted_parameter_tuple = merge(default_parameters,fitted_parameters)

    serialize(seasonal_parameters_path,seasonal_transmission_dist)
    serialize(pre_inf_behaviour_parameters_path, pre_inf_behaviour_parameters)

    post_inf_behaviour_parameters = fit_post_inf_behavioural_parameters(fitted_parameter_tuple) 
    serialize(post_inf_behaviour_parameters_path, post_inf_behaviour_parameters)
    
    

    fitted_parameters_with_post_inf_behaviour = (;
        deserialize(seasonal_parameters_path)...,
        deserialize(pre_inf_behaviour_parameters_path)...,
        deserialize(post_inf_behaviour_parameters_path)...
    )
    display(map(mode,fitted_parameters_with_post_inf_behaviour))
    return plot_fitting_posteriors("post_inf_fitting",fitted_parameters_with_post_inf_behaviour,default_parameters)
end 

function plot_max_posterior(fname,particles,parameters)
    samples = 5
    base_transmission = mode(particles.base_transmission_probability)
    p_tuple_without_vac = merge(parameters,
        (
            sim_length = 150,
            immunization_begin_day = 0,
            infection_introduction_day = 1,
            immunizing = false,
        )
    )
    new_params = merge(p_tuple_without_vac, (base_transmission_probability = base_transmission,))
    out = mean_solve(samples, new_params ,DebugRecorder)
    p = plot_model(nothing,[nothing],[out],new_params.infection_introduction_day,new_params.immunization_begin_day)
    savefig(p,"$fname.pdf")
    hist = StatsBase.fit(Histogram,particles.base_transmission_probability; nbins = 25)
    p = plot(hist;legend = false)
    savefig(p,"$(fname)_posterior.pdf")
end
198

199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
function plot_fitting_posteriors(fname,particles_tuple,parameters)
    p_tuple_adjust = merge(parameters,
        (
            sim_length = 500,
            I_0_fraction = 0.002,
            immunization_begin_day =60, 
            infection_introduction_day = 180,
            immunizing = true,
        )
    )
    out = mean_solve(5, merge(p_tuple_adjust,map(mode,particles_tuple)) ,DebugRecorder)
    p = plot_model(nothing,[nothing],[out],parameters.infection_introduction_day,parameters.immunization_begin_day)
    savefig(p, "$fname.pdf")
    
    plts = [plot() for i in 1:length(particles_tuple)]
    for (plt,(k,v)) in zip(plts,pairs(particles_tuple))
        hist = StatsBase.fit(Histogram,v; nbins = 40)
        plot!(plt,hist;legend = false,xlabel = k)            
    end
    p = plot(plts...; size = (1000,600))
    savefig(p,"$(fname)_posteriors.pdf")
    return out
end