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
using Agents: abmproperties, abmrng

@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 = StandardABM(PoorSoul, space; agent_step!, 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)
    return recover_or_die!(agent, model)
end

function migrate!(agent, model)
    pid = agent.pos
    d = DiscreteNonParametric(1:(model.C), model.migration_rates[pid, :])
    m = rand(abmrng(model), d)
    return 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(abmrng(model), d)
    n == 0 && return

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

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

function recover_or_die!(agent, model)
    return if agent.days_infected ≥ model.infection_period
        if rand(abmrng(model)) ≤ model.death_rate
            @a remove_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.0, 100.0), adata = to_collect)
agent sir_model with uuid 54cf095b 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 a7bb2f20 of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(5461, 2, 0, :S)
    agent 23986 with uuid 36550bce of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(23986, 7, 0, :S)
    agent 10467 with uuid 28849cc8 of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(10467, 3, 0, :S)
    agent 18565 with uuid 1b4e4626 of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(18565, 5, 0, :S)
    agent 26767 with uuid ac231f6c 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 54cf095b of type ABMAgent 
   custom properties:
   abm: 
StandardABM with 27416 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                2                 0          28540
   6 │     5                2                 0          28540
   7 │     6                3                 0          28540
   8 │     7                5                 0          28540
  ⋮  │   ⋮           ⋮                ⋮                ⋮
  95 │    94            27416                 0          27416
  96 │    95            27416                 0          27416
  97 │    96            27416                 0          27416
  98 │    97            27416                 0          27416
  99 │    98            27416                 0          27416
 100 │    99            27416                 0          27416
 101 │   100            27414                 2          27416
                                                86 rows omitted
   inner agents: 
    agent 5461 with uuid a7bb2f20 of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(5461, 4, 20, :I)
    agent 23986 with uuid 36550bce of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(23986, 3, 12, :I)
    agent 10467 with uuid 28849cc8 of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(10467, 2, 14, :I)
    agent 18565 with uuid 1b4e4626 of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(18565, 2, 13, :I)
    agent 26767 with uuid ac231f6c of type AAgent 
       custom properties:
       agent: 
Main.PoorSoul(26767, 1, 14, :I)
    27411 more agent(s) not shown ...

Print the data collected during the simulation:

println(m.df_agents)
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                2                 0          28540
   6 │     5                2                 0          28540
   7 │     6                3                 0          28540
   8 │     7                5                 0          28540
   9 │     8                9                 0          28540
  10 │     9               10                 0          28540
  11 │    10               19                 0          28540
  12 │    11               28                 0          28540
  13 │    12               40                 0          28540
  14 │    13               63                 0          28540
  15 │    14              116                 0          28540
  16 │    15              209                 0          28540
  17 │    16              330                 0          28540
  18 │    17              575                 0          28540
  19 │    18              944                 0          28540
  20 │    19             1563                 0          28540
  21 │    20             2658                 0          28540
  22 │    21             3410                 0          28540
  23 │    22             4652                 0          28540
  24 │    23             6735                 0          28540
  25 │    24            10114                 0          28540
  26 │    25            14925                 0          28540
  27 │    26            19545                 0          28540
  28 │    27            24133                 0          28540
  29 │    28            28383                 0          28540
  30 │    29            28540                 0          28540
  31 │    30            28540                 0          28540
  32 │    31            28540                 0          28540
  33 │    32            28540                 0          28540
  34 │    33            28540                 0          28540
  35 │    34            28540                 0          28540
  36 │    35            28540                 0          28540
  37 │    36            28540                 0          28540
  38 │    37            28539                 0          28539
  39 │    38            28539                 0          28539
  40 │    39            28539                 0          28539
  41 │    40            28537                 1          28538
  42 │    41            28536                 2          28538
  43 │    42            28534                 3          28537
  44 │    43            28532                 5          28537
  45 │    44            28526                 7          28533
  46 │    45            28518                14          28532
  47 │    46            28502                29          28531
  48 │    47            28472                56          28528
  49 │    48            28436                80          28516
  50 │    49            28381               115          28496
  51 │    50            28427                50          28477
  52 │    51            28348               108          28456
  53 │    52            28265               157          28422
  54 │    53            28089               264          28353
  55 │    54            27863               413          28276
  56 │    55            27988               202          28190
  57 │    56            27953               167          28120
  58 │    57            27890               110          28000
  59 │    58            27957                 4          27961
  60 │    59            27961                 0          27961
  61 │    60            27961                 0          27961
  62 │    61            27961                 0          27961
  63 │    62            27961                 0          27961
  64 │    63            27961                 0          27961
  65 │    64            27961                 0          27961
  66 │    65            27961                 0          27961
  67 │    66            27961                 0          27961
  68 │    67            27961                 0          27961
  69 │    68            27961                 0          27961
  70 │    69            27961                 0          27961
  71 │    70            27960                 1          27961
  72 │    71            27960                 1          27961
  73 │    72            27959                 2          27961
  74 │    73            27953                 7          27960
  75 │    74            27949                 8          27957
  76 │    75            27937                18          27955
  77 │    76            27926                27          27953
  78 │    77            27906                41          27947
  79 │    78            27865                73          27938
  80 │    79            27771               131          27902
  81 │    80            27830                51          27881
  82 │    81            27757               109          27866
  83 │    82            27658               172          27830
  84 │    83            27534               245          27779
  85 │    84            27331               373          27704
  86 │    85            27432               184          27616
  87 │    86            27383               159          27542
  88 │    87            27357                95          27452
  89 │    88            27409                 7          27416
  90 │    89            27416                 0          27416
  91 │    90            27416                 0          27416
  92 │    91            27416                 0          27416
  93 │    92            27416                 0          27416
  94 │    93            27416                 0          27416
  95 │    94            27416                 0          27416
  96 │    95            27416                 0          27416
  97 │    96            27416                 0          27416
  98 │    97            27416                 0          27416
  99 │    98            27416                 0          27416
 100 │    99            27416                 0          27416
 101 │   100            27414                 2          27416
draw(m)