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
 
@@ -79,9 +76,7 @@ Creates the `TimeDepMixingGraph` for our specific model.
 
 Assumes the simulation begins on Thursday arbitrarily.
 """
-function time_dep_mixing_graphs(len,base_network,demographics,index_vectors,ws_matrix_tuple,rest_matrix_tuple)
-
-    home_static_edges = WeightedGraph(base_network,demographics,contact_time_distributions.hh) #network with households and LTC homes
+function time_dep_mixing_graphs(len,home_static_edges,demographics,index_vectors,ws_matrix_tuple,rest_matrix_tuple)
 
     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)
@@ -119,17 +114,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`. 
@@ -153,62 +156,65 @@ 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
     function WeightedGraph(demographics::AbstractVector,index_vectors,mixing_matrix::M1, sampler_matrix::M2) where {M1,M2}
-        g = Graph(length(demographics))
-        weights_dict = Dictionary{GraphEdge,UInt8}()
-        assemble_graph!(g,weights_dict,index_vectors,mixing_matrix,sampler_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,
         )
+        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,
         )
+        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)
+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)
+function assemble_graph!(g,index_vectors,mixing_matrix,sampler_matrix)
     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])
+            stubs_a,stubs_b = 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])
+            stubs_a,stubs_b = 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)
+        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
 function fast_chung_lu_bipartite(pop_i,pop_j,mixing_dist_ij,mixing_dist_ji)
     num_degrees_ij = similar(pop_i)
@@ -221,7 +227,7 @@ function fast_chung_lu_bipartite(pop_i,pop_j,mixing_dist_ij,mixing_dist_ji)
         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)
+    return stubs_i,stubs_j
 end
 
 function fast_chung_lu(pop_i,mixing_dist)
@@ -235,11 +241,38 @@ function fast_chung_lu(pop_i,mixing_dist)
         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 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

@@ -79,7 +79,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)
@@ -98,7 +98,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,length(app_user_index)) #two weeks worth of values
         time_of_last_alert = fill(-1,length(app_user_index)) #time of last alert is negative if no alert has been recieved

CovidAlertVaccinationModel/src/ABM/solve.jl

@@ -21,9 +21,9 @@ Base.@propagate_inbounds @views function update_alert_durations!(t,modelsol) # B
         end
         total_weight_i = 0
         for mixing_graph in inf_network.graph_list[t]
-            for j in neighbors(mixing_graph,node)
+            for (j,w) in neighbors_and_weights(mixing_graph,node)
                 if app_user[j]
-                    total_weight_i+= get_weight(mixing_graph,GraphEdge(node,j))
+                    total_weight_i += w[]
                 end
             end
         end
@@ -61,11 +61,11 @@ Base.@propagate_inbounds @views function update_infection_state!(t,modelsol; rec
         end
         if  agent_status == Susceptible || agent_status == Immunized 
             for mixing_graph in inf_network.graph_list[t]
-                for j in neighbors(mixing_graph,i)    
+                for (j,w) in neighbors_and_weights(mixing_graph,i)   
                     if u_inf[j] == Infected && u_next_inf[i] != Infected
                         β_vec = (β_y,β_m,β_o)
                         α_vec = (α_y,α_m,α_o)
-                        if rand(Random.default_rng(Threads.threadid())) < contact_weight(β_vec[Int(agent_demo)],get_weight(mixing_graph,GraphEdge(i,j)))
+                        if rand(Random.default_rng(Threads.threadid())) < contact_weight(β_vec[Int(agent_demo)],w[])
                             if agent_status == Immunized && rand(Random.default_rng(Threads.threadid())) < 1- α_vec[Int(agent_demo)]
                                 agent_transition!(i, Immunized,Infected)
                                 output_data.daily_cases_by_age[Int(agent_demo),t]+=1
@@ -111,7 +111,7 @@ Base.@propagate_inbounds @views function update_vaccination_opinion_state!(t,mod
         (π_base_y*ζ,π_base_m*ζ,π_base_o*ζ)
         random_soc_network = sample(Random.default_rng(Threads.threadid()), soc_network.graph_list[t])
         if !isempty(neighbors(random_soc_network,i))
-                random_neighbour = sample(Random.default_rng(Threads.threadid()), neighbors(random_soc_network.g,i))
+                random_neighbour = sample(Random.default_rng(Threads.threadid()), neighbors(random_soc_network,i))
                 app_vac_payoff = 0.0    
                 if app_user[i] && time_of_last_alert[app_user_list[i]]>=0
                     output_data.daily_total_notified_agents[t] += 1
@@ -140,8 +140,8 @@ end
 function weighted_degree(t,node,network::TimeDepMixingGraph)
     weighted_degree = 0
     for g in network.graph_list[t]
-        for j in neighbors(g,node)
-            weighted_degree += get_weight(g,GraphEdge(node,j))
+        for (j,w) in neighbors_and_weights(mixing_graph,node)
+            weighted_degree += w[]
         end
     end
     return weighted_degree
@@ -151,8 +151,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

 module CovidAlertVaccinationModel
 using Intervals: Ending
+import Base.empty!
 using Base: Float64
 using LightGraphs
 using RandomNumbers.Xorshifts
@@ -28,7 +29,7 @@ using CSV
 import Pandas: read_csv
 using DataFrames
 using StaticArrays
-import LightGraphs.neighbors
+import LightGraphs:neighbors,add_edge!
 
 export intervalsmodel, hh, ws, rest, abm,multivariate_simulations