agents.jl 3.77 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
29
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.
"""
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
function complete_graph_from_households_composition(households_composition)
    total_household_pop = sum(sum.(households_composition))
    network = SimpleGraph(total_household_pop)
    vertex_pointer = 1
    for household in households_composition
        num_vertices = sum(household)
        for v in vertex_pointer:(vertex_pointer + num_vertices - 1)
            for w in vertex_pointer:(vertex_pointer + num_vertices - 1)
                if v != w
                    add_edge!(network,v,w)
                end
            end
        end
        vertex_pointer+=num_vertices
    end
    return network
end


Peter Jentsch's avatar
Peter Jentsch committed
49
50
51
52
"""
    generate_population(num_households::Int)

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

Peter Jentsch's avatar
Peter Jentsch committed
54
55
56
57
58
59
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
60
61
62
function generate_population(num_households)
    households_composition = sample_household_data(num_households)

63
    household_networks = complete_graph_from_households_composition(households_composition)
Peter Jentsch's avatar
Peter Jentsch committed
64
65
66
    
    #test the complete_graph_from_households_composition function against the mapreduce version to ensure it is correct.
    @c_assert all(adjacency_matrix(household_networks) .== adjacency_matrix( mapreduce(LightGraphs.complete_graph,blockdiag, sum.(households_composition); init = SimpleGraph())))
Peter Jentsch's avatar
Peter Jentsch committed
67

68
69
70
    households = map( l -> [fill(AgentDemographic(i),n) for (i,n) in enumerate(l)],households_composition) |>
    l -> reduce(vcat,l) |>
    l -> reduce(vcat,l)
Peter Jentsch's avatar
Peter Jentsch committed
71
72
    population_list = reduce(vcat,households)
    index_vectors = [findall(x -> x == AgentDemographic(i), population_list) for i in 1:(AgentDemographic.size-1)]
73
74

    
Peter Jentsch's avatar
Peter Jentsch committed
75
76
    return (;
        population_list,
77
        household_networks,
Peter Jentsch's avatar
Peter Jentsch committed
78
79
80
81
82
        index_vectors
    )
end


Peter Jentsch's avatar
Peter Jentsch committed
83
84
85
"""
Defines pretty printing methods so that `display(AgentStatus)` is more readable. 
"""
86
87
88
89
90
91
92
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")
93
94
    elseif status == Immunized
        print(io, "Vac")
95
96
97
    end
end

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