mixing_graphs.jl 8.4 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

"""
    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
39
40
41
function sample_mixing_edges!(weights_dict,sampler_matrix,demographics)
    indices = keys(weights_dict)
    for ind in indices
        e = ind
        i = e.a
        j = e.b
        weight = rand(Random.default_rng(Threads.threadid()),sampler_matrix[Int(demographics[j]),Int(demographics[i])])
        setindex!(weights_dict,weight,ind)
42
43
44
    end
end

Peter Jentsch's avatar
Peter Jentsch committed
45
46
47
48
49
"""
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
50
51
    remade_graphs::NTuple{N,G}

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

Peter Jentsch's avatar
docs    
Peter Jentsch committed
54
55
56
    resampled_graphs::NTuple{N,G}   

These are references to the graphs that get resampled everyday.
Peter Jentsch's avatar
Peter Jentsch committed
57
58
59
60


    graph_list::Vector{Vector{G}}

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

Peter Jentsch's avatar
Peter Jentsch committed
63
"""
Peter Jentsch's avatar
Peter Jentsch committed
64
65
66
struct TimeDepMixingGraph{G,T1,T2}
    remade_graphs::T1
    resampled_graphs::T2
67
    graph_list::Vector{Vector{G}}
Peter Jentsch's avatar
Peter Jentsch committed
68
69
70
    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,
71
72
73
74
75
            resampled_graphs,
            [copy(base_graph_list) for i in 1:len]
        )
    end
end
Peter Jentsch's avatar
Peter Jentsch committed
76
77
78
79
80
81

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

Assumes the simulation begins on Thursday arbitrarily.
"""
82
function time_dep_mixing_graphs(len,base_network,demographics,index_vectors,ws_matrix_tuple,rest_matrix_tuple)
83

Peter Jentsch's avatar
Peter Jentsch committed
84
    home_static_edges = WeightedGraph(base_network,demographics,contact_time_distributions.hh) #network with households and LTC homes
85

86
87
    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)
88
    ws_justonce_edges = WeightedGraph(demographics,index_vectors,ws_matrix_tuple.otherwise,contact_time_distributions.ws)
89

90
91
    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)
92
    rest_justonce_edges = WeightedGraph(demographics,index_vectors,rest_matrix_tuple.otherwise,contact_time_distributions.rest)
93
   
94
    inf_network_list = [home_static_edges,rest_static_edges,ws_justonce_edges,rest_justonce_edges] 
95
    soc_network_list = [home_static_edges,rest_static_edges,ws_static_edges]
Peter Jentsch's avatar
Peter Jentsch committed
96

97
    remade_graphs = (ws_justonce_edges,rest_justonce_edges)
Peter Jentsch's avatar
Peter Jentsch committed
98
99
100
101
    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)

102
103
104
105
106
    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
107
        if rand(Random.default_rng(Threads.threadid()))<5/7
108
109
110
111
112
            push!(l, ws_weekly_edges)
            push!(l, rest_weekly_edges)
        end
    end

Peter Jentsch's avatar
Peter Jentsch committed
113
    soc_mixing_graph = TimeDepMixingGraph(len,remade_graphs,resampled_graphs,soc_network_list)
114
115
116
    return infected_mixing_graph,soc_mixing_graph
end

Peter Jentsch's avatar
Peter Jentsch committed
117
118
119
"""
Completely remake all the graphs in `time_dep_mixing_graph.resampled_graphs`.
"""
Peter Jentsch's avatar
Peter Jentsch committed
120
function remake!(t,time_dep_mixing_graph,index_vectors,demographics)
121
122
123
    for wg in time_dep_mixing_graph.remade_graphs
        remake!(wg,index_vectors)
    end
Peter Jentsch's avatar
Peter Jentsch committed
124
    # display_degree(time_dep_mixing_graph.resampled_graphs[1])
125
126
    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
127
            sample_mixing_edges!(wg.weights_dict,wg.sampler_matrix,demographics)
128
129
        end
    end
Peter Jentsch's avatar
Peter Jentsch committed
130
    # display_degree(time_dep_mixing_graph.resampled_graphs[1])
Peter Jentsch's avatar
Peter Jentsch committed
131
end
Peter Jentsch's avatar
Peter Jentsch committed
132

Peter Jentsch's avatar
Peter Jentsch committed
133

Peter Jentsch's avatar
Peter Jentsch committed
134
"""
135
Weighted graph type. Stores the graph in `g`, and the weights and edges in `mixing_edges`. 
Peter Jentsch's avatar
docs    
Peter Jentsch committed
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
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
155
"""
Peter Jentsch's avatar
Peter Jentsch committed
156
struct WeightedGraph{G,M1,M2} 
Peter Jentsch's avatar
Peter Jentsch committed
157
    g::G
158
159
160
    weights_dict::Dictionary{GraphEdge,UInt8}
    mixing_matrix::M1
    sampler_matrix::M2
Peter Jentsch's avatar
Peter Jentsch committed
161
    function WeightedGraph(demographics::AbstractVector,index_vectors,mixing_matrix::M1, sampler_matrix::M2) where {M1,M2}
162
        g = Graph(length(demographics))
Peter Jentsch's avatar
Peter Jentsch committed
163
        weights_dict = Dictionary{GraphEdge,UInt8}()
Peter Jentsch's avatar
Peter Jentsch committed
164
        assemble_graph!(g,weights_dict,index_vectors,mixing_matrix,sampler_matrix)
165
        return new{typeof(g),M1,M2}(
Peter Jentsch's avatar
Peter Jentsch committed
166
            g,
167
168
            weights_dict,
            mixing_matrix,
Peter Jentsch's avatar
Peter Jentsch committed
169
            sampler_matrix,
Peter Jentsch's avatar
Peter Jentsch committed
170
171
        )
    end
Peter Jentsch's avatar
Peter Jentsch committed
172
    function WeightedGraph(g::G,demographics,sampler_matrix::M2) where {G<:SimpleGraph,M2}
Peter Jentsch's avatar
Peter Jentsch committed
173
        weights_dict = Dictionary{GraphEdge,UInt8}(;sizehint = ne(g))
174
        for e in edges(g)
Peter Jentsch's avatar
Peter Jentsch committed
175
176
177
178
            j = src(e)
            i = dst(e)
            weight = rand(Random.default_rng(Threads.threadid()),sampler_matrix[Int(demographics[j]),Int(demographics[i])])
            set!(weights_dict,GraphEdge(j,i),weight)
179
180
        end
        return new{typeof(g),Nothing,M2}(
Peter Jentsch's avatar
Peter Jentsch committed
181
            g,
182
183
            weights_dict,
            nothing,
Peter Jentsch's avatar
Peter Jentsch committed
184
            sampler_matrix,
Peter Jentsch's avatar
Peter Jentsch committed
185
        )
Peter Jentsch's avatar
Peter Jentsch committed
186
187
    end
end
188
189
190
191
192

function remake!(wg::WeightedGraph,index_vectors)
    empty!.(wg.g.fadjlist) #empty all the vector edgelists
    wg.g.ne = 0
    empty!(wg.weights_dict)
Peter Jentsch's avatar
Peter Jentsch committed
193
194
195
196
197
    assemble_graph!(wg.g,wg.weights_dict,index_vectors,wg.mixing_matrix,wg.sampler_matrix)
end


function assemble_graph!(g,weights_dict,index_vectors,mixing_matrix,sampler_matrix)
Peter Jentsch's avatar
Peter Jentsch committed
198
199
    for i in 1:3, j in 1:i #diagonal
        if i != j 
Peter Jentsch's avatar
Peter Jentsch committed
200
            edges = 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
201
        else #from one group to itself we need another algorithm
Peter Jentsch's avatar
Peter Jentsch committed
202
            edges = fast_chung_lu(index_vectors[i],mixing_matrix[i,i])
Peter Jentsch's avatar
Peter Jentsch committed
203
204
        end   
        for e in edges
Peter Jentsch's avatar
Peter Jentsch committed
205
206
            edge_weight_k = rand(Random.default_rng(Threads.threadid()),sampler_matrix[j,i])
            set!(weights_dict, e, edge_weight_k)
Peter Jentsch's avatar
Peter Jentsch committed
207
208
        end
    end
Peter Jentsch's avatar
Peter Jentsch committed
209
210
211
    for e in keys(weights_dict)
        add_edge!(g,e.a,e.b)
    end
212
end
Peter Jentsch's avatar
Peter Jentsch committed
213
214
215
216
217
218
219
220
221
222
223
224
225
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
    return GraphEdge.(stubs_i,stubs_j)
end
Peter Jentsch's avatar
Peter Jentsch committed
226

Peter Jentsch's avatar
Peter Jentsch committed
227
228
229
230
231
232
233
234
235
236
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)
237
    end
Peter Jentsch's avatar
Peter Jentsch committed
238
    return GraphEdge.(stubs_i,stubs_j)
Peter Jentsch's avatar
Peter Jentsch committed
239
240
241
242
243
244
245
246
247
end



neighbors(g::WeightedGraph,i) = neighbors(g.g,i)
get_weight(g::WeightedGraph,e) = g.weights_dict[e]
function Base.show(io::IO, g::WeightedGraph) 
    print(io, "WG $(ne(g.g))")
end