Commit 34a8709f authored by Peter Jentsch's avatar Peter Jentsch
Browse files

tweak tests

parent df40693d
......@@ -32,7 +32,7 @@ Fills i_to_j_contacts and j_to_i_contacts with degrees sampled from ij_dist and
Given `μz_i = mean(ij_dist)` and `μ_j = mean(ji_dist)`, these must satisfy `μ_i* length(i_to_j_contacts) == μ_j* length(j_to_i_contacts)`
"""
function generate_contact_vectors!(ij_dist,ji_dist,i_to_j_contacts::Vector{T}, j_to_i_contacts::Vector{T}) where T
function generate_contact_vectors!(ij_dist,ji_dist,i_to_j_contacts::AbstractVector{T}, j_to_i_contacts::AbstractVector{T}) where T
rand!(Random.default_rng(Threads.threadid()),ij_dist,i_to_j_contacts)
rand!(Random.default_rng(Threads.threadid()),ji_dist,j_to_i_contacts)
......
......@@ -67,52 +67,51 @@ Stores the weights used in the graph, so they can be easily resampled.
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!`.
"""
struct MixingEdges{M,M2}
total_edges::Int
contact_array::Matrix{Vector{GraphEdge}}
sample_cache::Matrix{Vector{UInt8}}
weights_dict::Dictionary{GraphEdge,UInt8}
sampler_matrix::M
mixing_matrix::Union{Nothing,M2}
function MixingEdges(total_edges,contact_array,sampler_matrix::M,mixing_matrix::M2) where {M,M2}
function MixingEdges(total_edges, contact_array,sampler_matrix::M,mixing_matrix::M2) where {M,M2}
sample_cache = map(v-> Vector{UInt8}(undef,length(v)),contact_array)
weights_dict = Dictionary{GraphEdge,UInt8}(;sizehint = total_edges)
me = new{M,M2}(total_edges,contact_array,sample_cache,weights_dict,sampler_matrix,mixing_matrix)
me = new{M,M2}(contact_array,sample_cache,weights_dict,sampler_matrix,mixing_matrix)
sample_mixing_edges!(me)
return me
end
end
using StrideArrays
"""
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.
"""
function create_mixing_edges(demographic_index_vectors,mixing_matrix,weights_distribution_matrix)
contact_array =[Vector{GraphEdge}() for i in 1:length(demographic_index_vectors),j in 1:length(demographic_index_vectors)]
function create_mixing_edges(index_vectors,mixing_matrix,weights_distribution_matrix)
contact_array =[Vector{GraphEdge}() for i in 1:length(index_vectors),j in 1:length(index_vectors)]
tot = 0
for i in 1:3, j in 1:i #diagonal
if i != j
num_degrees_ij = similar(demographic_index_vectors[i])
num_degrees_ji = similar(demographic_index_vectors[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 = Vector{Int}(undef,num_edges)
stubs_j = similar(stubs_i)
if num_edges>0
sample!(Random.default_rng(Threads.threadid()),demographic_index_vectors[i],Weights(num_degrees_ij./num_edges),stubs_i)
sample!(Random.default_rng(Threads.threadid()),demographic_index_vectors[j],Weights(num_degrees_ji./num_edges),stubs_j)
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(demographic_index_vectors[i])
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 = Vector{Int}(undef,num_edges)
stubs_j = similar(stubs_i)
if m>0
sample!(Random.default_rng(Threads.threadid()),demographic_index_vectors[i],Weights(num_degrees_ii./m),stubs_i)
sample!(Random.default_rng(Threads.threadid()),demographic_index_vectors[i],Weights(num_degrees_ii./m),stubs_j)
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)
......@@ -125,8 +124,8 @@ end
"""
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`.
"""
function create_mixing_edges(g::SimpleGraph,demographics,demographic_index_vectors,weights_distribution_matrix)
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)]
function create_mixing_edges(g::SimpleGraph,demographics,index_vectors,weights_distribution_matrix)
contact_array = [sizehint!(Vector{GraphEdge}(),length(index_vectors[i])*2) for i in 1:length(index_vectors),j in 1:length(index_vectors)]
for e in edges(g)
i = src(e)
j = dst(e)
......@@ -206,12 +205,12 @@ end
"""
Completely remake all the graphs in `time_dep_mixing_graph.resampled_graphs`.
"""
function remake!(t,time_dep_mixing_graph,demographic_index_vectors)
function remake!(t,time_dep_mixing_graph,index_vectors)
for weighted_graph in time_dep_mixing_graph.remade_graphs
empty!.(weighted_graph.g.fadjlist) #empty all the vector edgelists
weighted_graph.g.ne = 0
weighted_graph.mixing_edges = create_mixing_edges(demographic_index_vectors,weighted_graph.mixing_edges.mixing_matrix,weighted_graph.mixing_edges.sampler_matrix)
weighted_graph.mixing_edges = create_mixing_edges(index_vectors,weighted_graph.mixing_edges.mixing_matrix,weighted_graph.mixing_edges.sampler_matrix)
graph_from_mixing_edges(weighted_graph.g,weighted_graph.mixing_edges)
end
......@@ -243,11 +242,8 @@ 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,mixing_edges)
for i in 1:3, j in 1:i #diagonal
for k in eachindex(mixing_edges.contact_array[j,i])
e = mixing_edges.contact_array[j,i][k]
add_edge!(g,e.a,e.b)
end
for e in keys(mixing_edges.weights_dict)
add_edge!(g,e.a,e.b)
end
end
......@@ -258,8 +254,8 @@ My own weighted graph type. Stores the graph in `g`, and the weights and edges i
mutable struct WeightedGraph{G,M}
g::G
mixing_edges::M
function WeightedGraph(demographics,demographic_index_vectors,mixing_matrix,weights_distribution_matrix)
mixing_edges = create_mixing_edges(demographic_index_vectors,mixing_matrix,weights_distribution_matrix)
function WeightedGraph(demographics,index_vectors,mixing_matrix,weights_distribution_matrix)
mixing_edges = create_mixing_edges(index_vectors,mixing_matrix,weights_distribution_matrix)
g = Graph(length(demographics))
graph_from_mixing_edges(g,mixing_edges)
return new{typeof(g),typeof(mixing_edges)}(
......@@ -267,8 +263,8 @@ mutable struct WeightedGraph{G,M}
mixing_edges
)
end
function WeightedGraph(g::SimpleGraph,demographics,demographic_index_vectors,weights_distribution_matrix)
mixing_edges = create_mixing_edges(g, demographics, demographic_index_vectors,weights_distribution_matrix)
function WeightedGraph(g::SimpleGraph,demographics,index_vectors,weights_distribution_matrix)
mixing_edges = create_mixing_edges(g, demographics, index_vectors,weights_distribution_matrix)
return new{typeof(g),typeof(mixing_edges)}(
g,
mixing_edges,
......
......@@ -3,7 +3,7 @@ function get_parameters()#(0.0000,0.00048,0.0005,0.16,-1.30,-1.24,-0.8,0.35,0.35
sim_length = 500,
num_households = 5000,
I_0_fraction = 0.003,
β_y = 0.0011,
β_y = 0.001105,
β_m = 0.00061,
β_o = 0.04,
α_y = 0.4,
......
......@@ -23,7 +23,7 @@ using Random
println("obs final size: $final_size, target: $target_final_size")
println("obs preinf vac: $total_preinf_vaccination, target: $target_preinf_vac")
println("obs postinf vac: $total_postinf_vaccination,target: $target_postinf_vac")
@test all(abs.(final_size .- target_final_size) .< [40,100,25])
@test all(abs.(final_size .- target_final_size) .< [50,100,25])
@test all(abs.(total_preinf_vaccination .- target_preinf_vac) .< [50,100,25])
@test all(abs.(total_postinf_vaccination .- target_postinf_vac) .< [100,100,25])
end
......
......@@ -7,5 +7,5 @@ using BenchmarkTools
import StatsBase.mean
include("ABM/mixing_test.jl")
# include("ABM/output_test.jl")
include("ABM/output_test.jl")
include("IntervalsModel/intervals_model_test.jl")
\ No newline at end of file
No preview for this file type
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