convert to alias sampler

src/params/progression.jl

@@ -5,7 +5,7 @@ struct ProgressionParams
   dist_hospitalization_time::Gamma{Float64}
   dist_mild_recovery_time::Uniform{Float64}
   dist_death_time::LogNormal{Float64}
-  dist_severity_by_age::Vector{Vector{Distributions.Categorical{P, Ps} where {P<:Real, Ps<:AbstractVector{P}}}}
+  dist_severity_by_age::Matrix{AliasSampler}
   effectiveness_table::Matrix{Float64}
   hospitalization_time_ratio::Float64
 end
@@ -40,11 +40,11 @@ const hospitalization_time_sampler = AliasSampler(Int, hospitalization_time_prob
 const age_hospitalization_thresholds = Int[0, 40, 50, 60, 70, 80]
 
 
-function sample_severity(rng::AbstractRNG, age::Real, gender::Bool, immunity::ImmunityState, effectiveness_table::Matrix{Float64}, dist_severity_by_age::Vector{Vector{Distributions.Categorical{P, Ps} where {P<:Real, Ps<:AbstractVector{P}}}})
+function sample_severity(rng::AbstractRNG, age::Real, gender::Bool, immunity::ImmunityState, effectiveness_table::Matrix{Float64}, dist_severity_by_age::Matrix{AliasSampler})
   @assert age >= 0 "age should be non-negative"
-  idx = gender + 1 |> Int32
-  dist = dist_severity_by_age[idx][age < max_age_hosp ? age + 1 : max_age_hosp]
-  severity = rand(rng, dist) |> Severity
+  gender_int = gender + 1 |> UInt8
+  dist = dist_severity_by_age[min(age + 1, max_age_hosp), gender_int]
+  severity = asample(dist) |> Severity
   severity_int = severity |> UInt8
   #reduction severity for immunited subject with some probability
   if size(effectiveness_table,1) > 1
@@ -197,11 +197,16 @@ end
 
 function make_dist_severity_by_age(hospitalization_men_probs::Vector{T}, hospitalization_women_probs::Vector{T}, hospitalization_multiplier::Float64, death_multiplier::Float64) where T<:Real
   n = length(hospitalization_women_probs)
-  dist_severity_by_age = [Vector{Categorical}(undef, n),Vector{Categorical}(undef, n)]
+  hosp_man = hospitalization_men_probs * hospitalization_multiplier
+  hosp_woman = hospitalization_women_probs * hospitalization_multiplier
+  men_sampler = AliasSampler[]
+  women_sampler = AliasSampler[]
   for idx in 1:n
     group_ids = agegroup(age_hospitalization_thresholds, idx-1)
-    dist_severity_by_age[1][idx] = Categorical([0,  1-(hospitalization_men_probs[idx]*hospitalization_multiplier),  hospitalization_multiplier*hospitalization_men_probs[idx]*(1-critical_probs[group_ids]*death_multiplier),  hospitalization_multiplier*hospitalization_men_probs[idx]*critical_probs[group_ids]*death_multiplier])
-    dist_severity_by_age[2][idx] = Categorical([0,  1-(hospitalization_women_probs[idx]*hospitalization_multiplier),  hospitalization_multiplier*hospitalization_women_probs[idx]*(1-critical_probs[group_ids]*death_multiplier),  hospitalization_multiplier*hospitalization_women_probs[idx]*critical_probs[group_ids]*death_multiplier])
+    hosp_prob = [0, 1-hosp_man[idx], hosp_man[idx]*(1-critical_probs[group_ids]*death_multiplier), hosp_man[idx]critical_probs[group_ids]*death_multiplier] 
+    push!(men_sampler, AliasSampler(UInt8,hosp_prob))
+    hosp_prob = [0, 1-hosp_woman[idx], hosp_woman[idx]*(1-critical_probs[group_ids]*death_multiplier), hosp_woman[idx]critical_probs[group_ids]*death_multiplier] 
+    push!(women_sampler, AliasSampler(UInt8,hosp_prob))
   end
-  dist_severity_by_age
+  hcat(men_sampler,women_sampler)
 end