Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix optim example #1019

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Fix optim example #1019

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
318 changes: 114 additions & 204 deletions examples/optim.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,11 +11,26 @@
# to optimize the parameters of an epidemiological model (SIR). We explain this model
# in detail in [SIR model for the spread of COVID-19](@ref).

# For brevity here, we just import

# ```julia
# include("siroptim.jl") # From the examples directory
# ```
# For brevity here, we just import it and create a initialization function which
# will be used to optimize the parameters

using AgentsExampleZoo: sir

function model_init(x, seed)
C = 3
sir(;
C = C,
Ns = [50, 50, 50],
max_travel_rate = x[1],
death_rate = x[2],
β_det = [x[3] for _ in 1:C],
β_und = [x[4] for _ in 1:C],
infection_period = x[5],
reinfection_probability = x[6],
detection_time = x[7],
seed = seed
)
end

# which provides us a `model_initiation` helper function to build a SIR model,
# and an `agent_step!` function.
Expand All @@ -28,60 +43,40 @@
# is detected. The function returns an *objective*: this value takes the form one
# or more numbers, which the optimiser will attempt to minimize.

# ```julia
# using BlackBoxOptim, Random
# using Statistics: mean
#
# function cost(x)
# model = sir(;
# Ns = [500, 500, 500],
# migration_rate = x[1],
# death_rate = x[2],
# β_det = x[3],
# β_und = x[4],
# infection_period = x[5],
# reinfection_probability = x[6],
# detection_time = x[7],
# )
#
# infected_fraction(model) =
# count(a.status == :I for a in allagents(model)) / nagents(model)
#
# _, mdf = run!(
# model,
# 50;
# mdata = [infected_fraction],
# when_model = [50],
# replicates = 10,
# )
#
# return mean(mdf.infected_fraction)
# end
# ```
using BlackBoxOptim, Random
using Statistics: mean

infected_fraction(model) = count(a.status == :I for a in allagents(model)) / max(1, nagents(model))

function cost(x)
rng = Xoshiro(43)
_, mdf = ensemblerun!(
[model_init(x, rand(rng, 1:1000)) for _ in 1:10],
50;
mdata = [infected_fraction],
when_model = [50],
init = false
)
return mean(mdf.infected_fraction)
end

# This cost function runs our model 10 times for 50 days, then returns the average number
# of infected people. When we pass this function to an optimiser, we will effectively be
# asking for a set of parameters that can reduce the number of infected people to the
# lowest possible number.

# We can now test the function cost with some reasonable parameter values.
# ```julia
# Random.seed!(10)
#
# x0 = [
# 0.2, # migration_rate
# 0.1, # death_rate
# 0.05, # β_det
# 0.3, # β_und
# 10, # infection_period
# 0.1, # reinfection_probability
# 5, # detection_time
# ]
# cost(x0)
# ```
# ```@raw html
# <pre class="documenter-example-output">0.9059485530546623</pre>
# ```

x0 = [
0.2, # max_travel_rate
0.1, # death_rate
0.05, # β_det
0.3, # β_und
10, # infection_period
0.1, # reinfection_probability
5, # detection_time
]
cost(x0)

# With these initial values, 94% of the population is infected after the 50 day period.

Expand All @@ -92,44 +87,26 @@
# cutoff time in the event that certain parameter sets are unfeasible and cause our model
# to never converge to a solution.

# ```julia
# result = bboptimize(
# cost,
# SearchRange = [
# (0.0, 1.0),
# (0.0, 1.0),
# (0.0, 1.0),
# (0.0, 1.0),
# (7.0, 13.0),
# (0.0, 1.0),
# (2.0, 6.0),
# ],
# NumDimensions = 7,
# MaxTime = 20,
# )
# best_fitness(result)
# ```
# ```@raw html
# <pre class="documenter-example-output">0.0</pre>
# ```
result = bboptimize(
cost,
SearchRange = [
(0.0, 1.0),
(0.0, 1.0),
(0.0, 1.0),
(0.0, 1.0),
(7.0, 13.0),
(0.0, 1.0),
(2.0, 6.0),
],
MaxTime = 10,
)
best_fitness(result)

# With the new parameter values found in `result`, we find that the
# fraction of the infected population can be dropped down to 11%.
# These values of these parameters are now:
# ```julia
# best_candidate(result)
# ```
# ```@raw html
# <pre class="documenter-example-output">7-element Array{Float64,1}:
# 0.1545049978104396
# 0.886202142470518
# 0.8258299702140992
# 0.7411762981538305
# 9.172098752376595
# 0.17302035312870545
# 5.907046385323653
# </pre>
# ```

best_candidate(result)

# Unfortunately we've not given the optimiser information we probably needed to.
# Notice that the death rate is 96%, with reinfection quite low.
Expand All @@ -140,144 +117,77 @@
# Let's modify the cost function to also keep the mortality rate low.

# First, we'll run the model with our new-found parameters:
# ```julia
# x = best_candidate(result)
#
# Random.seed!(0)
#
# model = model_initiation(;
# Ns = [500, 500, 500],
# migration_rate = x[1],
# death_rate = x[2],
# β_det = x[3],
# β_und = x[4],
# infection_period = x[5],
# reinfection_probability = x[6],
# detection_time = x[7],
# )
#
# _, data =
# run!(model, agent_step!, 50; mdata = [nagents], when_model = [50], replicates = 10)
#
# mean(data.nagents)
# ```
# ```@raw html
# <pre class="documenter-example-output">2.0</pre>
# ```

x = best_candidate(result)

rng = Xoshiro(43)
models = [model_init(x, rand(rng, 1:1000)) for _ in 1:10]
ensemblerun!(models, 50)
mean(nagents.(models))

# About 10% of the population dies with these parameters over our 50 day window.

# We can define a multi-objective cost function that minimizes the number of infected and
# deaths by returning more than one value in our cost function.
# ```julia
# function cost_multi(x)
# model = model_initiation(;
# Ns = [500, 500, 500],
# migration_rate = x[1],
# death_rate = x[2],
# β_det = x[3],
# β_und = x[4],
# infection_period = x[5],
# reinfection_probability = x[6],
# detection_time = x[7],
# )
#
# initial_size = nagents(model)
#
# infected_fraction(model) =
# count(a.status == :I for a in allagents(model)) / nagents(model)
# n_fraction(model) = -1.0 * nagents(model) / initial_size
#
# mdata = [infected_fraction, n_fraction]
# _, data = run!(
# model,
# agent_step!,
# 50;
# mdata,
# when_model = [50],
# replicates = 10,
# )
#
# return mean(data[!, dataname(mdata[1])), mean(data[!, dataname(mdata[2]))
# end
# ```

function cost_multi(x)
model = model_init(x, 1)
initial_size = nagents(model)
n_fraction(model) = -1.0 * nagents(model) / initial_size
rng = Xoshiro(43)
mdata = [infected_fraction, n_fraction]
_, data = ensemblerun!(
[model_init(x, rand(rng, 1:1000)) for _ in 1:10],
50;
mdata,
when_model = [50],
init = false
)
return mean(data[!, dataname(mdata[1])]), mean(data[!, dataname(mdata[2])])
end

# Notice that our new objective `n_fraction` is negative. It would be simpler to
# state we'd like to 'maximise the living population', but the optimiser we're using
# here focuses on minimising objectives only, therefore we must 'minimise the number of
# agents dying.

# ```julia
# cost_multi(x0)
# ```
# ```@raw html
# <pre class="documenter-example-output">(0.9812286689419796, -0.7813333333333333)</pre>
# ```
cost_multi(x0)

# The cost of our initial parameter values is high: most of the population (96%) is
# infected and 22% die.

# Let's minimize this multi-objective cost function.
# There is more than one way to approach such an optimisation.
# Again, refer to the BlackBoxOptim.jl documentation for specifics.
# ```julia
# result = bboptimize(
# cost_multi,
# Method = :borg_moea,
# FitnessScheme = ParetoFitnessScheme{2}(is_minimizing = true),
# SearchRange = [
# (0.0, 1.0),
# (0.0, 1.0),
# (0.0, 1.0),
# (0.0, 1.0),
# (7.0, 13.0),
# (0.0, 1.0),
# (2.0, 6.0),
# ],
# NumDimensions = 7,
# MaxTime = 55,
# )
# best_fitness(result)
# ```
# ```@raw html
# <pre class="documenter-example-output">(0.0047011417058428475, -0.9926666666666668)</pre>
# ```

# ```julia
# best_candidate(result)
# ```
# ```@raw html
# <pre class="documenter-example-output">7-element Array{Float64,1}:
# 0.8798741355149663
# 0.6703698358420607
# 0.07093587652308599
# 0.07760264834010584
# 10.65213641721431
# 0.9911248984077646
# 5.869646301829334
# </pre>
# ```

# These parameters look better: about 0.3% of the population dies and 0.02% are infected:

# The algorithm managed to minimize the number of infected and deaths while still
# increasing death rate to 42%, reinfection probability to 53%, and migration rates to 33%.
# The most important change however, was decreasing the transmission rate when individuals
# are infected and undetected from 30% in our initial calculation, to 0.2%.

# Over a longer period of time than 50 days, that high death rate will take its toll though.

result = bboptimize(
cost_multi,
Method = :borg_moea,
FitnessScheme = ParetoFitnessScheme{2}(is_minimizing = true),
SearchRange = [
(0.0, 1.0),
(0.0, 1.0),
(0.0, 1.0),
(0.0, 1.0),
(7.0, 13.0),
(0.0, 1.0),
(2.0, 6.0),
],
MaxTime = 20,
)
best_fitness(result)

best_candidate(result)

# These parameters look better: about 3% of the population dies and 31% are infected:

# Over a longer period of time than 50 days, the high death rate will take its toll though.
# Let's reduce that rate and check the cost.

# ```julia
# x = best_candidate(result)
# x[2] = 0.02
# cost_multi(x)
# ```
# ```@raw html
# <pre class="documenter-example-output">(0.03933333333333333, -1.0)</pre>
# ```
x = best_candidate(result)
x[2] = 0.02
cost_multi(x)

# The fraction of infected increases to 0.04%. This is an interesting result:
# The fraction of infected decreases to 0.04%. This is an interesting result:
# since this virus model is not as deadly, the chances of re-infection increase.
# We now have a set of parameters to strive towards in the real world. Insights
# such as these assist us to enact countermeasures like social distancing to mitigate
Expand Down
Loading