mixing_graphs.jl 9.16 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
Define a hash on GraphEdge such that `hash(a,b) = hash(b,a)` (hash is commutative).
Peter Jentsch's avatar
Peter Jentsch committed
14
15
16

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

"""
    sample_mixing_graph!(mixing_graph::Graph)


Resample all the weights in `mixing_graph`
"""
Peter Jentsch's avatar
Peter Jentsch committed
34
35
36
37
38
function sample_mixing_edges!(wg,sampler_matrix,demographics)
    for k in 1:ne(wg)
        i = wg.srcs[k]
        j = wg.sinks[k]
        wg.weights[k].x = rand(Random.default_rng(Threads.threadid()),sampler_matrix[Int(demographics[j]),Int(demographics[i])])
39
40
41
    end
end

Peter Jentsch's avatar
Peter Jentsch committed
42
43
44
45
46
"""
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 

Peter Jentsch's avatar
docs    
Peter Jentsch committed
47
48
    remade_graphs::NTuple{N,G}

Peter Jentsch's avatar
Peter Jentsch committed
49
50
These are references to the graphs that get resampled everyday.

Peter Jentsch's avatar
docs    
Peter Jentsch committed
51
52
53
    resampled_graphs::NTuple{N,G}   

These are references to the graphs that get resampled everyday.
Peter Jentsch's avatar
Peter Jentsch committed
54
55
56
57


    graph_list::Vector{Vector{G}}

Peter Jentsch's avatar
docs    
Peter Jentsch committed
58
59
List of lists of graphs, one list for each day. 

Peter Jentsch's avatar
Peter Jentsch committed
60
"""
Peter Jentsch's avatar
Peter Jentsch committed
61
62
63
struct TimeDepMixingGraph{G,T1,T2}
    remade_graphs::T1
    resampled_graphs::T2
64
    graph_list::Vector{Vector{G}}
Peter Jentsch's avatar
Peter Jentsch committed
65
66
67
    function TimeDepMixingGraph(len,remade_graphs::T1,resampled_graphs::T2,base_graph_list::Vector{G}) where {G,T1,T2}
        return new{G,T1,T2}(
            remade_graphs,
68
69
70
71
72
            resampled_graphs,
            [copy(base_graph_list) for i in 1:len]
        )
    end
end
Peter Jentsch's avatar
Peter Jentsch committed
73
74
75
76
77
78

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

Assumes the simulation begins on Thursday arbitrarily.
"""
Peter Jentsch's avatar
Peter Jentsch committed
79
function time_dep_mixing_graphs(len,home_static_edges,demographics,index_vectors,ws_matrix_tuple,rest_matrix_tuple)
80

81
82
    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)
83
    ws_justonce_edges = WeightedGraph(demographics,index_vectors,ws_matrix_tuple.otherwise,contact_time_distributions.ws)
84

85
86
    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)
87
    rest_justonce_edges = WeightedGraph(demographics,index_vectors,rest_matrix_tuple.otherwise,contact_time_distributions.rest)
88
   
89
    inf_network_list = [home_static_edges,rest_static_edges,ws_justonce_edges,rest_justonce_edges] 
90
    soc_network_list = [home_static_edges,rest_static_edges,ws_static_edges]
Peter Jentsch's avatar
Peter Jentsch committed
91

92
    remade_graphs = (ws_justonce_edges,rest_justonce_edges)
Peter Jentsch's avatar
Peter Jentsch committed
93
94
95
96
    resampled_graphs = (home_static_edges,rest_static_edges,ws_static_edges,rest_weekly_edges,ws_weekly_edges)

    infected_mixing_graph = TimeDepMixingGraph(len,remade_graphs,resampled_graphs,inf_network_list)

97
98
99
100
101
    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
Peter Jentsch's avatar
Peter Jentsch committed
102
        if rand(Random.default_rng(Threads.threadid()))<5/7
103
104
105
106
107
            push!(l, ws_weekly_edges)
            push!(l, rest_weekly_edges)
        end
    end

Peter Jentsch's avatar
Peter Jentsch committed
108
    soc_mixing_graph = TimeDepMixingGraph(len,remade_graphs,resampled_graphs,soc_network_list)
109
110
111
    return infected_mixing_graph,soc_mixing_graph
end

Peter Jentsch's avatar
Peter Jentsch committed
112
113
114
"""
Completely remake all the graphs in `time_dep_mixing_graph.resampled_graphs`.
"""
Peter Jentsch's avatar
Peter Jentsch committed
115
function remake_all!(t,time_dep_mixing_graph,index_vectors,demographics)
116
    for wg in time_dep_mixing_graph.remade_graphs
Peter Jentsch's avatar
Peter Jentsch committed
117
118
        empty!(wg)
        assemble_graph!(wg,index_vectors,wg.mixing_matrix,wg.sampler_matrix)
119
    end
Peter Jentsch's avatar
Peter Jentsch committed
120
    # display_degree(time_dep_mixing_graph.resampled_graphs[1])
121
122
    for wg in time_dep_mixing_graph.resampled_graphs
        if wg in time_dep_mixing_graph.graph_list[t]
Peter Jentsch's avatar
Peter Jentsch committed
123
            sample_mixing_edges!(wg,wg.sampler_matrix,demographics)
124
125
        end
    end
Peter Jentsch's avatar
Peter Jentsch committed
126
    # display_degree(time_dep_mixing_graph.resampled_graphs[1])
Peter Jentsch's avatar
Peter Jentsch committed
127
end
Peter Jentsch's avatar
Peter Jentsch committed
128

Peter Jentsch's avatar
Peter Jentsch committed
129
130
131
132
133
134
135
struct Neighborhood
    neighbourhood::Vector{Int}
    weights::Vector{Base.RefValue{UInt8}}
    function Neighborhood()
        return new(Vector{Int}(),Vector{Base.RefValue{UInt8}}())
    end
end
Peter Jentsch's avatar
Peter Jentsch committed
136

Peter Jentsch's avatar
Peter Jentsch committed
137
"""
138
Weighted graph type. Stores the graph in `g`, and the weights and edges in `mixing_edges`. 
Peter Jentsch's avatar
docs    
Peter Jentsch committed
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
Fields

    g::SimpleGraph

Stores the actual graph structure

    weights_dict::Dictionary{GraphEdge,UInt8}

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

    mixing_matrix::M1

Matrix of distributions determining node degrees 


    sampler_matrix::M

Matrix of distributions determining the edge weights

Peter Jentsch's avatar
Peter Jentsch committed
158
"""
Peter Jentsch's avatar
Peter Jentsch committed
159
160
161
162
163
struct WeightedGraph{M1,M2} 
    srcs::Vector{Int}
    sinks::Vector{Int}
    weights::Vector{Base.RefValue{UInt8}}
    fadjlist::Vector{Neighborhood}
164
165
    mixing_matrix::M1
    sampler_matrix::M2
Peter Jentsch's avatar
Peter Jentsch committed
166
    function WeightedGraph(demographics::AbstractVector,index_vectors,mixing_matrix::M1, sampler_matrix::M2) where {M1,M2}
Peter Jentsch's avatar
Peter Jentsch committed
167
168
169
170
171
172
        srcs = Vector{Int}()
        sinks = Vector{Int}()
        weights = Vector{Base.RefValue{UInt8}}()
        fadjlist = [Neighborhood() for _ in 1:length(demographics)]
        wg = new{M1,M2}(
            srcs,sinks,weights,fadjlist,
173
            mixing_matrix,
Peter Jentsch's avatar
Peter Jentsch committed
174
            sampler_matrix,
Peter Jentsch's avatar
Peter Jentsch committed
175
        )
Peter Jentsch's avatar
Peter Jentsch committed
176
177
        assemble_graph!(wg,index_vectors,mixing_matrix,sampler_matrix)
        return wg
Peter Jentsch's avatar
Peter Jentsch committed
178
    end
Peter Jentsch's avatar
Peter Jentsch committed
179
180
181
182
183
184
185
    function WeightedGraph(nv,sampler_matrix::M2) where M2
        srcs = Vector{Int}()
        sinks = Vector{Int}()
        weights = Vector{Base.RefValue{UInt8}}()
        fadjlist = [Neighborhood() for _ in 1:nv]
        wg = new{Nothing,M2}(
            srcs,sinks,weights,fadjlist,
186
            nothing,
Peter Jentsch's avatar
Peter Jentsch committed
187
            sampler_matrix,
Peter Jentsch's avatar
Peter Jentsch committed
188
        )
Peter Jentsch's avatar
Peter Jentsch committed
189
        return wg
Peter Jentsch's avatar
Peter Jentsch committed
190
191
    end
end
Peter Jentsch's avatar
Peter Jentsch committed
192
193
194
195
196
197
198
function empty!(wg::WeightedGraph)
    for nbhd in wg.fadjlist    
        empty!(nbhd)
    end
    empty!(wg.srcs)
    empty!(wg.sinks)
    empty!(wg.weights)
Peter Jentsch's avatar
Peter Jentsch committed
199
200
end

Peter Jentsch's avatar
Peter Jentsch committed
201
202
203
204
function empty!(nbhd::Neighborhood)
    empty!(nbhd.neighbourhood) #empty all the vector edgelists
    empty!(nbhd.weights)
end
Peter Jentsch's avatar
Peter Jentsch committed
205

Peter Jentsch's avatar
Peter Jentsch committed
206
function assemble_graph!(g,index_vectors,mixing_matrix,sampler_matrix)
Peter Jentsch's avatar
Peter Jentsch committed
207
208
    for i in 1:3, j in 1:i #diagonal
        if i != j 
Peter Jentsch's avatar
Peter Jentsch committed
209
            stubs_a,stubs_b = fast_chung_lu_bipartite(index_vectors[i],index_vectors[j],mixing_matrix[i,j],mixing_matrix[j,i])
Peter Jentsch's avatar
Peter Jentsch committed
210
        else #from one group to itself we need another algorithm
Peter Jentsch's avatar
Peter Jentsch committed
211
            stubs_a,stubs_b = fast_chung_lu(index_vectors[i],mixing_matrix[i,i])
Peter Jentsch's avatar
Peter Jentsch committed
212
        end   
Peter Jentsch's avatar
Peter Jentsch committed
213
214
215
        for (a,b) in zip(stubs_a,stubs_b)
            edge_weight_k::UInt8 = rand(Random.default_rng(Threads.threadid()),sampler_matrix[j,i])
            add_edge!(g,a,b,edge_weight_k)
Peter Jentsch's avatar
Peter Jentsch committed
216
217
        end
    end
218
end
Peter Jentsch's avatar
Peter Jentsch committed
219
220
221
222
223
224
225
226
227
228
229
function fast_chung_lu_bipartite(pop_i,pop_j,mixing_dist_ij,mixing_dist_ji)
    num_degrees_ij = similar(pop_i)
    num_degrees_ji =  similar(pop_j)
    generate_contact_vectors!(mixing_dist_ij,mixing_dist_ji,num_degrees_ij,num_degrees_ji)
    num_edges = sum(num_degrees_ij)
    stubs_i = Vector{Int}(undef,num_edges)
    stubs_j = similar(stubs_i)
    if num_edges>0
        sample!(Random.default_rng(Threads.threadid()),pop_i,Weights(num_degrees_ij./num_edges),stubs_i)
        sample!(Random.default_rng(Threads.threadid()),pop_j,Weights(num_degrees_ji./num_edges),stubs_j)
    end
Peter Jentsch's avatar
Peter Jentsch committed
230
    return stubs_i,stubs_j
Peter Jentsch's avatar
Peter Jentsch committed
231
end
Peter Jentsch's avatar
Peter Jentsch committed
232

Peter Jentsch's avatar
Peter Jentsch committed
233
234
235
236
237
238
239
240
241
242
function fast_chung_lu(pop_i,mixing_dist)
    num_degrees_ii = similar(pop_i)
    generate_contact_vectors!(mixing_dist,num_degrees_ii)
    m = sum(num_degrees_ii)
    num_edges= div(m,2)
    stubs_i = Vector{Int}(undef,num_edges)
    stubs_j = similar(stubs_i)
    if m>0
        sample!(Random.default_rng(Threads.threadid()),pop_i,Weights(num_degrees_ii./m),stubs_i)
        sample!(Random.default_rng(Threads.threadid()),pop_i,Weights(num_degrees_ii./m),stubs_j)
243
    end
Peter Jentsch's avatar
Peter Jentsch committed
244
    return stubs_i,stubs_j
Peter Jentsch's avatar
Peter Jentsch committed
245
246
end

Peter Jentsch's avatar
Peter Jentsch committed
247
neighbors(g::WeightedGraph,i) = g.fadjlist[i].neighbourhood
Peter Jentsch's avatar
Peter Jentsch committed
248
function Base.show(io::IO, g::WeightedGraph) 
Peter Jentsch's avatar
Peter Jentsch committed
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
    print(io, "WG $(ne(g)))")
end 



function add_edge!(g::WeightedGraph{T}, a,b,weight) where {T}
    verts = vertices(g)
    (a in verts && b in verts) || return false  # edge out of bounds
    a_nbhd = g.fadjlist[a]
    (b in a_nbhd.neighbourhood) && return false  # edge already in graph
    weight_boxed = Ref(weight)

    push!(a_nbhd.neighbourhood, b)
    push!(a_nbhd.weights,weight_boxed)
    push!(g.srcs,a)
    push!(g.sinks,b)
    push!(g.weights,weight_boxed)

    a == b && return true  # selfloop
    
    b_nbhd = g.fadjlist[b]
    push!(b_nbhd.neighbourhood, a)
    push!(b_nbhd.weights,weight_boxed)
    
    return true  # edge successfully added
end
is_directed(g::WeightedGraph) = false
ne(g::WeightedGraph) = length(g.srcs)
vertices(g::WeightedGraph) = 1:length(g.fadjlist)
neighbors_and_weights(g::WeightedGraph,i) = zip(g.fadjlist[i].neighbourhood,g.fadjlist[i].weights)