writing docs to avoid marking

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)

CovidAlertVaccinationModel/src/ABM/agents.jl

 #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")

CovidAlertVaccinationModel/src/ABM/contact_vectors.jl

+#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)

CovidAlertVaccinationModel/src/ABM/mixing_distributions.jl

+
+"""
+    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

CovidAlertVaccinationModel/src/ABM/mixing_graphs.jl

 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