Even on this simple example Ark.jl surpasses a Vector{Symbol} by 2-3x.
utils:
using Random
function rate_to_probability(r::T, t::T) where {T<:AbstractFloat}
1 - exp(-r * t)
end
const rng = Xoshiro(42)
# Time domain parameters
const δt = 0.1
const nsteps = 400
const tf = nsteps * δt
const t = 0:δt:tf;
# System parameters
const β = 0.05
const c = 10.0
const γ = rate_to_probability(0.25, δt);
# Initial conditions
const N = 1_000_000
const I0 = 5;
Ark.jl version (on dev):
using Ark
abstract type HealthState end
struct S <: HealthState end
struct I <: HealthState end
struct R <: HealthState end;
function get_count(world, ::Type{T}) where {T<:HealthState}
count = 0
for (entities, ) in Query(world, (), with=(T, ))
count += length(entities)
end
return count
end
function run_sir()
# set up world, add initial entities
world = World(S,I,R)
new_entities!(world, N-I0, (S(),))
new_entities!(world, I0, (I(),))
# output
trajectory = zeros(Int, nsteps+1, length(subtypes(HealthState)))
trajectory[1, :] = [get_count(world, T) for T in (S, I, R)]
# vectors to contain Entities making state transitions on this step
i_to_r = Entity[]
s_to_i = Entity[]
# vector to contain uniform random numbers for state transitions
rands = Float64[]
# run sim
for t in 1:nsteps
# S->I
i = get_count(world, I)
foi = β * c * i/N
prob = rate_to_probability(foi, δt)
for (entities, ) in Query(world, (), with=(S, ))
resize!(rands, length(entities)); rand!(rng, rands)
@inbounds for i in eachindex(entities)
if rands[i] <= prob
push!(s_to_i, entities[i])
end
end
end
# I->R
for (entities, ) in Query(world, (), with=(I, ))
resize!(rands, length(entities)); rand!(rng, rands)
@inbounds for i in eachindex(entities)
if rands[i] <= γ
push!(i_to_r, entities[i])
end
end
end
# apply transitions
for entity in s_to_i
exchange_components!(world, entity, add = (I(),), remove = (S,))
end
for entity in i_to_r
exchange_components!(world, entity, add = (R(),), remove = (I,))
end
# record state
trajectory[t+1, :] = [get_count(world, T) for T in (S, I, R)]
# reuse vectors to avoid allocations
resize!(i_to_r, 0)
resize!(s_to_i, 0)
end
return trajectory
end
@time run_sir(); # 0.439895 seconds
Vector{Symbol}:
function get_counts(model)
count_S, count_I, count_R = 0, 0, 0
for agent in model
if agent == :S
count_S += 1
elseif agent == :I
count_I += 1
elseif agent == :R
count_R += 1
end
end
return [count_S, count_I, count_R]
end
function run_sir()
model = Symbol[]
for _ in 1:N-I0
push!(model, :S)
end
for _ in 1:I0
push!(model, :I)
end
trajectory = zeros(Int, nsteps+1, 3)
trajectory[1, :] = get_counts(model)
for t in 1:nsteps
i = get_counts(model)[2]
foi = β * c * i/N
prob = rate_to_probability(foi, δt)
@inbounds for i in eachindex(model)
if model[i] == :S && rand(rng) <= prob
model[i] = :I
elseif model[i] == :I && rand(rng) <= γ
model[i] = :R
end
end
trajectory[t+1, :] = get_counts(model)
end
return trajectory
end
@time run_sir(); # 1.317357 seconds
Optimized Vector{Symbol} → going more unreadable:
function get_counts(model)
count_S, count_I, count_R = 0, 0, 0
for agent in model
if agent == :S
count_S += 1
elseif agent == :I
count_I += 1
elseif agent == :R
count_R += 1
end
end
return [count_S, count_I, count_R]
end
function run_sir()
model = Symbol[]
for _ in 1:N-I0
push!(model, :S)
end
for _ in 1:I0
push!(model, :I)
end
trajectory = zeros(Int, nsteps+1, 3)
trajectory[1, :] = get_counts(model)
for t in 1:nsteps
i = sum(1 for agent in model if agent==:I)
foi = β * c * i/N
prob = rate_to_probability(foi, δt)
count_S, count_I, count_R = 0, 0, 0
@inbounds for i in eachindex(model)
if model[i] == :S
count_S += 1
if rand(rng) <= prob
model[i] = :I
end
elseif model[i] == :I
count_I += 1
if rand(rng) <= γ
model[i] = :R
end
else
count_R += 1
end
end
trajectory[t+1, :] = [count_S, count_I, count_R]
end
return trajectory
end
@time run_sir(); # 0.884714 seconds
