roughly working, not happy with mixing graph implementation

Manifest.toml

@@ -6,6 +6,11 @@ git-tree-sha1 = "345a14764e43fe927d6f5c250fe4c8e4664e6ee8"
 uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 version = "2.4.0"
 
+[[ArgCheck]]
+git-tree-sha1 = "dedbbb2ddb876f899585c4ec4433265e3017215a"
+uuid = "dce04be8-c92d-5529-be00-80e4d2c0e197"
+version = "2.1.0"
+
 [[ArgTools]]
 uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"
 
@@ -18,9 +23,21 @@ version = "0.0.2"
 [[Artifacts]]
 uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
 
+[[BangBang]]
+deps = ["Compat", "ConstructionBase", "Future", "InitialValues", "LinearAlgebra", "Requires", "Setfield", "Tables", "ZygoteRules"]
+git-tree-sha1 = "f42321255afc37da855b6cd9f2a1fc36c017ceee"
+uuid = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
+version = "0.3.29"
+
 [[Base64]]
 uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 
+[[BenchmarkTools]]
+deps = ["JSON", "Logging", "Printf", "Statistics", "UUIDs"]
+git-tree-sha1 = "9e62e66db34540a0c919d72172cc2f642ac71260"
+uuid = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+version = "0.5.0"
+
 [[Bzip2_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
 git-tree-sha1 = "c3598e525718abcc440f69cc6d5f60dda0a1b61e"
@@ -73,6 +90,16 @@ version = "3.25.0"
 deps = ["Artifacts", "Libdl"]
 uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
 
+[[CompositionsBase]]
+git-tree-sha1 = "f3955eb38944e5dd0fabf8ca1e267d94941d34a5"
+uuid = "a33af91c-f02d-484b-be07-31d278c5ca2b"
+version = "0.1.0"
+
+[[ConstructionBase]]
+git-tree-sha1 = "a2a6a5fea4d6f730ec4c18a76d27ec10e8ec1c50"
+uuid = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
+version = "1.0.0"
+
 [[Contour]]
 deps = ["StaticArrays"]
 git-tree-sha1 = "9f02045d934dc030edad45944ea80dbd1f0ebea7"
@@ -110,6 +137,11 @@ version = "1.0.0"
 deps = ["Printf"]
 uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"
 
+[[DefineSingletons]]
+git-tree-sha1 = "77b4ca280084423b728662fe040e5ff8819347c5"
+uuid = "244e2a9f-e319-4986-a169-4d1fe445cd52"
+version = "0.1.1"
+
 [[DelimitedFiles]]
 deps = ["Mmap"]
 uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"
@@ -250,6 +282,11 @@ git-tree-sha1 = "098e4d2c533924c921f9f9847274f2ad89e018b8"
 uuid = "83e8ac13-25f8-5344-8a64-a9f2b223428f"
 version = "0.5.0"
 
+[[InitialValues]]
+git-tree-sha1 = "26c8832afd63ac558b98a823265856670d898b6c"
+uuid = "22cec73e-a1b8-11e9-2c92-598750a2cf9c"
+version = "0.2.10"
+
 [[InteractiveUtils]]
 deps = ["Markdown"]
 uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
@@ -421,6 +458,12 @@ git-tree-sha1 = "e498ddeee6f9fdb4551ce855a46f54dbd900245f"
 uuid = "442fdcdd-2543-5da2-b0f3-8c86c306513e"
 version = "0.3.1"
 
+[[MicroCollections]]
+deps = ["BangBang", "Setfield"]
+git-tree-sha1 = "e991b6a9d38091c4a0d7cd051fcb57c05f98ac03"
+uuid = "128add7d-3638-4c79-886c-908ea0c25c34"
+version = "0.1.0"
+
 [[Missings]]
 deps = ["DataAPI"]
 git-tree-sha1 = "ed61674a0864832495ffe0a7e889c0da76b0f4c8"
@@ -438,6 +481,11 @@ git-tree-sha1 = "bfe47e760d60b82b66b61d2d44128b62e3a369fb"
 uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 version = "0.3.5"
 
+[[NamedTupleTools]]
+git-tree-sha1 = "63831dcea5e11db1c0925efe5ef5fc01d528c522"
+uuid = "d9ec5142-1e00-5aa0-9d6a-321866360f50"
+version = "0.13.7"
+
 [[NetworkLayout]]
 deps = ["GeometryBasics", "LinearAlgebra", "SparseArrays"]
 git-tree-sha1 = "01b5715cdd1b7c5d493c26cc05e4af663ba9a052"
@@ -586,6 +634,11 @@ git-tree-sha1 = "7b1d07f411bc8ddb7977ec7f377b97b158514fe0"
 uuid = "189a3867-3050-52da-a836-e630ba90ab69"
 version = "0.2.0"
 
+[[Referenceables]]
+git-tree-sha1 = "4a32b1dd124a846580608eb347d4337f873c2499"
+uuid = "42d2dcc6-99eb-4e98-b66c-637b7d73030e"
+version = "0.1.0"
+
 [[Requires]]
 deps = ["UUIDs"]
 git-tree-sha1 = "cfbac6c1ed70c002ec6361e7fd334f02820d6419"
@@ -616,6 +669,12 @@ version = "1.0.3"
 [[Serialization]]
 uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 
+[[Setfield]]
+deps = ["ConstructionBase", "Future", "MacroTools", "Requires"]
+git-tree-sha1 = "d5640fc570fb1b6c54512f0bd3853866bd298b3e"
+uuid = "efcf1570-3423-57d1-acb7-fd33fddbac46"
+version = "0.7.0"
+
 [[SharedArrays]]
 deps = ["Distributed", "Mmap", "Random", "Serialization"]
 uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"
@@ -651,6 +710,12 @@ git-tree-sha1 = "75394dbe2bd346beeed750fb02baa6445487b862"
 uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
 version = "1.2.1"
 
+[[SplittablesBase]]
+deps = ["Setfield", "Test"]
+git-tree-sha1 = "e6e5e266aa58c2e419d19e8eccbe6af55d5dc6e1"
+uuid = "171d559e-b47b-412a-8079-5efa626c420e"
+version = "0.1.11"
+
 [[StaticArrays]]
 deps = ["LinearAlgebra", "Random", "Statistics"]
 git-tree-sha1 = "9da72ed50e94dbff92036da395275ed114e04d49"
@@ -713,6 +778,18 @@ uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"
 deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
 uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
+[[ThreadsX]]
+deps = ["ArgCheck", "BangBang", "ConstructionBase", "InitialValues", "MicroCollections", "Referenceables", "Setfield", "SplittablesBase", "Transducers"]
+git-tree-sha1 = "269f5c1955c1194086cf6d2029aa4a0b4fb8018b"
+uuid = "ac1d9e8a-700a-412c-b207-f0111f4b6c0d"
+version = "0.1.7"
+
+[[Transducers]]
+deps = ["ArgCheck", "BangBang", "CompositionsBase", "DefineSingletons", "Distributed", "InitialValues", "Logging", "Markdown", "MicroCollections", "Requires", "Setfield", "SplittablesBase", "Tables"]
+git-tree-sha1 = "eb5d4ff48c25762319f1aae0e5b501394beaefd7"
+uuid = "28d57a85-8fef-5791-bfe6-a80928e7c999"
+version = "0.4.56"
+
 [[UUIDs]]
 deps = ["Random", "SHA"]
 uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
@@ -885,6 +962,12 @@ git-tree-sha1 = "6f1abcb0c44f184690912aa4b0ba861dd64f11b9"
 uuid = "3161d3a3-bdf6-5164-811a-617609db77b4"
 version = "1.4.5+2"
 
+[[ZygoteRules]]
+deps = ["MacroTools"]
+git-tree-sha1 = "9e7a1e8ca60b742e508a315c17eef5211e7fbfd7"
+uuid = "700de1a5-db45-46bc-99cf-38207098b444"
+version = "0.2.1"
+
 [[libass_jll]]
 deps = ["Artifacts", "Bzip2_jll", "FreeType2_jll", "FriBidi_jll", "JLLWrappers", "Libdl", "Pkg", "Zlib_jll"]
 git-tree-sha1 = "acc685bcf777b2202a904cdcb49ad34c2fa1880c"

Project.toml

@@ -4,16 +4,19 @@ authors = ["pjentsch <pjentsch@uwaterloo.ca> and contributors"]
 version = "0.1.0"
 
 [deps]
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
+NamedTupleTools = "d9ec5142-1e00-5aa0-9d6a-321866360f50"
 NetworkLayout = "46757867-2c16-5918-afeb-47bfcb05e46a"
 Pipe = "b98c9c47-44ae-5843-9183-064241ee97a0"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 RandomNumbers = "e6cf234a-135c-5ec9-84dd-332b85af5143"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+ThreadsX = "ac1d9e8a-700a-412c-b207-f0111f4b6c0d"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 
 [compat]

src/CovidAlertVaccinationModel.jl

 module CovidAlertVaccinationModel
-export get_u_0, solve, agents_step!,get_graph,plot_network,main,parse_cases_data,generate_population
+export get_u_0, solve, agents_step!,get_graph,plot_model,main,parse_cases_data,generate_population
 using LightGraphs
 using RandomNumbers.Xorshifts
+using Random
 using UnPack
 using Plots
 using DataFrames
@@ -12,15 +13,17 @@ using ThreadsX
 using DelimitedFiles
 using NamedTupleTools
 # using CUDA
+using NetworkLayout:Stress
 using NetworkLayout:SFDP
+import Base.rand
 
 
-
+include("types.jl")
+include("mixing_matrices.jl")
 include("data.jl")
 include("model.jl")
 include("plotting.jl")
-include("types.jl")
-
+include("graphs.jl")
 
 const population = 14.57e6
 
@@ -41,21 +44,25 @@ using BenchmarkTools
 function main()
 
     household_size_dist = Truncated(Poisson(2.60),0,Inf) #https://www.researchgate.net/publication/226081704_Household_size_and_the_Poisson_distribution
-    pop_list,network,index_vectors = generate_population(500,2,household_size_dist)
-    u_0 = get_u_0(length(pop_list))
-    params = get_parameters()
-    steps = 200
-    
+    # pop_list,network,index_vectors = generate_population(500,2,household_size_dist)
+
     r = Xoroshiro128Plus()
-    sol1 = solve!(u_0,network,params,steps,r,pop_list,index_vectors)
     # display(sol1
-    # pop_list,network,index_vectors = generate_population(500,2,household_size_dist)
-    # contacts = annealed_contacts(r,Young,index_vectors) 
+    pop_list,base_network,index_vectors = generate_population(r,5000,1,household_size_dist)
+    u_0 = get_u_0(length(pop_list))
+
+    # plot_model(network,[u_0],pop_list)
+    # networks = [deepcopy(base_network) for i in 1:5]
+    # for network in networks
+    #     generate_mixing_graph!(r,network,pop_list,index_vectors,workschool_mixing_matrix) 
+    #     generate_mixing_graph!(r,network,pop_list,index_vectors,rest_mixing_matrix) 
+    # end
+    params = get_parameters()
+    steps = 300
+    display(length(pop_list))
+    @btime solve!($u_0,$base_network,$params,$steps,$r,$pop_list,$index_vectors);
+    # plot_model(graphs,sol1,pop_list)
     # contact_time = contact_weight(r,0.1, Young,Elderly) 
 end
 
-
-
-
-
-end
+end
\ No newline at end of file

src/graphs.jl

+
+
+function generate_mixing_graph!(rng,g,population_list,index_vectors,mixing_matrix)
+    agents_young = population_list[index_vectors[1]]
+    # g = SimpleGraph(length(population_list))
+    for i in 1:5, j in 1:i #diagonal
+            agent_contact_dist = mixing_matrix[i,j]
+            i_to_j_contacts = rand(rng,agent_contact_dist,length(index_vectors[i]))
+            
+            j_to_i_contacts = rand(rng,agent_contact_dist,length(index_vectors[j]))
+            
+            equalize_degree_lists!(i_to_j_contacts,j_to_i_contacts)
+            random_bipartite_graph_fast_CL!(rng,g,index_vectors[i],index_vectors[j],i_to_j_contacts,j_to_i_contacts)
+            # display(g)
+    end
+    return g
+end
+function equalize_degree_lists!(l1,l2)
+    if sum(l1) > sum(l2)
+        map(i -> l1[i]-=1, sample(1:length(l1),sum(l1) - sum(l2)))
+    else
+        map(i -> l2[i]-=1, sample(1:length(l2),sum(l2) - sum(l1)))
+    end
+end
+function random_bipartite_graph_fast_CL!(rng,g,anodes,bnodes,aseq,bseq) #Aksoy, S. G., Kolda, T. G., & Pinar, A. (2017). Measuring and modeling bipartite graphs with community structure
+    lena = length(aseq)
+    lenb = length(bseq)
+    m = sum(aseq)
+    @assert sum(aseq) == sum(bseq) "degree sequences must have equal sum"
+    astubs = sample(rng,anodes,StatsBase.weights(aseq./m),m;replace = true)
+    bstubs = sample(rng,bnodes,StatsBase.weights(bseq./m),m;replace = true)
+    for k in 1:m
+        add_edge!(g,astubs[k],bstubs[k])
+    end
+    return g
+end
+
+function generate_population(rng,num_households,num_LTC,household_size_distribution)
+    household_sizes = Int.(rand(rng,household_size_distribution,num_households))
+    
+    total_household_population = sum(household_sizes) 
+    
+    # LTC_sizes = rand(LTC_distribution,num_LTC)
+
+    prob_adult_is_HCW = 0.01
+    prob_elderly_in_LTC = 0.05
+    agent_distribution_in_households = Categorical([ #find a better way to do this maybe
+        demographic_distribution[1],
+        demographic_distribution[2] * (1 - prob_adult_is_HCW),
+        demographic_distribution[2] * prob_adult_is_HCW,
+        demographic_distribution[3] * (1 - prob_elderly_in_LTC),
+        demographic_distribution[3] * prob_elderly_in_LTC
+    ])
+
+    households = Vector{Vector{AgentDemographic}}()
+    ltc_pop_list = Vector{AgentDemographic}()
+
+    for household_size in household_sizes
+        household = Vector{AgentDemographic}()
+        while length(household)<household_size
+            agent_sample = AgentDemographic(rand(rng,agent_distribution_in_households))
+            if agent_sample == ElderlyLTC
+                push!(ltc_pop_list,agent_sample)
+            else
+                push!(household,agent_sample)
+            end
+        end
+        push!(households,household)
+    end
+
+    ltc_pop = length(ltc_pop_list)
+    println("ltc: $ltc_pop")
+    ltc_home_size = Int(ceil(ltc_pop/num_LTC))
+
+
+    if ltc_home_size != 0
+        ltc_homes = [ltc_pop_list[i:min(i + ltc_home_size - 1,ltc_pop)] for i = 1:ltc_home_size:ltc_pop]
+    else 
+        ltc_homes = Vector{Vector{AgentDemographic}}()
+    end
+
+    ltc_networks = map(LightGraphs.complete_graph,length.(ltc_homes))
+    household_networks = map(LightGraphs.complete_graph, household_sizes)
+    # display(household_networks)
+    population_list = vcat(vcat(households...), vcat(ltc_homes...))
+    network = blockdiag(reduce(blockdiag,household_networks; init = SimpleGraph()), reduce(blockdiag,ltc_networks; init = SimpleGraph()))
+    # display(network)
+    index_vectors = [findall(x -> x == AgentDemographic(i), population_list) for i in 1:AgentDemographic.size+1]
+  
+    return (;
+        population_list,
+        network,
+        index_vectors
+    )
+end
\ No newline at end of file

src/mixing_matrices.jl

+
+const home_mixing_matrix = [
+    Poisson(1.749) Poisson(2.629) Poisson(0.648) Poisson(2.629) Poisson(0.648)
+    Poisson(0.726) Poisson(1.909) Poisson(0.357) Poisson(1.909) Poisson(0.357)
+    Poisson(0.022) Poisson(0.120) Poisson(0.677) Poisson(0.120) Poisson(0.677)
+    Poisson(0.726) Poisson(1.909) Poisson(0.357) Poisson(1.909) Poisson(0.357)
+    Poisson(0.022) Poisson(0.120) Poisson(0.677) Poisson(0.120) Poisson(0.677)
+]
+
+
+const workschool_mixing_matrix = [
+    ZeroGeometric(0.439,0.126) ZeroGeometric(0.414,0.230) ZeroPoisson(0.907,0.243) ZeroGeometric(0.414,0.230) ZeroPoisson(0.907,0.243)
+    ZeroGeometric(0.767,0.457) ZeroGeometric(0.786,0.106) ZeroPoisson(0.842,0.122) ZeroGeometric(0.786,0.106) ZeroPoisson(0.842,0.122)
+    ZeroPoisson(0.898,0.015)   ZeroPoisson(0.786,0.036)   ZeroPoisson(0.941,0.373) ZeroPoisson(0.786,0.036)   ZeroPoisson(0.941,0.373)
+    ZeroGeometric(0.767,0.457) ZeroGeometric(0.786,0.106) ZeroPoisson(0.842,0.122) ZeroGeometric(0.786,0.106) ZeroPoisson(0.842,0.122)
+    ZeroPoisson(0.898,0.015)   ZeroPoisson(0.786,0.036)   ZeroPoisson(0.941,0.373) ZeroPoisson(0.786,0.036)   ZeroPoisson(0.941,0.373)
+]
+
+
+
+const rest_mixing_matrix = [
+    Geometric(0.450) Geometric(0.784)  Poisson(0.155)  Geometric(0.784)  Poisson(0.155)
+    Poisson(0.850)  Geometric(0.269)  Poisson(0.741)  Geometric(0.269)  Poisson(0.741)
+    Poisson(0.217)    Poisson(0.783)    Poisson(0.667)  Poisson(0.783)    Poisson(0.667)
+    Poisson(0.850)    Geometric(0.269)  Poisson(0.741)  Geometric(0.269)  Poisson(0.741)
+    Poisson(0.217)    Poisson(0.783)    Poisson(0.667)  Poisson(0.783)    Poisson(0.667)
+]
+
+const contact_time_distribution_matrix = [Hypergeometric(5,2,4) for i in 1:AgentDemographic.size + 1, j in 1:AgentDemographic.size + 1]
+
+

src/model.jl

-
-function generate_population(num_households,num_LTC,household_size_distribution)
-    household_sizes = Int.(rand(household_size_distribution,num_households))
-    
-    total_household_population = sum(household_sizes) 
-    
-    # LTC_sizes = rand(LTC_distribution,num_LTC)
-
-    prob_adult_is_HCW = 0.01
-    prob_elderly_in_LTC = 0.05
-    agent_distribution_in_households = Categorical([ #find a better way to do this maybe
-        demographic_distribution[1],
-        demographic_distribution[2] * (1 - prob_adult_is_HCW),
-        demographic_distribution[2] * prob_adult_is_HCW,
-        demographic_distribution[3] * (1 - prob_elderly_in_LTC),
-        demographic_distribution[3] * prob_elderly_in_LTC
-    ])
-
-    households = Vector{Vector{AgentDemographic}}()
-    ltc_pop_list = Vector{AgentDemographic}()
-
-    for household_size in household_sizes
-        household = Vector{AgentDemographic}()
-        while length(household)<household_size
-            agent_sample = AgentDemographic(rand(agent_distribution_in_households))
-            if agent_sample == ElderlyLTC
-                push!(ltc_pop_list,agent_sample)
-            else
-                push!(household,agent_sample)
-            end
-        end
-        push!(households,household)
-    end
-
-    ltc_pop = length(ltc_pop_list)
-    println("ltc: $ltc_pop")
-    ltc_home_size = Int(ceil(ltc_pop/num_LTC))
-
-
-    if ltc_home_size != 0
-        ltc_homes = [ltc_pop_list[i:min(i + ltc_home_size - 1,ltc_pop)] for i = 1:ltc_home_size:ltc_pop]
-    else 
-        ltc_homes = Vector{Vector{AgentDemographic}}()
-    end
-
-    ltc_networks = map(LightGraphs.complete_graph,length.(ltc_homes))
-    household_networks = map(LightGraphs.complete_graph, household_sizes)
-
-    population_list = vcat(vcat(households...), vcat(ltc_homes...))
-    network = join(ThreadsX.reduce(join,household_networks; init = Graph()), ThreadsX.reduce(join,ltc_networks; init = Graph()))
-    index_vectors = [findall(x -> x == AgentDemographic(i), population_list) for i in 1:AgentDemographic.size+1]
-  
-    return (;
-        population_list,
-        network,
-        index_vectors
-    )
-end
-
-
-
-
-function get_graph()
-    return LightGraphs.grid([100,100]) 
-end
-
 function get_u_0(nodes)
     u_0 = fill(Susceptible,nodes)
     u_0[1] = Infected
@@ -71,15 +5,6 @@ function get_u_0(nodes)
     return u_0
 end
 
-
-const contact_distibution_matrix = [NegativeBinomial(2,0.5) for i in 1:AgentDemographic.size + 1, j in 1:AgentDemographic.size + 1]
-const contact_time_distribution_matrix = [Hypergeometric(5,2,4) for i in 1:AgentDemographic.size + 1, j in 1:AgentDemographic.size + 1]
-
-function annealed_contacts(rand_gen,agent,index_vectors) 
-    agent_contact_number_distributions = @view contact_distibution_matrix[Int(agent),:]
-    agent_contact_numbers = map(dst -> rand(rand_gen,dst),agent_contact_number_distributions)
-    return mapreduce(i -> sample(rand_gen,index_vectors[i],agent_contact_numbers[i]; replace = false), vcat, 1:length(agent_contact_numbers))
-end
 function contact_weight(rand_gen,p, agent_v,agent_w) 
     contact_time = rand(rand_gen,contact_time_distribution_matrix[Int(agent_v),Int(agent_w)])
     return 1 - (1-p)^contact_time
@@ -93,19 +18,10 @@ function agents_step!(u_next,u,population_list,graph,params,t,rand_gen,index_vec
         agent_demo = population_list[i]
         if agent_status == Susceptible
             for j in LightGraphs.neighbors(graph,i)
-                if u[j] == Infected && rand(rand_gen) < p #change
+                if u[j] == Infected && rand(rand_gen) < contact_weight(rand_gen,p, agent_demo,population_list[j])
                     u_next[i] = Infected
                 end
             end
-            for j in annealed_contacts(rand_gen,agent_demo,index_vectors)
-                if u[i] == Infected  && j != i
-                    contact_demo = population_list[j]
-                    p_n = contact_weight(rand_gen,p,agent_demo,contact_demo)
-                    if rand(rand_gen) < p_n
-                        u_next[i] = Infected
-                    end
-                end
-            end
         elseif agent_status == Infected
             if rand(rand_gen) < recovery_rate
                 u_next[i] = Recovered
@@ -113,23 +29,22 @@ function agents_step!(u_next,u,population_list,graph,params,t,rand_gen,index_vec
         end
     end
 end
+
+
 function solve!(u_0,graph,params,steps,rand_gen,population_list,index_vectors)
     solution = vcat([u_0], [fill(Susceptible,length(u_0)) for i in 1:steps])
+    graphs = [graph]
     for t in 1:steps
-        agents_step!(solution[t+1],solution[t],population_list,graph,params,t,rand_gen,index_vectors)
+        graph_t = deepcopy(graph)
+        generate_mixing_graph!(rand_gen,graph_t,population_list,index_vectors,workschool_mixing_matrix) 
+        generate_mixing_graph!(rand_gen,graph_t,population_list,index_vectors,rest_mixing_matrix) 
+        push!(graphs,graph_t)
+        agents_step!(solution[t+1],solution[t],population_list,graph_t,params,t,rand_gen,index_vectors)
     end
    
-    return solution
+    return solution,graphs
 end
-# function solve_2!(u_0,graph,params,steps,rand_gen)
-#     S_vec = findall(x -> x == Susceptible, u_0)
-#     solution = vcat((u_0,S_vec), [(fill(Susceptible,length(u_0)), deepcopy(S_vec)) for i in 1:steps])
-#     for t in 1:steps
-#         agents_step_2!(solution[t+1],solution[t],graph,params,t,rand_gen)
-#     end
-   
-#     return solution
-# end
+
 function get_parameters()
     params = (
         p = 0.9,

src/plotting.jl

 
-function plot_network(g,solution,weights) #it's shockingly hard to plot animated networks in julia so this function does that...
-    a = LightGraphs.adjacency_matrix(g) 
-    network_pts = SFDP.layout(a,SFDP.Point2f0;tol=0.1,C=1,K=1,iterations=10) # generate 2D layout
-    # display(network_pts)
+function plot_edges!(plt,g,network_pts)
     xpts = [i[1] for i in network_pts]
     ypts = [i[2] for i in network_pts]
-    plt = plot()
     for edge in edges(g)
         e = [src(edge), dst(edge)]
         plot!(plt,xpts[e],ypts[e], color = :black)
     end
+end
+
+function plot_nodes!(plt,network_pts,u_t)
+    for (i,var) in enumerate(instances(AgentStatus))
+        indices_of_var = findall(x -> x == var, u_t)
+        xpts = [i[1] for i in network_pts]
+        ypts = [i[2] for i in network_pts]
+        scatter!(plt,xpts[indices_of_var],ypts[indices_of_var]; color = color_palette[i], legend = false,markersize = 5)
+    end
+end
+function make_layout(g)
+    a = LightGraphs.adjacency_matrix(g) 
+    # network_pts = Stress.layout(a,Stress.Point2f0) # generate 2D layout
+    network_pts = SFDP.layout(a,SFDP.Point2f0,tol=0.1,C=1,K=1,iterations=10)
+    return network_pts
+end
+
+function plot_model(g,solution,population) 
+    plt = plot()
+    network_pts = make_layout(g)
+    plot_edges!(plt,g,network_pts)  
     anim = Animation()
     for (t,u_t) in enumerate(solution)
         frame_t = deepcopy(plt)
-        for (i,var) in enumerate(instances(AgentStatus))
-            indices_of_var = findall(x -> x == var, u_t)
-            scatter!(frame_t,xpts[indices_of_var],ypts[indices_of_var]; color = color_palette[i], legend = false,markersize = 5)
-        end
+        plot_nodes!(frame_t,network_pts,u_t)
         plot!(frame_t;title = "t = $t")
         frame(anim,frame_t)
     end
         
     gif(anim, "graphplot.gif", fps = 1.5)
+end
+function plot_model(g::AbstractVector,solution,population) 
+plt = plot()
+network_pts = make_layout(g[1])
+anim = Animation()
+for (t,(u_t,g_t)) in enumerate(zip(solution,g))
+    frame_t = plot()
+    plot_nodes!(frame_t,network_pts,u_t)
+    plot_edges!(frame_t,g_t,network_pts)  
+    plot!(frame_t;title = "t = $t")
+    frame(anim,frame_t)
+end
+    
+gif(anim, "graphplot.gif", fps = 1.5)
 end
\ No newline at end of file

src/types.jl

 #Agent type definitions
 
-
 @enum AgentDemographic begin
     Young = 1
     Adult
-    AdultHCW
     Elderly
+    AdultHCW
     ElderlyLTC
 end
 @enum AgentStatus begin
@@ -19,6 +18,29 @@ struct AgentList
     status::Vector{AgentStatus}
 end
 
+struct ZWDist{BaseDistType,T} <: Sampleable{Univariate,Discrete}
+    α::T
+    base_dist::BaseDistType
+end
+function Base.rand(rng::AbstractRNG, s::ZWDist)
+    if Base.rand(rng) < s.α
+        return 0
+    else
+        return Base.rand(rng,s.base_dist)
+    end
+end