Commit 53679241 authored by Peter Jentsch's avatar Peter Jentsch
Browse files

initial implementation

parent 40d24b5f
......@@ -29,10 +29,10 @@ This function wires together a graph such that each household is in a complete s
Also returns a vector of AgentDemographic representing each agent, defined by the household compositions.
"""
@views function complete_graph_from_households_composition(households_composition)
@views function complete_graph_from_households_composition(households_composition, sampler_matrix)
total_household_pop = sum(sum.(households_composition))
population_list = Vector{AgentDemographic}(undef,total_household_pop)
network = SimpleGraph(total_household_pop)
wg = WeightedGraph(total_household_pop,sampler_matrix)
vertex_pointer = 1
for household in households_composition
num_vertices = sum(household)
......@@ -43,12 +43,13 @@ Also returns a vector of AgentDemographic representing each agent, defined by th
end
for v in vertex_pointer:(vertex_pointer + num_vertices - 1), w in vertex_pointer:(vertex_pointer + num_vertices - 1)
if v != w
add_edge!(network,v,w)
weight::UInt8 = rand(Random.default_rng(Threads.threadid()),sampler_matrix[Int(population_list[v]),Int(population_list[w])])
add_edge!(wg,v,w,weight)
end
end
vertex_pointer+=num_vertices
end
return network,population_list
return wg,population_list
end
......@@ -65,7 +66,7 @@ index_vectors: Vector of 3 vectors, each of which contains the indexes of all th
"""
function generate_population(num_households)
households_composition = sample_household_data(num_households)
household_networks,population_list = complete_graph_from_households_composition(households_composition)
household_networks,population_list = complete_graph_from_households_composition(households_composition, contact_time_distributions.hh)
index_vectors = [findall(x -> x == AgentDemographic(i), population_list) for i in 1:(AgentDemographic.size-1)]
return (;
population_list,
......
......@@ -31,14 +31,11 @@ end
Resample all the weights in `mixing_graph`
"""
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)
function sample_mixing_edges!(wg,sampler_matrix,demographics)
for k in 1:ne(wg)
i = wg.srcs[k]
j = wg.sinks[k]
wg.weights[k].x = rand(Random.default_rng(Threads.threadid()),sampler_matrix[Int(demographics[j]),Int(demographics[i])])
end
end
......@@ -125,17 +122,25 @@ Completely remake all the graphs in `time_dep_mixing_graph.resampled_graphs`.
"""
function remake_all!(t,time_dep_mixing_graph,index_vectors,demographics)
for wg in time_dep_mixing_graph.remade_graphs
remake!(wg,index_vectors)
empty!(wg)
assemble_graph!(wg,index_vectors,wg.mixing_matrix,wg.sampler_matrix)
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.sampler_matrix,demographics)
sample_mixing_edges!(wg,wg.sampler_matrix,demographics)
end
end
# display_degree(time_dep_mixing_graph.resampled_graphs[1])
end
struct Neighborhood
neighbourhood::Vector{Int}
weights::Vector{Base.RefValue{UInt8}}
function Neighborhood()
return new(Vector{Int}(),Vector{Base.RefValue{UInt8}}())
end
end
"""
Weighted graph type. Stores the graph in `g`, and the weights and edges in `mixing_edges`.
......@@ -159,50 +164,55 @@ Matrix of distributions determining node degrees
Matrix of distributions determining the edge weights
"""
struct WeightedGraph{G,M1,M2}
g::G
weights_dict::Dictionary{GraphEdge,UInt8}
struct WeightedGraph{M1,M2}
srcs::Vector{Int}
sinks::Vector{Int}
weights::Vector{Base.RefValue{UInt8}}
fadjlist::Vector{Neighborhood}
mixing_matrix::M1
sampler_matrix::M2
degrees_matrix::Union{Nothing,Matrix{Vector{Int}}}
function WeightedGraph(demographics::AbstractVector,index_vectors,mixing_matrix::M1, sampler_matrix::M2) where {M1,M2}
g = Graph(length(demographics))
weights_dict = Dictionary{GraphEdge,UInt8}()
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)
return new{typeof(g),M1,M2}(
g,
weights_dict,
srcs = Vector{Int}()
sinks = Vector{Int}()
weights = Vector{Base.RefValue{UInt8}}()
fadjlist = [Neighborhood() for _ in 1:length(demographics)]
wg = new{M1,M2}(
srcs,sinks,weights,fadjlist,
mixing_matrix,
sampler_matrix,
degrees_matrix
)
assemble_graph!(wg,index_vectors,mixing_matrix,sampler_matrix)
return wg
end
function WeightedGraph(g::G,demographics,sampler_matrix::M2) where {G<:SimpleGraph,M2}
weights_dict = Dictionary{GraphEdge,UInt8}(;sizehint = ne(g))
for e in edges(g)
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
return new{typeof(g),Nothing,M2}(
g,
weights_dict,
function WeightedGraph(nv,sampler_matrix::M2) where M2
srcs = Vector{Int}()
sinks = Vector{Int}()
weights = Vector{Base.RefValue{UInt8}}()
fadjlist = [Neighborhood() for _ in 1:nv]
wg = new{Nothing,M2}(
srcs,sinks,weights,fadjlist,
nothing,
sampler_matrix,
nothing
)
return wg
end
end
function remake!(wg::WeightedGraph,index_vectors)
empty!.(wg.g.fadjlist) #empty all the vector edgelists
wg.g.ne = 0
empty!(wg.weights_dict)
assemble_graph!(wg.g,wg.weights_dict,index_vectors,wg.mixing_matrix,wg.sampler_matrix,wg.degrees_matrix; resample_degrees = false)
function empty!(wg::WeightedGraph)
for nbhd in wg.fadjlist
empty!(nbhd)
end
empty!(wg.srcs)
empty!(wg.sinks)
empty!(wg.weights)
end
function empty!(nbhd::Neighborhood)
empty!(nbhd.neighbourhood) #empty all the vector edgelists
empty!(nbhd.weights)
end
function assemble_graph!(g,weights_dict,index_vectors,mixing_matrix,sampler_matrix,degree_matrix; resample_degrees=true)
for i in 1:3, j in 1:i #diagonal
......@@ -217,14 +227,11 @@ function assemble_graph!(g,weights_dict,index_vectors,mixing_matrix,sampler_matr
end
edges = fast_chung_lu(degree_matrix[i,j],index_vectors[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)
for (a,b) in zip(stubs_a,stubs_b)
edge_weight_k::UInt8 = rand(Random.default_rng(Threads.threadid()),sampler_matrix[j,i])
add_edge!(g,a,b,edge_weight_k)
end
end
for e in keys(weights_dict)
add_edge!(g,e.a,e.b)
end
end
......@@ -239,7 +246,7 @@ function fast_chung_lu_bipartite(degrees_ij,degrees_ji,index_vectors_i,index_vec
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)
end
return GraphEdge.(stubs_i,stubs_j)
return stubs_i,stubs_j
end
function fast_chung_lu(degrees_ii,index_vectors_i)
......@@ -251,12 +258,38 @@ function fast_chung_lu(degrees_ii,index_vectors_i)
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)
end
return GraphEdge.(stubs_i,stubs_j)
return stubs_i,stubs_j
end
neighbors(g::WeightedGraph,i) = LightGraphs.neighbors(g.g,i)
get_weight(g::WeightedGraph,e) = g.weights_dict[e]
neighbors(g::WeightedGraph,i) = g.fadjlist[i].neighbourhood
function Base.show(io::IO, g::WeightedGraph)
print(io, "WG $(ne(g.g))")
end
\ No newline at end of file
print(io, "WG $(ne(g)))")
end
function add_edge!(g::WeightedGraph{T}, a,b,weight) where {T}
verts = vertices(g)
(a in verts && b in verts) || return false # edge out of bounds
a_nbhd = g.fadjlist[a]
(b in a_nbhd.neighbourhood) && return false # edge already in graph
weight_boxed = Ref(weight)
push!(a_nbhd.neighbourhood, b)
push!(a_nbhd.weights,weight_boxed)
push!(g.srcs,a)
push!(g.sinks,b)
push!(g.weights,weight_boxed)
a == b && return true # selfloop
b_nbhd = g.fadjlist[b]
push!(b_nbhd.neighbourhood, a)
push!(b_nbhd.weights,weight_boxed)
return true # edge successfully added
end
is_directed(g::WeightedGraph) = false
ne(g::WeightedGraph) = length(g.srcs)
vertices(g::WeightedGraph) = 1:length(g.fadjlist)
neighbors_and_weights(g::WeightedGraph,i) = zip(g.fadjlist[i].neighbourhood,g.fadjlist[i].weights)
\ No newline at end of file
......@@ -80,7 +80,7 @@ mutable struct ModelSolution{T,InfNet,SocNet,WSMixingDist,RestMixingDist,Recorde
immunization_countdown::Vector{Int}
output_data::RecorderType
@views function ModelSolution(sim_length,params::T,num_households; record_degrees = false) where T
demographics,base_network,index_vectors = generate_population(num_households)
demographics,household_network,index_vectors = generate_population(num_households)
nodes = length(demographics)
pop_sizes = length.(index_vectors)
ws_mixing_tuple_preshift = deepcopy(workschool_mixing)
......@@ -96,7 +96,7 @@ mutable struct ModelSolution{T,InfNet,SocNet,WSMixingDist,RestMixingDist,Recorde
u_0_inf,u_0_vac = get_u_0(nodes,params.vaccinator_prob)
infected_mixing_graph,soc_mixing_graph = time_dep_mixing_graphs(sim_length,base_network,demographics,index_vectors,ws_matrix_tuple,rest_matrix_tuple)
infected_mixing_graph,soc_mixing_graph = time_dep_mixing_graphs(sim_length,household_network,demographics,index_vectors,ws_matrix_tuple,rest_matrix_tuple)
covid_alert_notifications = zeros(Bool,14,nodes) #two weeks worth of values
time_of_last_alert = fill(-1,nodes) #time of last alert is negative if no alert has been recieved
......
......@@ -144,8 +144,8 @@ function sample_initial_nodes(nodes,graphs,I_0_fraction)
weighted_degrees = zeros(nodes)
for v in 1:nodes
for g in graphs
for w in neighbors(g,v)
weighted_degrees[v] += get_weight(g,GraphEdge(v,w))
for (j,w) in neighbors_and_weights(g,v)
weighted_degrees[v] += w[]
end
end
end
......
......@@ -25,7 +25,7 @@ using KissABC
using CSV
using DataFrames
using StaticArrays
import LightGraphs.neighbors
import LightGraphs:neighbors,add_edge!
export intervalsmodel, hh, ws, rest, abm,multivariate_simulations
......
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