diff --git a/examples/optim.jl b/examples/optim.jl
index 76434bdb81..372c3133d7 100644
--- a/examples/optim.jl
+++ b/examples/optim.jl
@@ -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.
@@ -28,36 +43,22 @@
# 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
@@ -65,23 +66,17 @@
# 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
-# 0.9059485530546623
-# ```
+
+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.
@@ -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
-# 0.0
-# ```
+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
-# 7-element Array{Float64,1}:
-# 0.1545049978104396
-# 0.886202142470518
-# 0.8258299702140992
-# 0.7411762981538305
-# 9.172098752376595
-# 0.17302035312870545
-# 5.907046385323653
-#
-# ```
+
+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.
@@ -140,79 +117,40 @@
# 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
-# 2.0
-# ```
+
+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
-# (0.9812286689419796, -0.7813333333333333)
-# ```
+cost_multi(x0)
# The cost of our initial parameter values is high: most of the population (96%) is
# infected and 22% die.
@@ -220,64 +158,36 @@
# 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
-# (0.0047011417058428475, -0.9926666666666668)
-# ```
-
-# ```julia
-# best_candidate(result)
-# ```
-# ```@raw html
-# 7-element Array{Float64,1}:
-# 0.8798741355149663
-# 0.6703698358420607
-# 0.07093587652308599
-# 0.07760264834010584
-# 10.65213641721431
-# 0.9911248984077646
-# 5.869646301829334
-#
-# ```
-
-# 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
-# (0.03933333333333333, -1.0)
-# ```
+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