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

84
    #weekly multiply durations by 1/5
85
    daily_shift(dur_dists) = shift_contact_distributions(dur_dists,1/5)
86
    #daily multiply by 1/2
87
    weekly_shift(dur_dists) = shift_contact_distributions(dur_dists,1/2)
88

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

91
    ws_static_edges = WeightedGraph(demographics,index_vectors,ws_matrix_tuple.daily,contact_time_distributions.ws)
92 93
    ws_weekly_edges = WeightedGraph(demographics,index_vectors,ws_matrix_tuple.twice_a_week,weekly_shift(contact_time_distributions.ws))
    ws_justonce_edges = WeightedGraph(demographics,index_vectors,ws_matrix_tuple.otherwise,daily_shift(contact_time_distributions.ws))
94

95
    rest_static_edges = WeightedGraph(demographics,index_vectors,rest_matrix_tuple.daily,contact_time_distributions.rest)
96 97
    rest_weekly_edges = WeightedGraph(demographics,index_vectors,rest_matrix_tuple.twice_a_week,weekly_shift(contact_time_distributions.rest))
    rest_justonce_edges = WeightedGraph(demographics,index_vectors,rest_matrix_tuple.otherwise,daily_shift(contact_time_distributions.rest))
98
   
99
 
Peter Jentsch's avatar
Peter Jentsch committed
100
    
101

102
    inf_network_list = [home_static_edges,rest_static_edges,ws_justonce_edges,rest_justonce_edges] 
103
    soc_network_list = [home_static_edges,rest_static_edges,ws_static_edges]
Peter Jentsch's avatar
Peter Jentsch committed
104

105
    remade_graphs = (ws_justonce_edges,rest_justonce_edges)
Peter Jentsch's avatar
Peter Jentsch committed
106 107 108 109
    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)

110 111 112 113 114
    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
115
        if rand(Random.default_rng(Threads.threadid()))<2/7
116 117 118 119 120
            push!(l, ws_weekly_edges)
            push!(l, rest_weekly_edges)
        end
    end

Peter Jentsch's avatar
Peter Jentsch committed
121
    soc_mixing_graph = TimeDepMixingGraph(len,remade_graphs,resampled_graphs,soc_network_list)
122 123 124
    return infected_mixing_graph,soc_mixing_graph
end

Peter Jentsch's avatar
Peter Jentsch committed
125 126 127
"""
Completely remake all the graphs in `time_dep_mixing_graph.resampled_graphs`.
"""
Peter Jentsch's avatar
Peter Jentsch committed
128
function remake_all!(t,time_dep_mixing_graph,index_vectors,demographics)
129 130 131
    for wg in time_dep_mixing_graph.remade_graphs
        remake!(wg,index_vectors)
    end
Peter Jentsch's avatar
Peter Jentsch committed
132
    # display_degree(time_dep_mixing_graph.resampled_graphs[1])
133 134
    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
135
            sample_mixing_edges!(wg.weights_dict,wg.sampler_matrix,demographics)
136 137
        end
    end
Peter Jentsch's avatar
Peter Jentsch committed
138
    # display_degree(time_dep_mixing_graph.resampled_graphs[1])
Peter Jentsch's avatar
Peter Jentsch committed
139
end
Peter Jentsch's avatar
Peter Jentsch committed
140

Peter Jentsch's avatar
Peter Jentsch committed
141

Peter Jentsch's avatar
Peter Jentsch committed
142
"""
143
Weighted graph type. Stores the graph in `g`, and the weights and edges in `mixing_edges`. 
Peter Jentsch's avatar
docs  
Peter Jentsch committed
144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162
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
163
"""
Peter Jentsch's avatar
Peter Jentsch committed
164
struct WeightedGraph{G,M1,M2} 
Peter Jentsch's avatar
Peter Jentsch committed
165
    g::G
166 167 168
    weights_dict::Dictionary{GraphEdge,UInt8}
    mixing_matrix::M1
    sampler_matrix::M2
169
    degrees_matrix::Union{Nothing,Matrix{Vector{Int}}}
Peter Jentsch's avatar
Peter Jentsch committed
170
    function WeightedGraph(demographics::AbstractVector,index_vectors,mixing_matrix::M1, sampler_matrix::M2) where {M1,M2}
171
        g = Graph(length(demographics))
Peter Jentsch's avatar
Peter Jentsch committed
172
        weights_dict = Dictionary{GraphEdge,UInt8}()
173 174
        degrees_matrix = [similar(index_vectors[i]) for i = 1:3, j = 1:3]
        assemble_graph!(g,weights_dict,index_vectors,mixing_matrix,sampler_matrix,degrees_matrix)
175
        return new{typeof(g),M1,M2}(
Peter Jentsch's avatar
Peter Jentsch committed
176
            g,
177 178
            weights_dict,
            mixing_matrix,
Peter Jentsch's avatar
Peter Jentsch committed
179
            sampler_matrix,
180
            degrees_matrix
Peter Jentsch's avatar
Peter Jentsch committed
181 182
        )
    end
Peter Jentsch's avatar
Peter Jentsch committed
183
    function WeightedGraph(g::G,demographics,sampler_matrix::M2) where {G<:SimpleGraph,M2}
Peter Jentsch's avatar
Peter Jentsch committed
184
        weights_dict = Dictionary{GraphEdge,UInt8}(;sizehint = ne(g))
185
        for e in edges(g)
Peter Jentsch's avatar
Peter Jentsch committed
186 187 188 189
            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)
190 191
        end
        return new{typeof(g),Nothing,M2}(
Peter Jentsch's avatar
Peter Jentsch committed
192
            g,
193 194
            weights_dict,
            nothing,
Peter Jentsch's avatar
Peter Jentsch committed
195
            sampler_matrix,
196
            nothing
Peter Jentsch's avatar
Peter Jentsch committed
197
        )
Peter Jentsch's avatar
Peter Jentsch committed
198 199
    end
end
200 201 202 203 204

function remake!(wg::WeightedGraph,index_vectors)
    empty!.(wg.g.fadjlist) #empty all the vector edgelists
    wg.g.ne = 0
    empty!(wg.weights_dict)
205
    assemble_graph!(wg.g,wg.weights_dict,index_vectors,wg.mixing_matrix,wg.sampler_matrix,wg.degrees_matrix; resample_degrees = false)
Peter Jentsch's avatar
Peter Jentsch committed
206 207 208
end


209
function assemble_graph!(g,weights_dict,index_vectors,mixing_matrix,sampler_matrix,degree_matrix; resample_degrees=true)
Peter Jentsch's avatar
Peter Jentsch committed
210 211
    for i in 1:3, j in 1:i #diagonal
        if i != j 
212 213 214 215
                if resample_degrees
                    generate_contact_vectors!(mixing_matrix[i,j],mixing_matrix[j,i],degree_matrix[i,j],degree_matrix[j,i])
                end
                edges = fast_chung_lu_bipartite(degree_matrix[i,j],degree_matrix[j,i],index_vectors[i],index_vectors[j])
Peter Jentsch's avatar
Peter Jentsch committed
216
        else #from one group to itself we need another algorithm
217 218 219 220
                if resample_degrees
                    generate_contact_vectors!(mixing_matrix[i,j],degree_matrix[i,j])
                end
                edges = fast_chung_lu(degree_matrix[i,j],index_vectors[i])       
Peter Jentsch's avatar
Peter Jentsch committed
221 222
        end   
        for e in edges
Peter Jentsch's avatar
Peter Jentsch committed
223 224
            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
225 226
        end
    end
Peter Jentsch's avatar
Peter Jentsch committed
227 228 229
    for e in keys(weights_dict)
        add_edge!(g,e.a,e.b)
    end
230
end
231 232 233 234 235

function fast_chung_lu_bipartite(degrees_ij,degrees_ji,index_vectors_i,index_vectors_j)
    m = sum(degrees_ij)
    @assert m == sum(degrees_ji)
    stubs_i = Vector{Int}(undef,m)
Peter Jentsch's avatar
Peter Jentsch committed
236
    stubs_j = similar(stubs_i)
237 238 239
    if m>0
        sample!(Random.default_rng(Threads.threadid()),index_vectors_i,Weights(degrees_ij./m),stubs_i)
        sample!(Random.default_rng(Threads.threadid()),index_vectors_j,Weights(degrees_ji./m),stubs_j)
Peter Jentsch's avatar
Peter Jentsch committed
240 241 242
    end
    return GraphEdge.(stubs_i,stubs_j)
end
Peter Jentsch's avatar
Peter Jentsch committed
243

244 245
function fast_chung_lu(degrees_ii,index_vectors_i)
    m = sum(degrees_ii)
Peter Jentsch's avatar
Peter Jentsch committed
246 247 248 249
    num_edges= div(m,2)
    stubs_i = Vector{Int}(undef,num_edges)
    stubs_j = similar(stubs_i)
    if m>0
250 251
        sample!(Random.default_rng(Threads.threadid()),index_vectors_i,Weights(degrees_ii./m),stubs_i)
        sample!(Random.default_rng(Threads.threadid()),index_vectors_i,Weights(degrees_ii./m),stubs_j)
252
    end
Peter Jentsch's avatar
Peter Jentsch committed
253
    return GraphEdge.(stubs_i,stubs_j)
Peter Jentsch's avatar
Peter Jentsch committed
254 255
end

256

257
neighbors(g::WeightedGraph,i) = LightGraphs.neighbors(g.g,i)
Peter Jentsch's avatar
Peter Jentsch committed
258 259 260 261
get_weight(g::WeightedGraph,e) = g.weights_dict[e]
function Base.show(io::IO, g::WeightedGraph) 
    print(io, "WG $(ne(g.g))")
end