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

optionally record degrees

parent ac3f9efc
......@@ -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`.
"""
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)
return model_sol
end
......@@ -78,7 +78,7 @@ mutable struct ModelSolution{T,InfNet,SocNet,WSMixingDist,RestMixingDist,Recorde
rest_matrix_tuple::RestMixingDist
immunization_countdown::Vector{Int}
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)
nodes = length(demographics)
pop_sizes = length.(index_vectors)
......@@ -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]
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),
typeof(soc_mixing_graph),typeof(ws_matrix_tuple),
typeof(rest_matrix_tuple), typeof(output_data)}(
......
......@@ -6,6 +6,7 @@ import OnlineStats.fit!
Recorder should store everything we might want to know about the model output.
"""
struct Recorder{ElType,ArrT1,ArrT2,ArrT3,StatAccumulator}
recorded_status_totals::ArrT1
daily_cases_by_age::ArrT3
total_vaccinators::ArrT2
......@@ -18,11 +19,13 @@ struct Recorder{ElType,ArrT1,ArrT2,ArrT3,StatAccumulator}
mean_time_since_last_notification::ArrT2
daily_immunized_by_age::ArrT3
daily_unvac_cases_by_age::ArrT3
record_degrees_flag::Bool
avg_weighted_degree_of_vaccinators::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]
state_totals = @LArray totals (S = (1,:),I = (2,:),R = (3,:), V = (4,:))
total_vaccinators = [copy(val) for j in 1:sim_length]
......@@ -41,7 +44,7 @@ struct Recorder{ElType,ArrT1,ArrT2,ArrT3,StatAccumulator}
final_size_by_age = [copy(val) for i 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)}(
state_totals,
......@@ -56,6 +59,7 @@ struct Recorder{ElType,ArrT1,ArrT2,ArrT3,StatAccumulator}
mean_time_since_last_notification,
daily_immunized_by_age,
daily_unvac_by_age,
record_degrees,
avg_weighted_degree_of_vaccinators,
avg_weighted_degree
)
......
......@@ -84,17 +84,18 @@ Base.@propagate_inbounds @views function update_infection_state!(t,modelsol; rec
agent_transition!(i, Infected,Recovered)
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
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)
agent_transition!(i, Susceptible,Immunized)
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
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
......
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