mixing_graphs.jl 10.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
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
    for i in 1:size(mixing_edges.contact_array)[1], j in 1:i  #diagonal     
Peter Jentsch's avatar
Peter Jentsch committed
37
        rand!(Random.default_rng(Threads.threadid()), mixing_edges.sampler_matrix[j,i],mixing_edges.sample_cache[j,i])
38
39
        for k in 1:length(mixing_edges.contact_array[j,i])
            edge_weight_k = mixing_edges.sample_cache[j,i][k]
40
            
41
            set!(mixing_edges.weights_dict, mixing_edges.contact_array[j,i][k], edge_weight_k)
42
        end
43
44
45
    end
end

Peter Jentsch's avatar
Peter Jentsch committed
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
"""
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!`.
"""
71
struct MixingEdges{M}
Peter Jentsch's avatar
Peter Jentsch committed
72
    total_edges::Int
73
    contact_array::Matrix{Vector{GraphEdge}}
74
    sample_cache::Matrix{Vector{UInt8}}
75
    weights_dict::Dictionary{GraphEdge,UInt8}
76
    sampler_matrix::M
77
    function MixingEdges(total_edges,contact_array,sampler_matrix::Matrix{M}) where M<:Sampleable{Univariate, Discrete}
78
        sample_cache = map(v-> Vector{UInt8}(undef,length(v)),contact_array)
79
        weights_dict = Dictionary{GraphEdge,UInt8}(;sizehint = total_edges)
80
81
82
        new{typeof(sampler_matrix)}(total_edges,contact_array,sample_cache,weights_dict,sampler_matrix)
    end
end
Peter Jentsch's avatar
Peter Jentsch committed
83

Peter Jentsch's avatar
Peter Jentsch committed
84
85
86
"""
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. 
"""
87
function create_mixing_edges(demographic_index_vectors,mixing_matrix,weights_distribution_matrix)
88
    contact_array =[Vector{GraphEdge}() for i in 1:length(demographic_index_vectors),j in 1:length(demographic_index_vectors)]
89
90
    tot = 0
    for i in 1:size(mixing_matrix)[1], j in 1:i  #diagonal
91
92
        num_degrees_ij = similar(demographic_index_vectors[i])
        num_degrees_ji = similar(demographic_index_vectors[j])
93
94
95
96
97
        generate_contact_vectors!(mixing_matrix[i,j],mixing_matrix[j,i],num_degrees_ij,num_degrees_ji)
        m = sum(num_degrees_ij)
        stubs_i = Vector{Int}(undef,m)
        stubs_j = Vector{Int}(undef,m)
        if m>0
Peter Jentsch's avatar
Peter Jentsch committed
98
99
            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)
100
101
            tot += m
        end    
102
103

        contact_array[j,i] = GraphEdge.(stubs_i,stubs_j)
104

Peter Jentsch's avatar
Peter Jentsch committed
105
    end
106
107
108
    return MixingEdges(tot,contact_array,weights_distribution_matrix)
end

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

"""
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`.
"""
113
function create_mixing_edges(g::SimpleGraph,demographics,demographic_index_vectors,weights_distribution_matrix) 
114
    contact_array = [sizehint!(Vector{GraphEdge}(),length(demographic_index_vectors[i])*2) for i in 1:length(demographic_index_vectors),j in 1:length(demographic_index_vectors)]
115
    for e in edges(g)
116
117
        i = src(e)
        j = dst(e)
118
        push!(contact_array[Int(demographics[i]),Int(demographics[j])], GraphEdge(i,j))
119
    end
120
    return MixingEdges(ne(g),contact_array,weights_distribution_matrix)
121
end
122

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

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}}

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

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

Assumes the simulation begins on Thursday arbitrarily.
"""
154
function time_dep_mixing_graphs(len,base_network,demographics,index_vectors,ws_matrix_tuple,rest_matrix_tuple)
Peter Jentsch's avatar
Peter Jentsch committed
155
    home_static_edges = WeightedGraph(base_network,demographics,index_vectors,contact_time_distributions.hh) #network with households and LTC homes
156
   
157
158
159
    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)
160
    
161
162
163
    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)
164
   
165
166
167
    inf_network_list = [home_static_edges,rest_static_edges] 
    soc_network_list = [home_static_edges,rest_static_edges,ws_static_edges]
    
168
169
    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)
170
171
172
173
174
175
    # 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
Peter Jentsch's avatar
Peter Jentsch committed
176
        if rand(Random.default_rng(Threads.threadid()))<5/7
177
178
179
180
181
182
183
184
185
186
            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
187
188
189
"""
Completely remake all the graphs in `time_dep_mixing_graph.resampled_graphs`.
"""
190
191
function remake!(time_dep_mixing_graph,demographic_index_vectors)
    for (weighted_graph,mixing_matrix) in time_dep_mixing_graph.resampled_graphs
192
193
        empty!.(weighted_graph.g.fadjlist) #empty all the vector edgelists
        weighted_graph.g.ne = 0
194
        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
195
        graph_from_mixing_edges(weighted_graph.g,weighted_graph.mixing_edges)
196
197
    end
end
Peter Jentsch's avatar
Peter Jentsch committed
198

Peter Jentsch's avatar
Peter Jentsch committed
199
200
201
"""
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
202
203
function graph_from_mixing_edges(g,mixing_edges)
    for i in 1:size(mixing_edges.contact_array)[1], j in 1:i  #diagonal
204
205
        for k in 1:length(mixing_edges.contact_array[i,j])
            e = mixing_edges.contact_array[i,j][k]
206
207
208
209
            success = add_edge!(g,e.a,e.b)
            # if !success
            #     display((has_edge(g,e.a,e.b)))
            # end
Peter Jentsch's avatar
Peter Jentsch committed
210
211
212
        end
    end
end
Peter Jentsch's avatar
Peter Jentsch committed
213
214
215
216
217

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

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