Commit c705d7f6 authored by Peter Jentsch's avatar Peter Jentsch
Browse files

small changes

parent 5e1ad5b9
......@@ -8,7 +8,9 @@ using DataFrames
using Distributions
using StatsBase
using Dates
using ThreadsX
using DelimitedFiles
using NamedTupleTools
# using CUDA
using NetworkLayout:SFDP
......@@ -37,17 +39,19 @@ default(framestyle = :box)
using BenchmarkTools
function main()
graph = get_graph()
u_0 = get_u_0(nv(graph))
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
r = Xoroshiro128Plus()
# @btime sol1 = solve!($u_0,$graph,$params,$steps,$r)
# display(sol1)
household_size_dist = Truncated(Poisson(2.60),0,Inf) #https://www.researchgate.net/publication/226081704_Household_size_and_the_Poisson_distribution
pop = generate_population(2,1,household_size_dist)
return pop
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)
# contact_time = contact_weight(r,0.1, Young,Elderly)
end
......
function generate_population(num_households,num_LTC,household_size_distribution)
household_sizes = rand(household_size_distribution,num_households)
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 = [ #find a better way to do this
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],
demographic_distribution[3] * 0.0
]
display(agent_distribution_in_households)
agents_demographics = map(household_sizes) do household_size
return AgentDemographic.(rand(Categorical(agent_distribution_in_households),Int(household_size)))
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
display(agents_demographics)
end
function get_graph(population)
num_households = 100
household_sizes = rand(household_size_distribution,num_households)
households = map(nodes -> LightGraphs.complete_graph(nodes),household_sizes)
num_households = 100
long_term_care_sizes = rand(household_size_distribution,num_households)
households = map(nodes -> LightGraphs.complete_graph(nodes),long_term_care_sizes)
ltc_pop = length(ltc_pop_list)
println("ltc: $ltc_pop")
ltc_home_size = Int(ceil(ltc_pop/num_LTC))
return LightGraphs.grid([100,100])
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
......@@ -43,27 +72,51 @@ function get_u_0(nodes)
end
function agents_step!(u_next,u,graph,params,t,rand_gen)
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
end
function agents_step!(u_next,u,population_list,graph,params,t,rand_gen,index_vectors)
@unpack p, recovery_rate = params
u_next .= u
@inbounds for (v,agent) in enumerate(u)
if agent == Susceptible
for w in LightGraphs.neighbors(graph,v)
if u[w] == Infected && rand(rand_gen) < p
u_next[v] = Infected
for i in 1:length(u)
agent_status = u[i]
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
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 == Infected
elseif agent_status == Infected
if rand(rand_gen) < recovery_rate
u_next[v] = Recovered
u_next[i] = Recovered
end
end
end
end
function solve!(u_0,graph,params,steps,rand_gen)
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])
for t in 1:steps
agents_step!(solution[t+1],solution[t],graph,params,t,rand_gen)
agents_step!(solution[t+1],solution[t],population_list,graph,params,t,rand_gen,index_vectors)
end
return solution
......
......@@ -9,7 +9,7 @@
ElderlyLTC
end
@enum AgentStatus begin
Susceptible
Susceptible = 1
Infected
Recovered
end
......@@ -34,10 +34,12 @@ function Base.show(io::IO, status::AgentDemographic)
print(io, "<20")
elseif status == Adult
print(io, "20-60")
elseif status == Adult
print(io, "20-60,HCW")
elseif status == AdultHCW
print(io, "20-60 HCW")
elseif status == Elderly
print(io, ">60")
elseif status == ElderlyLTC
print(io, ">60 LTC")
end
end
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment