diff --git a/CovidAlertVaccinationModel/src/ABM/abm.jl b/CovidAlertVaccinationModel/src/ABM/abm.jl
index f63f50b8d1f0a724b6bf9b5607c6d5b5eaf0bc0a..687088da65734ace551ccdb4810b05899f4d6e6f 100644
--- a/CovidAlertVaccinationModel/src/ABM/abm.jl
+++ b/CovidAlertVaccinationModel/src/ABM/abm.jl
@@ -6,7 +6,11 @@ const age_bins = [(0.0, 25.0),(25.0,65.0),(65.0,Inf)]
default(dpi = 300)
default(framestyle = :box)
+"""
+ bench()
+Runs the model with 5k households for 500 timesteps.
+"""
function bench()
Random.seed!(RNG,1)
steps = 500
@@ -15,6 +19,9 @@ function bench()
output = solve!(model_sol,recording )
end
+"""
+The only function we export. I just put whatever I am currently working on into this function to make the model easier to test from the REPL.
+"""
function abm()
b1 = @benchmark bench() seconds = 20
display(b1)
diff --git a/CovidAlertVaccinationModel/src/ABM/agents.jl b/CovidAlertVaccinationModel/src/ABM/agents.jl
index 822fb2fe31d405ca2058b188a2f49d50302eee15..5b830aa895bd7730e29e75cee059a7f757e8ac30 100644
--- a/CovidAlertVaccinationModel/src/ABM/agents.jl
+++ b/CovidAlertVaccinationModel/src/ABM/agents.jl
@@ -1,9 +1,17 @@
#agent type definitions
+
+"""
+Define the possible agent demographic groups. Note that these can be converted to/from integers (starting at 1).
+"""
@enum AgentDemographic begin
Young = 1
Adult
Elderly
end
+
+"""
+Define the possible agent infection statuses. Note that these can be converted to/from integers (starting at 1).
+"""
@enum AgentStatus begin
Susceptible = 1
Infected
@@ -11,7 +19,14 @@ end
Immunized
end
+"""
+
+ complete_graph_from_households_composition(households_composition::AbstractVector{AbstractVector{Int}})
+Generates a complete graph from a vector of household compositions. Each household composition is a 3 element vectors (one for each demographic group) of integers where each element describes the number of the corresponding demographic group present in that household.
+
+This function wires together a graph such that each household is in a complete subgraph. It is much faster than the `mapreduce` solution used previously.
+"""
function complete_graph_from_households_composition(households_composition)
total_household_pop = sum(sum.(households_composition))
network = SimpleGraph(total_household_pop)
@@ -31,13 +46,24 @@ function complete_graph_from_households_composition(households_composition)
end
+"""
+ generate_population(num_households::Int)
+
+Returns a namedtuple with fields `population_list`, `household_networks` and `index_vectors`.
-#generate population with households distributed according to household_size_distribution
+population_list: Vector of AgentDemographic where each entry is the demographic of a person in the simulation. The length is equal to the total number of people in the model.
+
+household_networks: SimpleGraph with all agents such that each household is in a complete subgraph.
+
+index_vectors: Vector of 3 vectors, each of which contains the indexes of all the agents in a particular demographic. If `population_list` is a mapping from agents to demographics, these are mappings from demographics to agents.
+"""
function generate_population(num_households)
households_composition = sample_household_data(num_households)
household_networks = complete_graph_from_households_composition(households_composition)
- # @assert all(adjacency_matrix(household_networks) .== adjacency_matrix( mapreduce(LightGraphs.complete_graph,blockdiag, sum.(households_composition); init = SimpleGraph())))
+
+ #test the complete_graph_from_households_composition function against the mapreduce version to ensure it is correct.
+ @c_assert all(adjacency_matrix(household_networks) .== adjacency_matrix( mapreduce(LightGraphs.complete_graph,blockdiag, sum.(households_composition); init = SimpleGraph())))
households = map( l -> [fill(AgentDemographic(i),n) for (i,n) in enumerate(l)],households_composition) |>
l -> reduce(vcat,l) |>
@@ -54,8 +80,9 @@ function generate_population(num_households)
end
-
-# func(a) = println("agent")
+"""
+Defines pretty printing methods so that `display(AgentStatus)` is more readable.
+"""
function Base.show(io::IO, status::AgentStatus)
if status == Susceptible
print(io, "S")
@@ -68,6 +95,9 @@ function Base.show(io::IO, status::AgentStatus)
end
end
+"""
+Defines pretty printing methods so that `display(AgentDemographic)` is more readable.
+"""
function Base.show(io::IO, status::AgentDemographic)
if status == Young
print(io, "<25")
diff --git a/CovidAlertVaccinationModel/src/ABM/contact_vectors.jl b/CovidAlertVaccinationModel/src/ABM/contact_vectors.jl
index 91ebac2a9dc1139396ca4890977b8a36404c02f4..013389efde5782cf1216764d5057c4584066fbec 100644
--- a/CovidAlertVaccinationModel/src/ABM/contact_vectors.jl
+++ b/CovidAlertVaccinationModel/src/ABM/contact_vectors.jl
@@ -1,4 +1,14 @@
+#here are the functions for
+
+
+
+
+"""
+Called in the inner loop of `generate_contact_vectors!`, this function removes the degree corresponding to kth entry of index_lists, replaces it with another randomly sampled degree, and returns `csum` adjusted to the new samples.
+
+We repeat this function until `csum == 0`.
+"""
@inline function reindex!(k,csum,index_list_i,index_list_j,j_to_i_contacts,i_to_j_contacts,sample_list_i,sample_list_j)
i_index = index_list_i[k]
j_index = index_list_j[k]
@@ -7,6 +17,15 @@
j_to_i_contacts[j_index] = sample_list_j[k]
return csum
end
+
+
+"""
+ generate_contact_vectors!(ij_dist<:Distribution,ji_dist<:Distribution,i_to_j_contacts::Vector{T}, j_to_i_contacts::Vector{T})
+
+Fills i_to_j_contacts and j_to_i_contacts with degrees sampled from ij_dist and ji_dist, such that i_to_j_contacts and j_to_i_contacts have equal sum.
+
+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
rand!(RNG,ij_dist,i_to_j_contacts)
rand!(RNG,ji_dist,j_to_i_contacts)
diff --git a/CovidAlertVaccinationModel/src/ABM/mixing_distributions.jl b/CovidAlertVaccinationModel/src/ABM/mixing_distributions.jl
index 8b4ad25750ae22a5cbba2342a939d09c259007e2..4247add4d92bac13e5795f3ff5c058ed37572fff 100644
--- a/CovidAlertVaccinationModel/src/ABM/mixing_distributions.jl
+++ b/CovidAlertVaccinationModel/src/ABM/mixing_distributions.jl
@@ -1,3 +1,9 @@
+
+"""
+ adjust_distributions_mean(distribution_matrix::AbstractArray{T},mean_shift_percentage::Number)
+
+Adjust the means of the distributions in `distribution_matrix` by `mean_shift_percentage`.
+"""
function adjust_distributions_mean(distribution_matrix::AbstractArray{T},mean_shift_percentage) where T<:Distribution
return map(distribution_matrix) do dist
new_mean = mean(dist)*(1 + mean_shift_percentage)
@@ -5,6 +11,11 @@ function adjust_distributions_mean(distribution_matrix::AbstractArray{T},mean_sh
end
end
+"""
+ symmetrize_means(population_sizes::Vector{Int},mixing_matrix::Matrix{<:Distribution})
+
+Symmetrize `mixing_matrix` such that it satisifes the conditions needed for `generate_contact_vectors!(..)`, given `population_sizes`.
+"""
function symmetrize_means(population_sizes,mixing_matrix::Matrix{<:Distribution})
mixing_means = mean.(mixing_matrix)
symmetrized_means = zeros((length(population_sizes),length(population_sizes)))
@@ -15,11 +26,17 @@ function symmetrize_means(population_sizes,mixing_matrix::Matrix{<:Distribution}
return map(((dst,sym_mean),)->from_mean(typeof(dst),sym_mean), zip(mixing_matrix,symmetrized_means))
end
+"""
+This is a namedtuple with entries `hh,ws,rest` which define matrices of Poisson distributions, with parameters given by IntervalsModel.
-
+See `load_contact_time_distributions()`.
+"""
const contact_time_distributions = load_contact_time_distributions()
+```
+Given alphas given in `unemployment_matrix` computes the zero weight for each pair of demographic contacts given individual zero weights.
+```
function alpha_matrix(alphas)
M = zeros(length(alphas),length(alphas))
for i in 1:length(alphas), j in 1:length(alphas)
@@ -36,6 +53,9 @@ const unemployment_matrix = alpha_matrix(
)
)
+"""
+Sample initial_workschool_mixing_matrix, which is the workschool distributions symmetrized for the full Canadian population, rather than subsets (as used in the ABM). This is used in IntervalsModel.
+"""
@views function ws_sample(age)
return rand.(initial_workschool_mixing_matrix[age,:]) * (rand(RNG) < (5/7))
end
diff --git a/CovidAlertVaccinationModel/src/ABM/mixing_graphs.jl b/CovidAlertVaccinationModel/src/ABM/mixing_graphs.jl
index 87d12b0db53193e9b09ace19009d63a8f73995cf..d3fe2329f46ccc611c6dae0b17d5b37e0fd43a34 100644
--- a/CovidAlertVaccinationModel/src/ABM/mixing_graphs.jl
+++ b/CovidAlertVaccinationModel/src/ABM/mixing_graphs.jl
@@ -1,21 +1,41 @@
using Dictionaries
+
+"""
+A type that stores a pair of nodes representing an edge in the graph.
+
+We need to define a custom type for these so we can define a hash function on graph edges, in order to more efficiently use them in hashmaps (dictionaries)
+"""
struct GraphEdge
a::Int
b::Int
end
+
+
+"""
+Define a hash on GraphEdge such that ``hash(a,b) = hash(b,a)`` (hash is commutative).
+
+This is helpful because then we only need to store (a,b) in the graph edges weights dictionary, rather than both (a,b) and (b,a).
+"""
function Base.hash(e::GraphEdge)
return hash(minmax(e.a,e.b))
end
+
+"""
+Define symmetric (?) edge equality, matches the hash function.
+"""
function Base.isequal(e1::GraphEdge,e2::GraphEdge)
return isequal(minmax(e1.a,e1.b),minmax(e2.a,e2.b))
end
-function sample_mixing_graph!(mixing_graph,population_demographics)
+
+"""
+ sample_mixing_graph!(mixing_graph::Graph)
+
+
+Resample all the weights in `mixing_graph`
+"""
+function sample_mixing_graph!(mixing_graph)
mixing_edges = mixing_graph.mixing_edges
- # display( mixing_edges.sampler_matrix)
- for i in 1:size(mixing_edges.contact_array)[1], j in 1:i #diagonal
- demo_i = Int(population_demographics[i])
- demo_j = Int(population_demographics[j])
-
+ for i in 1:size(mixing_edges.contact_array)[1], j in 1:i #diagonal
rand!(RNG, mixing_edges.sampler_matrix[demo_i,demo_j],mixing_edges.sample_cache[i,j])
for k in 1:length(mixing_edges.contact_array[i,j][1])
kth_node_i = mixing_edges.contact_array[i,j][1][k]
@@ -26,7 +46,31 @@ function sample_mixing_graph!(mixing_graph,population_demographics)
end
end
-#A type that defines a matrix of vectors, such that the ijth vector contains the nodes to connect to the jith vector, ala. Chung-Lu bipartite
+"""
+A type that defines the graph edges such that they can be easily resampled and then re-added to the graph. Created with `create_mixing_edges(..)`.
+
+#Fields
+
+ total_edges::Int
+
+Total number of edges in the struct.
+
+ contact_array::Matrix{Tuple{Vector{Int},Vector{Int}}}
+
+A matrix of vector pairs, such that one node is in the first vector and the other is in the second. The position of the vector pairs in the matrix corresponds to the distribution in sampler_matrix used to sample them. There is probably a better way to represent these, I did this so they could be sampled fast. We only use the upper triangle of this but Julia lacks a good Symmetric matrix type.
+
+ sample_cache::Matrix{Vector{Int}}
+
+The structure that pre-allocates the space for the weights of the edges in contact array. This is filled with weights and then the weights are added to weights_dict. We only use the upper triangle of this but Julia lacks a good Symmetric matrix type.
+
+ weights_dict::Dictionary{GraphEdge,UInt8}
+
+Stores the weights used in the graph, so they can be easily resampled.
+
+ sampler_matrix::M
+
+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}
total_edges::Int
contact_array::Matrix{Tuple{Vector{Int},Vector{Int}}}
@@ -40,6 +84,9 @@ struct MixingEdges{M}
end
end
+"""
+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{Int}(),Vector{Int}()) for i in 1:length(demographic_index_vectors),j in 1:length(demographic_index_vectors)]
tot = 0
@@ -61,6 +108,10 @@ function create_mixing_edges(demographic_index_vectors,mixing_matrix,weights_dis
return MixingEdges(tot,contact_array,weights_distribution_matrix)
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 = [(Vector{Int}(),Vector{Int}()) for i in 1:length(demographic_index_vectors),j in 1:length(demographic_index_vectors)]
for e in edges(g)
@@ -73,6 +124,21 @@ function create_mixing_edges(g::SimpleGraph,demographics,demographic_index_vecto
return MixingEdges(ne(g),contact_array,weights_distribution_matrix)
end
+"""
+
+Stores the full time dependent mixing graph for the model. I think this might be a weird abstraction for this idea but it works fine.
+
+#Fields
+
+These are references to the graphs that get resampled everyday.
+
+ resampled_graphs::NTuple{N,G}
+
+List of lists of graphs, one list for each day.
+
+ graph_list::Vector{Vector{G}}
+
+"""
struct TimeDepMixingGraph{N,G}
resampled_graphs::NTuple{N,G}
graph_list::Vector{Vector{G}}
@@ -83,6 +149,12 @@ struct TimeDepMixingGraph{N,G}
)
end
end
+
+"""
+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,index_vectors,contact_time_distributions.hh) #network with households and LTC homes
ws_static_edges = WeightedGraph(demographics,index_vectors,ws_matrix_tuple.daily,contact_time_distributions.ws)
@@ -113,6 +185,11 @@ function time_dep_mixing_graphs(len,base_network,demographics,index_vectors,ws_m
return infected_mixing_graph,soc_mixing_graph
end
+"""
+Completely remake all the graphs in `time_dep_mixing_graph.resampled_graphs`.
+
+This is a huge bottleneck for the performance of the model but I doubt it can be improved much more.
+"""
function remake!(time_dep_mixing_graph,demographic_index_vectors,mixing_matrix)
for weighted_graph in time_dep_mixing_graph.resampled_graphs
empty!.(weighted_graph.g.fadjlist) #empty all the vector edgelists
@@ -122,6 +199,9 @@ function remake!(time_dep_mixing_graph,demographic_index_vectors,mixing_matrix)
end
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:size(mixing_edges.contact_array)[1], j in 1:i #diagonal
for k in 1:length(mixing_edges.contact_array[i,j][1])
@@ -129,7 +209,11 @@ function graph_from_mixing_edges(g,mixing_edges)
end
end
end
-#defining my own weighted graph type cause we need to be able to resample the edge weights in a very particular way
+
+"""
+My own weighted graph type. Stores the graph in `g`, and the weights and edges in `mixing_edges`.
+
+"""
mutable struct WeightedGraph{G,M}
g::G
mixing_edges::M