initial implementation

CovidAlertVaccinationModel/src/ABM/agents.jl

@@ -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,

CovidAlertVaccinationModel/src/ABM/mixing_graphs.jl

@@ -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

CovidAlertVaccinationModel/src/ABM/model_setup.jl

@@ -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

CovidAlertVaccinationModel/src/ABM/solve.jl

@@ -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

CovidAlertVaccinationModel/src/CovidAlertVaccinationModel.jl

@@ -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