Commit 6cfb5d27 authored by Peter Jentsch's avatar Peter Jentsch
Browse files

optionally record degrees

parent ac3f9efc
...@@ -34,7 +34,7 @@ end ...@@ -34,7 +34,7 @@ end
Run the model with given parameter tuple and output recorder. See `get_parameters` for list of parameters. See `output.jl` for the list of recorders. Currently just 'DebugRecorder`. Run the model with given parameter tuple and output recorder. See `get_parameters` for list of parameters. See `output.jl` for the list of recorders. Currently just 'DebugRecorder`.
""" """
function abm(parameters) function abm(parameters)
model_sol = ModelSolution(parameters.sim_length,parameters,parameters.num_households) model_sol = ModelSolution(parameters.sim_length,parameters,parameters.num_households; record_degrees = false)
solve!(model_sol) solve!(model_sol)
return model_sol return model_sol
end end
...@@ -78,7 +78,7 @@ mutable struct ModelSolution{T,InfNet,SocNet,WSMixingDist,RestMixingDist,Recorde ...@@ -78,7 +78,7 @@ mutable struct ModelSolution{T,InfNet,SocNet,WSMixingDist,RestMixingDist,Recorde
rest_matrix_tuple::RestMixingDist rest_matrix_tuple::RestMixingDist
immunization_countdown::Vector{Int} immunization_countdown::Vector{Int}
output_data::RecorderType output_data::RecorderType
@views function ModelSolution(sim_length,params::T,num_households) where T @views function ModelSolution(sim_length,params::T,num_households; record_degrees = false) where T
demographics,base_network,index_vectors = generate_population(num_households) demographics,base_network,index_vectors = generate_population(num_households)
nodes = length(demographics) nodes = length(demographics)
pop_sizes = length.(index_vectors) pop_sizes = length.(index_vectors)
...@@ -105,7 +105,8 @@ mutable struct ModelSolution{T,InfNet,SocNet,WSMixingDist,RestMixingDist,Recorde ...@@ -105,7 +105,8 @@ mutable struct ModelSolution{T,InfNet,SocNet,WSMixingDist,RestMixingDist,Recorde
status_totals = [count(==(AgentStatus(i)), u_0_inf) for i in 1:AgentStatus.size] status_totals = [count(==(AgentStatus(i)), u_0_inf) for i in 1:AgentStatus.size]
immunization_countdown = fill(-1, nodes) #immunization countdown is negative if not counting down immunization_countdown = fill(-1, nodes) #immunization countdown is negative if not counting down
output_data = Recorder(0,sim_length) output_data = Recorder(0,sim_length; record_degrees)
return new{T,typeof(infected_mixing_graph), return new{T,typeof(infected_mixing_graph),
typeof(soc_mixing_graph),typeof(ws_matrix_tuple), typeof(soc_mixing_graph),typeof(ws_matrix_tuple),
typeof(rest_matrix_tuple), typeof(output_data)}( typeof(rest_matrix_tuple), typeof(output_data)}(
......
...@@ -6,6 +6,7 @@ import OnlineStats.fit! ...@@ -6,6 +6,7 @@ import OnlineStats.fit!
Recorder should store everything we might want to know about the model output. Recorder should store everything we might want to know about the model output.
""" """
struct Recorder{ElType,ArrT1,ArrT2,ArrT3,StatAccumulator} struct Recorder{ElType,ArrT1,ArrT2,ArrT3,StatAccumulator}
recorded_status_totals::ArrT1 recorded_status_totals::ArrT1
daily_cases_by_age::ArrT3 daily_cases_by_age::ArrT3
total_vaccinators::ArrT2 total_vaccinators::ArrT2
...@@ -18,11 +19,13 @@ struct Recorder{ElType,ArrT1,ArrT2,ArrT3,StatAccumulator} ...@@ -18,11 +19,13 @@ struct Recorder{ElType,ArrT1,ArrT2,ArrT3,StatAccumulator}
mean_time_since_last_notification::ArrT2 mean_time_since_last_notification::ArrT2
daily_immunized_by_age::ArrT3 daily_immunized_by_age::ArrT3
daily_unvac_cases_by_age::ArrT3 daily_unvac_cases_by_age::ArrT3
record_degrees_flag::Bool
avg_weighted_degree_of_vaccinators::Vector{StatAccumulator} avg_weighted_degree_of_vaccinators::Vector{StatAccumulator}
avg_weighted_degree::Vector{StatAccumulator} avg_weighted_degree::Vector{StatAccumulator}
function Recorder(val::T,sim_length) where T function Recorder(val::T,sim_length; record_degrees = false) where T
totals = [copy(val) for i in 1:4, j in 1:sim_length] totals = [copy(val) for i in 1:4, j in 1:sim_length]
state_totals = @LArray totals (S = (1,:),I = (2,:),R = (3,:), V = (4,:)) state_totals = @LArray totals (S = (1,:),I = (2,:),R = (3,:), V = (4,:))
total_vaccinators = [copy(val) for j in 1:sim_length] total_vaccinators = [copy(val) for j in 1:sim_length]
...@@ -41,7 +44,7 @@ struct Recorder{ElType,ArrT1,ArrT2,ArrT3,StatAccumulator} ...@@ -41,7 +44,7 @@ struct Recorder{ElType,ArrT1,ArrT2,ArrT3,StatAccumulator}
final_size_by_age = [copy(val) for i in 1:3] final_size_by_age = [copy(val) for i in 1:3]
avg_weighted_degree_of_vaccinators = [Variance() for _ in 1:3] avg_weighted_degree_of_vaccinators = [Variance() for _ in 1:3]
avg_weighted_degree = Variance() avg_weighted_degree = [Variance() for _ in 1:3]
return new{T,typeof(state_totals),typeof(total_vaccinators),typeof(daily_immunized_by_age),eltype(avg_weighted_degree)}( return new{T,typeof(state_totals),typeof(total_vaccinators),typeof(daily_immunized_by_age),eltype(avg_weighted_degree)}(
state_totals, state_totals,
...@@ -56,6 +59,7 @@ struct Recorder{ElType,ArrT1,ArrT2,ArrT3,StatAccumulator} ...@@ -56,6 +59,7 @@ struct Recorder{ElType,ArrT1,ArrT2,ArrT3,StatAccumulator}
mean_time_since_last_notification, mean_time_since_last_notification,
daily_immunized_by_age, daily_immunized_by_age,
daily_unvac_by_age, daily_unvac_by_age,
record_degrees,
avg_weighted_degree_of_vaccinators, avg_weighted_degree_of_vaccinators,
avg_weighted_degree avg_weighted_degree
) )
......
...@@ -84,17 +84,18 @@ Base.@propagate_inbounds @views function update_infection_state!(t,modelsol; rec ...@@ -84,17 +84,18 @@ Base.@propagate_inbounds @views function update_infection_state!(t,modelsol; rec
agent_transition!(i, Infected,Recovered) agent_transition!(i, Infected,Recovered)
end end
end end
weighted_degree_of_i::Int = record_degrees ? weighted_degree(t,i,inf_network) : 0
weighted_degree_of_i::Int = output_data.record_degrees_flag ? weighted_degree(t,i,inf_network) : 0
if immunization_countdown[i] == 0 if immunization_countdown[i] == 0
output_data.daily_immunized_by_age[Int(agent_demo),t] += 1 output_data.daily_immunized_by_age[Int(agent_demo),t] += 1
fit!(output_data.avg_weighted_degree_of_vaccinators[Int(agent_demo)],weighted_degree_of_i) fit!(output_data.avg_weighted_degree_of_vaccinators[Int(agent_demo)],weighted_degree_of_i)
agent_transition!(i, Susceptible,Immunized) agent_transition!(i, Susceptible,Immunized)
elseif immunization_countdown[i]>0 elseif immunization_countdown[i]>0
fit!(output_data.avg_weighted_degree,weighted_degree_of_i) fit!(output_data.avg_weighted_degree[Int(agent_demo)],weighted_degree_of_i)
immunization_countdown[i] -= 1 immunization_countdown[i] -= 1
else else
fit!(output_data.avg_weighted_degree,weighted_degree_of_i) fit!(output_data.avg_weighted_degree[Int(agent_demo)],weighted_degree_of_i)
end end
end end
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