agents.jl 3.89 KB
Newer Older
1
#agent type definitions
Peter Jentsch's avatar
Peter Jentsch committed
2
3
4
5

"""
Define the possible agent demographic groups. Note that these can be converted to/from integers (starting at 1).
"""
6
7
8
9
10
@enum AgentDemographic begin
    Young = 1
    Adult
    Elderly
end
Peter Jentsch's avatar
Peter Jentsch committed
11
12
13
14

"""
Define the possible agent infection statuses. Note that these can be converted to/from integers (starting at 1).
"""
15
@enum AgentStatus begin
Peter Jentsch's avatar
Peter Jentsch committed
16
    Susceptible = 1
17
18
    Infected
    Recovered
Peter Jentsch's avatar
Peter Jentsch committed
19
    Immunized
20
21
end

Peter Jentsch's avatar
Peter Jentsch committed
22
23
24
"""

    complete_graph_from_households_composition(households_composition::AbstractVector{AbstractVector{Int}})
25

Peter Jentsch's avatar
Peter Jentsch committed
26
27
28
Generates a complete graph from a vector of household compositions. Each household composition is a 3 element vectors (one for each demographic group) of integers where each element describes the number of the corresponding demographic group present in that household.

This function wires together a graph such that each household is in a complete subgraph. It is much faster than the `mapreduce` solution used previously.
29
30

Also returns a vector of AgentDemographic representing each agent, defined by the household compositions. 
Peter Jentsch's avatar
Peter Jentsch committed
31
"""
Peter Jentsch's avatar
Peter Jentsch committed
32
@views function complete_graph_from_households_composition(households_composition, sampler_matrix)
33
    total_household_pop = sum(sum.(households_composition))
34
    population_list = Vector{AgentDemographic}(undef,total_household_pop)
Peter Jentsch's avatar
Peter Jentsch committed
35
    wg = WeightedGraph(total_household_pop,sampler_matrix)
36
37
38
    vertex_pointer = 1
    for household in households_composition
        num_vertices = sum(household)
39
40
41
42
43
44
45
        household_pointer = 0
        for (k,size) in enumerate(household)
            @inbounds population_list[vertex_pointer+household_pointer:vertex_pointer+household_pointer+size-1] .= AgentDemographic(k)
            household_pointer += size
        end
        for v in vertex_pointer:(vertex_pointer + num_vertices - 1),  w in vertex_pointer:(vertex_pointer + num_vertices - 1)
            if v != w
Peter Jentsch's avatar
Peter Jentsch committed
46
47
                weight::UInt8 = rand(Random.default_rng(Threads.threadid()),sampler_matrix[Int(population_list[v]),Int(population_list[w])])
                add_edge!(wg,v,w,weight)
48
49
50
51
            end
        end
        vertex_pointer+=num_vertices
    end
Peter Jentsch's avatar
Peter Jentsch committed
52
    return wg,population_list
53
54
55
end


Peter Jentsch's avatar
Peter Jentsch committed
56
57
58
59
"""
    generate_population(num_households::Int)

Returns a namedtuple with fields `population_list`, `household_networks` and `index_vectors`.
60

Peter Jentsch's avatar
Peter Jentsch committed
61
62
63
64
65
66
population_list: Vector of AgentDemographic where each entry is the demographic of a person in the simulation. The length is equal to the total number of people in the model.

household_networks: SimpleGraph with all agents such that each household is in a complete subgraph.

index_vectors: Vector of 3 vectors, each of which contains the indexes of all the agents in a particular demographic. If `population_list` is a mapping from agents to demographics, these are mappings from demographics to agents.
"""
Peter Jentsch's avatar
Peter Jentsch committed
67
68
function generate_population(num_households)
    households_composition = sample_household_data(num_households)
Peter Jentsch's avatar
Peter Jentsch committed
69
    household_networks,population_list = complete_graph_from_households_composition(households_composition, contact_time_distributions.hh)
Peter Jentsch's avatar
Peter Jentsch committed
70
71
72
    index_vectors = [findall(x -> x == AgentDemographic(i), population_list) for i in 1:(AgentDemographic.size-1)]
    return (;
        population_list,
73
        household_networks,
Peter Jentsch's avatar
Peter Jentsch committed
74
75
76
77
78
        index_vectors
    )
end


Peter Jentsch's avatar
Peter Jentsch committed
79
"""
80
Defines pretty printing methods so that `display(s::AgentStatus)` is more readable. 
Peter Jentsch's avatar
Peter Jentsch committed
81
"""
82
83
84
85
86
87
88
function Base.show(io::IO, status::AgentStatus) 
    if status == Susceptible
        print(io, "S")
    elseif status == Infected
        print(io, "I")
    elseif status == Recovered
        print(io, "R")
89
90
    elseif status == Immunized
        print(io, "Vac")
91
92
93
    end
end

Peter Jentsch's avatar
Peter Jentsch committed
94
"""
95
Defines pretty printing methods so that `display(s::AgentDemographic)` is more readable. 
Peter Jentsch's avatar
Peter Jentsch committed
96
"""
97
98
function Base.show(io::IO, status::AgentDemographic) 
    if status == Young
99
        print(io, "<25")
100
    elseif status == Adult
101
        print(io, "20-65")
102
    elseif status == Elderly
103
        print(io, ">65")
104
105
    end
end