simplify more

CovidAlertVaccinationModel/src/ABM/mixing_graphs.jl

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

CovidAlertVaccinationModel/src/ABM/model_setup.jl

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

CovidAlertVaccinationModel/src/ABM/solve.jl

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

CovidAlertVaccinationModel/src/utils.jl

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

CovidAlertVaccinationModel/test/ABM/mixing_test.jl

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

scratch_code/alloc_test2.jl

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