mixing_graphs.jl 10.5 KB
Newer Older
1
using Dictionaries
Peter Jentsch's avatar
Peter Jentsch committed
2
3
4
5
6
"""
A type that stores a pair of nodes representing an edge in the graph. 

We need to define a custom type for these so we can define a hash function on graph edges, in order to more efficiently use them in hashmaps (dictionaries)
"""
7
8
9
10
struct GraphEdge
    a::Int
    b::Int
end
Peter Jentsch's avatar
Peter Jentsch committed
11
12
13
14
15
16

"""
Define a hash on GraphEdge such that ``hash(a,b) = hash(b,a)`` (hash is commutative).

This is helpful because then we only need to store (a,b) in the graph edges weights dictionary, rather than both (a,b) and (b,a).
"""
17
18
19
function Base.hash(e::GraphEdge)
    return hash(minmax(e.a,e.b))
end
Peter Jentsch's avatar
Peter Jentsch committed
20
21

"""
22
Define symmetric edge equality, matches the hash function.
Peter Jentsch's avatar
Peter Jentsch committed
23
"""
24
25
26
function Base.isequal(e1::GraphEdge,e2::GraphEdge)
    return isequal(minmax(e1.a,e1.b),minmax(e2.a,e2.b))
end
Peter Jentsch's avatar
Peter Jentsch committed
27
28
29
30
31
32
33
34

"""
    sample_mixing_graph!(mixing_graph::Graph)


Resample all the weights in `mixing_graph`
"""
function sample_mixing_graph!(mixing_graph)
35
    mixing_edges = mixing_graph.mixing_edges
36
37
38
    # display(length.(mixing_edges.contact_array))
    # display(length.(mixing_edges.sample_cache))
    for i in 1:size(mixing_edges.contact_array)[1], j in 1:i  #diagonal     
39
        rand!(Random.default_rng(Threads.threadid()), mixing_edges.sampler_matrix[j,i],mixing_edges.sample_cache[j,i])
40
41
42
        for k in 1:length(mixing_edges.contact_array[j,i])
            edge_weight_k = mixing_edges.sample_cache[j,i][k]
            set!(mixing_edges.weights_dict, mixing_edges.contact_array[j,i][k], edge_weight_k)
43
        end
44
45
46
    end
end

Peter Jentsch's avatar
Peter Jentsch committed
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
"""
A type that defines the graph edges such that they can be easily resampled and then re-added to the graph. Created with `create_mixing_edges(..)`.

#Fields

    total_edges::Int

Total number of edges in the struct.    

    contact_array::Matrix{Tuple{Vector{Int},Vector{Int}}}

A matrix of vector pairs, such that one node is in the first vector and the other is in the second. The position of the vector pairs in the matrix corresponds to the distribution in sampler_matrix used to sample them. There is probably a better way to represent these, I did this so they could be sampled fast. We only use the upper triangle of this but Julia lacks a good Symmetric matrix type.

    sample_cache::Matrix{Vector{Int}}

The structure that pre-allocates the space for the weights of the edges in contact array. This is filled with weights and then the weights are added to weights_dict. We only use the upper triangle of this but Julia lacks a good Symmetric matrix type.

    weights_dict::Dictionary{GraphEdge,UInt8}

Stores the weights used in the graph, so they can be easily resampled.    

    sampler_matrix::M

This is the matrix of distributions, from which the edge weights are sampled. Specifically, weights for edges in `contact_array[i,j]` come from the distribution in `sampler_matrix[i,j]`, and are placed into `sample_cache[i,j]`.  We only use the upper triangle of this but Julia lacks a good Symmetric matrix type. See `sample_mixing_graph!`.
"""
72
struct MixingEdges{M}
Peter Jentsch's avatar
Peter Jentsch committed
73
    total_edges::Int
74
    contact_array::Matrix{Vector{GraphEdge}}
75
76
    sample_cache::Matrix{Vector{Int}}
    weights_dict::Dictionary{GraphEdge,UInt8}
77
    sampler_matrix::M
78
    function MixingEdges(total_edges,contact_array,sampler_matrix::Matrix{M}) where M<:Sampleable{Univariate, Discrete}
79
        sample_cache = map(v-> Vector{Int}(undef,length(v)),contact_array)
80
        weights_dict = Dictionary{GraphEdge,UInt8}(;sizehint = total_edges)
81
82
83
        new{typeof(sampler_matrix)}(total_edges,contact_array,sample_cache,weights_dict,sampler_matrix)
    end
end
Peter Jentsch's avatar
Peter Jentsch committed
84

Peter Jentsch's avatar
Peter Jentsch committed
85
86
87
"""
This function constructs the `MixingEdges` type for rest and WS graphs. Calls `generate_contact_vectors!(..)` to get the degree distributions for a bipartite graph for each pair of demographics, and then adds edges to MixingEdges.contact_array with the Chung-Lu approach. 
"""
88
function create_mixing_edges(demographic_index_vectors,mixing_matrix,weights_distribution_matrix)
89
    contact_array =[Vector{GraphEdge}() for i in 1:length(demographic_index_vectors),j in 1:length(demographic_index_vectors)]
90
91
92
93
94
    tot = 0
    for i in 1:size(mixing_matrix)[1], j in 1:i  #diagonal
        num_degrees_ij = zeros(Int,length(demographic_index_vectors[i]))
        num_degrees_ji = zeros(Int,length(demographic_index_vectors[j]))
        generate_contact_vectors!(mixing_matrix[i,j],mixing_matrix[j,i],num_degrees_ij,num_degrees_ji)
Peter Jentsch's avatar
Peter Jentsch committed
95

96
97
98
99
        m = sum(num_degrees_ij)
        stubs_i = Vector{Int}(undef,m)
        stubs_j = Vector{Int}(undef,m)
        if m>0
100
101
            sample!(Random.default_rng(Threads.threadid()),demographic_index_vectors[i],Weights(num_degrees_ij./m),stubs_i)
            sample!(Random.default_rng(Threads.threadid()),demographic_index_vectors[j],Weights(num_degrees_ji./m),stubs_j)
102
103
            tot += m
        end    
104
105

        contact_array[j,i] = GraphEdge.(stubs_i,stubs_j)
Peter Jentsch's avatar
Peter Jentsch committed
106
    end
107
108
109
    return MixingEdges(tot,contact_array,weights_distribution_matrix)
end

Peter Jentsch's avatar
Peter Jentsch committed
110
111
112
113

"""
This function constructs the `MixingEdges` type for the home graphs. Simply adds edges to MixingEdges from the graph `g`, since we already create that in `generate_population`.
"""
114
function create_mixing_edges(g::SimpleGraph,demographics,demographic_index_vectors,weights_distribution_matrix) 
115
116
    contact_array = [sizehint!(Vector{GraphEdge}(),ne(g)) for i in 1:length(demographic_index_vectors),j in 1:length(demographic_index_vectors)]
    for e in edges(g)
117
118
        i = src(e)
        j = dst(e)
119
        push!(contact_array[Int(demographics[i]),Int(demographics[j])], GraphEdge(i,j))
120
    end
121
    sample_cache = map(v-> Vector{Int}(undef,length(v)),contact_array)
122
    return MixingEdges(ne(g),contact_array,weights_distribution_matrix)
123
end
124

Peter Jentsch's avatar
Peter Jentsch committed
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
"""

Stores the full time dependent mixing graph for the model. I think this might be a weird abstraction for this idea but it works fine.

#Fields 

These are references to the graphs that get resampled everyday.

    resampled_graphs::NTuple{N,G}

List of lists of graphs, one list for each day. 

    graph_list::Vector{Vector{G}}

"""
140
141
struct TimeDepMixingGraph{N,G,GNT}
    resampled_graphs::NTuple{N,GNT}
142
    graph_list::Vector{Vector{G}}
143
144
    function TimeDepMixingGraph(len,resampled_graphs::NTuple{N,GNT},base_graph_list::Vector{G}) where {GNT,G,N}
        return new{N,G,GNT}(
145
146
147
148
149
            resampled_graphs,
            [copy(base_graph_list) for i in 1:len]
        )
    end
end
Peter Jentsch's avatar
Peter Jentsch committed
150
151
152
153
154
155

"""
Creates the `TimeDepMixingGraph` for our specific model. 

Assumes the simulation begins on Thursday arbitrarily.
"""
156
function time_dep_mixing_graphs(len,base_network,demographics,index_vectors,ws_matrix_tuple,rest_matrix_tuple)
Peter Jentsch's avatar
Peter Jentsch committed
157
    home_static_edges = WeightedGraph(base_network,demographics,index_vectors,contact_time_distributions.hh) #network with households and LTC homes
158
   
159
160
161
    ws_static_edges = WeightedGraph(demographics,index_vectors,ws_matrix_tuple.daily,contact_time_distributions.ws)
    ws_weekly_edges = WeightedGraph(demographics,index_vectors,ws_matrix_tuple.twice_a_week,contact_time_distributions.ws)
    ws_daily_edges = WeightedGraph(demographics,index_vectors,ws_matrix_tuple.otherwise,contact_time_distributions.ws)
162
    
163
164
165
    rest_static_edges = WeightedGraph(demographics,index_vectors,rest_matrix_tuple.daily,contact_time_distributions.rest)
    rest_weekly_edges = WeightedGraph(demographics,index_vectors,rest_matrix_tuple.twice_a_week,contact_time_distributions.rest)
    rest_daily_edges = WeightedGraph(demographics,index_vectors,rest_matrix_tuple.otherwise,contact_time_distributions.rest)
166
   
167
168
169
    inf_network_list = [home_static_edges,rest_static_edges] 
    soc_network_list = [home_static_edges,rest_static_edges,ws_static_edges]
    
170
171
    infected_mixing_graph = TimeDepMixingGraph(len,((ws_daily_edges,ws_matrix_tuple.daily),(rest_daily_edges,rest_matrix_tuple.daily)),inf_network_list)
    soc_mixing_graph = TimeDepMixingGraph(len,((ws_daily_edges,ws_matrix_tuple.daily),(rest_daily_edges,rest_matrix_tuple.daily)),soc_network_list)
172
173
174
175
176
177
    # display(infected_mixing_graph.graph_list)
    for (t,l) in enumerate(infected_mixing_graph.graph_list)
        day_of_week = mod(t,7)
        if !(day_of_week == 3 || day_of_week == 4) #simulation begins on thursday I guess
            push!(l, ws_static_edges)
        end
178
        if rand(Random.default_rng(Threads.threadid()))<5/7
179
180
181
182
183
184
185
186
187
188
            push!(l, ws_weekly_edges)
            push!(l, rest_weekly_edges)
        end
        push!(l,ws_daily_edges)
        push!(l,rest_daily_edges)
    end

    return infected_mixing_graph,soc_mixing_graph
end

Peter Jentsch's avatar
Peter Jentsch committed
189
190
191
"""
Completely remake all the graphs in `time_dep_mixing_graph.resampled_graphs`.
"""
192
193
function remake!(time_dep_mixing_graph,demographic_index_vectors)
    for (weighted_graph,mixing_matrix) in time_dep_mixing_graph.resampled_graphs
194
195
        empty!.(weighted_graph.g.fadjlist) #empty all the vector edgelists
        weighted_graph.g.ne = 0
196
        weighted_graph.mixing_edges = create_mixing_edges(demographic_index_vectors,mixing_matrix,weighted_graph.mixing_edges.sampler_matrix)
Peter Jentsch's avatar
Peter Jentsch committed
197
        graph_from_mixing_edges(weighted_graph.g,weighted_graph.mixing_edges)
198
199
    end
end
Peter Jentsch's avatar
Peter Jentsch committed
200

Peter Jentsch's avatar
Peter Jentsch committed
201
202
203
"""
Add the edges defined by MixingEdges to the actual graph G. Another big bottleneck since adjancency lists don't add edges super efficiently, and there are a ton of them. 
"""
Peter Jentsch's avatar
Peter Jentsch committed
204
205
function graph_from_mixing_edges(g,mixing_edges)
    for i in 1:size(mixing_edges.contact_array)[1], j in 1:i  #diagonal
206
207
208
        for k in 1:length(mixing_edges.contact_array[i,j])
            e = mixing_edges.contact_array[i,j][k]
            add_edge!(g,e.a,e.b)
Peter Jentsch's avatar
Peter Jentsch committed
209
210
211
        end
    end
end
Peter Jentsch's avatar
Peter Jentsch committed
212
213
214
215
216

"""
My own weighted graph type. Stores the graph in `g`, and the weights and edges in `mixing_edges`. 

"""
217
mutable struct WeightedGraph{G,M} 
Peter Jentsch's avatar
Peter Jentsch committed
218
    g::G
219
    mixing_edges::M
220
    function WeightedGraph(demographics,demographic_index_vectors,mixing_matrix,weights_distribution_matrix)
221
        mixing_edges = create_mixing_edges(demographic_index_vectors,mixing_matrix,weights_distribution_matrix)
222
        g = Graph(length(demographics))
Peter Jentsch's avatar
Peter Jentsch committed
223
        graph_from_mixing_edges(g,mixing_edges)
224
        return new{typeof(g),typeof(mixing_edges)}(
Peter Jentsch's avatar
Peter Jentsch committed
225
            g,
226
            mixing_edges
Peter Jentsch's avatar
Peter Jentsch committed
227
228
        )
    end
Peter Jentsch's avatar
Peter Jentsch committed
229
    function WeightedGraph(g::SimpleGraph,demographics,demographic_index_vectors,weights_distribution_matrix)
230
231
        mixing_edges = create_mixing_edges(g, demographics, demographic_index_vectors,weights_distribution_matrix)
        weights_dict = Dict{GraphEdge,UInt8}()
232
        sample_cache = map(v-> Vector{Int}(undef,length(v)),mixing_edges.contact_array)
233
        return new{typeof(g),typeof(mixing_edges)}(
Peter Jentsch's avatar
Peter Jentsch committed
234
            g,
Peter Jentsch's avatar
Peter Jentsch committed
235
            mixing_edges,
Peter Jentsch's avatar
Peter Jentsch committed
236
        )
Peter Jentsch's avatar
Peter Jentsch committed
237
238
    end
end
Peter Jentsch's avatar
Peter Jentsch committed
239
neighbors(g::WeightedGraph,i) = neighbors(g.g,i)
240
get_weight(g::WeightedGraph,e) = g.mixing_edges.weights_dict[e]
241
function Base.show(io::IO, g::WeightedGraph) 
Peter Jentsch's avatar
Peter Jentsch committed
242
    print(io, "WG $(ne(g.g))")
243
end