ridzikowski/new-immunity

benchmark/benchmark.jl

@@ -42,7 +42,10 @@ rng = Random.MersenneTwister(0);
 
   forward_detection_delay=1.75,
   backward_detection_delay=1.75,
-  testing_time=3.0
+  testing_time=3.0,
+  hospitalization_time_ratio = 0.5,
+  hospitalization_multiplier=1.5,
+  death_multiplier = 0.5
 );
 
 mutable struct Callback
@@ -59,6 +62,13 @@ end
 @time MocosSim.reset!(state, MersenneTwister(0))
 @time MocosSim.InstantOutsideCases(;num_infections=100)(state, params)
 
+immunization_thresholds = Int32[0, 18, 40]
+immunization_table = Float32[0.25 0.00;
+                             0.50 0.25;
+                             0.75 0.50]
+immunization_previously_infected  = Float32[0.24, 0.24, 0.24]                          
+MocosSim.immunize!(state, params, immunization_thresholds, immunization_table,immunization_previously_infected)
+
 @info "warm-up"
 cb = Callback(50, 0)
 @time MocosSim.simulate!(state, params, cb)

src/enums.jl

@@ -33,8 +33,10 @@ const NUM_STRAINS = length(instances(StrainKind)) - 1
   NullImmunity = 0
   NoImmunity
   NaturalImmunity
-  VecVacImmunity
-  MRNAVacImmunity
+  VacImmunity
+  VacNaturalImmunity
+  BoostImmunity
+  BoostNaturalImmunity
 end
 
 const NUM_IMMUNITIES = length(instances(ImmunityState)) - 1

src/event.jl

@@ -64,6 +64,7 @@ struct Event
   Event(::Val{DetectionEvent}, time::Real, subject::Integer, detectionkind::DetectionKind) = new(time, subject, 0, DetectionEvent, UInt8(detectionkind), NullStrain)
   Event(::Val{ImmunizationEvent}, time::Real, subject::Integer, new_immunity::ImmunityState) = new(time, subject, 0, ImmunizationEvent, UInt8(new_immunity), NullStrain)
   Event(::Val{ScreeningEvent}, time::Real) = new(time, 0, 0, ScreeningEvent, 0, NullStrain)
+  Event(::Val{BecomeInfectiousEvent}, time::Real, subject::Integer, strain::StrainKind) = new(time, subject, 0, BecomeInfectiousEvent , 0, strain)
 end
 
 time(event::Event) = event.time

src/event_execution.jl

@@ -68,7 +68,7 @@ function execute!(::Val{OutsideInfectionEvent}, state::SimState, params::SimPara
 
   incubation_time = progression.incubation_time
   @assert !ismissing(incubation_time)
-  event = Event(Val(BecomeInfectiousEvent), time(event) + incubation_time, subject(event))
+  event = Event(Val(BecomeInfectiousEvent), time(event) + incubation_time, subject(event), event.strain)
   push!(state.queue, event)
   return true
 end
@@ -119,7 +119,7 @@ function execute!(::Val{TransmissionEvent}, state::SimState, params::SimParams,
   incubation_time = progression.incubation_time
   @assert !ismissing(incubation_time)
   push!(state.queue,
-    Event(Val(BecomeInfectiousEvent), time(event) + incubation_time, subject(event))
+    Event(Val(BecomeInfectiousEvent), time(event) + incubation_time, subject(event), event.strain)
   )
   return true
 end
@@ -137,6 +137,7 @@ function execute!(::Val{BecomeInfectiousEvent}, state::SimState, params::SimPara
   progression = progressionof(state, subject_id)
 
   severity = progression.severity
+  setstrain!(state, subject_id, event.strain)
 
   infected_time = time(event) - progression.incubation_time
   if Asymptomatic == severity
@@ -354,7 +355,7 @@ function execute!(::Val{QuarantinedEvent}, state::SimState, params::SimParams, e
   end
 
   if extension(event)
-    @assert Undetected != detected(state, subject_id)
+    #@assert Undetected != detected(state, subject_id)
     setfreedom!(state, subject_id, HomeQuarantine)
   elseif Free == freedom_state
     @assert !is_already_quarantined "quarantined cannot be free, but $subject_id is"
@@ -410,7 +411,7 @@ end
 
 function execute!(::Val{ImmunizationEvent}, state::SimState, params::SimParams, event::Event)::Bool
   subject_id = subject(event)
-  new_immunity = immunitystate(event)
+  new_immunity = immunitykind(event)
   setimmunity!(state, subject_id, new_immunity)
   return true
 end
@@ -435,7 +436,9 @@ function quarantinehousehold!(state::SimState, params::SimParams, subject_id::In
     if !include_subject && (member == subject_id)
       continue
     end
-
+    if rand(state.rng) > params.household_params.quarantine_prob && member != subject_id
+      continue
+    end
     member_freedom = freedom(state, member)
 
     if (Hospitalized == member_freedom) || (Released == member_freedom)
@@ -449,8 +452,11 @@ end
 
 function tracehousehold!(state::SimState, params::SimParams, subject_id::Integer; trace_household_connections::Bool)
   for member in householdof(params, subject_id)
-     backtrace!(state, params, member, trace_household_connections=trace_household_connections)
-     forwardtrace!(state, params, member, trace_household_connections=trace_household_connections)
+    if rand(state.rng) > params.household_params.trace_prob && member != subject_id
+      continue
+    end
+    backtrace!(state, params, member, trace_household_connections=trace_household_connections)
+    forwardtrace!(state, params, member, trace_household_connections=trace_household_connections)
   end
   nothing
 end

src/params/households.jl

@@ -1,3 +1,8 @@
+Base.@kwdef struct HouseholdParams
+  quarantine_prob::Float64 = 0.5
+  trace_prob::Float64 = 0.5
+end
+
 function make_household_ptrs!(
     ptrs::AbstractVector{Tuple{Ti,Ti}},
     household_indices::AbstractVector{T} where T<:Integer

src/params/immunity.jl

@@ -9,18 +9,20 @@ function make_infectivity_table(;base_multiplier::Real=1.0, british_multiplier::
   immunity
 end
 
-function make_susceptibility_table()
+function make_susceptibility_table(;delta_susceptibility::Vector{T}, omicron_susceptibility::Vector{T}) where T <: Real
   #each column is distinct StrainKind
   #each row is distinct ImmunityState
 
   table = @SMatrix [
-    #ChineseStrain  BritishStrain DeltaStrain OmicronStrain
-    1.00            1.00          1.00        1.00;         # NoImmunity
-    0.01            0.01          0.10        0.60;         # NaturalImmunity
-    0.10            0.10          0.50        0.75;         # VecVacImmunity
-    0.03            0.10          0.30        0.75;         # MRNAVacImmunity
+    #ChineseStrain  BritishStrain DeltaStrain
+    1.00   1.00;
+    0.01   0.01;
+    0.10   0.10;
+    0.10   0.10;
+    0.03   0.10;
+    0.03   0.10;
   ]
-
+  table = hcat(table, delta_susceptibility, omicron_susceptibility)
   @assert all( 0 .<= table .<= 1)
   table
 end
@@ -33,6 +35,33 @@ end
 
 straininfectivity(table::StrainInfectivityTable, strain::StrainKind) = table[UInt(strain)]
 susceptibility(table::StrainSusceptibilityTable, immunity::ImmunityState, strain::StrainKind) = table[UInt(immunity), UInt(strain)]
+immunited(immunity::ImmunityState) = immunity != NoImmunity
+vaccinated(immunity::ImmunityState) = immunity == VacImmunity || immunity == VacNaturalImmunity || immunity == BoostImmunity || immunity == BoostNaturalImmunity
+boostered(immunity::ImmunityState) = immunity == BoostImmunity || immunity == BoostNaturalImmunity
+previnfected(immunity::ImmunityState) = immunity == NaturalImmunity || immunity == VacNaturalImmunity || BoostNaturalImmunity
+
+
+function immunize!(state::AbstractSimState, params::AbstractSimParams, immunization_thresholds::Vector{Int32},immunization_table::Matrix{Float32}, previously_infected::Vector{Float32})::Nothing
+  @info "Immunizing"
+  N = numindividuals(state)
+  @assert length(previously_infected) == 3
+  Samplers = AliasSampler[]
+  for gr in 1:length(immunization_thresholds)
+    no_immunity = 1-immunization_table[gr,1]
+    full_vaccinated = immunization_table[gr,1]-immunization_table[gr,2]
+    boostered = immunization_table[gr,2]
+    immunization_prob = [no_immunity*(1-previously_infected[1]), no_immunity*previously_infected[1], full_vaccinated*(1-previously_infected[2]), full_vaccinated*previously_infected[2], boostered*(1-previously_infected[3]), boostered*previously_infected[3]]
+    push!(Samplers, MocosSim.AliasSampler(UInt8,immunization_prob))
+  end
+  for id in 1:N
+    age = params.ages[id]
+    group_ids = agegroup(immunization_thresholds,age)
+    immunity_int = asample(Samplers[group_ids])
+    immunity = immunity_int |> ImmunityState
+    setimmunity!(state, id, immunity)
+  end
+  nothing
+end
 
 function immunize!(state::SimState, immunization::ImmunizationOrder; enqueue::Bool)::Nothing
   @info "Immunizing"

src/params/modulations.jl

@@ -24,6 +24,30 @@ function evalmodulation(f::TanhModulation, state::AbstractSimState, ::AbstractSi
   tanh_modulation(fear, f.loc, f.scale, 1.0, f.limit_value)
 end
 
+struct TwoTanhModulations <: InfectionModulation
+  weight_detected::Float64
+  weight_deaths::Float64
+  weight_days::Float64
+  loc1::Float64
+  scale1::Float64
+  limit_value1::Float64
+  loc2::Float64
+  scale2::Float64
+  limit_value2::Float64
+
+  TwoTanhModulations(;weight_detected::Real=0, weight_deaths::Real=0, weight_days::Real=0, loc1::Real=0, scale1::Real=0, limit_value1::Real=0, loc2::Real=0, scale2::Real=0, limit_value2::Real=0) =
+    0 <= limit_value2 <= limit_value1 <= 1 && loc1 < loc2 ? new(weight_detected, weight_deaths, weight_days, loc1, scale1, limit_value1, loc2, scale2, limit_value2) : error("initial_value1 and initial_value2 must be from 0 to 1 and initial_value2 must be not larger than initial_value1, got ($initial_value1, $initial_value2), also loc2 must be after loc1, got ($loc1, $loc2)")
+end
+
+function evalmodulation(f::TwoTanhModulations, state::AbstractSimState, ::AbstractSimParams)::Float64
+  num_days = time(state)
+  num_detected = numdetected(state)
+  num_deaths = numdead(state)
+  fear = num_detected * f.weight_detected + num_deaths * f.weight_deaths + num_days * f.weight_days
+  modulation1 = tanh_modulation(fear, f.loc1, f.scale1, 1.0, f.limit_value1)
+  tanh_modulation(fear, f.loc2, f.scale2, modulation1, f.limit_value2)
+end
+
 struct IncreasingTanhModulation <: InfectionModulation
   weight_detected::Float64
   weight_deaths::Float64
@@ -45,6 +69,31 @@ function evalmodulation(f::IncreasingTanhModulation, state::AbstractSimState, ::
   tanh_modulation(fear, f.loc, f.scale, f.initial_value, 1.0)
 end
 
+struct IncreasingTwoTanhModulations <: InfectionModulation
+  weight_detected::Float64
+  weight_deaths::Float64
+  weight_days::Float64
+  loc1::Float64
+  scale1::Float64
+  initial_value1::Float64
+  loc2::Float64
+  scale2::Float64
+  initial_value2::Float64
+
+  IncreasingTwoTanhModulations(;weight_detected::Real=0, weight_deaths::Real=0, weight_days::Real=0, loc1::Real=0, scale1::Real=0, initial_value1::Real=0, loc2::Real=0, scale2::Real=0, initial_value2::Real=0) =
+    0 <= initial_value1 <= initial_value2 <= 1 && loc1 < loc2 ? new(weight_detected, weight_deaths, weight_days, loc1, scale1, initial_value1, loc2, scale2, initial_value2) : error("initial_value1 and initial_value2 must be from 0 to 1 and initial_value2 must be not smaller than initial_value1, got ($initial_value1, $initial_value2), also loc2 must be after loc1, got ($loc1, $loc2)")
+end
+
+function evalmodulation(f::IncreasingTwoTanhModulations, state::AbstractSimState, ::AbstractSimParams)::Float64
+  num_days = time(state)
+  num_detected = numdetected(state)
+  num_deaths = numdead(state)
+
+  fear = num_detected * f.weight_detected + num_deaths * f.weight_deaths + num_days * f.weight_days
+  modulation1 = tanh_modulation(fear, f.loc1, f.scale1, f.initial_value1, f.initial_value2)
+  tanh_modulation(fear, f.loc2, f.scale2, modulation1, 1.0)
+end
+
 evalmodulation(modulation::Nothing, state::AbstractSimState, params::AbstractSimParams)::Float64 = 1.0
 
 function infectionsuccess(modulation::InfectionModulation, state::AbstractSimState, params::AbstractSimParams, event::Event)::Bool
@@ -77,7 +126,9 @@ end
 # This all to avoid using @eval and others
 const modulations = Dict{String, Type{T} where T}(
     "TanhModulation" => TanhModulation,
-    "IncreasingTanhModulation" => IncreasingTanhModulation
+    "TwoTanhModulations" => TwoTanhModulations,
+    "IncreasingTanhModulation" => IncreasingTanhModulation,
+    "IncreasingTwoTanhModulations" => IncreasingTwoTanhModulations
 )
 
 make_infection_modulation(::Nothing) = nothing