Agents.jl Integration

We instantiate an agent-based SIR model based on Agents.jl: SIR model for the spread of COVID-19 make use of a SIR model constructor from an Agents.jl' SIR model for the spread of COVID-19, and then we simulate the model using AlgebraicAgents.jl

SIR Model in Agents.jl

To begin with, we define the Agents.jl model:

# SIR model for the spread of COVID-19
# taken from https://juliadynamics.github.io/Agents.jl/stable/examples/sir/
using AlgebraicAgents
using Agents, Random
using Agents.DataFrames, Agents.Graphs
using Distributions: Poisson, DiscreteNonParametric
using DrWatson: @dict
using Plots

@agent struct PoorSoul(GraphAgent)
    days_infected::Int  ## number of days since is infected
    status::Symbol  ## 1: S, 2: I, 3:R
end

Let's provide the constructors:

function model_initiation(;
                          Ns,
                          migration_rates,
                          β_und,
                          β_det,
                          infection_period = 30,
                          reinfection_probability = 0.05,
                          detection_time = 14,
                          death_rate = 0.02,
                          Is = [zeros(Int, length(Ns) - 1)..., 1],
                          seed = 0)
    rng = MersenneTwister(seed)
    @assert length(Ns)==
            length(Is)==
            length(β_und)==
            length(β_det)==
            size(migration_rates, 1) "length of Ns, Is, and B, and number of rows/columns in migration_rates should be the same "
    @assert size(migration_rates, 1)==size(migration_rates, 2) "migration_rates rates should be a square matrix"

    C = length(Ns)
    # normalize migration_rates
    migration_rates_sum = sum(migration_rates, dims = 2)
    for c in 1:C
        migration_rates[c, :] ./= migration_rates_sum[c]
    end

    properties = @dict(Ns,
                       Is,
                       β_und,
                       β_det,
                       β_det,
                       migration_rates,
                       infection_period,
                       infection_period,
                       reinfection_probability,
                       detection_time,
                       C,
                       death_rate)
    space = GraphSpace(complete_digraph(C))
    model = ABM(PoorSoul, space; properties, rng, model_step! = identity)

    # Add initial individuals
    for city in 1:C, n in 1:Ns[city]
        ind = add_agent!(city, model, 0, :S) ## Susceptible
    end
    # add infected individuals
    for city in 1:C
        inds = ids_in_position(city, model)
        for n in 1:Is[city]
            agent = model[inds[n]]
            agent.status = :I ## Infected
            agent.days_infected = 1
        end
    end
    return model
end

using LinearAlgebra: diagind

function create_params(;
                       C,
                       max_travel_rate,
                       infection_period = 30,
                       reinfection_probability = 0.05,
                       detection_time = 14,
                       death_rate = 0.02,
                       Is = [zeros(Int, C - 1)..., 1],
                       seed = 19)
    Random.seed!(seed)
    Ns = rand(50:5000, C)
    β_und = rand(0.3:0.02:0.6, C)
    β_det = β_und ./ 10

    Random.seed!(seed)
    migration_rates = zeros(C, C)
    for c in 1:C
        for c2 in 1:C
            migration_rates[c, c2] = (Ns[c] + Ns[c2]) / Ns[c]
        end
    end
    maxM = maximum(migration_rates)
    migration_rates = (migration_rates .* max_travel_rate) ./ maxM
    migration_rates[diagind(migration_rates)] .= 1.0

    params = @dict(Ns,
                   β_und,
                   β_det,
                   migration_rates,
                   infection_period,
                   reinfection_probability,
                   detection_time,
                   death_rate,
                   Is)

    return params
end

create_params (generic function with 1 method)

It remains to provide the SIR stepping functions:

function agent_step!(agent, model)
    @get_model model
    extract_agent(model, agent)
    migrate!(agent, model)
    transmit!(agent, model)
    update!(agent, model)
    recover_or_die!(agent, model)
end

function migrate!(agent, model)
    pid = agent.pos
    d = DiscreteNonParametric(1:(model.C), model.migration_rates[pid, :])
    m = rand(model.rng, d)
    if m ≠ pid
        move_agent!(agent, m, model)
    end
end

function transmit!(agent, model)
    agent.status == :S && return
    rate = if agent.days_infected < model.detection_time
        model.β_und[agent.pos]
    else
        model.β_det[agent.pos]
    end

    d = Poisson(rate)
    n = rand(model.rng, d)
    n == 0 && return

    for contactID in ids_in_position(agent, model)
        contact = model[contactID]
        if contact.status == :S ||
           (contact.status == :R && rand(model.rng) ≤ model.reinfection_probability)
            contact.status = :I
            n -= 1
            n == 0 && return
        end
    end
end

update!(agent, model) = agent.status == :I && (agent.days_infected += 1)

function recover_or_die!(agent, model)
    if agent.days_infected ≥ model.infection_period
        if rand(model.rng) ≤ model.death_rate
            @a kill_agent!(agent, model)
        else
            agent.status = :R
            agent.days_infected = 0
        end
    end
end

recover_or_die! (generic function with 1 method)

That's it!

Simulation Using AlgebraicAgents.jl

We instantiate a sample ABM model:

# create a sample agent based model
params = create_params(C = 8, max_travel_rate = 0.01)
abm = model_initiation(; params...)

StandardABM with 28540 agents of type PoorSoul
 agents container: Dict
 space: GraphSpace with 8 positions and 56 edges
 scheduler: fastest
 properties: Is, death_rate, infection_period, β_und, Ns, migration_rates, detection_time, reinfection_probability, β_det, C

Let's specify what data to collect:

infected(x) = count(i == :I for i in x)
recovered(x) = count(i == :R for i in x)
to_collect = [(:status, f) for f in (infected, recovered, length)]

3-element Vector{Tuple{Symbol, Function}}:
 (:status, Main.infected)
 (:status, Main.recovered)
 (:status, length)

We wrap the model as an agent:

m = ABMAgent("sir_model", abm; tspan=(0., 100.), adata=to_collect)

agent sir_model with uuid 30943f02 of type ABMAgent 
   custom properties:
   abm: 
StandardABM with 28540 agents of type PoorSoul
 agents container: Dict
 space: GraphSpace with 8 positions and 56 edges
 scheduler: fastest
 properties: migration_rates, Is, infection_period, __aagent__, detection_time, reinfection_probability, Ns, β_det, death_rate, β_und, C
   df_agents: 
0×0 DataFrame
   df_model: 
0×0 DataFrame
   inner agents: 
    agent 5461 with uuid cf5e1333 of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(5461, 2, 0, :S)
    agent 23986 with uuid 97381e00 of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(23986, 7, 0, :S)
    agent 10467 with uuid 1939fbea of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(10467, 3, 0, :S)
    agent 18565 with uuid 801dea46 of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(18565, 5, 0, :S)
    agent 26767 with uuid fe91cb9c of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(26767, 7, 0, :S)
    28535 more agent(s) not shown ...

And we simulate the dynamics:

simulate(m)

agent sir_model with uuid 30943f02 of type ABMAgent 
   custom properties:
   abm: 
StandardABM with 28540 agents of type PoorSoul
 agents container: Dict
 space: GraphSpace with 8 positions and 56 edges
 scheduler: fastest
 properties: migration_rates, Is, infection_period, __aagent__, detection_time, reinfection_probability, Ns, β_det, death_rate, β_und, C
   df_agents: 
0×0 DataFrame
   df_model: 
101×4 DataFrame
 Row │ time   infected_status  recovered_status  length_status
     │ Int64  Int64            Int64             Int64
─────┼─────────────────────────────────────────────────────────
   1 │     0                1                 0          28540
   2 │     1                1                 0          28540
   3 │     2                1                 0          28540
   4 │     3                1                 0          28540
   5 │     4                1                 0          28540
   6 │     5                1                 0          28540
   7 │     6                1                 0          28540
   8 │     7                1                 0          28540
  ⋮  │   ⋮           ⋮                ⋮                ⋮
  95 │    94                1                 0          28540
  96 │    95                1                 0          28540
  97 │    96                1                 0          28540
  98 │    97                1                 0          28540
  99 │    98                1                 0          28540
 100 │    99                1                 0          28540
 101 │   100                1                 0          28540
                                                86 rows omitted
   inner agents: 
    agent 5461 with uuid cf5e1333 of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(5461, 2, 0, :S)
    agent 23986 with uuid 97381e00 of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(23986, 7, 0, :S)
    agent 10467 with uuid 1939fbea of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(10467, 3, 0, :S)
    agent 18565 with uuid 801dea46 of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(18565, 5, 0, :S)
    agent 26767 with uuid fe91cb9c of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(26767, 7, 0, :S)
    28535 more agent(s) not shown ...

draw(m)