mixing_graphs.jl 7.99 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
using Dictionaries
struct GraphEdge
    a::Int
    b::Int
end

function Base.hash(e::GraphEdge)
    # display(e)
    return hash(minmax(e.a,e.b))
end
function Base.isequal(e1::GraphEdge,e2::GraphEdge)
    return isequal(minmax(e1.a,e1.b),minmax(e2.a,e2.b))
end
function sample_mixing_graph!(mixing_graph,population_demographics)
    mixing_edges = mixing_graph.mixing_edges
    for i in 1:size(mixing_edges.contact_array)[1], j in 1:i  #diagonal
            demo_i = Int(population_demographics[i])
            demo_j = Int(population_demographics[j])
            rand!(RNG, mixing_edges.weights_distribution_matrix[demo_i,demo_j],mixing_edges.sample_cache[i,j])
            for k in 1:length(mixing_edges.contact_array[i,j][1])
                kth_node_i = mixing_edges.contact_array[i,j][1][k]
                kth_node_j = mixing_edges.contact_array[i,j][2][k]
                edge_weight_k = mixing_edges.sample_cache[i,j][k]
                set!(mixing_edges.weights_dict,GraphEdge(kth_node_i,kth_node_j), edge_weight_k)
            end
    end
end

Peter Jentsch's avatar
Peter Jentsch committed
29
#A type that defines a matrix of vectors, such that the ijth vector contains the nodes to connect to the jith vector, ala. Chung-Lu bipartite
30
struct MixingEdges{M}
Peter Jentsch's avatar
Peter Jentsch committed
31
    total_edges::Int
Peter Jentsch's avatar
Peter Jentsch committed
32
    contact_array::Matrix{Tuple{Vector{Int},Vector{Int}}}
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
    sample_cache::Matrix{Vector{Int}}
    weights_dict::Dictionary{GraphEdge,UInt8}
    weights_distribution_matrix::M
    function MixingEdges(total_edges,contact_array,sampler_matrix::Matrix{M}) where M<:Sampleable{Univariate, Discrete}
        sample_cache = map(v-> Vector{Int}(undef,length(v[1])),contact_array)
        weights_dict = Dictionary{GraphEdge,UInt8}()
        new{typeof(sampler_matrix)}(total_edges,contact_array,sample_cache,weights_dict,sampler_matrix)
    end
    function MixingEdges(total_edges,contact_array,weights_distribution_matrix::Matrix{M}) where M<:UnivariateDistribution
        sample_cache = map(v-> Vector{Int}(undef,length(v[1])),contact_array)
        weights_dict = Dictionary{GraphEdge,UInt8}()
        sampler_matrix =  map(m -> Distributions.PoissonCountSampler.(m),weights_distribution_matrix)
        new{typeof(sampler_matrix)}(total_edges,contact_array,sample_cache,weights_dict,sampler_matrix)
    end
    
end
Peter Jentsch's avatar
Peter Jentsch committed
49

50
51
52
53
54
55
56
function create_mixing_edges(demographic_index_vectors,mixing_matrix,weights_distribution_matrix)
    contact_array =[(Vector{Int}(),Vector{Int}()) for i in 1:length(demographic_index_vectors),j in 1:length(demographic_index_vectors)]
    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
57

58
59
60
61
62
63
64
65
66
67
        m = sum(num_degrees_ij)
        stubs_i = Vector{Int}(undef,m)
        stubs_j = Vector{Int}(undef,m)
        if m>0
            # @assert m == sum(num_degrees_ji)
            sample!(RNG,demographic_index_vectors[i],Weights(num_degrees_ij./m),stubs_i)
            sample!(RNG,demographic_index_vectors[j],Weights(num_degrees_ji./m),stubs_j)
            tot += m
        end    
        contact_array[j,i] = (stubs_i,stubs_j)
Peter Jentsch's avatar
Peter Jentsch committed
68
    end
69
70
71
72
73
74
75
76
77
78
    return MixingEdges(tot,contact_array,weights_distribution_matrix)
end

function create_mixing_edges(g::SimpleGraph,demographics,demographic_index_vectors,weights_distribution_matrix) 
    contact_array = [(Vector{Int}(),Vector{Int}()) for i in 1:length(demographic_index_vectors),j in 1:length(demographic_index_vectors)]
    for e in edges(g)    
        i = src(e)
        j = dst(e)
        push!(contact_array[Int(demographics[j]),Int(demographics[i])][1],i)
        push!(contact_array[Int(demographics[j]),Int(demographics[i])][2],j)
79
    end
80
81
    sample_cache = map(v-> Vector{Int}(undef,length(v[1])),contact_array)
    return MixingEdges(ne(g),contact_array,weights_distribution_matrix)
82
end
83

84
85
86
87
88
89
90
91
92
93
94
struct TimeDepMixingGraph{N,G}
    resampled_graphs::NTuple{N,G}
    graph_list::Vector{Vector{G}}
    function TimeDepMixingGraph(len,resampled_graphs::NTuple{N,G},base_graph_list::Vector{G}) where {G,N}
        return new{N,G}(
            resampled_graphs,
            [copy(base_graph_list) for i in 1:len]
        )
    end
end
function time_dep_mixing_graphs(len,base_network,demographics,index_vectors,ws_matrix_tuple,rest_matrix_tuple)
Peter Jentsch's avatar
Peter Jentsch committed
95
    home_static_edges = WeightedGraph(base_network,demographics,index_vectors,contact_time_distributions.hh) #network with households and LTC homes
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
    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)
    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)
    inf_network_list = [home_static_edges,rest_static_edges] 
    soc_network_list = [home_static_edges,rest_static_edges,ws_static_edges]
    
    infected_mixing_graph = TimeDepMixingGraph(len,(ws_daily_edges,rest_daily_edges),inf_network_list)
    soc_mixing_graph = TimeDepMixingGraph(len,(ws_daily_edges,rest_daily_edges),soc_network_list)
    # 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
        if rand(RNG)<5/7
            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

function remake!(time_dep_mixing_graph,demographic_index_vectors,mixing_matrix)
    for weighted_graph in time_dep_mixing_graph.resampled_graphs
        empty!.(weighted_graph.g.fadjlist) #empty all the vector edgelists
        weighted_graph.g.ne = 0
128
        weighted_graph.mixing_edges = create_mixing_edges(demographic_index_vectors,mixing_matrix,weighted_graph.mixing_edges.weights_distribution_matrix)
Peter Jentsch's avatar
Peter Jentsch committed
129
        graph_from_mixing_edges(weighted_graph.g,weighted_graph.mixing_edges)
130
131
    end
end
Peter Jentsch's avatar
Peter Jentsch committed
132
133
134
135
136
137
138
139

function graph_from_mixing_edges(g,mixing_edges)
    for i in 1:size(mixing_edges.contact_array)[1], j in 1:i  #diagonal
        for k in 1:length(mixing_edges.contact_array[i,j][1])
            add_edge!(g,mixing_edges.contact_array[i,j][1][k],mixing_edges.contact_array[i,j][2][k])
        end
    end
end
140
#defining my own weighted graph type cause we need to be able to resample the edge weights in a very particular way
141
mutable struct WeightedGraph{G,M} 
Peter Jentsch's avatar
Peter Jentsch committed
142
    g::G
143
    mixing_edges::M
144
    function WeightedGraph(demographics,demographic_index_vectors,mixing_matrix,weights_distribution_matrix)
145
        mixing_edges = create_mixing_edges(demographic_index_vectors,mixing_matrix,weights_distribution_matrix)
146
        g = Graph(length(demographics))
Peter Jentsch's avatar
Peter Jentsch committed
147
        graph_from_mixing_edges(g,mixing_edges)
148
        return new{typeof(g),typeof(mixing_edges)}(
Peter Jentsch's avatar
Peter Jentsch committed
149
            g,
150
            mixing_edges
Peter Jentsch's avatar
Peter Jentsch committed
151
152
        )
    end
Peter Jentsch's avatar
Peter Jentsch committed
153
    function WeightedGraph(g::SimpleGraph,demographics,demographic_index_vectors,weights_distribution_matrix)
154
155
        mixing_edges = create_mixing_edges(g, demographics, demographic_index_vectors,weights_distribution_matrix)
        weights_dict = Dict{GraphEdge,UInt8}()
Peter Jentsch's avatar
Peter Jentsch committed
156
        sample_cache = map(v-> Vector{Int}(undef,length(v[1])),mixing_edges.contact_array)
157
        return new{typeof(g),typeof(mixing_edges)}(
Peter Jentsch's avatar
Peter Jentsch committed
158
            g,
Peter Jentsch's avatar
Peter Jentsch committed
159
            mixing_edges,
Peter Jentsch's avatar
Peter Jentsch committed
160
        )
Peter Jentsch's avatar
Peter Jentsch committed
161
162
    end
end
163
164

get_weight(g::WeightedGraph,e) = g.mixing_edges.weights_dict[e]
165
function Base.show(io::IO, g::WeightedGraph) 
Peter Jentsch's avatar
Peter Jentsch committed
166
    print(io, "WG $(ne(g.g))")
167
end