Commit 239cdbbb authored by Peter Jentsch's avatar Peter Jentsch
Browse files

simplify more

parent 078cd4e2
......@@ -31,13 +31,14 @@ end
Resample all the weights in `mixing_graph`
"""
function sample_mixing_edges!(weights_dict,contact_array,sample_cache,sampler_matrix)
for i in 1:3, j in 1:i #diagonal
rand!(Random.default_rng(Threads.threadid()),sampler_matrix[j,i],sample_cache[j,i])
for k in 1:length(contact_array[j,i])
edge_weight_k = sample_cache[j,i][k]
set!(weights_dict, contact_array[j,i][k], edge_weight_k)
end
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)
end
end
......@@ -116,43 +117,22 @@ end
"""
Completely remake all the graphs in `time_dep_mixing_graph.resampled_graphs`.
"""
function remake!(t,time_dep_mixing_graph,index_vectors)
function remake!(t,time_dep_mixing_graph,index_vectors,demographics)
for wg in time_dep_mixing_graph.remade_graphs
remake!(wg,index_vectors)
end
# display_degree(time_dep_mixing_graph.resampled_graphs[1])
for wg in time_dep_mixing_graph.resampled_graphs
if wg in time_dep_mixing_graph.graph_list[t]
sample_mixing_edges!(wg.weights_dict,wg.contact_array,wg.sample_cache,wg.sampler_matrix)
end
end
end
function display_weighted_degree(g)
weighted_degree = 0.0
for node in vertices(g.g)
for j in neighbors(g,node)
weighted_degree += get_weight(g,GraphEdge(node,j))
sample_mixing_edges!(wg.weights_dict,wg.sampler_matrix,demographics)
end
end
display(weighted_degree)
end
function display_degree(g)
weighted_degree = 0.0
for node in vertices(g.g)
for j in neighbors(g,node)
weighted_degree += 1
end
end
display(weighted_degree)
end
"""
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.
"""
function graph_from_mixing_edges(g,weights_dict)
for e in keys(weights_dict)
add_edge!(g,e.a,e.b)
end
# display_degree(time_dep_mixing_graph.resampled_graphs[1])
end
"""
Weighted graph type. Stores the graph in `g`, and the weights and edges in `mixing_edges`.
Fields
......@@ -185,48 +165,44 @@ Matrix of distributions determining the edge weights
"""
mutable struct WeightedGraph{G,M1,M2}
g::G
contact_array::Matrix{Vector{GraphEdge}}
sample_cache::Matrix{Vector{UInt8}}
weights_dict::Dictionary{GraphEdge,UInt8}
mixing_matrix::M1
sampler_matrix::M2
function WeightedGraph(demographics::AbstractVector,index_vectors,mixing_matrix::M1, weights_distribution_matrix::M2) where {M1,M2}
contact_array =[Vector{GraphEdge}() for i in 1:length(index_vectors),j in 1:length(index_vectors)]
function WeightedGraph(demographics::AbstractVector,index_vectors,mixing_matrix::M1, sampler_matrix::M2) where {M1,M2}
g = Graph(length(demographics))
total_edges = fast_chung_lu!(g,index_vectors,contact_array,mixing_matrix)
sample_cache = map(v-> Vector{UInt8}(undef,length(v)),contact_array)
weights_dict = Dictionary{GraphEdge,UInt8}(;sizehint = total_edges)
sample_mixing_edges!(weights_dict,contact_array,sample_cache,weights_distribution_matrix)
graph_from_mixing_edges(g,weights_dict)
weights_dict = Dictionary{GraphEdge,UInt8}()
for i in 1:3, j in 1:i #diagonal
if i != j
edges = fast_chung_lu_bipartite(index_vectors[i],index_vectors[j],mixing_matrix[i,j],mixing_matrix[j,i])
else #from one group to itself we need another algorithm
edges = fast_chung_lu(index_vectors[i],mixing_matrix[i,i])
end
for e in edges
edge_weight_k = rand(Random.default_rng(Threads.threadid()),sampler_matrix[j,i])
set!(weights_dict, e, edge_weight_k)
add_edge!(g,e.a,e.b)
end
end
return new{typeof(g),M1,M2}(
g,
contact_array,
sample_cache,
weights_dict,
mixing_matrix,
weights_distribution_matrix,
sampler_matrix,
)
end
function WeightedGraph(g::G,demographics,index_vectors,weights_distribution_matrix::M2) where {G<:SimpleGraph,M2}
contact_array = [sizehint!(Vector{GraphEdge}(),length(index_vectors[i])*2) for i in 1:length(index_vectors),j in 1:length(index_vectors)]
function WeightedGraph(g::G,demographics,index_vectors,sampler_matrix::M2) where {G<:SimpleGraph,M2}
weights_dict = Dictionary{GraphEdge,UInt8}(;sizehint = ne(g))
for e in edges(g)
i = src(e)
j = dst(e)
push!(contact_array[Int(demographics[i]),Int(demographics[j])], GraphEdge(i,j))
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)
end
sample_cache = map(v-> Vector{UInt8}(undef,length(v)),contact_array)
weights_dict = Dictionary{GraphEdge,UInt8}(;sizehint = ne(g))
sample_mixing_edges!(weights_dict,contact_array,sample_cache,weights_distribution_matrix)
return new{typeof(g),Nothing,M2}(
g,
contact_array,
sample_cache,
weights_dict,
nothing,
weights_distribution_matrix,
sampler_matrix,
)
end
end
......@@ -234,15 +210,19 @@ end
function remake!(wg::WeightedGraph,index_vectors)
empty!.(wg.g.fadjlist) #empty all the vector edgelists
wg.g.ne = 0
fast_chung_lu!(wg,index_vectors,wg.contact_array,wg.mixing_matrix)
for i in eachindex(wg.contact_array)
wg.sample_cache[i] = Vector{UInt8}(undef,length(wg.contact_array[i]))
end
empty!(wg.weights_dict)
sample_mixing_edges!(wg.weights_dict,wg.contact_array,wg.sample_cache,wg.sampler_matrix)
graph_from_mixing_edges(wg.g,wg.weights_dict)
for i in 1:3, j in 1:i #diagonal
if i != j
edges = fast_chung_lu_bipartite(index_vectors[i],index_vectors[j],wg.mixing_matrix[i,j],wg.mixing_matrix[j,i])
else #from one group to itself we need another algorithm
edges = fast_chung_lu(index_vectors[i],wg.mixing_matrix[i,i])
end
for e in edges
edge_weight_k = rand(Random.default_rng(Threads.threadid()),wg.sampler_matrix[j,i])
set!(wg.weights_dict, e, edge_weight_k)
add_edge!(wg. g,e.a,e.b)
end
end
end
......@@ -252,38 +232,29 @@ function Base.show(io::IO, g::WeightedGraph)
print(io, "WG $(ne(g.g))")
end
function fast_chung_lu!(wg,index_vectors,contact_array,mixing_matrix)
tot = 0
for i in 1:3, j in 1:i #diagonal
if i != j
num_degrees_ij = similar(index_vectors[i])
num_degrees_ji = similar(index_vectors[j])
generate_contact_vectors!(mixing_matrix[i,j],mixing_matrix[j,i],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()),index_vectors[i],Weights(num_degrees_ij./num_edges),stubs_i)
sample!(Random.default_rng(Threads.threadid()),index_vectors[j],Weights(num_degrees_ji./num_edges),stubs_j)
tot += num_edges
end
contact_array[j,i] = GraphEdge.(stubs_i,stubs_j)
else #from one group to itself we need another algorithm
num_degrees_ii = similar(index_vectors[i])
generate_contact_vectors!(mixing_matrix[i,j],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()),index_vectors[i],Weights(num_degrees_ii./m),stubs_i)
sample!(Random.default_rng(Threads.threadid()),index_vectors[i],Weights(num_degrees_ii./m),stubs_j)
tot += num_edges
end
contact_array[j,i] = GraphEdge.(stubs_i,stubs_j)
end
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
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)
end
return tot
return GraphEdge.(stubs_i,stubs_j)
end
\ No newline at end of file
......@@ -15,7 +15,7 @@ function get_parameters()#(0.0000,0.00048,0.0005,0.16,-1.30,-1.24,-0.8,0.35,0.35
π_base_o = -0.95,
η = 0.0,
κ = 0.0,
ω = 0.0057,
ω = 0.0056,
ω_en = 0.00,
Γ = 1/7,
ξ = 5.0,
......
......@@ -150,7 +150,7 @@ function solve!(modelsol,recordings...)
#this also resamples the soc network weights since they point to the same objects, but those are never used
if t>1
remake!(t,modelsol.inf_network,modelsol.index_vectors)
remake!(t,modelsol.inf_network,modelsol.index_vectors,modelsol.demographics)
end
if t>modelsol.params.infection_introduction_day
......
......@@ -39,3 +39,21 @@ function from_mean_workschool(_,μ)
end
function display_weighted_degree(g)
weighted_degree = 0.0
for node in vertices(g.g)
for j in neighbors(g,node)
weighted_degree += get_weight(g,GraphEdge(node,j))
end
end
display(weighted_degree)
end
function display_degree(g)
weighted_degree = 0.0
for node in vertices(g.g)
for j in neighbors(g,node)
weighted_degree += 1
end
end
display(weighted_degree)
end
\ No newline at end of file
......@@ -143,7 +143,7 @@ end
end
end
for i in eachindex(mixing_dist)
@test mean(mixing_dist[i]) expected_dist_mean atol = 0.01
@test mean(mixing_dist[i]) expected_dist_mean[i] atol = 0.05
@test mean(mixing_dist[i]) mean(dist[i]) atol = 0.2
end
end
......
......@@ -43,12 +43,10 @@ function generate_contact_vectors!(ij_dist,ji_dist,i_to_j_contacts::AbstractVect
index_list_j = MArray{Tuple{10,},Int,1,10}(undef)
sample_list_i = MArray{Tuple{10,},Int,1,10}(undef)
sample_list_j = MArray{Tuple{10,},Int,1,10}(undef)
while csum != 0
GC.@preserve index_list_i index_list_j begin
sample!(Random.default_rng(Threads.threadid()),1:l_i,index_list_i)
sample!(Random.default_rng(Threads.threadid()),1:l_j,index_list_j)
end
sample!(Random.default_rng(Threads.threadid()),1:l_i,index_list_i)
sample!(Random.default_rng(Threads.threadid()),1:l_j,index_list_j)
rand!(Random.default_rng(Threads.threadid()),ij_dist,sample_list_i)
rand!(Random.default_rng(Threads.threadid()),ji_dist,sample_list_j)
@inbounds for i = 1:inner_iter
......@@ -64,27 +62,24 @@ function alloc_test()
model = ModelSolution(1, get_parameters(), 5000)
@unpack demographics, index_vectors, rest_matrix_tuple,ws_matrix_tuple = model
dist = ws_matrix_tuple.daily
contact_array = [sizehint!(Vector{GraphEdge}(),length(index_vectors[i])*2) for i in 1:length(index_vectors),j in 1:length(index_vectors)]
fast_chung_lu!(index_vectors,contact_array,dist)
contact_array = [Vector{GraphEdge}() for i in 1:length(index_vectors),j in 1:length(index_vectors)]
@btime fast_chung_lu!($index_vectors,$contact_array,$dist)
return contact_array
end
function fast_chung_lu!(index_vectors,contact_array,mixing_matrix)
tot = 0
num_degrees_alloc_ij = Array{Int}(undef,maximum(length.(index_vectors)))
num_degrees_alloc_ji = similar(num_degrees_alloc_ij)
stubs_alloc_i = similar(num_degrees_alloc_ij)
stubs_alloc_j = similar(num_degrees_alloc_ij)
num_degrees_allocs = similar.(index_vectors)
for i in 1:3, j in 1:i #diagonal
if i != j
num_degrees_ij = view(num_degrees_alloc_ij,1:length(index_vectors[i]))
num_degrees_ji = view(num_degrees_alloc_ji,1:length(index_vectors[j]))
generate_contact_vectors!(mixing_matrix[i,j],mixing_matrix[j,i],num_degrees_ij,num_degrees_ji)
num_degrees_ij = view(num_degrees_allocs[i],1:length(index_vectors[i]))
num_degrees_ji = view(num_degrees_allocs[j],1:length(index_vectors[j]))
generate_contact_vectors!(mixing_matrix[i,j],mixing_matrix[j,i],(num_degrees_ij),num_degrees_ji)
num_edges = sum(num_degrees_ij)
stubs_i = view(stubs_alloc_i,1:num_edges)
stubs_j = view(stubs_alloc_i,1:num_edges)
stubs_i = Vector{Int}(undef,num_edges)
stubs_j = similar(stubs_i)
if num_edges>0
sample!(Random.default_rng(Threads.threadid()),index_vectors[i],Weights(num_degrees_ij./num_edges),stubs_i)
sample!(Random.default_rng(Threads.threadid()),index_vectors[j],Weights(num_degrees_ji./num_edges),stubs_j)
......@@ -96,4 +91,4 @@ function fast_chung_lu!(index_vectors,contact_array,mixing_matrix)
end
return tot
end
@btime alloc_test()
\ No newline at end of file
alloc_test()
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment