From 0e32d50890e7191b8d105280f998d2b271e316fa Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 8 Mar 2024 11:28:26 +0000 Subject: [PATCH 01/57] first pass at AR model --- EpiAware/src/EpiLatentModels/randomwalk.jl | 57 ++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/EpiAware/src/EpiLatentModels/randomwalk.jl b/EpiAware/src/EpiLatentModels/randomwalk.jl index 5f47bd4b1..a4166bb47 100644 --- a/EpiAware/src/EpiLatentModels/randomwalk.jl +++ b/EpiAware/src/EpiLatentModels/randomwalk.jl @@ -1,3 +1,8 @@ +function generate_latent(latent_model::AbstractLatentModel, n) + @info "No concrete implementation for generate_latent is defined." + return nothing +end + struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel init_prior::D std_prior::S @@ -20,3 +25,55 @@ end end return rw, (; σ_RW, rw_init) end + +struct AR <: AbstractLatentModel + """A distribution, a vector of distributions, or a product of distributions representing the prior distributions of the damping factors.""" + damp_prior::Union{Distribution, Vector{<:Distribution}, Product} + + """A distribution representing the prior distribution of the variance.""" + var_prior::Distribution + + """A distribution, a vector of distributions, or a product of distributions representing the prior distributions of the initial values.""" + init_prior::Union{Distribution, Vector{<:Distribution}, Product} + + """ + The order of the AR process, determined by the length of `damp_prior`. + """ + p::Int + + function AR(damp_prior::Union{Distribution, Vector{<:Distribution}, Product}, + var_prior::Distribution, + init_prior::Union{Distribution, Vector{<:Distribution}, Product}) + p = length(damp_prior) + @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process" + return new(damp_prior, var_prior, init_prior, p) + end +end + +function default_ar_priors() + return (:damp_prior => [truncated(Normal(0.0, 0.05), 0.0, 1)], + :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), + :init_prior => Normal()) |> Dict +end + +@model function generate_latent(latent_model::AR, n) + p = latent_model.p + ϵ_t ~ MvNormal(ones(n - p)) + σ²_AR ~ latent_model.var_prior + ar_init ~ latent_model.init_prior + damp_AR ~ latent_model.damp_prior + σ_AR = sqrt(σ²_AR) + + @assert n>p "n must be longer than p" + + # Initialize the AR series with the initial values + ar = Vector{Float64}(undef, n) + ar[1:p] = ar_init + + # Generate the rest of the AR series + for t in (p + 1):n + ar[t] = damp_AR' * ar[(t - p):(t - 1)] + σ_AR * ϵ_t[t - p] + end + + return ar, (; σ_AR, ar_init, damp_AR) +end From 57d4fcf4d52cd73b973f29ff4209e6bbaa074d2c Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 8 Mar 2024 12:14:45 +0000 Subject: [PATCH 02/57] move files around and add differencing infrastructure --- EpiAware/src/EpiAware.jl | 2 - EpiAware/src/EpiLatentModels/randomwalk.jl | 54 +++++++++++++++++----- 2 files changed, 42 insertions(+), 14 deletions(-) diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl index 3d5a9cbbf..d382a44f9 100644 --- a/EpiAware/src/EpiAware.jl +++ b/EpiAware/src/EpiAware.jl @@ -74,5 +74,3 @@ export make_epi_aware include("docstrings.jl") include("make_epi_aware.jl") - -end diff --git a/EpiAware/src/EpiLatentModels/randomwalk.jl b/EpiAware/src/EpiLatentModels/randomwalk.jl index a4166bb47..d63920f22 100644 --- a/EpiAware/src/EpiLatentModels/randomwalk.jl +++ b/EpiAware/src/EpiLatentModels/randomwalk.jl @@ -1,8 +1,3 @@ -function generate_latent(latent_model::AbstractLatentModel, n) - @info "No concrete implementation for generate_latent is defined." - return nothing -end - struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel init_prior::D std_prior::S @@ -26,25 +21,30 @@ end return rw, (; σ_RW, rw_init) end +# Define the Priors type alias +const Priors = Union{Distribution, Vector{<:Distribution}, Product} + struct AR <: AbstractLatentModel - """A distribution, a vector of distributions, or a product of distributions representing the prior distributions of the damping factors.""" - damp_prior::Union{Distribution, Vector{<:Distribution}, Product} + """A distribution representing the prior distribution of the damping factors.""" + damp_prior::Priors """A distribution representing the prior distribution of the variance.""" var_prior::Distribution - """A distribution, a vector of distributions, or a product of distributions representing the prior distributions of the initial values.""" - init_prior::Union{Distribution, Vector{<:Distribution}, Product} + """A distribution representing the prior distribution of the initial values.""" + init_prior::Priors """ The order of the AR process, determined by the length of `damp_prior`. """ p::Int - function AR(damp_prior::Union{Distribution, Vector{<:Distribution}, Product}, - var_prior::Distribution, - init_prior::Union{Distribution, Vector{<:Distribution}, Product}) + function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors) p = length(damp_prior) + return new(damp_prior, var_prior, init_prior, p) + end + + function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors, p::Int) @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process" return new(damp_prior, var_prior, init_prior, p) end @@ -77,3 +77,33 @@ end return ar, (; σ_AR, ar_init, damp_AR) end + +struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel + model::T + d::Int + init_prior::Priors + + function DiffLatentModel(model::T, init_prior::Priors) + d = length(init_prior) + return new(model, d, init_prior) + end + + function DiffLatentModel(model::T, d::Int, init_prior::Priors) + @assert d>0 "d must be greater than 0" + @assert length(init_prior)==d "Length of init_prior must be equal to d" + return new(model, d, init_prior) + end +end + +@model function generate_latent(latent_model::DiffLatentModel, n) + d = latent_model.d + @assert n>d "n must be longer than d" + init_latent ~ latent_model.init_prior + + @submodel diff_latent, diff_latent_aux = generate_latent(latent_model.model, n - d) + + latent = vcat(init_latent, diff_latent) |> + cumsum + # Return the reconstructed series and the parameters + return latent, (; init_latent, diff_latent_aux...) +end From 47613aa4610e1425e5b32ed25c2c4f714c190276 Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 8 Mar 2024 12:17:49 +0000 Subject: [PATCH 03/57] remove scaffold comment --- EpiAware/src/EpiLatentModels/randomwalk.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/EpiAware/src/EpiLatentModels/randomwalk.jl b/EpiAware/src/EpiLatentModels/randomwalk.jl index d63920f22..65760cba0 100644 --- a/EpiAware/src/EpiLatentModels/randomwalk.jl +++ b/EpiAware/src/EpiLatentModels/randomwalk.jl @@ -104,6 +104,6 @@ end latent = vcat(init_latent, diff_latent) |> cumsum - # Return the reconstructed series and the parameters + return latent, (; init_latent, diff_latent_aux...) end From 105d8b4093a04558a2144195b74cf47297dff499 Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 8 Mar 2024 12:23:15 +0000 Subject: [PATCH 04/57] change arg order --- EpiAware/src/EpiLatentModels/randomwalk.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/EpiAware/src/EpiLatentModels/randomwalk.jl b/EpiAware/src/EpiLatentModels/randomwalk.jl index 65760cba0..98efa7484 100644 --- a/EpiAware/src/EpiLatentModels/randomwalk.jl +++ b/EpiAware/src/EpiLatentModels/randomwalk.jl @@ -80,18 +80,18 @@ end struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel model::T - d::Int init_prior::Priors + d::Int function DiffLatentModel(model::T, init_prior::Priors) d = length(init_prior) - return new(model, d, init_prior) + return new(model, init_prior, d) end function DiffLatentModel(model::T, d::Int, init_prior::Priors) @assert d>0 "d must be greater than 0" @assert length(init_prior)==d "Length of init_prior must be equal to d" - return new(model, d, init_prior) + return new(model, init_prior, d) end end From 34635c7fe5065f19135c68f4835bd880943c0fc1 Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 8 Mar 2024 12:41:45 +0000 Subject: [PATCH 05/57] clean up use of new --- EpiAware/src/EpiLatentModels/randomwalk.jl | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/EpiAware/src/EpiLatentModels/randomwalk.jl b/EpiAware/src/EpiLatentModels/randomwalk.jl index 98efa7484..ac2a13b62 100644 --- a/EpiAware/src/EpiLatentModels/randomwalk.jl +++ b/EpiAware/src/EpiLatentModels/randomwalk.jl @@ -41,7 +41,7 @@ struct AR <: AbstractLatentModel function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors) p = length(damp_prior) - return new(damp_prior, var_prior, init_prior, p) + return AR(damp_prior, var_prior, init_prior, p) end function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors, p::Int) @@ -83,15 +83,15 @@ struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel init_prior::Priors d::Int - function DiffLatentModel(model::T, init_prior::Priors) + function DiffLatentModel(model::AbstractModel, init_prior::Priors) d = length(init_prior) - return new(model, init_prior, d) + return DiffLatentModel(model, init_prior, d) # Add the missing type parameter to the new function call end - function DiffLatentModel(model::T, d::Int, init_prior::Priors) + function DiffLatentModel(model::AbstractModel, init_prior::Priors, d::Int) @assert d>0 "d must be greater than 0" @assert length(init_prior)==d "Length of init_prior must be equal to d" - return new(model, init_prior, d) + return new{T <: AbstractModel}(model, init_prior, d) # Add the missing type parameter to the new function call end end From 24b29b8bd052d5465eff7bc739c06c3ccb27a577 Mon Sep 17 00:00:00 2001 From: Sam Date: Mon, 11 Mar 2024 10:20:10 +0000 Subject: [PATCH 06/57] reorg based on main --- .../src/EpiLatentModels/autoregressive.jl | 53 +++++++++++ .../EpiLatentModels/differencedlatentmodel.jl | 29 +++++++ EpiAware/src/EpiLatentModels/randomwalk.jl | 87 ------------------- EpiAware/src/latentmodels/latentmodels.jl | 5 ++ 4 files changed, 87 insertions(+), 87 deletions(-) create mode 100644 EpiAware/src/EpiLatentModels/autoregressive.jl create mode 100644 EpiAware/src/EpiLatentModels/differencedlatentmodel.jl create mode 100644 EpiAware/src/latentmodels/latentmodels.jl diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl new file mode 100644 index 000000000..3c2fc7648 --- /dev/null +++ b/EpiAware/src/EpiLatentModels/autoregressive.jl @@ -0,0 +1,53 @@ +struct AR <: AbstractLatentModel + """A distribution representing the prior distribution of the damping factors.""" + damp_prior::Priors + + """A distribution representing the prior distribution of the variance.""" + var_prior::Distribution + + """A distribution representing the prior distribution of the initial values.""" + init_prior::Priors + + """ + The order of the AR process, determined by the length of `damp_prior`. + """ + p::Int + + function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors) + p = length(damp_prior) + return AR(damp_prior, var_prior, init_prior, p) + end + + function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors, p::Int) + @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process" + return new(damp_prior, var_prior, init_prior, p) + end +end + +function default_ar_priors() + return (:damp_prior => [truncated(Normal(0.0, 0.05), 0.0, 1)], + :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), + :init_prior => Normal()) |> Dict +end + +@model function generate_latent(latent_model::AR, n) + p = latent_model.p + ϵ_t ~ MvNormal(ones(n - p)) + σ²_AR ~ latent_model.var_prior + ar_init ~ latent_model.init_prior + damp_AR ~ latent_model.damp_prior + σ_AR = sqrt(σ²_AR) + + @assert n>p "n must be longer than p" + + # Initialize the AR series with the initial values + ar = Vector{Float64}(undef, n) + ar[1:p] = ar_init + + # Generate the rest of the AR series + for t in (p + 1):n + ar[t] = damp_AR' * ar[(t - p):(t - 1)] + σ_AR * ϵ_t[t - p] + end + + return ar, (; σ_AR, ar_init, damp_AR) +end diff --git a/EpiAware/src/EpiLatentModels/differencedlatentmodel.jl b/EpiAware/src/EpiLatentModels/differencedlatentmodel.jl new file mode 100644 index 000000000..1d38e6052 --- /dev/null +++ b/EpiAware/src/EpiLatentModels/differencedlatentmodel.jl @@ -0,0 +1,29 @@ +struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel + model::T + init_prior::Priors + d::Int + + function DiffLatentModel(model::AbstractModel, init_prior::Priors) + d = length(init_prior) + return DiffLatentModel(model, init_prior, d) # Add the missing type parameter to the new function call + end + + function DiffLatentModel(model::AbstractModel, init_prior::Priors, d::Int) + @assert d>0 "d must be greater than 0" + @assert length(init_prior)==d "Length of init_prior must be equal to d" + return new{T <: AbstractModel}(model, init_prior, d) # Add the missing type parameter to the new function call + end +end + +@model function generate_latent(latent_model::DiffLatentModel, n) + d = latent_model.d + @assert n>d "n must be longer than d" + init_latent ~ latent_model.init_prior + + @submodel diff_latent, diff_latent_aux = generate_latent(latent_model.model, n - d) + + latent = vcat(init_latent, diff_latent) |> + cumsum + + return latent, (; init_latent, diff_latent_aux...) +end diff --git a/EpiAware/src/EpiLatentModels/randomwalk.jl b/EpiAware/src/EpiLatentModels/randomwalk.jl index ac2a13b62..5f47bd4b1 100644 --- a/EpiAware/src/EpiLatentModels/randomwalk.jl +++ b/EpiAware/src/EpiLatentModels/randomwalk.jl @@ -20,90 +20,3 @@ end end return rw, (; σ_RW, rw_init) end - -# Define the Priors type alias -const Priors = Union{Distribution, Vector{<:Distribution}, Product} - -struct AR <: AbstractLatentModel - """A distribution representing the prior distribution of the damping factors.""" - damp_prior::Priors - - """A distribution representing the prior distribution of the variance.""" - var_prior::Distribution - - """A distribution representing the prior distribution of the initial values.""" - init_prior::Priors - - """ - The order of the AR process, determined by the length of `damp_prior`. - """ - p::Int - - function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors) - p = length(damp_prior) - return AR(damp_prior, var_prior, init_prior, p) - end - - function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors, p::Int) - @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process" - return new(damp_prior, var_prior, init_prior, p) - end -end - -function default_ar_priors() - return (:damp_prior => [truncated(Normal(0.0, 0.05), 0.0, 1)], - :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), - :init_prior => Normal()) |> Dict -end - -@model function generate_latent(latent_model::AR, n) - p = latent_model.p - ϵ_t ~ MvNormal(ones(n - p)) - σ²_AR ~ latent_model.var_prior - ar_init ~ latent_model.init_prior - damp_AR ~ latent_model.damp_prior - σ_AR = sqrt(σ²_AR) - - @assert n>p "n must be longer than p" - - # Initialize the AR series with the initial values - ar = Vector{Float64}(undef, n) - ar[1:p] = ar_init - - # Generate the rest of the AR series - for t in (p + 1):n - ar[t] = damp_AR' * ar[(t - p):(t - 1)] + σ_AR * ϵ_t[t - p] - end - - return ar, (; σ_AR, ar_init, damp_AR) -end - -struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel - model::T - init_prior::Priors - d::Int - - function DiffLatentModel(model::AbstractModel, init_prior::Priors) - d = length(init_prior) - return DiffLatentModel(model, init_prior, d) # Add the missing type parameter to the new function call - end - - function DiffLatentModel(model::AbstractModel, init_prior::Priors, d::Int) - @assert d>0 "d must be greater than 0" - @assert length(init_prior)==d "Length of init_prior must be equal to d" - return new{T <: AbstractModel}(model, init_prior, d) # Add the missing type parameter to the new function call - end -end - -@model function generate_latent(latent_model::DiffLatentModel, n) - d = latent_model.d - @assert n>d "n must be longer than d" - init_latent ~ latent_model.init_prior - - @submodel diff_latent, diff_latent_aux = generate_latent(latent_model.model, n - d) - - latent = vcat(init_latent, diff_latent) |> - cumsum - - return latent, (; init_latent, diff_latent_aux...) -end diff --git a/EpiAware/src/latentmodels/latentmodels.jl b/EpiAware/src/latentmodels/latentmodels.jl new file mode 100644 index 000000000..7a2757966 --- /dev/null +++ b/EpiAware/src/latentmodels/latentmodels.jl @@ -0,0 +1,5 @@ +# Define the Priors type alias +const Priors = Union{Distribution, Vector{<:Distribution}, Product} + +include("randomwalk.jl") +include("autoregressive.jl") From cb53aeaeef9ef4484944968f26c806888aa4dea6 Mon Sep 17 00:00:00 2001 From: Sam Date: Mon, 11 Mar 2024 11:57:24 +0000 Subject: [PATCH 07/57] improve testing --- EpiAware/src/EpiAware.jl | 6 ++ EpiAware/test/test_autoregressive.jl | 69 +++++++++++++++++++ ...st_latent-models.jl => test_randomwalk.jl} | 0 3 files changed, 75 insertions(+) create mode 100644 EpiAware/test/test_autoregressive.jl rename EpiAware/test/{test_latent-models.jl => test_randomwalk.jl} (100%) diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl index d382a44f9..1a76d0f71 100644 --- a/EpiAware/src/EpiAware.jl +++ b/EpiAware/src/EpiAware.jl @@ -36,9 +36,15 @@ include("EpiAwareBase/EpiAwareBase.jl") using .EpiAwareBase +<<<<<<< HEAD export AbstractModel, AbstractLatentModel, AbstractEpiModel, AbstractObservationModel, generate_latent, generate_latent_infs, generate_observations +======= +# Exported types +export EpiData, Renewal, ExpGrowthRate, DirectInfections, RandomWalk, + DelayObservations, AR, DiffLatentModel +>>>>>>> 9fb97d9 (improve testing) include("EpiAwareUtils/EpiAwareUtils.jl") using .EpiAwareUtils diff --git a/EpiAware/test/test_autoregressive.jl b/EpiAware/test/test_autoregressive.jl new file mode 100644 index 000000000..08de1d924 --- /dev/null +++ b/EpiAware/test/test_autoregressive.jl @@ -0,0 +1,69 @@ +@testitem "Testing default_ar_priors" begin + using Distributions + + @testset "damp_prior" begin + priors = EpiAware.default_ar_priors() + damp = rand(priors[:damp_prior][1]) + @test 0.0 <= damp <= 1.0 + end + + @testset "var_prior" begin + priors = EpiAware.default_ar_priors() + var_AR = rand(priors[:var_prior]) + @test var_AR >= 0.0 + end + + @testset "init_prior" begin + priors = EpiAware.default_ar_priors() + init_ar_value = rand(priors[:init_prior]) + @test typeof(init_ar_value) == Float64 + end +end + +@testitem "Testing AR constructor" begin + using Distributions + + damp_prior = [truncated(Normal(0.0, 0.05), 0.0, 1)] + var_prior = truncated(Normal(0.0, 0.05), 0.0, Inf) + init_prior = Normal() + ar_process = EpiAware.AR(damp_prior, var_prior, init_prior) + + @test ar_process.damp_prior == damp_prior + @test ar_process.var_prior == var_prior + @test ar_process.init_prior == init_prior + @test ar_process.p == 1 +end + +@testitem "Testing AR process against theoretical properties" begin + using DynamicPPL, Turing + using HypothesisTests: ExactOneSampleKSTest, pvalue + using Distributions + + ar_model = EpiAware.AR(EpiAware.default_ar_priors()[:damp_prior], + EpiAware.default_ar_priors()[:var_prior], + EpiAware.default_ar_priors()[:init_prior] + ) + n = 1000 + damp = [0.1] + σ²_AR = 1.0 + ar_init = [0.0] + + model = EpiAware.generate_latent(ar_model, n) + fixed_model = fix(model, (σ²_AR = σ²_AR, damp_AR = damp, ar_init = ar_init)) + + n_samples = 100 + samples = sample(fixed_model, Prior(), n_samples) |> + chn -> mapreduce(vcat, generated_quantities(fixed_model, chn)) do gen + gen[1] + end + + theoretical_mean = 0.0 + theoretical_var = σ²_AR / (1 - damp[1]^2) + + @test isapprox(mean(samples), theoretical_mean, atol = 0.1) + @test isapprox(var(samples), theoretical_var, atol = 0.2) + + ks_test_pval = ExactOneSampleKSTest( + samples, Normal(theoretical_mean, sqrt(theoretical_var))) |> pvalue + @test ks_test_pval > 1e-6 +end diff --git a/EpiAware/test/test_latent-models.jl b/EpiAware/test/test_randomwalk.jl similarity index 100% rename from EpiAware/test/test_latent-models.jl rename to EpiAware/test/test_randomwalk.jl From b54918ef9cd888c1fb17be413034215aa61c6435 Mon Sep 17 00:00:00 2001 From: Sam Date: Mon, 11 Mar 2024 13:21:35 +0000 Subject: [PATCH 08/57] catching out of date export --- EpiAware/src/EpiAware.jl | 6 ------ 1 file changed, 6 deletions(-) diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl index 1a76d0f71..d382a44f9 100644 --- a/EpiAware/src/EpiAware.jl +++ b/EpiAware/src/EpiAware.jl @@ -36,15 +36,9 @@ include("EpiAwareBase/EpiAwareBase.jl") using .EpiAwareBase -<<<<<<< HEAD export AbstractModel, AbstractLatentModel, AbstractEpiModel, AbstractObservationModel, generate_latent, generate_latent_infs, generate_observations -======= -# Exported types -export EpiData, Renewal, ExpGrowthRate, DirectInfections, RandomWalk, - DelayObservations, AR, DiffLatentModel ->>>>>>> 9fb97d9 (improve testing) include("EpiAwareUtils/EpiAwareUtils.jl") using .EpiAwareUtils From 4cd830779eb76299376d1b2405f94b7071369d29 Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 12 Mar 2024 17:17:44 +0000 Subject: [PATCH 09/57] update to use parametric types and use @kwdef to simplify construction --- .../src/EpiLatentModels/autoregressive.jl | 28 ++++--------------- 1 file changed, 6 insertions(+), 22 deletions(-) diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl index 3c2fc7648..da19fd54c 100644 --- a/EpiAware/src/EpiLatentModels/autoregressive.jl +++ b/EpiAware/src/EpiLatentModels/autoregressive.jl @@ -1,33 +1,17 @@ -struct AR <: AbstractLatentModel +@kwdef struct AR{D <: Priors, S <: Distribution, I <: Priors, P <: Int} """A distribution representing the prior distribution of the damping factors.""" - damp_prior::Priors + damp_prior::D = [truncated(Normal(0.0, 0.05), 0.0, 1)] """A distribution representing the prior distribution of the variance.""" - var_prior::Distribution + var_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf) """A distribution representing the prior distribution of the initial values.""" - init_prior::Priors + init_prior::I = Normal() """ The order of the AR process, determined by the length of `damp_prior`. """ - p::Int - - function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors) - p = length(damp_prior) - return AR(damp_prior, var_prior, init_prior, p) - end - - function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors, p::Int) - @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process" - return new(damp_prior, var_prior, init_prior, p) - end -end - -function default_ar_priors() - return (:damp_prior => [truncated(Normal(0.0, 0.05), 0.0, 1)], - :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), - :init_prior => Normal()) |> Dict + p::P = length(damp_prior) end @model function generate_latent(latent_model::AR, n) @@ -38,7 +22,7 @@ end damp_AR ~ latent_model.damp_prior σ_AR = sqrt(σ²_AR) - @assert n>p "n must be longer than p" + @assert n>p "n must be longer than latent_model.p" # Initialize the AR series with the initial values ar = Vector{Float64}(undef, n) From 3142f200d27ece30860a5a756e9da93f0172a1fe Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 12 Mar 2024 18:22:17 +0000 Subject: [PATCH 10/57] update to use parameters.jl and drop p calc in struct as don't work with having defaults --- EpiAware/src/EpiAware.jl | 2 ++ .../src/EpiLatentModels/autoregressive.jl | 24 ++++++------------- 2 files changed, 9 insertions(+), 17 deletions(-) diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl index d382a44f9..3d5a9cbbf 100644 --- a/EpiAware/src/EpiAware.jl +++ b/EpiAware/src/EpiAware.jl @@ -74,3 +74,5 @@ export make_epi_aware include("docstrings.jl") include("make_epi_aware.jl") + +end diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl index da19fd54c..075be9c6c 100644 --- a/EpiAware/src/EpiLatentModels/autoregressive.jl +++ b/EpiAware/src/EpiLatentModels/autoregressive.jl @@ -1,28 +1,18 @@ -@kwdef struct AR{D <: Priors, S <: Distribution, I <: Priors, P <: Int} - """A distribution representing the prior distribution of the damping factors.""" - damp_prior::D = [truncated(Normal(0.0, 0.05), 0.0, 1)] - - """A distribution representing the prior distribution of the variance.""" - var_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf) - - """A distribution representing the prior distribution of the initial values.""" +@with_kw struct AR{D <: Priors, S <: Distribution, I <: Priors} <: AbstractLatentModel + damp_prior::D = truncated(Normal(0.0, 0.05), 0.0, 1) + std_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf) init_prior::I = Normal() - - """ - The order of the AR process, determined by the length of `damp_prior`. - """ - p::P = length(damp_prior) + @assert length(damp_prior)==length(init_prior) "damp_prior and init_prior must have the same length" end @model function generate_latent(latent_model::AR, n) - p = latent_model.p + p = length(damp_prior) ϵ_t ~ MvNormal(ones(n - p)) - σ²_AR ~ latent_model.var_prior + σ_AR ~ latent_model.std_prior ar_init ~ latent_model.init_prior damp_AR ~ latent_model.damp_prior - σ_AR = sqrt(σ²_AR) - @assert n>p "n must be longer than latent_model.p" + @assert n>p "n must be longer than order of the autoregressive process" # Initialize the AR series with the initial values ar = Vector{Float64}(undef, n) From 668b7eea0dc7f864fd2b2b35c2341827e2afdbab Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 12 Mar 2024 18:31:24 +0000 Subject: [PATCH 11/57] work through tests --- .../src/EpiLatentModels/autoregressive.jl | 2 +- EpiAware/test/test_autoregressive.jl | 53 ++++++++----------- 2 files changed, 24 insertions(+), 31 deletions(-) diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl index 075be9c6c..828eccd0b 100644 --- a/EpiAware/src/EpiLatentModels/autoregressive.jl +++ b/EpiAware/src/EpiLatentModels/autoregressive.jl @@ -6,7 +6,7 @@ end @model function generate_latent(latent_model::AR, n) - p = length(damp_prior) + p = length(latent_model.damp_prior) ϵ_t ~ MvNormal(ones(n - p)) σ_AR ~ latent_model.std_prior ar_init ~ latent_model.init_prior diff --git a/EpiAware/test/test_autoregressive.jl b/EpiAware/test/test_autoregressive.jl index 08de1d924..0c3090640 100644 --- a/EpiAware/test/test_autoregressive.jl +++ b/EpiAware/test/test_autoregressive.jl @@ -1,55 +1,48 @@ -@testitem "Testing default_ar_priors" begin +@testitem "Testing AR constructor" begin using Distributions + damp_prior = [truncated(Normal(0.0, 0.05), 0.0, 1)] + std_prior = truncated(Normal(0.0, 0.05), 0.0, Inf) + init_prior = Normal() + ar_process = EpiAware.AR(damp_prior, std_prior, init_prior) + + @test ar_process.damp_prior == damp_prior + @test ar_process.std_prior == std_prior + @test ar_process.init_prior == init_prior +end + +@testitem "Test AR defaults" begin + using Distributions + ar = EpiAware.AR() @testset "damp_prior" begin - priors = EpiAware.default_ar_priors() - damp = rand(priors[:damp_prior][1]) + damp = rand(ar.damp_prior) @test 0.0 <= damp <= 1.0 end - @testset "var_prior" begin - priors = EpiAware.default_ar_priors() - var_AR = rand(priors[:var_prior]) - @test var_AR >= 0.0 + @testset "std_prior" begin + std_AR = rand(ar.std_prior) + @test std_AR >= 0.0 end @testset "init_prior" begin - priors = EpiAware.default_ar_priors() - init_ar_value = rand(priors[:init_prior]) + init_ar_value = rand(ar.init_prior) @test typeof(init_ar_value) == Float64 end end -@testitem "Testing AR constructor" begin - using Distributions - - damp_prior = [truncated(Normal(0.0, 0.05), 0.0, 1)] - var_prior = truncated(Normal(0.0, 0.05), 0.0, Inf) - init_prior = Normal() - ar_process = EpiAware.AR(damp_prior, var_prior, init_prior) - - @test ar_process.damp_prior == damp_prior - @test ar_process.var_prior == var_prior - @test ar_process.init_prior == init_prior - @test ar_process.p == 1 -end - @testitem "Testing AR process against theoretical properties" begin using DynamicPPL, Turing using HypothesisTests: ExactOneSampleKSTest, pvalue using Distributions - ar_model = EpiAware.AR(EpiAware.default_ar_priors()[:damp_prior], - EpiAware.default_ar_priors()[:var_prior], - EpiAware.default_ar_priors()[:init_prior] - ) + ar_model = EpiAware.AR() n = 1000 damp = [0.1] - σ²_AR = 1.0 + σ_AR = 1.0 ar_init = [0.0] model = EpiAware.generate_latent(ar_model, n) - fixed_model = fix(model, (σ²_AR = σ²_AR, damp_AR = damp, ar_init = ar_init)) + fixed_model = fix(model, (σ_AR = σ_AR, damp_AR = damp, ar_init = ar_init)) n_samples = 100 samples = sample(fixed_model, Prior(), n_samples) |> @@ -58,7 +51,7 @@ end end theoretical_mean = 0.0 - theoretical_var = σ²_AR / (1 - damp[1]^2) + theoretical_var = σ_AR^2 / (1 - damp[1]^2) @test isapprox(mean(samples), theoretical_mean, atol = 0.1) @test isapprox(var(samples), theoretical_var, atol = 0.2) From 50c5673e6d284f7cd3d207329148e27165e9c47c Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 12 Mar 2024 18:40:49 +0000 Subject: [PATCH 12/57] check tests and add basic docs --- .../src/EpiLatentModels/autoregressive.jl | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl index 828eccd0b..a6f86b996 100644 --- a/EpiAware/src/EpiLatentModels/autoregressive.jl +++ b/EpiAware/src/EpiLatentModels/autoregressive.jl @@ -1,3 +1,11 @@ +""" +The autoregressive (AR) model struct. + +# Fields +- `damp_prior::D`: The prior distribution for the damping factor. +- `std_prior::S`: The prior distribution for the standard deviation. +- `init_prior::I`: The prior distribution for the initial values. +""" @with_kw struct AR{D <: Priors, S <: Distribution, I <: Priors} <: AbstractLatentModel damp_prior::D = truncated(Normal(0.0, 0.05), 0.0, 1) std_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf) @@ -5,6 +13,21 @@ @assert length(damp_prior)==length(init_prior) "damp_prior and init_prior must have the same length" end +""" +Generate a latent AR series. + +# Arguments +- `latent_model::AR`: The AR model. +- `n::Int`: The length of the AR series. + +# Returns +- `ar::Vector{Float64}`: The generated AR series. +- `params::NamedTuple`: A named tuple containing the generated parameters (`σ_AR`, `ar_init`, `damp_AR`). + +# Notes +- The length of `damp_prior` and `init_prior` must be the same. +- `n` must be longer than the order of the autoregressive process. +""" @model function generate_latent(latent_model::AR, n) p = length(latent_model.damp_prior) ϵ_t ~ MvNormal(ones(n - p)) From 030193f0a44a6535c94ea2eb85fe3cde681de6a1 Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 12 Mar 2024 21:54:23 +0000 Subject: [PATCH 13/57] tweak docs --- EpiAware/src/EpiLatentModels/autoregressive.jl | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl index a6f86b996..04927cc9f 100644 --- a/EpiAware/src/EpiLatentModels/autoregressive.jl +++ b/EpiAware/src/EpiLatentModels/autoregressive.jl @@ -1,11 +1,6 @@ -""" +@doc raw" The autoregressive (AR) model struct. - -# Fields -- `damp_prior::D`: The prior distribution for the damping factor. -- `std_prior::S`: The prior distribution for the standard deviation. -- `init_prior::I`: The prior distribution for the initial values. -""" +" @with_kw struct AR{D <: Priors, S <: Distribution, I <: Priors} <: AbstractLatentModel damp_prior::D = truncated(Normal(0.0, 0.05), 0.0, 1) std_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf) @@ -13,10 +8,11 @@ The autoregressive (AR) model struct. @assert length(damp_prior)==length(init_prior) "damp_prior and init_prior must have the same length" end -""" +@doc raw" Generate a latent AR series. # Arguments + - `latent_model::AR`: The AR model. - `n::Int`: The length of the AR series. @@ -27,7 +23,7 @@ Generate a latent AR series. # Notes - The length of `damp_prior` and `init_prior` must be the same. - `n` must be longer than the order of the autoregressive process. -""" +" @model function generate_latent(latent_model::AR, n) p = length(latent_model.damp_prior) ϵ_t ~ MvNormal(ones(n - p)) From e11a2e9d61ef2cf138cb9cfa544c6fea515a4477 Mon Sep 17 00:00:00 2001 From: Sam Date: Thu, 14 Mar 2024 21:47:37 +0000 Subject: [PATCH 14/57] clean up merge issues --- .../src/EpiLatentModels/autoregressive.jl | 2 +- .../EpiLatentModels/differencedlatentmodel.jl | 29 ------------------- 2 files changed, 1 insertion(+), 30 deletions(-) delete mode 100644 EpiAware/src/EpiLatentModels/differencedlatentmodel.jl diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl index 04927cc9f..18fa7037e 100644 --- a/EpiAware/src/EpiLatentModels/autoregressive.jl +++ b/EpiAware/src/EpiLatentModels/autoregressive.jl @@ -24,7 +24,7 @@ Generate a latent AR series. - The length of `damp_prior` and `init_prior` must be the same. - `n` must be longer than the order of the autoregressive process. " -@model function generate_latent(latent_model::AR, n) +@model function EpiAwareBase.generate_latent(latent_model::AR, n) p = length(latent_model.damp_prior) ϵ_t ~ MvNormal(ones(n - p)) σ_AR ~ latent_model.std_prior diff --git a/EpiAware/src/EpiLatentModels/differencedlatentmodel.jl b/EpiAware/src/EpiLatentModels/differencedlatentmodel.jl deleted file mode 100644 index 1d38e6052..000000000 --- a/EpiAware/src/EpiLatentModels/differencedlatentmodel.jl +++ /dev/null @@ -1,29 +0,0 @@ -struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel - model::T - init_prior::Priors - d::Int - - function DiffLatentModel(model::AbstractModel, init_prior::Priors) - d = length(init_prior) - return DiffLatentModel(model, init_prior, d) # Add the missing type parameter to the new function call - end - - function DiffLatentModel(model::AbstractModel, init_prior::Priors, d::Int) - @assert d>0 "d must be greater than 0" - @assert length(init_prior)==d "Length of init_prior must be equal to d" - return new{T <: AbstractModel}(model, init_prior, d) # Add the missing type parameter to the new function call - end -end - -@model function generate_latent(latent_model::DiffLatentModel, n) - d = latent_model.d - @assert n>d "n must be longer than d" - init_latent ~ latent_model.init_prior - - @submodel diff_latent, diff_latent_aux = generate_latent(latent_model.model, n - d) - - latent = vcat(init_latent, diff_latent) |> - cumsum - - return latent, (; init_latent, diff_latent_aux...) -end From 9a172f99f0570348e82d40aabb35e1457974c6e4 Mon Sep 17 00:00:00 2001 From: Sam Date: Thu, 14 Mar 2024 23:14:01 +0000 Subject: [PATCH 15/57] reimplement AR --- EpiAware/src/EpiAware.jl | 2 +- .../src/EpiLatentModels/EpiLatentModels.jl | 5 +-- .../src/EpiLatentModels/autoregressive.jl | 35 +++++++++++++++---- EpiAware/src/EpiLatentModels/randomwalk.jl | 11 ++---- EpiAware/src/EpiLatentModels/utils.jl | 6 ++++ EpiAware/test/test_randomwalk.jl | 19 +++++----- 6 files changed, 51 insertions(+), 27 deletions(-) create mode 100644 EpiAware/src/EpiLatentModels/utils.jl diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl index 3d5a9cbbf..fdf21fe48 100644 --- a/EpiAware/src/EpiAware.jl +++ b/EpiAware/src/EpiAware.jl @@ -48,7 +48,7 @@ export spread_draws, scan, create_discrete_pmf include("EpiLatentModels/EpiLatentModels.jl") using .EpiLatentModels -export RandomWalk, default_rw_priors +export RandomWalk, AR include("EpiInfModels/EpiInfModels.jl") using .EpiInfModels diff --git a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl index 9d22427d2..57cb74c3d 100644 --- a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl +++ b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl @@ -8,8 +8,9 @@ using ..EpiAwareBase using Turing, Distributions, DocStringExtensions -export RandomWalk, default_rw_priors +export RandomWalk, AR include("randomwalk.jl") - +include("autoregressive.jl") +include("utils.jl") end diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl index 18fa7037e..a08e99b2d 100644 --- a/EpiAware/src/EpiLatentModels/autoregressive.jl +++ b/EpiAware/src/EpiLatentModels/autoregressive.jl @@ -1,11 +1,34 @@ @doc raw" The autoregressive (AR) model struct. " -@with_kw struct AR{D <: Priors, S <: Distribution, I <: Priors} <: AbstractLatentModel - damp_prior::D = truncated(Normal(0.0, 0.05), 0.0, 1) - std_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf) - init_prior::I = Normal() - @assert length(damp_prior)==length(init_prior) "damp_prior and init_prior must have the same length" +struct AR{D <: Distribution, S <: Distribution, I <: Distribution, P <: Int} <: AbstractLatentModel + damp_prior::D, + std_prior::S, + init_prior::I, + p::P + + function AR(damp_prior::Distribution, std_prior::Distribution; init_prior::Distribution; p::Int = 1) + + damp_priors = fill(damp_prior, p) + init_priors = fill(init_prior, p) + return AR(; damp_priors = damp_priors, std_prior = std_prior, init_priors = init_priors, p = p) + end + + function AR(; damp_priors::Vector{D} = [truncated(Normal(0.0, 0.05))], + std_prior::Distribution = truncated(Normal(0.0, 0.05), 0.0, Inf), + init_priors::Vector{I} = [Normal()]) where {D <: Distribution, I <: Distribution} + p = length(damp_priors) + damp_prior = _expand_dist(damp_prior) + init_prior = _expand_dist(init_priors) + return AR(damp_prior, std_prior, init_prior, p) + end + + function AR(damp_prior::Distribution, std_prior::Distribution, init_prior::Distribution, p::Int) + @assert p > 0 "p must be greater than 0" + @assert length(damp_prior) == length(init_prior) "damp_prior and init_prior must have the same length" + @assert p == length(damp_prior) "p must be equal to the length of damp_prior" + new{typeof(damp_prior), typeof(std_prior), typeof(init_prior), typeof(p)}(damp_prior, std_prior, init_prior, p) + end end @doc raw" @@ -25,7 +48,7 @@ Generate a latent AR series. - `n` must be longer than the order of the autoregressive process. " @model function EpiAwareBase.generate_latent(latent_model::AR, n) - p = length(latent_model.damp_prior) + p = latent_model.p ϵ_t ~ MvNormal(ones(n - p)) σ_AR ~ latent_model.std_prior ar_init ~ latent_model.init_prior diff --git a/EpiAware/src/EpiLatentModels/randomwalk.jl b/EpiAware/src/EpiLatentModels/randomwalk.jl index 5f47bd4b1..e0c1c2567 100644 --- a/EpiAware/src/EpiLatentModels/randomwalk.jl +++ b/EpiAware/src/EpiLatentModels/randomwalk.jl @@ -1,11 +1,6 @@ -struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel - init_prior::D - std_prior::S -end - -function default_rw_priors() - return (:var_RW_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), - :init_rw_value_prior => Normal()) |> Dict +@kwdef struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel + init_prior::D = truncated(Normal(0.0, 0.05), 0.0, Inf) + std_prior::S = Normal() end @model function EpiAwareBase.generate_latent(latent_model::RandomWalk, n) diff --git a/EpiAware/src/EpiLatentModels/utils.jl b/EpiAware/src/EpiLatentModels/utils.jl new file mode 100644 index 000000000..c4f1558ce --- /dev/null +++ b/EpiAware/src/EpiLatentModels/utils.jl @@ -0,0 +1,6 @@ +function _expand_dist(dist::Vector{D} where {D <: Distribution}) + d = length(dist) + product_dist = all(first(dist) .== dist) ? + filldist(first(dist), d) : arraydist(dist) + return product_dist +end diff --git a/EpiAware/test/test_randomwalk.jl b/EpiAware/test/test_randomwalk.jl index e34f90dff..f661d94ca 100644 --- a/EpiAware/test/test_randomwalk.jl +++ b/EpiAware/test/test_randomwalk.jl @@ -4,7 +4,6 @@ using HypothesisTests: ExactOneSampleKSTest, pvalue n = 5 - priors = default_rw_priors() rw_process = RandomWalk(Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) model = generate_latent(rw_process, n) @@ -18,17 +17,17 @@ ks_test_pval = ExactOneSampleKSTest(samples_day_5, Normal(0.0, sqrt(5))) |> pvalue @test ks_test_pval > 1e-6 #Very unlikely to fail if the model is correctly implemented end -@testitem "Testing default_rw_priors" begin - @testset "var_RW_prior" begin - priors = default_rw_priors() - var_RW = rand(priors[:var_RW_prior]) - @test var_RW >= 0.0 +@testitem "Testing default RW priors" begin + @testset "std_prior" begin + priors = RandomWalk() + std_rw = rand(priors.std_prior]) + @test std_rw >= 0.0 end - @testset "init_rw_value_prior" begin - priors = default_rw_priors() - init_rw_value = rand(priors[:init_rw_value_prior]) - @test typeof(init_rw_value) == Float64 + @testset "init_prior" begin + priors = RandomWalk() + init_value = rand(priors.init_prior) + @test typeof(init_value) == Float64 end end @testset "Testing RandomWalk constructor" begin From e21e6fce6cb467ebf2958aec2dc1077bd21b8173 Mon Sep 17 00:00:00 2001 From: Sam Date: Thu, 14 Mar 2024 23:15:03 +0000 Subject: [PATCH 16/57] reimplement AR --- .../src/EpiLatentModels/autoregressive.jl | 27 +++++++++++-------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl index a08e99b2d..c2e9eb2c2 100644 --- a/EpiAware/src/EpiLatentModels/autoregressive.jl +++ b/EpiAware/src/EpiLatentModels/autoregressive.jl @@ -1,33 +1,38 @@ @doc raw" The autoregressive (AR) model struct. " -struct AR{D <: Distribution, S <: Distribution, I <: Distribution, P <: Int} <: AbstractLatentModel +struct AR{D <: Distribution, S <: Distribution, I <: Distribution, P <: Int} <: + AbstractLatentModel damp_prior::D, std_prior::S, init_prior::I, p::P - function AR(damp_prior::Distribution, std_prior::Distribution; init_prior::Distribution; p::Int = 1) - + function AR(damp_prior::Distribution, std_prior::Distribution; + init_prior::Distribution; p::Int = 1) damp_priors = fill(damp_prior, p) init_priors = fill(init_prior, p) - return AR(; damp_priors = damp_priors, std_prior = std_prior, init_priors = init_priors, p = p) + return AR(; damp_priors = damp_priors, std_prior = std_prior, + init_priors = init_priors, p = p) end function AR(; damp_priors::Vector{D} = [truncated(Normal(0.0, 0.05))], - std_prior::Distribution = truncated(Normal(0.0, 0.05), 0.0, Inf), - init_priors::Vector{I} = [Normal()]) where {D <: Distribution, I <: Distribution} + std_prior::Distribution = truncated(Normal(0.0, 0.05), 0.0, Inf), + init_priors::Vector{I} = [Normal()]) where { + D <: Distribution, I <: Distribution} p = length(damp_priors) damp_prior = _expand_dist(damp_prior) init_prior = _expand_dist(init_priors) return AR(damp_prior, std_prior, init_prior, p) end - function AR(damp_prior::Distribution, std_prior::Distribution, init_prior::Distribution, p::Int) - @assert p > 0 "p must be greater than 0" - @assert length(damp_prior) == length(init_prior) "damp_prior and init_prior must have the same length" - @assert p == length(damp_prior) "p must be equal to the length of damp_prior" - new{typeof(damp_prior), typeof(std_prior), typeof(init_prior), typeof(p)}(damp_prior, std_prior, init_prior, p) + function AR(damp_prior::Distribution, std_prior::Distribution, + init_prior::Distribution, p::Int) + @assert p>0 "p must be greater than 0" + @assert length(damp_prior)==length(init_prior) "damp_prior and init_prior must have the same length" + @assert p==length(damp_prior) "p must be equal to the length of damp_prior" + new{typeof(damp_prior), typeof(std_prior), typeof(init_prior), typeof(p)}( + damp_prior, std_prior, init_prior, p) end end From e34e952c6582f3bd584d32224df4d5cec9132bad Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 15 Mar 2024 09:15:07 +0000 Subject: [PATCH 17/57] fix AR --- .../src/EpiLatentModels/autoregressive.jl | 22 +++++++++++-------- EpiAware/test/test_randomwalk.jl | 2 +- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl index c2e9eb2c2..8e1063a46 100644 --- a/EpiAware/src/EpiLatentModels/autoregressive.jl +++ b/EpiAware/src/EpiLatentModels/autoregressive.jl @@ -1,19 +1,22 @@ @doc raw" The autoregressive (AR) model struct. " -struct AR{D <: Distribution, S <: Distribution, I <: Distribution, P <: Int} <: +struct AR{D <: Sampleable, S <: Sampleable, I <: Sampleable, P <: Int} <: AbstractLatentModel - damp_prior::D, - std_prior::S, - init_prior::I, + "Prior distribution for the damping coefficients." + damp_prior::D + "Prior distribution for the standard deviation." + std_prior::S + "Prior distribution for the initial conditions" + init_prior::I + "Order of the AR model." p::P - - function AR(damp_prior::Distribution, std_prior::Distribution; - init_prior::Distribution; p::Int = 1) + function AR(damp_prior::Distribution, std_prior::Distribution, + init_prior::Distribution; p::Int) damp_priors = fill(damp_prior, p) init_priors = fill(init_prior, p) return AR(; damp_priors = damp_priors, std_prior = std_prior, - init_priors = init_priors, p = p) + init_priors = init_priors) end function AR(; damp_priors::Vector{D} = [truncated(Normal(0.0, 0.05))], @@ -32,7 +35,8 @@ struct AR{D <: Distribution, S <: Distribution, I <: Distribution, P <: Int} <: @assert length(damp_prior)==length(init_prior) "damp_prior and init_prior must have the same length" @assert p==length(damp_prior) "p must be equal to the length of damp_prior" new{typeof(damp_prior), typeof(std_prior), typeof(init_prior), typeof(p)}( - damp_prior, std_prior, init_prior, p) + damp_prior, std_prior, init_prior, p + ) end end diff --git a/EpiAware/test/test_randomwalk.jl b/EpiAware/test/test_randomwalk.jl index f661d94ca..ab9591311 100644 --- a/EpiAware/test/test_randomwalk.jl +++ b/EpiAware/test/test_randomwalk.jl @@ -20,7 +20,7 @@ end @testitem "Testing default RW priors" begin @testset "std_prior" begin priors = RandomWalk() - std_rw = rand(priors.std_prior]) + std_rw = rand(priors.std_prior) @test std_rw >= 0.0 end From 8f5d20334463e310173470660f78ce677a30028b Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 15 Mar 2024 09:16:31 +0000 Subject: [PATCH 18/57] fix current AR tests --- EpiAware/test/test_autoregressive.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/EpiAware/test/test_autoregressive.jl b/EpiAware/test/test_autoregressive.jl index 0c3090640..87201e12f 100644 --- a/EpiAware/test/test_autoregressive.jl +++ b/EpiAware/test/test_autoregressive.jl @@ -1,7 +1,7 @@ @testitem "Testing AR constructor" begin using Distributions - damp_prior = [truncated(Normal(0.0, 0.05), 0.0, 1)] + damp_prior = truncated(Normal(0.0, 0.05), 0.0, 1) std_prior = truncated(Normal(0.0, 0.05), 0.0, Inf) init_prior = Normal() ar_process = EpiAware.AR(damp_prior, std_prior, init_prior) From 4164ec4b01c8c417dfbbf74a49d4170cd3371f5e Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 15 Mar 2024 09:22:24 +0000 Subject: [PATCH 19/57] current tests passing --- EpiAware/src/EpiLatentModels/autoregressive.jl | 4 ++-- EpiAware/test/test_autoregressive.jl | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl index 8e1063a46..b7ac5bb29 100644 --- a/EpiAware/src/EpiLatentModels/autoregressive.jl +++ b/EpiAware/src/EpiLatentModels/autoregressive.jl @@ -12,7 +12,7 @@ struct AR{D <: Sampleable, S <: Sampleable, I <: Sampleable, P <: Int} <: "Order of the AR model." p::P function AR(damp_prior::Distribution, std_prior::Distribution, - init_prior::Distribution; p::Int) + init_prior::Distribution; p::Int = 1) damp_priors = fill(damp_prior, p) init_priors = fill(init_prior, p) return AR(; damp_priors = damp_priors, std_prior = std_prior, @@ -24,7 +24,7 @@ struct AR{D <: Sampleable, S <: Sampleable, I <: Sampleable, P <: Int} <: init_priors::Vector{I} = [Normal()]) where { D <: Distribution, I <: Distribution} p = length(damp_priors) - damp_prior = _expand_dist(damp_prior) + damp_prior = _expand_dist(damp_priors) init_prior = _expand_dist(init_priors) return AR(damp_prior, std_prior, init_prior, p) end diff --git a/EpiAware/test/test_autoregressive.jl b/EpiAware/test/test_autoregressive.jl index 87201e12f..4bbfbb101 100644 --- a/EpiAware/test/test_autoregressive.jl +++ b/EpiAware/test/test_autoregressive.jl @@ -1,14 +1,14 @@ @testitem "Testing AR constructor" begin - using Distributions + using Distributions, Turing damp_prior = truncated(Normal(0.0, 0.05), 0.0, 1) std_prior = truncated(Normal(0.0, 0.05), 0.0, Inf) init_prior = Normal() ar_process = EpiAware.AR(damp_prior, std_prior, init_prior) - @test ar_process.damp_prior == damp_prior + @test ar_process.damp_prior == filldist(damp_prior, 1) @test ar_process.std_prior == std_prior - @test ar_process.init_prior == init_prior + @test ar_process.init_prior == filldist(init_prior, 1) end @testitem "Test AR defaults" begin @@ -16,7 +16,7 @@ end ar = EpiAware.AR() @testset "damp_prior" begin damp = rand(ar.damp_prior) - @test 0.0 <= damp <= 1.0 + @test 0.0 <= damp[1] <= 1.0 end @testset "std_prior" begin @@ -26,7 +26,7 @@ end @testset "init_prior" begin init_ar_value = rand(ar.init_prior) - @test typeof(init_ar_value) == Float64 + @test typeof(init_ar_value[1]) == Float64 end end From 54676d5eb482157a4d4d594a705afb18b0e13c1b Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 15 Mar 2024 09:34:08 +0000 Subject: [PATCH 20/57] copy doc strings everywhere --- EpiAware/src/EpiAwareBase/EpiAwareBase.jl | 1 + EpiAware/src/EpiAwareBase/docstrings.jl | 24 +++++++++++++++++++ EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl | 1 + EpiAware/src/EpiAwareUtils/docstrings.jl | 24 +++++++++++++++++++ EpiAware/src/EpiInfModels/EpiInfModels.jl | 1 + EpiAware/src/EpiInfModels/docstrings.jl | 24 +++++++++++++++++++ EpiAware/src/EpiInference/EpiInference.jl | 1 + EpiAware/src/EpiInference/docstrings.jl | 24 +++++++++++++++++++ .../src/EpiLatentModels/EpiLatentModels.jl | 1 + .../src/EpiLatentModels/autoregressive.jl | 17 +++++++++++++ EpiAware/src/EpiLatentModels/docstrings.jl | 24 +++++++++++++++++++ EpiAware/src/EpiObsModels/EpiObsModels.jl | 1 + EpiAware/src/EpiObsModels/docstrings.jl | 24 +++++++++++++++++++ EpiAware/src/latentmodels/latentmodels.jl | 5 ---- 14 files changed, 167 insertions(+), 5 deletions(-) create mode 100644 EpiAware/src/EpiAwareBase/docstrings.jl create mode 100644 EpiAware/src/EpiAwareUtils/docstrings.jl create mode 100644 EpiAware/src/EpiInfModels/docstrings.jl create mode 100644 EpiAware/src/EpiInference/docstrings.jl create mode 100644 EpiAware/src/EpiLatentModels/docstrings.jl create mode 100644 EpiAware/src/EpiObsModels/docstrings.jl delete mode 100644 EpiAware/src/latentmodels/latentmodels.jl diff --git a/EpiAware/src/EpiAwareBase/EpiAwareBase.jl b/EpiAware/src/EpiAwareBase/EpiAwareBase.jl index 4b81cfaf5..2d803a58f 100644 --- a/EpiAware/src/EpiAwareBase/EpiAwareBase.jl +++ b/EpiAware/src/EpiAwareBase/EpiAwareBase.jl @@ -10,6 +10,7 @@ export AbstractModel, AbstractEpiModel, AbstractLatentModel, AbstractObservationModel, generate_latent, generate_latent_infs, generate_observations +include("docstrings.jl") include("types.jl") include("functions.jl") diff --git a/EpiAware/src/EpiAwareBase/docstrings.jl b/EpiAware/src/EpiAwareBase/docstrings.jl new file mode 100644 index 000000000..5cde01b87 --- /dev/null +++ b/EpiAware/src/EpiAwareBase/docstrings.jl @@ -0,0 +1,24 @@ +@template (FUNCTIONS, METHODS, MACROS) = """ + $(TYPEDSIGNATURES) + $(DOCSTRING) + """ + +@template (TYPES) = """ + $(TYPEDEF) + $(DOCSTRING) + + --- + ## Fields + $(TYPEDFIELDS) + """ + +@template MODULES = """ +$(DOCSTRING) + +--- +## Exports +$(EXPORTS) +--- +## Imports +$(IMPORTS) +""" diff --git a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl index 1b20aa9ef..c53e128e5 100644 --- a/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl +++ b/EpiAware/src/EpiAwareUtils/EpiAwareUtils.jl @@ -14,6 +14,7 @@ using DocStringExtensions, QuadGK export scan, spread_draws, create_discrete_pmf +include("docstrings.jl") include("prior-tools.jl") include("distributions.jl") include("scan.jl") diff --git a/EpiAware/src/EpiAwareUtils/docstrings.jl b/EpiAware/src/EpiAwareUtils/docstrings.jl new file mode 100644 index 000000000..5cde01b87 --- /dev/null +++ b/EpiAware/src/EpiAwareUtils/docstrings.jl @@ -0,0 +1,24 @@ +@template (FUNCTIONS, METHODS, MACROS) = """ + $(TYPEDSIGNATURES) + $(DOCSTRING) + """ + +@template (TYPES) = """ + $(TYPEDEF) + $(DOCSTRING) + + --- + ## Fields + $(TYPEDFIELDS) + """ + +@template MODULES = """ +$(DOCSTRING) + +--- +## Exports +$(EXPORTS) +--- +## Imports +$(IMPORTS) +""" diff --git a/EpiAware/src/EpiInfModels/EpiInfModels.jl b/EpiAware/src/EpiInfModels/EpiInfModels.jl index 3a77c39ce..23c6e294e 100644 --- a/EpiAware/src/EpiInfModels/EpiInfModels.jl +++ b/EpiAware/src/EpiInfModels/EpiInfModels.jl @@ -13,6 +13,7 @@ using Turing, Distributions, DocStringExtensions, LinearAlgebra export EpiData, DirectInfections, ExpGrowthRate, Renewal, R_to_r, r_to_R +include("docstrings.jl") include("epidata.jl") include("directinfections.jl") include("expgrowthrate.jl") diff --git a/EpiAware/src/EpiInfModels/docstrings.jl b/EpiAware/src/EpiInfModels/docstrings.jl new file mode 100644 index 000000000..5cde01b87 --- /dev/null +++ b/EpiAware/src/EpiInfModels/docstrings.jl @@ -0,0 +1,24 @@ +@template (FUNCTIONS, METHODS, MACROS) = """ + $(TYPEDSIGNATURES) + $(DOCSTRING) + """ + +@template (TYPES) = """ + $(TYPEDEF) + $(DOCSTRING) + + --- + ## Fields + $(TYPEDFIELDS) + """ + +@template MODULES = """ +$(DOCSTRING) + +--- +## Exports +$(EXPORTS) +--- +## Imports +$(IMPORTS) +""" diff --git a/EpiAware/src/EpiInference/EpiInference.jl b/EpiAware/src/EpiInference/EpiInference.jl index 210a733c1..743fde9d4 100644 --- a/EpiAware/src/EpiInference/EpiInference.jl +++ b/EpiAware/src/EpiInference/EpiInference.jl @@ -10,6 +10,7 @@ using DynamicPPL, DocStringExtensions export manypathfinder +include("docstrings.jl") include("manypathfinder.jl") end diff --git a/EpiAware/src/EpiInference/docstrings.jl b/EpiAware/src/EpiInference/docstrings.jl new file mode 100644 index 000000000..5cde01b87 --- /dev/null +++ b/EpiAware/src/EpiInference/docstrings.jl @@ -0,0 +1,24 @@ +@template (FUNCTIONS, METHODS, MACROS) = """ + $(TYPEDSIGNATURES) + $(DOCSTRING) + """ + +@template (TYPES) = """ + $(TYPEDEF) + $(DOCSTRING) + + --- + ## Fields + $(TYPEDFIELDS) + """ + +@template MODULES = """ +$(DOCSTRING) + +--- +## Exports +$(EXPORTS) +--- +## Imports +$(IMPORTS) +""" diff --git a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl index 57cb74c3d..71b317641 100644 --- a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl +++ b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl @@ -10,6 +10,7 @@ using Turing, Distributions, DocStringExtensions export RandomWalk, AR +include("docstrings.jl") include("randomwalk.jl") include("autoregressive.jl") include("utils.jl") diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl index b7ac5bb29..924a835eb 100644 --- a/EpiAware/src/EpiLatentModels/autoregressive.jl +++ b/EpiAware/src/EpiLatentModels/autoregressive.jl @@ -1,5 +1,22 @@ @doc raw" The autoregressive (AR) model struct. + +# Constructors +- `AR(damp_prior::Distribution, std_prior::Distribution, init_prior::Distribution; p::Int = 1)`: Constructs an AR model with the specified prior distributions for damping coefficients, standard deviation, and initial conditions. The order of the AR model can also be specified. + +- `AR(; damp_priors::Vector{D} = [truncated(Normal(0.0, 0.05))], std_prior::Distribution = truncated(Normal(0.0, 0.05), 0.0, Inf), init_priors::Vector{I} = [Normal()]) where {D <: Distribution, I <: Distribution}`: Constructs an AR model with the specified prior distributions for damping coefficients, standard deviation, and initial conditions. The order of the AR model is determined by the length of the `damp_priors` vector. + +- `AR(damp_prior::Distribution, std_prior::Distribution, init_prior::Distribution, p::Int)`: Constructs an AR model with the specified prior distributions for damping coefficients, standard deviation, and initial conditions. The order of the AR model is explicitly specified. + +# Examples + +```julia +using Distributions +using EpiAware +ar = AR() +ar_model = generate_latent(ar, 10) +rand(ar_model) +``` " struct AR{D <: Sampleable, S <: Sampleable, I <: Sampleable, P <: Int} <: AbstractLatentModel diff --git a/EpiAware/src/EpiLatentModels/docstrings.jl b/EpiAware/src/EpiLatentModels/docstrings.jl new file mode 100644 index 000000000..5cde01b87 --- /dev/null +++ b/EpiAware/src/EpiLatentModels/docstrings.jl @@ -0,0 +1,24 @@ +@template (FUNCTIONS, METHODS, MACROS) = """ + $(TYPEDSIGNATURES) + $(DOCSTRING) + """ + +@template (TYPES) = """ + $(TYPEDEF) + $(DOCSTRING) + + --- + ## Fields + $(TYPEDFIELDS) + """ + +@template MODULES = """ +$(DOCSTRING) + +--- +## Exports +$(EXPORTS) +--- +## Imports +$(IMPORTS) +""" diff --git a/EpiAware/src/EpiObsModels/EpiObsModels.jl b/EpiAware/src/EpiObsModels/EpiObsModels.jl index 7c0925063..2173b5c8b 100644 --- a/EpiAware/src/EpiObsModels/EpiObsModels.jl +++ b/EpiAware/src/EpiObsModels/EpiObsModels.jl @@ -12,6 +12,7 @@ using Turing, Distributions, DocStringExtensions, SparseArrays export DelayObservations, default_delay_obs_priors +include("docstrings.jl") include("delayobservations.jl") include("utils.jl") diff --git a/EpiAware/src/EpiObsModels/docstrings.jl b/EpiAware/src/EpiObsModels/docstrings.jl new file mode 100644 index 000000000..5cde01b87 --- /dev/null +++ b/EpiAware/src/EpiObsModels/docstrings.jl @@ -0,0 +1,24 @@ +@template (FUNCTIONS, METHODS, MACROS) = """ + $(TYPEDSIGNATURES) + $(DOCSTRING) + """ + +@template (TYPES) = """ + $(TYPEDEF) + $(DOCSTRING) + + --- + ## Fields + $(TYPEDFIELDS) + """ + +@template MODULES = """ +$(DOCSTRING) + +--- +## Exports +$(EXPORTS) +--- +## Imports +$(IMPORTS) +""" diff --git a/EpiAware/src/latentmodels/latentmodels.jl b/EpiAware/src/latentmodels/latentmodels.jl deleted file mode 100644 index 7a2757966..000000000 --- a/EpiAware/src/latentmodels/latentmodels.jl +++ /dev/null @@ -1,5 +0,0 @@ -# Define the Priors type alias -const Priors = Union{Distribution, Vector{<:Distribution}, Product} - -include("randomwalk.jl") -include("autoregressive.jl") From 83b2022324ec00bf7d9a5b5c82ab631e98dd983a Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 15 Mar 2024 09:35:24 +0000 Subject: [PATCH 21/57] reorder tests for RandomWalk --- EpiAware/src/EpiLatentModels/randomwalk.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/EpiAware/src/EpiLatentModels/randomwalk.jl b/EpiAware/src/EpiLatentModels/randomwalk.jl index e0c1c2567..03911d717 100644 --- a/EpiAware/src/EpiLatentModels/randomwalk.jl +++ b/EpiAware/src/EpiLatentModels/randomwalk.jl @@ -1,6 +1,6 @@ @kwdef struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel - init_prior::D = truncated(Normal(0.0, 0.05), 0.0, Inf) - std_prior::S = Normal() + init_prior::D = Normal() + std_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf) end @model function EpiAwareBase.generate_latent(latent_model::RandomWalk, n) From 42364fe015fc69430851291ba4b21c0eeb52d32c Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 15 Mar 2024 09:45:06 +0000 Subject: [PATCH 22/57] add AR(2) --- EpiAware/test/test_autoregressive.jl | 36 ++++++++++++++++++++++++---- 1 file changed, 32 insertions(+), 4 deletions(-) diff --git a/EpiAware/test/test_autoregressive.jl b/EpiAware/test/test_autoregressive.jl index 4bbfbb101..7ef7c97fc 100644 --- a/EpiAware/test/test_autoregressive.jl +++ b/EpiAware/test/test_autoregressive.jl @@ -4,7 +4,7 @@ damp_prior = truncated(Normal(0.0, 0.05), 0.0, 1) std_prior = truncated(Normal(0.0, 0.05), 0.0, Inf) init_prior = Normal() - ar_process = EpiAware.AR(damp_prior, std_prior, init_prior) + ar_process = AR(damp_prior, std_prior, init_prior) @test ar_process.damp_prior == filldist(damp_prior, 1) @test ar_process.std_prior == std_prior @@ -13,7 +13,7 @@ end @testitem "Test AR defaults" begin using Distributions - ar = EpiAware.AR() + ar = AR() @testset "damp_prior" begin damp = rand(ar.damp_prior) @test 0.0 <= damp[1] <= 1.0 @@ -30,18 +30,46 @@ end end end +@testitem "Test AR(2)" begin + using Distributions + ar = AR( + damp_priors = [truncated(Normal(0.0, 0.05), 0.0, 1), + truncated(Normal(0.0, 0.05), 0.0, 1)], + std_prior = truncated(Normal(0.0, 0.05), 0.0, Inf), + init_priors = [Normal(), Normal()] + ) + @testset "damp_prior" begin + damp = rand(ar.damp_prior) + for i in 1:2 + @test 0.0 <= damp[i] <= 1.0 + end + end + + @testset "std_prior" begin + std_AR = rand(ar.std_prior) + @test std_AR >= 0.0 + end + + @testset "init_prior" begin + init_ar_value = rand(ar.init_prior) + for i in 1:2 + @test typeof(init_ar_value[i]) == Float64 + end + end +end + @testitem "Testing AR process against theoretical properties" begin using DynamicPPL, Turing using HypothesisTests: ExactOneSampleKSTest, pvalue using Distributions - ar_model = EpiAware.AR() + ar_model = AR() n = 1000 damp = [0.1] σ_AR = 1.0 ar_init = [0.0] - model = EpiAware.generate_latent(ar_model, n) + model = generate_latent(ar_model, n) fixed_model = fix(model, (σ_AR = σ_AR, damp_AR = damp, ar_init = ar_init)) n_samples = 100 From 03c7c5ba14071ee935c712f0942153799c44671b Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 15 Mar 2024 10:25:02 +0000 Subject: [PATCH 23/57] update ar default --- EpiAware/src/EpiLatentModels/autoregressive.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl index 924a835eb..82ecd1c63 100644 --- a/EpiAware/src/EpiLatentModels/autoregressive.jl +++ b/EpiAware/src/EpiLatentModels/autoregressive.jl @@ -36,7 +36,7 @@ struct AR{D <: Sampleable, S <: Sampleable, I <: Sampleable, P <: Int} <: init_priors = init_priors) end - function AR(; damp_priors::Vector{D} = [truncated(Normal(0.0, 0.05))], + function AR(; damp_priors::Vector{D} = [truncated(Normal(0.0, 0.05), 0, 1)], std_prior::Distribution = truncated(Normal(0.0, 0.05), 0.0, Inf), init_priors::Vector{I} = [Normal()]) where { D <: Distribution, I <: Distribution} From 455c36b26f8b9b41ef0021a2bcd4f3be9a33441b Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 8 Mar 2024 11:28:26 +0000 Subject: [PATCH 24/57] first pass at AR model --- EpiAware/src/latentmodels/randomwalk.jl | 79 +++++++++++++++++++++++++ 1 file changed, 79 insertions(+) create mode 100644 EpiAware/src/latentmodels/randomwalk.jl diff --git a/EpiAware/src/latentmodels/randomwalk.jl b/EpiAware/src/latentmodels/randomwalk.jl new file mode 100644 index 000000000..acab8d398 --- /dev/null +++ b/EpiAware/src/latentmodels/randomwalk.jl @@ -0,0 +1,79 @@ +function generate_latent(latent_model::AbstractLatentModel, n) + @info "No concrete implementation for generate_latent is defined." + return nothing +end + +struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel + init_prior::D + std_prior::S +end + +function default_rw_priors() + return (:var_RW_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), + :init_rw_value_prior => Normal()) |> Dict +end + +@model function generate_latent(latent_model::RandomWalk, n) + ϵ_t ~ MvNormal(ones(n)) + σ_RW ~ latent_model.std_prior + rw_init ~ latent_model.init_prior + rw = Vector{eltype(ϵ_t)}(undef, n) + + rw[1] = rw_init + σ_RW * ϵ_t[1] + for t in 2:n + rw[t] = rw[t - 1] + σ_RW * ϵ_t[t] + end + return rw, (; σ_RW, rw_init) +end + +struct AR <: AbstractLatentModel + """A distribution, a vector of distributions, or a product of distributions representing the prior distributions of the damping factors.""" + damp_prior::Union{Distribution, Vector{<:Distribution}, Product} + + """A distribution representing the prior distribution of the variance.""" + var_prior::Distribution + + """A distribution, a vector of distributions, or a product of distributions representing the prior distributions of the initial values.""" + init_prior::Union{Distribution, Vector{<:Distribution}, Product} + + """ + The order of the AR process, determined by the length of `damp_prior`. + """ + p::Int + + function AR(damp_prior::Union{Distribution, Vector{<:Distribution}, Product}, + var_prior::Distribution, + init_prior::Union{Distribution, Vector{<:Distribution}, Product}) + p = length(damp_prior) + @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process" + return new(damp_prior, var_prior, init_prior, p) + end +end + +function default_ar_priors() + return (:damp_prior => [truncated(Normal(0.0, 0.05), 0.0, 1)], + :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), + :init_prior => Normal()) |> Dict +end + +@model function generate_latent(latent_model::AR, n) + p = latent_model.p + ϵ_t ~ MvNormal(ones(n - p)) + σ²_AR ~ latent_model.var_prior + ar_init ~ latent_model.init_prior + damp_AR ~ latent_model.damp_prior + σ_AR = sqrt(σ²_AR) + + @assert n>p "n must be longer than p" + + # Initialize the AR series with the initial values + ar = Vector{Float64}(undef, n) + ar[1:p] = ar_init + + # Generate the rest of the AR series + for t in (p + 1):n + ar[t] = damp_AR' * ar[(t - p):(t - 1)] + σ_AR * ϵ_t[t - p] + end + + return ar, (; σ_AR, ar_init, damp_AR) +end From b8668f222cb9a13b3cbae406791d65cf0b497fd7 Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 8 Mar 2024 12:14:45 +0000 Subject: [PATCH 25/57] move files around and add differencing infrastructure --- .../src/EpiObsModels/delayobservations.jl | 4 + EpiAware/src/latentmodels/randomwalk.jl | 79 ------------------- 2 files changed, 4 insertions(+), 79 deletions(-) delete mode 100644 EpiAware/src/latentmodels/randomwalk.jl diff --git a/EpiAware/src/EpiObsModels/delayobservations.jl b/EpiAware/src/EpiObsModels/delayobservations.jl index b7ffbb157..2a1e67f91 100644 --- a/EpiAware/src/EpiObsModels/delayobservations.jl +++ b/EpiAware/src/EpiObsModels/delayobservations.jl @@ -30,7 +30,11 @@ function default_delay_obs_priors() Normal(0, 0.1 * sqrt(pi) / sqrt(2)), 0.0, Inf),) |> Dict end +<<<<<<< HEAD:EpiAware/src/EpiObsModels/delayobservations.jl @model function EpiAwareBase.generate_observations(observation_model::DelayObservations, +======= +@model function generate_observations(observation_model::DelayObservations, +>>>>>>> c8ad6fd (move files around and add differencing infrastructure):EpiAware/src/observationmodels/delayobservations.jl y_t, I_t; pos_shift) diff --git a/EpiAware/src/latentmodels/randomwalk.jl b/EpiAware/src/latentmodels/randomwalk.jl deleted file mode 100644 index acab8d398..000000000 --- a/EpiAware/src/latentmodels/randomwalk.jl +++ /dev/null @@ -1,79 +0,0 @@ -function generate_latent(latent_model::AbstractLatentModel, n) - @info "No concrete implementation for generate_latent is defined." - return nothing -end - -struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel - init_prior::D - std_prior::S -end - -function default_rw_priors() - return (:var_RW_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), - :init_rw_value_prior => Normal()) |> Dict -end - -@model function generate_latent(latent_model::RandomWalk, n) - ϵ_t ~ MvNormal(ones(n)) - σ_RW ~ latent_model.std_prior - rw_init ~ latent_model.init_prior - rw = Vector{eltype(ϵ_t)}(undef, n) - - rw[1] = rw_init + σ_RW * ϵ_t[1] - for t in 2:n - rw[t] = rw[t - 1] + σ_RW * ϵ_t[t] - end - return rw, (; σ_RW, rw_init) -end - -struct AR <: AbstractLatentModel - """A distribution, a vector of distributions, or a product of distributions representing the prior distributions of the damping factors.""" - damp_prior::Union{Distribution, Vector{<:Distribution}, Product} - - """A distribution representing the prior distribution of the variance.""" - var_prior::Distribution - - """A distribution, a vector of distributions, or a product of distributions representing the prior distributions of the initial values.""" - init_prior::Union{Distribution, Vector{<:Distribution}, Product} - - """ - The order of the AR process, determined by the length of `damp_prior`. - """ - p::Int - - function AR(damp_prior::Union{Distribution, Vector{<:Distribution}, Product}, - var_prior::Distribution, - init_prior::Union{Distribution, Vector{<:Distribution}, Product}) - p = length(damp_prior) - @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process" - return new(damp_prior, var_prior, init_prior, p) - end -end - -function default_ar_priors() - return (:damp_prior => [truncated(Normal(0.0, 0.05), 0.0, 1)], - :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), - :init_prior => Normal()) |> Dict -end - -@model function generate_latent(latent_model::AR, n) - p = latent_model.p - ϵ_t ~ MvNormal(ones(n - p)) - σ²_AR ~ latent_model.var_prior - ar_init ~ latent_model.init_prior - damp_AR ~ latent_model.damp_prior - σ_AR = sqrt(σ²_AR) - - @assert n>p "n must be longer than p" - - # Initialize the AR series with the initial values - ar = Vector{Float64}(undef, n) - ar[1:p] = ar_init - - # Generate the rest of the AR series - for t in (p + 1):n - ar[t] = damp_AR' * ar[(t - p):(t - 1)] + σ_AR * ϵ_t[t - p] - end - - return ar, (; σ_AR, ar_init, damp_AR) -end From d2e0e7275f124a8bdb165a98e1b9f7daa040a01a Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 8 Mar 2024 12:17:49 +0000 Subject: [PATCH 26/57] remove scaffold comment --- EpiAware/src/EpiObsModels/delayobservations.jl | 4 ---- 1 file changed, 4 deletions(-) diff --git a/EpiAware/src/EpiObsModels/delayobservations.jl b/EpiAware/src/EpiObsModels/delayobservations.jl index 2a1e67f91..b7ffbb157 100644 --- a/EpiAware/src/EpiObsModels/delayobservations.jl +++ b/EpiAware/src/EpiObsModels/delayobservations.jl @@ -30,11 +30,7 @@ function default_delay_obs_priors() Normal(0, 0.1 * sqrt(pi) / sqrt(2)), 0.0, Inf),) |> Dict end -<<<<<<< HEAD:EpiAware/src/EpiObsModels/delayobservations.jl @model function EpiAwareBase.generate_observations(observation_model::DelayObservations, -======= -@model function generate_observations(observation_model::DelayObservations, ->>>>>>> c8ad6fd (move files around and add differencing infrastructure):EpiAware/src/observationmodels/delayobservations.jl y_t, I_t; pos_shift) From f51c17f2998d42a454f966c1ca6c050eabe6c05f Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 8 Mar 2024 12:41:45 +0000 Subject: [PATCH 27/57] clean up use of new --- EpiAware/src/latentmodels/randomwalk.jl | 109 ++++++++++++++++++++++++ 1 file changed, 109 insertions(+) create mode 100644 EpiAware/src/latentmodels/randomwalk.jl diff --git a/EpiAware/src/latentmodels/randomwalk.jl b/EpiAware/src/latentmodels/randomwalk.jl new file mode 100644 index 000000000..9592f970e --- /dev/null +++ b/EpiAware/src/latentmodels/randomwalk.jl @@ -0,0 +1,109 @@ +struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel + init_prior::D + std_prior::S +end + +function default_rw_priors() + return (:var_RW_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), + :init_rw_value_prior => Normal()) |> Dict +end + +@model function generate_latent(latent_model::RandomWalk, n) + ϵ_t ~ MvNormal(ones(n)) + σ_RW ~ latent_model.std_prior + rw_init ~ latent_model.init_prior + rw = Vector{eltype(ϵ_t)}(undef, n) + + rw[1] = rw_init + σ_RW * ϵ_t[1] + for t in 2:n + rw[t] = rw[t - 1] + σ_RW * ϵ_t[t] + end + return rw, (; σ_RW, rw_init) +end + +# Define the Priors type alias +const Priors = Union{Distribution, Vector{<:Distribution}, Product} + +struct AR <: AbstractLatentModel + """A distribution representing the prior distribution of the damping factors.""" + damp_prior::Priors + + """A distribution representing the prior distribution of the variance.""" + var_prior::Distribution + + """A distribution representing the prior distribution of the initial values.""" + init_prior::Priors + + """ + The order of the AR process, determined by the length of `damp_prior`. + """ + p::Int + + function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors) + p = length(damp_prior) + return AR(damp_prior, var_prior, init_prior, p) + end + + function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors, p::Int) + @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process" + return new(damp_prior, var_prior, init_prior, p) + end +end + +function default_ar_priors() + return (:damp_prior => [truncated(Normal(0.0, 0.05), 0.0, 1)], + :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), + :init_prior => Normal()) |> Dict +end + +@model function generate_latent(latent_model::AR, n) + p = latent_model.p + ϵ_t ~ MvNormal(ones(n - p)) + σ²_AR ~ latent_model.var_prior + ar_init ~ latent_model.init_prior + damp_AR ~ latent_model.damp_prior + σ_AR = sqrt(σ²_AR) + + @assert n>p "n must be longer than p" + + # Initialize the AR series with the initial values + ar = Vector{Float64}(undef, n) + ar[1:p] = ar_init + + # Generate the rest of the AR series + for t in (p + 1):n + ar[t] = damp_AR' * ar[(t - p):(t - 1)] + σ_AR * ϵ_t[t - p] + end + + return ar, (; σ_AR, ar_init, damp_AR) +end + +struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel + model::T + init_prior::Priors + d::Int + + function DiffLatentModel(model::AbstractModel, init_prior::Priors) + d = length(init_prior) + return DiffLatentModel(model, init_prior, d) # Add the missing type parameter to the new function call + end + + function DiffLatentModel(model::AbstractModel, init_prior::Priors, d::Int) + @assert d>0 "d must be greater than 0" + @assert length(init_prior)==d "Length of init_prior must be equal to d" + return new{T <: AbstractModel}(model, init_prior, d) # Add the missing type parameter to the new function call + end +end + +@model function generate_latent(latent_model::DiffLatentModel, n) + d = latent_model.d + @assert n>d "n must be longer than d" + init_latent ~ latent_model.init_prior + + @submodel diff_latent, diff_latent_aux = generate_latent(latent_model.model, n - d) + + latent = vcat(init_latent, diff_latent) |> + cumsum + + return latent, (; init_latent, diff_latent_aux...) +end From 3220e4fe1fa184ed64a3fd9f0b731ca8d3ee9e89 Mon Sep 17 00:00:00 2001 From: Sam Date: Mon, 11 Mar 2024 10:20:10 +0000 Subject: [PATCH 28/57] reorg based on main --- EpiAware/src/latentmodels/autoregressive.jl | 53 +++++++++++ .../latentmodels/differencedlatentmodel.jl | 29 +++++++ EpiAware/src/latentmodels/latentmodels.jl | 5 ++ EpiAware/src/latentmodels/randomwalk.jl | 87 ------------------- 4 files changed, 87 insertions(+), 87 deletions(-) create mode 100644 EpiAware/src/latentmodels/autoregressive.jl create mode 100644 EpiAware/src/latentmodels/differencedlatentmodel.jl create mode 100644 EpiAware/src/latentmodels/latentmodels.jl diff --git a/EpiAware/src/latentmodels/autoregressive.jl b/EpiAware/src/latentmodels/autoregressive.jl new file mode 100644 index 000000000..3c2fc7648 --- /dev/null +++ b/EpiAware/src/latentmodels/autoregressive.jl @@ -0,0 +1,53 @@ +struct AR <: AbstractLatentModel + """A distribution representing the prior distribution of the damping factors.""" + damp_prior::Priors + + """A distribution representing the prior distribution of the variance.""" + var_prior::Distribution + + """A distribution representing the prior distribution of the initial values.""" + init_prior::Priors + + """ + The order of the AR process, determined by the length of `damp_prior`. + """ + p::Int + + function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors) + p = length(damp_prior) + return AR(damp_prior, var_prior, init_prior, p) + end + + function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors, p::Int) + @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process" + return new(damp_prior, var_prior, init_prior, p) + end +end + +function default_ar_priors() + return (:damp_prior => [truncated(Normal(0.0, 0.05), 0.0, 1)], + :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), + :init_prior => Normal()) |> Dict +end + +@model function generate_latent(latent_model::AR, n) + p = latent_model.p + ϵ_t ~ MvNormal(ones(n - p)) + σ²_AR ~ latent_model.var_prior + ar_init ~ latent_model.init_prior + damp_AR ~ latent_model.damp_prior + σ_AR = sqrt(σ²_AR) + + @assert n>p "n must be longer than p" + + # Initialize the AR series with the initial values + ar = Vector{Float64}(undef, n) + ar[1:p] = ar_init + + # Generate the rest of the AR series + for t in (p + 1):n + ar[t] = damp_AR' * ar[(t - p):(t - 1)] + σ_AR * ϵ_t[t - p] + end + + return ar, (; σ_AR, ar_init, damp_AR) +end diff --git a/EpiAware/src/latentmodels/differencedlatentmodel.jl b/EpiAware/src/latentmodels/differencedlatentmodel.jl new file mode 100644 index 000000000..1d38e6052 --- /dev/null +++ b/EpiAware/src/latentmodels/differencedlatentmodel.jl @@ -0,0 +1,29 @@ +struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel + model::T + init_prior::Priors + d::Int + + function DiffLatentModel(model::AbstractModel, init_prior::Priors) + d = length(init_prior) + return DiffLatentModel(model, init_prior, d) # Add the missing type parameter to the new function call + end + + function DiffLatentModel(model::AbstractModel, init_prior::Priors, d::Int) + @assert d>0 "d must be greater than 0" + @assert length(init_prior)==d "Length of init_prior must be equal to d" + return new{T <: AbstractModel}(model, init_prior, d) # Add the missing type parameter to the new function call + end +end + +@model function generate_latent(latent_model::DiffLatentModel, n) + d = latent_model.d + @assert n>d "n must be longer than d" + init_latent ~ latent_model.init_prior + + @submodel diff_latent, diff_latent_aux = generate_latent(latent_model.model, n - d) + + latent = vcat(init_latent, diff_latent) |> + cumsum + + return latent, (; init_latent, diff_latent_aux...) +end diff --git a/EpiAware/src/latentmodels/latentmodels.jl b/EpiAware/src/latentmodels/latentmodels.jl new file mode 100644 index 000000000..7a2757966 --- /dev/null +++ b/EpiAware/src/latentmodels/latentmodels.jl @@ -0,0 +1,5 @@ +# Define the Priors type alias +const Priors = Union{Distribution, Vector{<:Distribution}, Product} + +include("randomwalk.jl") +include("autoregressive.jl") diff --git a/EpiAware/src/latentmodels/randomwalk.jl b/EpiAware/src/latentmodels/randomwalk.jl index 9592f970e..1dda224b2 100644 --- a/EpiAware/src/latentmodels/randomwalk.jl +++ b/EpiAware/src/latentmodels/randomwalk.jl @@ -20,90 +20,3 @@ end end return rw, (; σ_RW, rw_init) end - -# Define the Priors type alias -const Priors = Union{Distribution, Vector{<:Distribution}, Product} - -struct AR <: AbstractLatentModel - """A distribution representing the prior distribution of the damping factors.""" - damp_prior::Priors - - """A distribution representing the prior distribution of the variance.""" - var_prior::Distribution - - """A distribution representing the prior distribution of the initial values.""" - init_prior::Priors - - """ - The order of the AR process, determined by the length of `damp_prior`. - """ - p::Int - - function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors) - p = length(damp_prior) - return AR(damp_prior, var_prior, init_prior, p) - end - - function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors, p::Int) - @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process" - return new(damp_prior, var_prior, init_prior, p) - end -end - -function default_ar_priors() - return (:damp_prior => [truncated(Normal(0.0, 0.05), 0.0, 1)], - :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), - :init_prior => Normal()) |> Dict -end - -@model function generate_latent(latent_model::AR, n) - p = latent_model.p - ϵ_t ~ MvNormal(ones(n - p)) - σ²_AR ~ latent_model.var_prior - ar_init ~ latent_model.init_prior - damp_AR ~ latent_model.damp_prior - σ_AR = sqrt(σ²_AR) - - @assert n>p "n must be longer than p" - - # Initialize the AR series with the initial values - ar = Vector{Float64}(undef, n) - ar[1:p] = ar_init - - # Generate the rest of the AR series - for t in (p + 1):n - ar[t] = damp_AR' * ar[(t - p):(t - 1)] + σ_AR * ϵ_t[t - p] - end - - return ar, (; σ_AR, ar_init, damp_AR) -end - -struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel - model::T - init_prior::Priors - d::Int - - function DiffLatentModel(model::AbstractModel, init_prior::Priors) - d = length(init_prior) - return DiffLatentModel(model, init_prior, d) # Add the missing type parameter to the new function call - end - - function DiffLatentModel(model::AbstractModel, init_prior::Priors, d::Int) - @assert d>0 "d must be greater than 0" - @assert length(init_prior)==d "Length of init_prior must be equal to d" - return new{T <: AbstractModel}(model, init_prior, d) # Add the missing type parameter to the new function call - end -end - -@model function generate_latent(latent_model::DiffLatentModel, n) - d = latent_model.d - @assert n>d "n must be longer than d" - init_latent ~ latent_model.init_prior - - @submodel diff_latent, diff_latent_aux = generate_latent(latent_model.model, n - d) - - latent = vcat(init_latent, diff_latent) |> - cumsum - - return latent, (; init_latent, diff_latent_aux...) -end From 2724d963ce2e10a265986cc3129bebf64220d53b Mon Sep 17 00:00:00 2001 From: Sam Date: Mon, 11 Mar 2024 11:58:05 +0000 Subject: [PATCH 29/57] remove differenced model for its own PR --- .../latentmodels/differencedlatentmodel.jl | 29 ------------------- 1 file changed, 29 deletions(-) delete mode 100644 EpiAware/src/latentmodels/differencedlatentmodel.jl diff --git a/EpiAware/src/latentmodels/differencedlatentmodel.jl b/EpiAware/src/latentmodels/differencedlatentmodel.jl deleted file mode 100644 index 1d38e6052..000000000 --- a/EpiAware/src/latentmodels/differencedlatentmodel.jl +++ /dev/null @@ -1,29 +0,0 @@ -struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel - model::T - init_prior::Priors - d::Int - - function DiffLatentModel(model::AbstractModel, init_prior::Priors) - d = length(init_prior) - return DiffLatentModel(model, init_prior, d) # Add the missing type parameter to the new function call - end - - function DiffLatentModel(model::AbstractModel, init_prior::Priors, d::Int) - @assert d>0 "d must be greater than 0" - @assert length(init_prior)==d "Length of init_prior must be equal to d" - return new{T <: AbstractModel}(model, init_prior, d) # Add the missing type parameter to the new function call - end -end - -@model function generate_latent(latent_model::DiffLatentModel, n) - d = latent_model.d - @assert n>d "n must be longer than d" - init_latent ~ latent_model.init_prior - - @submodel diff_latent, diff_latent_aux = generate_latent(latent_model.model, n - d) - - latent = vcat(init_latent, diff_latent) |> - cumsum - - return latent, (; init_latent, diff_latent_aux...) -end From 04e7cebcae3e546e55ceffe39df5d1bbb2eb4c70 Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 12 Mar 2024 17:17:44 +0000 Subject: [PATCH 30/57] update to use parametric types and use @kwdef to simplify construction --- EpiAware/src/latentmodels/autoregressive.jl | 28 +++++---------------- 1 file changed, 6 insertions(+), 22 deletions(-) diff --git a/EpiAware/src/latentmodels/autoregressive.jl b/EpiAware/src/latentmodels/autoregressive.jl index 3c2fc7648..da19fd54c 100644 --- a/EpiAware/src/latentmodels/autoregressive.jl +++ b/EpiAware/src/latentmodels/autoregressive.jl @@ -1,33 +1,17 @@ -struct AR <: AbstractLatentModel +@kwdef struct AR{D <: Priors, S <: Distribution, I <: Priors, P <: Int} """A distribution representing the prior distribution of the damping factors.""" - damp_prior::Priors + damp_prior::D = [truncated(Normal(0.0, 0.05), 0.0, 1)] """A distribution representing the prior distribution of the variance.""" - var_prior::Distribution + var_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf) """A distribution representing the prior distribution of the initial values.""" - init_prior::Priors + init_prior::I = Normal() """ The order of the AR process, determined by the length of `damp_prior`. """ - p::Int - - function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors) - p = length(damp_prior) - return AR(damp_prior, var_prior, init_prior, p) - end - - function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors, p::Int) - @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process" - return new(damp_prior, var_prior, init_prior, p) - end -end - -function default_ar_priors() - return (:damp_prior => [truncated(Normal(0.0, 0.05), 0.0, 1)], - :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), - :init_prior => Normal()) |> Dict + p::P = length(damp_prior) end @model function generate_latent(latent_model::AR, n) @@ -38,7 +22,7 @@ end damp_AR ~ latent_model.damp_prior σ_AR = sqrt(σ²_AR) - @assert n>p "n must be longer than p" + @assert n>p "n must be longer than latent_model.p" # Initialize the AR series with the initial values ar = Vector{Float64}(undef, n) From 1d76b4586fe3593e321672bca85d509f66ac0612 Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 12 Mar 2024 18:22:17 +0000 Subject: [PATCH 31/57] update to use parameters.jl and drop p calc in struct as don't work with having defaults --- EpiAware/src/latentmodels/autoregressive.jl | 24 ++++++--------------- 1 file changed, 7 insertions(+), 17 deletions(-) diff --git a/EpiAware/src/latentmodels/autoregressive.jl b/EpiAware/src/latentmodels/autoregressive.jl index da19fd54c..075be9c6c 100644 --- a/EpiAware/src/latentmodels/autoregressive.jl +++ b/EpiAware/src/latentmodels/autoregressive.jl @@ -1,28 +1,18 @@ -@kwdef struct AR{D <: Priors, S <: Distribution, I <: Priors, P <: Int} - """A distribution representing the prior distribution of the damping factors.""" - damp_prior::D = [truncated(Normal(0.0, 0.05), 0.0, 1)] - - """A distribution representing the prior distribution of the variance.""" - var_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf) - - """A distribution representing the prior distribution of the initial values.""" +@with_kw struct AR{D <: Priors, S <: Distribution, I <: Priors} <: AbstractLatentModel + damp_prior::D = truncated(Normal(0.0, 0.05), 0.0, 1) + std_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf) init_prior::I = Normal() - - """ - The order of the AR process, determined by the length of `damp_prior`. - """ - p::P = length(damp_prior) + @assert length(damp_prior)==length(init_prior) "damp_prior and init_prior must have the same length" end @model function generate_latent(latent_model::AR, n) - p = latent_model.p + p = length(damp_prior) ϵ_t ~ MvNormal(ones(n - p)) - σ²_AR ~ latent_model.var_prior + σ_AR ~ latent_model.std_prior ar_init ~ latent_model.init_prior damp_AR ~ latent_model.damp_prior - σ_AR = sqrt(σ²_AR) - @assert n>p "n must be longer than latent_model.p" + @assert n>p "n must be longer than order of the autoregressive process" # Initialize the AR series with the initial values ar = Vector{Float64}(undef, n) From 3d2fec1086ab420e7b0aa43ee182148673e655c8 Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 12 Mar 2024 18:31:24 +0000 Subject: [PATCH 32/57] work through tests --- EpiAware/src/latentmodels/autoregressive.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/EpiAware/src/latentmodels/autoregressive.jl b/EpiAware/src/latentmodels/autoregressive.jl index 075be9c6c..828eccd0b 100644 --- a/EpiAware/src/latentmodels/autoregressive.jl +++ b/EpiAware/src/latentmodels/autoregressive.jl @@ -6,7 +6,7 @@ end @model function generate_latent(latent_model::AR, n) - p = length(damp_prior) + p = length(latent_model.damp_prior) ϵ_t ~ MvNormal(ones(n - p)) σ_AR ~ latent_model.std_prior ar_init ~ latent_model.init_prior From 5936ae4510051e0c725aebf56dc363b75173e06f Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 12 Mar 2024 18:40:49 +0000 Subject: [PATCH 33/57] check tests and add basic docs --- EpiAware/src/latentmodels/autoregressive.jl | 23 +++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/EpiAware/src/latentmodels/autoregressive.jl b/EpiAware/src/latentmodels/autoregressive.jl index 828eccd0b..a6f86b996 100644 --- a/EpiAware/src/latentmodels/autoregressive.jl +++ b/EpiAware/src/latentmodels/autoregressive.jl @@ -1,3 +1,11 @@ +""" +The autoregressive (AR) model struct. + +# Fields +- `damp_prior::D`: The prior distribution for the damping factor. +- `std_prior::S`: The prior distribution for the standard deviation. +- `init_prior::I`: The prior distribution for the initial values. +""" @with_kw struct AR{D <: Priors, S <: Distribution, I <: Priors} <: AbstractLatentModel damp_prior::D = truncated(Normal(0.0, 0.05), 0.0, 1) std_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf) @@ -5,6 +13,21 @@ @assert length(damp_prior)==length(init_prior) "damp_prior and init_prior must have the same length" end +""" +Generate a latent AR series. + +# Arguments +- `latent_model::AR`: The AR model. +- `n::Int`: The length of the AR series. + +# Returns +- `ar::Vector{Float64}`: The generated AR series. +- `params::NamedTuple`: A named tuple containing the generated parameters (`σ_AR`, `ar_init`, `damp_AR`). + +# Notes +- The length of `damp_prior` and `init_prior` must be the same. +- `n` must be longer than the order of the autoregressive process. +""" @model function generate_latent(latent_model::AR, n) p = length(latent_model.damp_prior) ϵ_t ~ MvNormal(ones(n - p)) From a4d2cc795b20c84c5e31e256e0195fb1a7f25321 Mon Sep 17 00:00:00 2001 From: Sam Date: Tue, 12 Mar 2024 21:54:23 +0000 Subject: [PATCH 34/57] tweak docs --- EpiAware/src/latentmodels/autoregressive.jl | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/EpiAware/src/latentmodels/autoregressive.jl b/EpiAware/src/latentmodels/autoregressive.jl index a6f86b996..04927cc9f 100644 --- a/EpiAware/src/latentmodels/autoregressive.jl +++ b/EpiAware/src/latentmodels/autoregressive.jl @@ -1,11 +1,6 @@ -""" +@doc raw" The autoregressive (AR) model struct. - -# Fields -- `damp_prior::D`: The prior distribution for the damping factor. -- `std_prior::S`: The prior distribution for the standard deviation. -- `init_prior::I`: The prior distribution for the initial values. -""" +" @with_kw struct AR{D <: Priors, S <: Distribution, I <: Priors} <: AbstractLatentModel damp_prior::D = truncated(Normal(0.0, 0.05), 0.0, 1) std_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf) @@ -13,10 +8,11 @@ The autoregressive (AR) model struct. @assert length(damp_prior)==length(init_prior) "damp_prior and init_prior must have the same length" end -""" +@doc raw" Generate a latent AR series. # Arguments + - `latent_model::AR`: The AR model. - `n::Int`: The length of the AR series. @@ -27,7 +23,7 @@ Generate a latent AR series. # Notes - The length of `damp_prior` and `init_prior` must be the same. - `n` must be longer than the order of the autoregressive process. -""" +" @model function generate_latent(latent_model::AR, n) p = length(latent_model.damp_prior) ϵ_t ~ MvNormal(ones(n - p)) From cebe72713a9d2b94f5beb69eb4bce9e3f5dcb09b Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 8 Mar 2024 11:28:26 +0000 Subject: [PATCH 35/57] first pass at AR model --- EpiAware/src/latentmodels/randomwalk.jl | 57 +++++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/EpiAware/src/latentmodels/randomwalk.jl b/EpiAware/src/latentmodels/randomwalk.jl index 1dda224b2..acab8d398 100644 --- a/EpiAware/src/latentmodels/randomwalk.jl +++ b/EpiAware/src/latentmodels/randomwalk.jl @@ -1,3 +1,8 @@ +function generate_latent(latent_model::AbstractLatentModel, n) + @info "No concrete implementation for generate_latent is defined." + return nothing +end + struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel init_prior::D std_prior::S @@ -20,3 +25,55 @@ end end return rw, (; σ_RW, rw_init) end + +struct AR <: AbstractLatentModel + """A distribution, a vector of distributions, or a product of distributions representing the prior distributions of the damping factors.""" + damp_prior::Union{Distribution, Vector{<:Distribution}, Product} + + """A distribution representing the prior distribution of the variance.""" + var_prior::Distribution + + """A distribution, a vector of distributions, or a product of distributions representing the prior distributions of the initial values.""" + init_prior::Union{Distribution, Vector{<:Distribution}, Product} + + """ + The order of the AR process, determined by the length of `damp_prior`. + """ + p::Int + + function AR(damp_prior::Union{Distribution, Vector{<:Distribution}, Product}, + var_prior::Distribution, + init_prior::Union{Distribution, Vector{<:Distribution}, Product}) + p = length(damp_prior) + @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process" + return new(damp_prior, var_prior, init_prior, p) + end +end + +function default_ar_priors() + return (:damp_prior => [truncated(Normal(0.0, 0.05), 0.0, 1)], + :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), + :init_prior => Normal()) |> Dict +end + +@model function generate_latent(latent_model::AR, n) + p = latent_model.p + ϵ_t ~ MvNormal(ones(n - p)) + σ²_AR ~ latent_model.var_prior + ar_init ~ latent_model.init_prior + damp_AR ~ latent_model.damp_prior + σ_AR = sqrt(σ²_AR) + + @assert n>p "n must be longer than p" + + # Initialize the AR series with the initial values + ar = Vector{Float64}(undef, n) + ar[1:p] = ar_init + + # Generate the rest of the AR series + for t in (p + 1):n + ar[t] = damp_AR' * ar[(t - p):(t - 1)] + σ_AR * ϵ_t[t - p] + end + + return ar, (; σ_AR, ar_init, damp_AR) +end From b9c853ce75f65a55cd341ce771702f561a1c3f96 Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 8 Mar 2024 12:14:45 +0000 Subject: [PATCH 36/57] move files around and add differencing infrastructure --- EpiAware/src/abstract-functions.jl | 17 ++++++++ EpiAware/src/latentmodels/randomwalk.jl | 54 +++++++++++++++++++------ 2 files changed, 59 insertions(+), 12 deletions(-) create mode 100644 EpiAware/src/abstract-functions.jl diff --git a/EpiAware/src/abstract-functions.jl b/EpiAware/src/abstract-functions.jl new file mode 100644 index 000000000..a996ac92b --- /dev/null +++ b/EpiAware/src/abstract-functions.jl @@ -0,0 +1,17 @@ +function generate_latent(latent_model::AbstractLatentModel, n) + @info "No concrete implementation for generate_latent is defined." + return nothing +end + +function generate_latent_infs(epi_model::AbstractEpiModel, latent_model) + @info "No concrete implementation for `generate_latent_infs` is defined." + return nothing +end + +function generate_observations(observation_model::AbstractObservationModel, + y_t, + I_t; + pos_shift) + @info "No concrete implementation for generate_observations is defined." + return nothing +end diff --git a/EpiAware/src/latentmodels/randomwalk.jl b/EpiAware/src/latentmodels/randomwalk.jl index acab8d398..ad5f697b3 100644 --- a/EpiAware/src/latentmodels/randomwalk.jl +++ b/EpiAware/src/latentmodels/randomwalk.jl @@ -1,8 +1,3 @@ -function generate_latent(latent_model::AbstractLatentModel, n) - @info "No concrete implementation for generate_latent is defined." - return nothing -end - struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel init_prior::D std_prior::S @@ -26,25 +21,30 @@ end return rw, (; σ_RW, rw_init) end +# Define the Priors type alias +const Priors = Union{Distribution, Vector{<:Distribution}, Product} + struct AR <: AbstractLatentModel - """A distribution, a vector of distributions, or a product of distributions representing the prior distributions of the damping factors.""" - damp_prior::Union{Distribution, Vector{<:Distribution}, Product} + """A distribution representing the prior distribution of the damping factors.""" + damp_prior::Priors """A distribution representing the prior distribution of the variance.""" var_prior::Distribution - """A distribution, a vector of distributions, or a product of distributions representing the prior distributions of the initial values.""" - init_prior::Union{Distribution, Vector{<:Distribution}, Product} + """A distribution representing the prior distribution of the initial values.""" + init_prior::Priors """ The order of the AR process, determined by the length of `damp_prior`. """ p::Int - function AR(damp_prior::Union{Distribution, Vector{<:Distribution}, Product}, - var_prior::Distribution, - init_prior::Union{Distribution, Vector{<:Distribution}, Product}) + function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors) p = length(damp_prior) + return new(damp_prior, var_prior, init_prior, p) + end + + function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors, p::Int) @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process" return new(damp_prior, var_prior, init_prior, p) end @@ -77,3 +77,33 @@ end return ar, (; σ_AR, ar_init, damp_AR) end + +struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel + model::T + d::Int + init_prior::Priors + + function DiffLatentModel(model::T, init_prior::Priors) + d = length(init_prior) + return new(model, d, init_prior) + end + + function DiffLatentModel(model::T, d::Int, init_prior::Priors) + @assert d>0 "d must be greater than 0" + @assert length(init_prior)==d "Length of init_prior must be equal to d" + return new(model, d, init_prior) + end +end + +@model function generate_latent(latent_model::DiffLatentModel, n) + d = latent_model.d + @assert n>d "n must be longer than d" + init_latent ~ latent_model.init_prior + + @submodel diff_latent, diff_latent_aux = generate_latent(latent_model.model, n - d) + + latent = vcat(init_latent, diff_latent) |> + cumsum + # Return the reconstructed series and the parameters + return latent, (; init_latent, diff_latent_aux...) +end From 49b78751ec6eec9b3884f41e26cc756ebf52b877 Mon Sep 17 00:00:00 2001 From: Sam Date: Mon, 11 Mar 2024 11:58:53 +0000 Subject: [PATCH 37/57] add differenced latent model --- EpiAware/src/latentmodels/difflatentmodel.jl | 29 ++++++++++++++++++++ 1 file changed, 29 insertions(+) create mode 100644 EpiAware/src/latentmodels/difflatentmodel.jl diff --git a/EpiAware/src/latentmodels/difflatentmodel.jl b/EpiAware/src/latentmodels/difflatentmodel.jl new file mode 100644 index 000000000..1d38e6052 --- /dev/null +++ b/EpiAware/src/latentmodels/difflatentmodel.jl @@ -0,0 +1,29 @@ +struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel + model::T + init_prior::Priors + d::Int + + function DiffLatentModel(model::AbstractModel, init_prior::Priors) + d = length(init_prior) + return DiffLatentModel(model, init_prior, d) # Add the missing type parameter to the new function call + end + + function DiffLatentModel(model::AbstractModel, init_prior::Priors, d::Int) + @assert d>0 "d must be greater than 0" + @assert length(init_prior)==d "Length of init_prior must be equal to d" + return new{T <: AbstractModel}(model, init_prior, d) # Add the missing type parameter to the new function call + end +end + +@model function generate_latent(latent_model::DiffLatentModel, n) + d = latent_model.d + @assert n>d "n must be longer than d" + init_latent ~ latent_model.init_prior + + @submodel diff_latent, diff_latent_aux = generate_latent(latent_model.model, n - d) + + latent = vcat(init_latent, diff_latent) |> + cumsum + + return latent, (; init_latent, diff_latent_aux...) +end From 486c25c8d8683cd2867e5251511320ffde5408b9 Mon Sep 17 00:00:00 2001 From: Sam Date: Mon, 11 Mar 2024 13:56:33 +0000 Subject: [PATCH 38/57] add default diff latent priors --- EpiAware/src/latentmodels/difflatentmodel.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/EpiAware/src/latentmodels/difflatentmodel.jl b/EpiAware/src/latentmodels/difflatentmodel.jl index 1d38e6052..8687ba075 100644 --- a/EpiAware/src/latentmodels/difflatentmodel.jl +++ b/EpiAware/src/latentmodels/difflatentmodel.jl @@ -15,6 +15,10 @@ struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel end end +function default_diff_latent_priors(d::Int) + return (init_prior = [Normal(0.0, 1.0) for i in 1:d],) +end + @model function generate_latent(latent_model::DiffLatentModel, n) d = latent_model.d @assert n>d "n must be longer than d" From e50db06050f9733f434960804666b533caccf3f3 Mon Sep 17 00:00:00 2001 From: Sam Date: Mon, 11 Mar 2024 14:28:34 +0000 Subject: [PATCH 39/57] clean up comments --- EpiAware/src/latentmodels/difflatentmodel.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/EpiAware/src/latentmodels/difflatentmodel.jl b/EpiAware/src/latentmodels/difflatentmodel.jl index 8687ba075..14386a427 100644 --- a/EpiAware/src/latentmodels/difflatentmodel.jl +++ b/EpiAware/src/latentmodels/difflatentmodel.jl @@ -5,13 +5,13 @@ struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel function DiffLatentModel(model::AbstractModel, init_prior::Priors) d = length(init_prior) - return DiffLatentModel(model, init_prior, d) # Add the missing type parameter to the new function call + return DiffLatentModel(model, init_prior, d) end function DiffLatentModel(model::AbstractModel, init_prior::Priors, d::Int) @assert d>0 "d must be greater than 0" @assert length(init_prior)==d "Length of init_prior must be equal to d" - return new{T <: AbstractModel}(model, init_prior, d) # Add the missing type parameter to the new function call + return new{T <: AbstractModel}(model, init_prior, d) end end From 6d23f5227ee9fe9cbeb5501c0f5fa8cde287b945 Mon Sep 17 00:00:00 2001 From: Sam Date: Mon, 11 Mar 2024 18:43:37 +0000 Subject: [PATCH 40/57] fix simple example --- EpiAware/src/latentmodels/difflatentmodel.jl | 21 ++++-- EpiAware/src/latentmodels/latentmodels.jl | 1 + EpiAware/test/test_difflatentmodel.jl | 79 ++++++++++++++++++++ 3 files changed, 94 insertions(+), 7 deletions(-) create mode 100644 EpiAware/test/test_difflatentmodel.jl diff --git a/EpiAware/src/latentmodels/difflatentmodel.jl b/EpiAware/src/latentmodels/difflatentmodel.jl index 14386a427..579441822 100644 --- a/EpiAware/src/latentmodels/difflatentmodel.jl +++ b/EpiAware/src/latentmodels/difflatentmodel.jl @@ -1,5 +1,5 @@ -struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel - model::T +struct DiffLatentModel <: AbstractLatentModel + model::AbstractModel init_prior::Priors d::Int @@ -11,7 +11,7 @@ struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel function DiffLatentModel(model::AbstractModel, init_prior::Priors, d::Int) @assert d>0 "d must be greater than 0" @assert length(init_prior)==d "Length of init_prior must be equal to d" - return new{T <: AbstractModel}(model, init_prior, d) + return new(model, init_prior, d) end end @@ -22,12 +22,19 @@ end @model function generate_latent(latent_model::DiffLatentModel, n) d = latent_model.d @assert n>d "n must be longer than d" - init_latent ~ latent_model.init_prior + latent_init ~ latent_model.init_prior @submodel diff_latent, diff_latent_aux = generate_latent(latent_model.model, n - d) - latent = vcat(init_latent, diff_latent) |> - cumsum + return _combine_diff(latent_init, diff_latent, d), (; latent_init, diff_latent_aux...) +end + +function _combine_diff(init, diff, d) + combined = vcat(init, diff) + + for i in 1:d + combined = cumsum(combined) + end - return latent, (; init_latent, diff_latent_aux...) + return combined end diff --git a/EpiAware/src/latentmodels/latentmodels.jl b/EpiAware/src/latentmodels/latentmodels.jl index 7a2757966..8f5e56cfa 100644 --- a/EpiAware/src/latentmodels/latentmodels.jl +++ b/EpiAware/src/latentmodels/latentmodels.jl @@ -3,3 +3,4 @@ const Priors = Union{Distribution, Vector{<:Distribution}, Product} include("randomwalk.jl") include("autoregressive.jl") +include("difflatentmodel.jl") diff --git a/EpiAware/test/test_difflatentmodel.jl b/EpiAware/test/test_difflatentmodel.jl new file mode 100644 index 000000000..7f8f8ecc2 --- /dev/null +++ b/EpiAware/test/test_difflatentmodel.jl @@ -0,0 +1,79 @@ +@testitem "Testing default_diff_latent_priors" begin + using Distributions + + d = 3 + priors = EpiAware.default_diff_latent_priors(d) + + @test length(priors[:init_prior]) == d + for prior in priors[:init_prior] + @test prior isa Normal + @test prior.μ == 0.0 + @test prior.σ == 1.0 + end +end + +@testitem "Testing DiffLatentModel constructor" begin + using Distributions + + model = EpiAware.RandomWalk(Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) + init_prior = [Normal(0.0, 1.0), Normal(1.0, 2.0)] + diff_model = EpiAware.DiffLatentModel(model, init_prior) + + @test diff_model.model == model + @test diff_model.init_prior == init_prior + @test diff_model.d == 2 +end + +@testitem "Testing DiffLatentModel process" begin + using DynamicPPL, Turing + using Distributions + + n = 100 + d = 2 + model = EpiAware.RandomWalk(Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) + init_prior = [Normal(0.0, 1.0), Normal(1.0, 2.0)] + diff_model = EpiAware.DiffLatentModel(model, init_prior) + + latent_model = EpiAware.generate_latent(diff_model, n) + fixed_model = fix( + latent_model, (latent_init = [0.0, 1.0], σ²_RW = 1.0, init_rw_value = 0.0)) + + n_samples = 100 + samples = sample(fixed_model, Prior(), n_samples) |> + chn -> mapreduce(vcat, generated_quantities(fixed_model, chn)) do gen + gen[1] + end + + @test size(samples) == (n, n_samples) + @test all(samples[1, :] .≈ 0.0) + @test all(samples[2, :] .≈ 1.0) + @test all(samples[3, :] .≈ samples[2, :] .+ samples[1, :]) +end + +@testitem "Testing DiffLatentModel with AR process" begin + using DynamicPPL, Turing + using Distributions + + n = 100 + d = 2 + model = EpiAware.AR(EpiAware.default_ar_priors()[:damp_prior], + EpiAware.default_ar_priors()[:var_prior], + EpiAware.default_ar_priors()[:init_prior]) + init_prior = [Normal(0.0, 1.0), Normal(1.0, 2.0)] + diff_model = EpiAware.DiffLatentModel(model, init_prior) + + latent_model = EpiAware.generate_latent(diff_model, n) + fixed_model = fix(latent_model, + (latant_init = [0.0, 1.0], σ²_AR = 1.0, damp_AR = [0.8], ar_init = [0.0])) + + n_samples = 100 + samples = sample(fixed_model, Prior(), n_samples) |> + chn -> mapreduce(vcat, generated_quantities(fixed_model, chn)) do gen + gen[1] + end + + @test size(samples) == (n, n_samples) + @test all(samples[1, :] .≈ 0.0) + @test all(samples[2, :] .≈ 1.0) + @test all(samples[3, :] .≈ samples[2, :] .+ samples[1, :]) +end From 3d389fb3bd47053e955789e6c3af3037442bb112 Mon Sep 17 00:00:00 2001 From: Sam Date: Mon, 11 Mar 2024 19:01:41 +0000 Subject: [PATCH 41/57] improve draft tests --- EpiAware/test/test_difflatentmodel.jl | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/EpiAware/test/test_difflatentmodel.jl b/EpiAware/test/test_difflatentmodel.jl index 7f8f8ecc2..d0fad9aab 100644 --- a/EpiAware/test/test_difflatentmodel.jl +++ b/EpiAware/test/test_difflatentmodel.jl @@ -36,7 +36,7 @@ end latent_model = EpiAware.generate_latent(diff_model, n) fixed_model = fix( - latent_model, (latent_init = [0.0, 1.0], σ²_RW = 1.0, init_rw_value = 0.0)) + latent_model, (σ²_RW = 0, rw_init = 0.0)) n_samples = 100 samples = sample(fixed_model, Prior(), n_samples) |> @@ -47,7 +47,6 @@ end @test size(samples) == (n, n_samples) @test all(samples[1, :] .≈ 0.0) @test all(samples[2, :] .≈ 1.0) - @test all(samples[3, :] .≈ samples[2, :] .+ samples[1, :]) end @testitem "Testing DiffLatentModel with AR process" begin From d2a967fda0130d2fa1628af9e6be3eec98e616da Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Wed, 13 Mar 2024 14:37:49 +0000 Subject: [PATCH 42/57] Add normal util --- EpiAware/src/EpiAwareUtils/distributions.jl | 116 ++++++++++++++++++++ 1 file changed, 116 insertions(+) diff --git a/EpiAware/src/EpiAwareUtils/distributions.jl b/EpiAware/src/EpiAwareUtils/distributions.jl index 0fe1037ea..81a72ca50 100644 --- a/EpiAware/src/EpiAwareUtils/distributions.jl +++ b/EpiAware/src/EpiAwareUtils/distributions.jl @@ -73,3 +73,119 @@ function create_discrete_pmf(dist::Distribution; Δd = 1.0, D) ts .|> (t -> ∫F(dist, t, Δd)) |> diff |> p -> p ./ sum(p) end + +""" +Compute the negative moment generating function (MGF) for a given rate `r` and weights `w`. + +# Arguments +- `r`: The rate parameter. +- `w`: An abstract vector of weights. + +# Returns +The value of the negative MGF. + +""" +function neg_MGF(r, w::AbstractVector) + return sum([w[i] * exp(-r * i) for i in 1:length(w)]) +end + +function dneg_MGF_dr(r, w::AbstractVector) + return -sum([w[i] * i * exp(-r * i) for i in 1:length(w)]) +end + +""" +This function computes an approximation to the exponential growth rate `r` +given the reproductive ratio `R₀` and the discretized generation interval `w` with +discretized interval width `Δd`. This is based on the implicit solution of + +```math +G(r) - {1 \\over R_0} = 0. +``` + +where + +```math +G(r) = \\sum_{i=1}^n w_i e^{-r i}. +``` + +is the negative moment generating function (MGF) of the generation interval distribution. + +The two step approximation is based on: + 1. Direct solution of implicit equation for a small `r` approximation. + 2. Improving the approximation using Newton's method for a fixed number of steps `newton_steps`. + +Returns: +- The approximate value of `r`. +""" +function R_to_r(R₀, w::Vector{T}; newton_steps = 2, Δd = 1.0) where {T <: AbstractFloat} + mean_gen_time = dot(w, 1:length(w)) * Δd + # Small r approximation as initial guess + r_approx = (R₀ - 1) / (R₀ * mean_gen_time) + # Newton's method + for _ in 1:newton_steps + r_approx -= (R₀ * neg_MGF(r_approx, w) - 1) / (R₀ * dneg_MGF_dr(r_approx, w)) + end + return r_approx +end + +function R_to_r(R₀, epi_model::AbstractEpiModel; newton_steps = 2, Δd = 1.0) + R_to_r(R₀, epi_model.data.gen_int; newton_steps = newton_steps, Δd = Δd) +end + +""" + r_to_R(r, w) + +Compute the reproductive ratio given exponential growth rate `r` + and discretized generation interval `w`. + +# Arguments +- `r`: The exponential growth rate. +- `w`: discretized generation interval. + +# Returns +- The reproductive ratio. +""" +function r_to_R(r, w::AbstractVector) + return 1 / neg_MGF(r, w::AbstractVector) +end + +""" +Compute the mean-cluster factor negative binomial distribution. + +# Arguments +- `μ`: The mean of the distribution. +- `α`: The clustering factor parameter. + +# Returns +A `NegativeBinomial` distribution object. +""" +function NegativeBinomialMeanClust(μ, α) + ex_σ² = (α * μ^2) + 1e-6 + p = μ / (μ + ex_σ² + 1e-6) + r = μ^2 / ex_σ² + return NegativeBinomial(r, p) +end + +""" +Create a half-normal prior distribution with the specified mean. + +# Arguments +- `prior_mean::AbstractFloat`: The mean of the prior distribution. + +# Returns +- `Truncated{Normal}`: The half-normal prior distribution. + +""" +function _make_halfnormal_prior(prior_mean::AbstractFloat) + return truncated(Normal(0.0, prior_mean * sqrt(pi) / sqrt(2)), 0.0, Inf) +end + +""" +Add two normal distributions `x` and `y` together. + +# Returns +- `Normal`: The sum of `x` and `y` as a normal distribution. +""" +function _add_normals(x::Normal, y::Normal) + return Normal(x.μ + y.μ, sqrt(x.σ^2 + y.σ^2)) +end From 259909c9405a74cdd6d2ca674245873fc4db94eb Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Wed, 13 Mar 2024 14:38:43 +0000 Subject: [PATCH 43/57] Change DiffLatentModel constructor --- EpiAware/src/latentmodels/difflatentmodel.jl | 38 ++++++++++++++------ 1 file changed, 28 insertions(+), 10 deletions(-) diff --git a/EpiAware/src/latentmodels/difflatentmodel.jl b/EpiAware/src/latentmodels/difflatentmodel.jl index 579441822..8eca2ca29 100644 --- a/EpiAware/src/latentmodels/difflatentmodel.jl +++ b/EpiAware/src/latentmodels/difflatentmodel.jl @@ -1,17 +1,35 @@ -struct DiffLatentModel <: AbstractLatentModel - model::AbstractModel - init_prior::Priors +""" +The `DiffLatentModel` struct represents a differential latent model that is a subtype of `AbstractLatentModel`. + +## Fields +- `model::M`: The underlying latent model. +- `init_prior::P`: The initial priors for the latent variables. +- `d::Int`: The dimensionality of the latent variables. + +## Constructors +- `DiffLatentModel(;latentmodel, init_priors::Vector{<:Distribution})`: Constructs a `DiffLatentModel` object with the given latent model and initial priors. +- `DiffLatentModel(;latentmodel, init_prior_distribution::Distribution, d::Int)`: Constructs a `DiffLatentModel` object with the given latent model, initial prior distribution, and dimensionality. + +""" +struct DiffLatentModel{M, P} <: AbstractLatentModel + "Underlying latent model for the differenced process" + model::M + "The prior distribution for the initial latent variables." + init_prior::P + "Number of times differenced." d::Int - function DiffLatentModel(model::AbstractModel, init_prior::Priors) - d = length(init_prior) - return DiffLatentModel(model, init_prior, d) + function DiffLatentModel(latentmodel; init_prior_distribution::Distribution, d::Int) + init_priors = fill(init_prior_distribution, d) + return DiffLatentModel(; latentmodel = latentmodel, init_priors = init_priors) end - function DiffLatentModel(model::AbstractModel, init_prior::Priors, d::Int) - @assert d>0 "d must be greater than 0" - @assert length(init_prior)==d "Length of init_prior must be equal to d" - return new(model, init_prior, d) + function DiffLatentModel(; + latentmodel, init_priors::Vector{D} where {D <: Distribution}) + d = length(init_priors) + init_prior = all(first(init_priors) .== init_priors) ? + filldist(first(init_priors), d) : arraydist(init_priors) + return new{typeof(latentmodel), typeof(init_prior)}(latentmodel, init_prior, d) end end From 326256833f3e096f21ea995726ab1bc1f96f8640 Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Wed, 13 Mar 2024 14:39:03 +0000 Subject: [PATCH 44/57] distribution tests for DiffLatentModel on randomwalk --- EpiAware/test/test_difflatentmodel.jl | 59 ++++++++++++++++++--------- 1 file changed, 40 insertions(+), 19 deletions(-) diff --git a/EpiAware/test/test_difflatentmodel.jl b/EpiAware/test/test_difflatentmodel.jl index d0fad9aab..f01c37157 100644 --- a/EpiAware/test/test_difflatentmodel.jl +++ b/EpiAware/test/test_difflatentmodel.jl @@ -13,43 +13,67 @@ end @testitem "Testing DiffLatentModel constructor" begin - using Distributions + using Distributions, Turing model = EpiAware.RandomWalk(Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) - init_prior = [Normal(0.0, 1.0), Normal(1.0, 2.0)] - diff_model = EpiAware.DiffLatentModel(model, init_prior) + @testset "Testing DiffLatentModel with vector of priors" begin + init_priors = [Normal(0.0, 1.0), Normal(1.0, 2.0)] + diff_model = EpiAware.DiffLatentModel(; + latentmodel = model, init_priors = init_priors) + + @test diff_model.model == model + @test diff_model.init_prior == arraydist(init_priors) + @test diff_model.d == 2 + end + + @testset "Testing DiffLatentModel with single prior and d" begin + d = 3 + init_prior = Normal() + diff_model = EpiAware.DiffLatentModel( + model; init_prior_distribution = init_prior, d = d) - @test diff_model.model == model - @test diff_model.init_prior == init_prior - @test diff_model.d == 2 + @test diff_model.model == model + @test diff_model.init_prior == filldist(init_prior, d) + @test diff_model.d == d + end end @testitem "Testing DiffLatentModel process" begin using DynamicPPL, Turing using Distributions + using HypothesisTests: ExactOneSampleKSTest, pvalue + using EpiAware: _add_normals n = 100 d = 2 model = EpiAware.RandomWalk(Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) - init_prior = [Normal(0.0, 1.0), Normal(1.0, 2.0)] - diff_model = EpiAware.DiffLatentModel(model, init_prior) + init_priors = [Normal(0.0, 1.0), Normal(1.0, 2.0)] + diff_model = EpiAware.DiffLatentModel(latentmodel = model, init_priors = init_priors) latent_model = EpiAware.generate_latent(diff_model, n) fixed_model = fix( latent_model, (σ²_RW = 0, rw_init = 0.0)) - n_samples = 100 + n_samples = 2000 samples = sample(fixed_model, Prior(), n_samples) |> - chn -> mapreduce(vcat, generated_quantities(fixed_model, chn)) do gen + chn -> mapreduce(hcat, generated_quantities(fixed_model, chn)) do gen gen[1] end + #Because of the recursive d-times cumsum to undifference the process, + #The distribution of the second day should be d lots of first day init distribution + #Plus day two distribution + day2_dist = foldl((x, y) -> _add_normals(x, init_priors[1]), 1:d, init = init_priors[2]) + + ks_test_pval_day1 = ExactOneSampleKSTest(samples[1, :], init_priors[1]) |> pvalue + ks_test_pval_day2 = ExactOneSampleKSTest(samples[2, :], day2_dist) |> pvalue + @test size(samples) == (n, n_samples) - @test all(samples[1, :] .≈ 0.0) - @test all(samples[2, :] .≈ 1.0) + @test ks_test_pval_day1 > 1e-6 #Very unlikely to fail if the model is correctly implemented + @test ks_test_pval_day2 > 1e-6 #Very unlikely to fail if the model is correctly implemented end -@testitem "Testing DiffLatentModel with AR process" begin +@testitem "Testing DiffLatentModel runs with AR process" begin using DynamicPPL, Turing using Distributions @@ -58,8 +82,8 @@ end model = EpiAware.AR(EpiAware.default_ar_priors()[:damp_prior], EpiAware.default_ar_priors()[:var_prior], EpiAware.default_ar_priors()[:init_prior]) - init_prior = [Normal(0.0, 1.0), Normal(1.0, 2.0)] - diff_model = EpiAware.DiffLatentModel(model, init_prior) + init_priors = [Normal(0.0, 1.0), Normal(1.0, 2.0)] + diff_model = EpiAware.DiffLatentModel(latentmodel = model, init_priors = init_priors) latent_model = EpiAware.generate_latent(diff_model, n) fixed_model = fix(latent_model, @@ -67,12 +91,9 @@ end n_samples = 100 samples = sample(fixed_model, Prior(), n_samples) |> - chn -> mapreduce(vcat, generated_quantities(fixed_model, chn)) do gen + chn -> mapreduce(hcat, generated_quantities(fixed_model, chn)) do gen gen[1] end @test size(samples) == (n, n_samples) - @test all(samples[1, :] .≈ 0.0) - @test all(samples[2, :] .≈ 1.0) - @test all(samples[3, :] .≈ samples[2, :] .+ samples[1, :]) end From 80405ae84662ee5f967660b3081de55773e819a3 Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Thu, 14 Mar 2024 09:49:27 +0000 Subject: [PATCH 45/57] doc strings for DiffLatentModel --- EpiAware/src/latentmodels/difflatentmodel.jl | 94 +++++++++++++++++--- 1 file changed, 80 insertions(+), 14 deletions(-) diff --git a/EpiAware/src/latentmodels/difflatentmodel.jl b/EpiAware/src/latentmodels/difflatentmodel.jl index 8eca2ca29..d7b23561f 100644 --- a/EpiAware/src/latentmodels/difflatentmodel.jl +++ b/EpiAware/src/latentmodels/difflatentmodel.jl @@ -1,17 +1,82 @@ -""" -The `DiffLatentModel` struct represents a differential latent model that is a subtype of `AbstractLatentModel`. -## Fields -- `model::M`: The underlying latent model. -- `init_prior::P`: The initial priors for the latent variables. -- `d::Int`: The dimensionality of the latent variables. +@doc raw""" +Model the latent process as a `d`-fold differenced version of another process. + +## Mathematical specification + +Let ``\Delta`` be the differencing operator. If ``\tilde{Z}_t`` is a realisation of +undifferenced latent model supplied to `DiffLatentModel`, then the differenced process is +given by, + +```math +\Delta^{(d)} Z_t = \tilde{Z}_t, \quad t = d+1, \ldots. +``` + +We can recover ``Z_t`` by applying the inverse differencing operator ``\Delta^{-1}``, which +corresponds to the cumulative sum operator `cumsum` in Julia, `d`-times. The `d` initial +terms ``Z_1, \ldots, Z_d`` are inferred. ## Constructors -- `DiffLatentModel(;latentmodel, init_priors::Vector{<:Distribution})`: Constructs a `DiffLatentModel` object with the given latent model and initial priors. -- `DiffLatentModel(;latentmodel, init_prior_distribution::Distribution, d::Int)`: Constructs a `DiffLatentModel` object with the given latent model, initial prior distribution, and dimensionality. + +- `DiffLatentModel(latentmodel, init_prior_distribution::Distribution; d::Int)` + Constructs a `DiffLatentModel` for `d`-fold differencing with `latentmodel` as the + undifferenced latent process. All initial terms have common prior + `init_prior_distribution`. +- `DiffLatentModel(;undiffmodel, init_priors::Vector{D} where {D <: Distribution})` + Constructs a `DiffLatentModel` for `d`-fold differencing with `latentmodel` as the + undifferenced latent process. The `d` initial terms have priors given by the vector + `init_priors`, therefore `length(init_priors)` sets `d`. + +## Example usage with `generate_latent` + +`generate_latent` can be used to construct a `Turing` model for the differenced latent +process. In this example, the underlying undifferenced process is a `RandomWalk` model. + + +First, we construct a `RandomWalk` struct with an initial value prior and a step size +standard deviation prior. + +```julia +using Distributions, EpiAware +rw = RandomWalk(Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) +``` + +Then, we can use `DiffLatentModel` to construct a `DiffLatentModel` for `d`-fold +differenced process with `rw` as the undifferenced latent process. + +We have two constructor options for `DiffLatentModel`. The first option is to supply a +common prior distribution for the initial terms and specify `d` as follows: + +```julia +diff_model = DiffLatentModel(rw, Normal(); d = 2) +``` + +Or we can supply a vector of priors for the initial terms and `d` is inferred as follows: +```julia +diff_model2 = DiffLatentModel(;undiffmodel = rw, init_priors = [Normal(), Normal()]) +``` + +Then, we can use `generate_latent` to construct a Turing model for the differenced latent +process generating a length `n` process, + +```julia +# Construct a Turing model +n = 100 +difference_mdl = generate_latent(diff_model, n) +``` + +Now we can use the `Turing` PPL API to sample underlying parameters and generate the +unobserved latent process. + +```julia +#Sample random parameters from prior +θ = rand(difference_mdl) +#Get a sampled latent process as a generated quantity from the model +Z_t = generated_quantities(difference_mdl, θ) +``` """ -struct DiffLatentModel{M, P} <: AbstractLatentModel +struct DiffLatentModel{M <: AbstractLatentModel, P} <: AbstractLatentModel "Underlying latent model for the differenced process" model::M "The prior distribution for the initial latent variables." @@ -19,17 +84,18 @@ struct DiffLatentModel{M, P} <: AbstractLatentModel "Number of times differenced." d::Int - function DiffLatentModel(latentmodel; init_prior_distribution::Distribution, d::Int) + function DiffLatentModel( + undiffmodel::AbstractLatentModel, init_prior_distribution::Distribution; d::Int) init_priors = fill(init_prior_distribution, d) - return DiffLatentModel(; latentmodel = latentmodel, init_priors = init_priors) + return DiffLatentModel(; undiffmodel = undiffmodel, init_priors = init_priors) end - function DiffLatentModel(; - latentmodel, init_priors::Vector{D} where {D <: Distribution}) + function DiffLatentModel(; undiffmodel::AbstractLatentModel, + init_priors::Vector{D} where {D <: Distribution}) d = length(init_priors) init_prior = all(first(init_priors) .== init_priors) ? filldist(first(init_priors), d) : arraydist(init_priors) - return new{typeof(latentmodel), typeof(init_prior)}(latentmodel, init_prior, d) + return new{typeof(undiffmodel), typeof(init_prior)}(undiffmodel, init_prior, d) end end From b21b9207a4a4d86571a8b6981d84229b2ee283fb Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Thu, 14 Mar 2024 09:51:08 +0000 Subject: [PATCH 46/57] remove unnecessary default prior creation --- EpiAware/src/latentmodels/difflatentmodel.jl | 4 ---- 1 file changed, 4 deletions(-) diff --git a/EpiAware/src/latentmodels/difflatentmodel.jl b/EpiAware/src/latentmodels/difflatentmodel.jl index d7b23561f..55d5772a2 100644 --- a/EpiAware/src/latentmodels/difflatentmodel.jl +++ b/EpiAware/src/latentmodels/difflatentmodel.jl @@ -99,10 +99,6 @@ struct DiffLatentModel{M <: AbstractLatentModel, P} <: AbstractLatentModel end end -function default_diff_latent_priors(d::Int) - return (init_prior = [Normal(0.0, 1.0) for i in 1:d],) -end - @model function generate_latent(latent_model::DiffLatentModel, n) d = latent_model.d @assert n>d "n must be longer than d" From b7c5abab438fe1042156e1d068dddca45786c441 Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Thu, 14 Mar 2024 09:51:50 +0000 Subject: [PATCH 47/57] style change to use _ for unused loop var --- EpiAware/src/latentmodels/difflatentmodel.jl | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/EpiAware/src/latentmodels/difflatentmodel.jl b/EpiAware/src/latentmodels/difflatentmodel.jl index 55d5772a2..dfb3b45cf 100644 --- a/EpiAware/src/latentmodels/difflatentmodel.jl +++ b/EpiAware/src/latentmodels/difflatentmodel.jl @@ -111,10 +111,8 @@ end function _combine_diff(init, diff, d) combined = vcat(init, diff) - - for i in 1:d + for _ in 1:d combined = cumsum(combined) end - return combined end From 60a673df7f45391dfd1557bbb122e831c71a2fda Mon Sep 17 00:00:00 2001 From: Samuel Brand Date: Thu, 14 Mar 2024 10:00:27 +0000 Subject: [PATCH 48/57] doc string for generated_latent method --- EpiAware/src/latentmodels/difflatentmodel.jl | 70 +++++++++++++++++++- 1 file changed, 69 insertions(+), 1 deletion(-) diff --git a/EpiAware/src/latentmodels/difflatentmodel.jl b/EpiAware/src/latentmodels/difflatentmodel.jl index dfb3b45cf..91833a71c 100644 --- a/EpiAware/src/latentmodels/difflatentmodel.jl +++ b/EpiAware/src/latentmodels/difflatentmodel.jl @@ -72,7 +72,8 @@ unobserved latent process. #Sample random parameters from prior θ = rand(difference_mdl) #Get a sampled latent process as a generated quantity from the model -Z_t = generated_quantities(difference_mdl, θ) +(Z_t, _) = generated_quantities(difference_mdl, θ) +Z_t ``` """ @@ -99,6 +100,73 @@ struct DiffLatentModel{M <: AbstractLatentModel, P} <: AbstractLatentModel end end +""" +Generate a `Turing` model for `n`-step latent process ``Z_t`` using a differenced +latent model defined by `latent_model`. + +# Arguments +- `latent_model::DiffLatentModel`: The differential latent model. +- `n`: The length of the latent variables. + +## Turing model specifications + +### Sampled random variables +- `latent_init`: The initial latent process variables. +- Other random variables defined by `model<:AbstractLatentModel` field of the undifferenced + model. + +### Generated quantities +- A tuple containing the generated latent process as its first argument + and a `NamedTuple` of sampled auxiliary variables as second argument. + +# Example usage with `DiffLatentModel` model constructor + +`generate_latent` can be used to construct a `Turing` model for the differenced latent +process. In this example, the underlying undifferenced process is a `RandomWalk` model. + +First, we construct a `RandomWalk` struct with an initial value prior and a step size +standard deviation prior. + +```julia +using Distributions, EpiAware +rw = RandomWalk(Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) +``` + +Then, we can use `DiffLatentModel` to construct a `DiffLatentModel` for `d`-fold +differenced process with `rw` as the undifferenced latent process. + +We have two constructor options for `DiffLatentModel`. The first option is to supply a +common prior distribution for the initial terms and specify `d` as follows: + +```julia +diff_model = DiffLatentModel(rw, Normal(); d = 2) +``` + +Or we can supply a vector of priors for the initial terms and `d` is inferred as follows: +```julia +diff_model2 = DiffLatentModel(;undiffmodel = rw, init_priors = [Normal(), Normal()]) +``` + +Then, we can use `generate_latent` to construct a Turing model for the differenced latent +process generating a length `n` process, + +```julia +# Construct a Turing model +n = 100 +difference_mdl = generate_latent(diff_model, n) +``` + +Now we can use the `Turing` PPL API to sample underlying parameters and generate the +unobserved latent process. + +```julia +#Sample random parameters from prior +θ = rand(difference_mdl) +#Get a sampled latent process as a generated quantity from the model +(Z_t, _) = generated_quantities(difference_mdl, θ) +Z_t +``` +""" @model function generate_latent(latent_model::DiffLatentModel, n) d = latent_model.d @assert n>d "n must be longer than d" From 8b1e20a81af88b78e98fab742e13f5c44530341f Mon Sep 17 00:00:00 2001 From: Sam Date: Thu, 14 Mar 2024 11:39:58 +0000 Subject: [PATCH 49/57] clean up further merging issues --- EpiAware/src/latentmodels/randomwalk.jl | 87 ------------------------- 1 file changed, 87 deletions(-) diff --git a/EpiAware/src/latentmodels/randomwalk.jl b/EpiAware/src/latentmodels/randomwalk.jl index ad5f697b3..1dda224b2 100644 --- a/EpiAware/src/latentmodels/randomwalk.jl +++ b/EpiAware/src/latentmodels/randomwalk.jl @@ -20,90 +20,3 @@ end end return rw, (; σ_RW, rw_init) end - -# Define the Priors type alias -const Priors = Union{Distribution, Vector{<:Distribution}, Product} - -struct AR <: AbstractLatentModel - """A distribution representing the prior distribution of the damping factors.""" - damp_prior::Priors - - """A distribution representing the prior distribution of the variance.""" - var_prior::Distribution - - """A distribution representing the prior distribution of the initial values.""" - init_prior::Priors - - """ - The order of the AR process, determined by the length of `damp_prior`. - """ - p::Int - - function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors) - p = length(damp_prior) - return new(damp_prior, var_prior, init_prior, p) - end - - function AR(damp_prior::Priors, var_prior::Distribution, init_prior::Priors, p::Int) - @assert length(init_prior)==p "Dimension of init_prior must be equal to the order of the AR process" - return new(damp_prior, var_prior, init_prior, p) - end -end - -function default_ar_priors() - return (:damp_prior => [truncated(Normal(0.0, 0.05), 0.0, 1)], - :var_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), - :init_prior => Normal()) |> Dict -end - -@model function generate_latent(latent_model::AR, n) - p = latent_model.p - ϵ_t ~ MvNormal(ones(n - p)) - σ²_AR ~ latent_model.var_prior - ar_init ~ latent_model.init_prior - damp_AR ~ latent_model.damp_prior - σ_AR = sqrt(σ²_AR) - - @assert n>p "n must be longer than p" - - # Initialize the AR series with the initial values - ar = Vector{Float64}(undef, n) - ar[1:p] = ar_init - - # Generate the rest of the AR series - for t in (p + 1):n - ar[t] = damp_AR' * ar[(t - p):(t - 1)] + σ_AR * ϵ_t[t - p] - end - - return ar, (; σ_AR, ar_init, damp_AR) -end - -struct DiffLatentModel{T <: AbstractModel} <: AbstractLatentModel - model::T - d::Int - init_prior::Priors - - function DiffLatentModel(model::T, init_prior::Priors) - d = length(init_prior) - return new(model, d, init_prior) - end - - function DiffLatentModel(model::T, d::Int, init_prior::Priors) - @assert d>0 "d must be greater than 0" - @assert length(init_prior)==d "Length of init_prior must be equal to d" - return new(model, d, init_prior) - end -end - -@model function generate_latent(latent_model::DiffLatentModel, n) - d = latent_model.d - @assert n>d "n must be longer than d" - init_latent ~ latent_model.init_prior - - @submodel diff_latent, diff_latent_aux = generate_latent(latent_model.model, n - d) - - latent = vcat(init_latent, diff_latent) |> - cumsum - # Return the reconstructed series and the parameters - return latent, (; init_latent, diff_latent_aux...) -end From bab886ee166ecb330e6c4de9d53469e9a6291eb1 Mon Sep 17 00:00:00 2001 From: Sam Date: Thu, 14 Mar 2024 11:45:54 +0000 Subject: [PATCH 50/57] clena out utils duplicate --- EpiAware/src/EpiAwareUtils/scan.jl | 34 ------------------------------ 1 file changed, 34 deletions(-) delete mode 100644 EpiAware/src/EpiAwareUtils/scan.jl diff --git a/EpiAware/src/EpiAwareUtils/scan.jl b/EpiAware/src/EpiAwareUtils/scan.jl deleted file mode 100644 index fbd76c5a6..000000000 --- a/EpiAware/src/EpiAwareUtils/scan.jl +++ /dev/null @@ -1,34 +0,0 @@ -""" -Apply `f` to each element of `xs` and accumulate the results. - -`f` must be a [callable](https://docs.julialang.org/en/v1/manual/methods/#Function-like-objects) - on a sub-type of `AbstractModel`. - -### Design note -`scan` is being restricted to `AbstractModel` sub-types to ensure: - 1. That compiler specialization is [activated](https://docs.julialang.org/en/v1/manual/performance-tips/#Be-aware-of-when-Julia-avoids-specializing) - 2. Also avoids potential compiler [overhead](https://docs.julialang.org/en/v1/devdocs/functions/#compiler-efficiency-issues) - from specialisation on `f<: Function`. - - - -# Arguments -- `f`: A callable/functor that takes two arguments, `carry` and `x`, and returns a new - `carry` and a result `y`. -- `init`: The initial value for the `carry` variable. -- `xs`: An iterable collection of elements. - -# Returns -- `ys`: An array containing the results of applying `f` to each element of `xs`. -- `carry`: The final value of the `carry` variable after processing all elements of `xs`. - -""" -function scan(f::F, init, xs) where {F <: EpiAwareBase.AbstractModel} - carry = init - ys = similar(xs) - for (i, x) in enumerate(xs) - carry, y = f(carry, x) - ys[i] = y - end - return ys, carry -end From a3439123fbd230240bc7ab4bfcba4b5c56b01839 Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 15 Mar 2024 10:44:30 +0000 Subject: [PATCH 51/57] update to new module structure and utils --- EpiAware/src/EpiAwareUtils/distributions.jl | 10 ---- .../src/EpiLatentModels/EpiLatentModels.jl | 1 + .../difflatentmodel.jl | 19 +++++--- EpiAware/src/abstract-functions.jl | 17 ------- EpiAware/src/latentmodels/autoregressive.jl | 46 ------------------- EpiAware/src/latentmodels/latentmodels.jl | 6 --- EpiAware/src/latentmodels/randomwalk.jl | 22 --------- EpiAware/test/test_difflatentmodel.jl | 10 ++++ 8 files changed, 23 insertions(+), 108 deletions(-) rename EpiAware/src/{latentmodels => EpiLatentModels}/difflatentmodel.jl (90%) delete mode 100644 EpiAware/src/abstract-functions.jl delete mode 100644 EpiAware/src/latentmodels/autoregressive.jl delete mode 100644 EpiAware/src/latentmodels/latentmodels.jl delete mode 100644 EpiAware/src/latentmodels/randomwalk.jl diff --git a/EpiAware/src/EpiAwareUtils/distributions.jl b/EpiAware/src/EpiAwareUtils/distributions.jl index 81a72ca50..8a2baac39 100644 --- a/EpiAware/src/EpiAwareUtils/distributions.jl +++ b/EpiAware/src/EpiAwareUtils/distributions.jl @@ -179,13 +179,3 @@ Create a half-normal prior distribution with the specified mean. function _make_halfnormal_prior(prior_mean::AbstractFloat) return truncated(Normal(0.0, prior_mean * sqrt(pi) / sqrt(2)), 0.0, Inf) end - -""" -Add two normal distributions `x` and `y` together. - -# Returns -- `Normal`: The sum of `x` and `y` as a normal distribution. -""" -function _add_normals(x::Normal, y::Normal) - return Normal(x.μ + y.μ, sqrt(x.σ^2 + y.σ^2)) -end diff --git a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl index 71b317641..788cf9963 100644 --- a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl +++ b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl @@ -13,5 +13,6 @@ export RandomWalk, AR include("docstrings.jl") include("randomwalk.jl") include("autoregressive.jl") +include("difflatentmodel.jl") include("utils.jl") end diff --git a/EpiAware/src/latentmodels/difflatentmodel.jl b/EpiAware/src/EpiLatentModels/difflatentmodel.jl similarity index 90% rename from EpiAware/src/latentmodels/difflatentmodel.jl rename to EpiAware/src/EpiLatentModels/difflatentmodel.jl index 91833a71c..ccebbde6a 100644 --- a/EpiAware/src/latentmodels/difflatentmodel.jl +++ b/EpiAware/src/EpiLatentModels/difflatentmodel.jl @@ -86,17 +86,22 @@ struct DiffLatentModel{M <: AbstractLatentModel, P} <: AbstractLatentModel d::Int function DiffLatentModel( - undiffmodel::AbstractLatentModel, init_prior_distribution::Distribution; d::Int) - init_priors = fill(init_prior_distribution, d) - return DiffLatentModel(; undiffmodel = undiffmodel, init_priors = init_priors) + model::AbstractLatentModel, init_prior::Distribution; d::Int) + init_priors = fill(init_prior, d) + return DiffLatentModel(; model = model, init_priors = init_priors) end - function DiffLatentModel(; undiffmodel::AbstractLatentModel, + function DiffLatentModel(; model::AbstractLatentModel, init_priors::Vector{D} where {D <: Distribution}) d = length(init_priors) - init_prior = all(first(init_priors) .== init_priors) ? - filldist(first(init_priors), d) : arraydist(init_priors) - return new{typeof(undiffmodel), typeof(init_prior)}(undiffmodel, init_prior, d) + init_prior = _expand_dist(init_priors) + return AR(model, init_prior, d) + end + + function DiffLatentModel(model::AbstractLatentModel, init_prior::Distribution, d::Int) + @assert d>0 "d must be greater than 0" + @assert d==length(init_prior) "d must be equal to the length of init_prior" + new{typeof(model), typeof(init_prior)}(model, init_prior, d) end end diff --git a/EpiAware/src/abstract-functions.jl b/EpiAware/src/abstract-functions.jl deleted file mode 100644 index a996ac92b..000000000 --- a/EpiAware/src/abstract-functions.jl +++ /dev/null @@ -1,17 +0,0 @@ -function generate_latent(latent_model::AbstractLatentModel, n) - @info "No concrete implementation for generate_latent is defined." - return nothing -end - -function generate_latent_infs(epi_model::AbstractEpiModel, latent_model) - @info "No concrete implementation for `generate_latent_infs` is defined." - return nothing -end - -function generate_observations(observation_model::AbstractObservationModel, - y_t, - I_t; - pos_shift) - @info "No concrete implementation for generate_observations is defined." - return nothing -end diff --git a/EpiAware/src/latentmodels/autoregressive.jl b/EpiAware/src/latentmodels/autoregressive.jl deleted file mode 100644 index 04927cc9f..000000000 --- a/EpiAware/src/latentmodels/autoregressive.jl +++ /dev/null @@ -1,46 +0,0 @@ -@doc raw" -The autoregressive (AR) model struct. -" -@with_kw struct AR{D <: Priors, S <: Distribution, I <: Priors} <: AbstractLatentModel - damp_prior::D = truncated(Normal(0.0, 0.05), 0.0, 1) - std_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf) - init_prior::I = Normal() - @assert length(damp_prior)==length(init_prior) "damp_prior and init_prior must have the same length" -end - -@doc raw" -Generate a latent AR series. - -# Arguments - -- `latent_model::AR`: The AR model. -- `n::Int`: The length of the AR series. - -# Returns -- `ar::Vector{Float64}`: The generated AR series. -- `params::NamedTuple`: A named tuple containing the generated parameters (`σ_AR`, `ar_init`, `damp_AR`). - -# Notes -- The length of `damp_prior` and `init_prior` must be the same. -- `n` must be longer than the order of the autoregressive process. -" -@model function generate_latent(latent_model::AR, n) - p = length(latent_model.damp_prior) - ϵ_t ~ MvNormal(ones(n - p)) - σ_AR ~ latent_model.std_prior - ar_init ~ latent_model.init_prior - damp_AR ~ latent_model.damp_prior - - @assert n>p "n must be longer than order of the autoregressive process" - - # Initialize the AR series with the initial values - ar = Vector{Float64}(undef, n) - ar[1:p] = ar_init - - # Generate the rest of the AR series - for t in (p + 1):n - ar[t] = damp_AR' * ar[(t - p):(t - 1)] + σ_AR * ϵ_t[t - p] - end - - return ar, (; σ_AR, ar_init, damp_AR) -end diff --git a/EpiAware/src/latentmodels/latentmodels.jl b/EpiAware/src/latentmodels/latentmodels.jl deleted file mode 100644 index 8f5e56cfa..000000000 --- a/EpiAware/src/latentmodels/latentmodels.jl +++ /dev/null @@ -1,6 +0,0 @@ -# Define the Priors type alias -const Priors = Union{Distribution, Vector{<:Distribution}, Product} - -include("randomwalk.jl") -include("autoregressive.jl") -include("difflatentmodel.jl") diff --git a/EpiAware/src/latentmodels/randomwalk.jl b/EpiAware/src/latentmodels/randomwalk.jl deleted file mode 100644 index 1dda224b2..000000000 --- a/EpiAware/src/latentmodels/randomwalk.jl +++ /dev/null @@ -1,22 +0,0 @@ -struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractLatentModel - init_prior::D - std_prior::S -end - -function default_rw_priors() - return (:var_RW_prior => truncated(Normal(0.0, 0.05), 0.0, Inf), - :init_rw_value_prior => Normal()) |> Dict -end - -@model function generate_latent(latent_model::RandomWalk, n) - ϵ_t ~ MvNormal(ones(n)) - σ_RW ~ latent_model.std_prior - rw_init ~ latent_model.init_prior - rw = Vector{eltype(ϵ_t)}(undef, n) - - rw[1] = rw_init + σ_RW * ϵ_t[1] - for t in 2:n - rw[t] = rw[t - 1] + σ_RW * ϵ_t[t] - end - return rw, (; σ_RW, rw_init) -end diff --git a/EpiAware/test/test_difflatentmodel.jl b/EpiAware/test/test_difflatentmodel.jl index f01c37157..b2bd17d6e 100644 --- a/EpiAware/test/test_difflatentmodel.jl +++ b/EpiAware/test/test_difflatentmodel.jl @@ -62,6 +62,16 @@ end #Because of the recursive d-times cumsum to undifference the process, #The distribution of the second day should be d lots of first day init distribution + """ +Add two normal distributions `x` and `y` together. + +# Returns +- `Normal`: The sum of `x` and `y` as a normal distribution. +""" + function _add_normals(x::Normal, y::Normal) + return Normal(x.μ + y.μ, sqrt(x.σ^2 + y.σ^2)) + end + #Plus day two distribution day2_dist = foldl((x, y) -> _add_normals(x, init_priors[1]), 1:d, init = init_priors[2]) From 5db0d93f9169d03791d543f5e8ed0a8c7ab2e1d9 Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 15 Mar 2024 10:48:54 +0000 Subject: [PATCH 52/57] add back in scan.jl --- EpiAware/src/EpiAwareUtils/scan.jl | 34 ++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) create mode 100644 EpiAware/src/EpiAwareUtils/scan.jl diff --git a/EpiAware/src/EpiAwareUtils/scan.jl b/EpiAware/src/EpiAwareUtils/scan.jl new file mode 100644 index 000000000..fbd76c5a6 --- /dev/null +++ b/EpiAware/src/EpiAwareUtils/scan.jl @@ -0,0 +1,34 @@ +""" +Apply `f` to each element of `xs` and accumulate the results. + +`f` must be a [callable](https://docs.julialang.org/en/v1/manual/methods/#Function-like-objects) + on a sub-type of `AbstractModel`. + +### Design note +`scan` is being restricted to `AbstractModel` sub-types to ensure: + 1. That compiler specialization is [activated](https://docs.julialang.org/en/v1/manual/performance-tips/#Be-aware-of-when-Julia-avoids-specializing) + 2. Also avoids potential compiler [overhead](https://docs.julialang.org/en/v1/devdocs/functions/#compiler-efficiency-issues) + from specialisation on `f<: Function`. + + + +# Arguments +- `f`: A callable/functor that takes two arguments, `carry` and `x`, and returns a new + `carry` and a result `y`. +- `init`: The initial value for the `carry` variable. +- `xs`: An iterable collection of elements. + +# Returns +- `ys`: An array containing the results of applying `f` to each element of `xs`. +- `carry`: The final value of the `carry` variable after processing all elements of `xs`. + +""" +function scan(f::F, init, xs) where {F <: EpiAwareBase.AbstractModel} + carry = init + ys = similar(xs) + for (i, x) in enumerate(xs) + carry, y = f(carry, x) + ys[i] = y + end + return ys, carry +end From cbc892b165af4052bd147cc95ee7ca4dc20505ba Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 15 Mar 2024 10:51:21 +0000 Subject: [PATCH 53/57] clean up more leakage from main/git hygiene --- EpiAware/src/EpiAwareUtils/distributions.jl | 106 -------------------- 1 file changed, 106 deletions(-) diff --git a/EpiAware/src/EpiAwareUtils/distributions.jl b/EpiAware/src/EpiAwareUtils/distributions.jl index 8a2baac39..0fe1037ea 100644 --- a/EpiAware/src/EpiAwareUtils/distributions.jl +++ b/EpiAware/src/EpiAwareUtils/distributions.jl @@ -73,109 +73,3 @@ function create_discrete_pmf(dist::Distribution; Δd = 1.0, D) ts .|> (t -> ∫F(dist, t, Δd)) |> diff |> p -> p ./ sum(p) end - -""" -Compute the negative moment generating function (MGF) for a given rate `r` and weights `w`. - -# Arguments -- `r`: The rate parameter. -- `w`: An abstract vector of weights. - -# Returns -The value of the negative MGF. - -""" -function neg_MGF(r, w::AbstractVector) - return sum([w[i] * exp(-r * i) for i in 1:length(w)]) -end - -function dneg_MGF_dr(r, w::AbstractVector) - return -sum([w[i] * i * exp(-r * i) for i in 1:length(w)]) -end - -""" -This function computes an approximation to the exponential growth rate `r` -given the reproductive ratio `R₀` and the discretized generation interval `w` with -discretized interval width `Δd`. This is based on the implicit solution of - -```math -G(r) - {1 \\over R_0} = 0. -``` - -where - -```math -G(r) = \\sum_{i=1}^n w_i e^{-r i}. -``` - -is the negative moment generating function (MGF) of the generation interval distribution. - -The two step approximation is based on: - 1. Direct solution of implicit equation for a small `r` approximation. - 2. Improving the approximation using Newton's method for a fixed number of steps `newton_steps`. - -Returns: -- The approximate value of `r`. -""" -function R_to_r(R₀, w::Vector{T}; newton_steps = 2, Δd = 1.0) where {T <: AbstractFloat} - mean_gen_time = dot(w, 1:length(w)) * Δd - # Small r approximation as initial guess - r_approx = (R₀ - 1) / (R₀ * mean_gen_time) - # Newton's method - for _ in 1:newton_steps - r_approx -= (R₀ * neg_MGF(r_approx, w) - 1) / (R₀ * dneg_MGF_dr(r_approx, w)) - end - return r_approx -end - -function R_to_r(R₀, epi_model::AbstractEpiModel; newton_steps = 2, Δd = 1.0) - R_to_r(R₀, epi_model.data.gen_int; newton_steps = newton_steps, Δd = Δd) -end - -""" - r_to_R(r, w) - -Compute the reproductive ratio given exponential growth rate `r` - and discretized generation interval `w`. - -# Arguments -- `r`: The exponential growth rate. -- `w`: discretized generation interval. - -# Returns -- The reproductive ratio. -""" -function r_to_R(r, w::AbstractVector) - return 1 / neg_MGF(r, w::AbstractVector) -end - -""" -Compute the mean-cluster factor negative binomial distribution. - -# Arguments -- `μ`: The mean of the distribution. -- `α`: The clustering factor parameter. - -# Returns -A `NegativeBinomial` distribution object. -""" -function NegativeBinomialMeanClust(μ, α) - ex_σ² = (α * μ^2) + 1e-6 - p = μ / (μ + ex_σ² + 1e-6) - r = μ^2 / ex_σ² - return NegativeBinomial(r, p) -end - -""" -Create a half-normal prior distribution with the specified mean. - -# Arguments -- `prior_mean::AbstractFloat`: The mean of the prior distribution. - -# Returns -- `Truncated{Normal}`: The half-normal prior distribution. - -""" -function _make_halfnormal_prior(prior_mean::AbstractFloat) - return truncated(Normal(0.0, prior_mean * sqrt(pi) / sqrt(2)), 0.0, Inf) -end From 9b379ed98ba96aa2920c9deac1c9e17edf7546dc Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 15 Mar 2024 10:59:41 +0000 Subject: [PATCH 54/57] add defaults to difflatentmodel --- EpiAware/src/EpiLatentModels/difflatentmodel.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/EpiAware/src/EpiLatentModels/difflatentmodel.jl b/EpiAware/src/EpiLatentModels/difflatentmodel.jl index ccebbde6a..2784b05e5 100644 --- a/EpiAware/src/EpiLatentModels/difflatentmodel.jl +++ b/EpiAware/src/EpiLatentModels/difflatentmodel.jl @@ -92,7 +92,7 @@ struct DiffLatentModel{M <: AbstractLatentModel, P} <: AbstractLatentModel end function DiffLatentModel(; model::AbstractLatentModel, - init_priors::Vector{D} where {D <: Distribution}) + init_priors::Vector{D} where {D <: Distribution} = [Normal()]) d = length(init_priors) init_prior = _expand_dist(init_priors) return AR(model, init_prior, d) From 45815224d8a428ddc9e218c986ab54fd902cdd8a Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 15 Mar 2024 11:16:20 +0000 Subject: [PATCH 55/57] get first difflatent test passing --- EpiAware/src/EpiAware.jl | 2 +- .../src/EpiLatentModels/EpiLatentModels.jl | 2 +- .../src/EpiLatentModels/difflatentmodel.jl | 2 +- EpiAware/test/test_difflatentmodel.jl | 22 +++---------------- 4 files changed, 6 insertions(+), 22 deletions(-) diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl index fdf21fe48..378ae7fff 100644 --- a/EpiAware/src/EpiAware.jl +++ b/EpiAware/src/EpiAware.jl @@ -48,7 +48,7 @@ export spread_draws, scan, create_discrete_pmf include("EpiLatentModels/EpiLatentModels.jl") using .EpiLatentModels -export RandomWalk, AR +export RandomWalk, AR, DiffLatentModel include("EpiInfModels/EpiInfModels.jl") using .EpiInfModels diff --git a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl index 788cf9963..024a00fcc 100644 --- a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl +++ b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl @@ -8,7 +8,7 @@ using ..EpiAwareBase using Turing, Distributions, DocStringExtensions -export RandomWalk, AR +export RandomWalk, AR, DiffLatentModel include("docstrings.jl") include("randomwalk.jl") diff --git a/EpiAware/src/EpiLatentModels/difflatentmodel.jl b/EpiAware/src/EpiLatentModels/difflatentmodel.jl index 2784b05e5..f72c569ba 100644 --- a/EpiAware/src/EpiLatentModels/difflatentmodel.jl +++ b/EpiAware/src/EpiLatentModels/difflatentmodel.jl @@ -95,7 +95,7 @@ struct DiffLatentModel{M <: AbstractLatentModel, P} <: AbstractLatentModel init_priors::Vector{D} where {D <: Distribution} = [Normal()]) d = length(init_priors) init_prior = _expand_dist(init_priors) - return AR(model, init_prior, d) + return DiffLatentModel(model, init_prior, d) end function DiffLatentModel(model::AbstractLatentModel, init_prior::Distribution, d::Int) diff --git a/EpiAware/test/test_difflatentmodel.jl b/EpiAware/test/test_difflatentmodel.jl index b2bd17d6e..7b2b4a5bc 100644 --- a/EpiAware/test/test_difflatentmodel.jl +++ b/EpiAware/test/test_difflatentmodel.jl @@ -1,25 +1,10 @@ -@testitem "Testing default_diff_latent_priors" begin - using Distributions - - d = 3 - priors = EpiAware.default_diff_latent_priors(d) - - @test length(priors[:init_prior]) == d - for prior in priors[:init_prior] - @test prior isa Normal - @test prior.μ == 0.0 - @test prior.σ == 1.0 - end -end - @testitem "Testing DiffLatentModel constructor" begin using Distributions, Turing - model = EpiAware.RandomWalk(Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) + model = RandomWalk() @testset "Testing DiffLatentModel with vector of priors" begin init_priors = [Normal(0.0, 1.0), Normal(1.0, 2.0)] - diff_model = EpiAware.DiffLatentModel(; - latentmodel = model, init_priors = init_priors) + diff_model = DiffLatentModel(; model = model, init_priors = init_priors) @test diff_model.model == model @test diff_model.init_prior == arraydist(init_priors) @@ -29,8 +14,7 @@ end @testset "Testing DiffLatentModel with single prior and d" begin d = 3 init_prior = Normal() - diff_model = EpiAware.DiffLatentModel( - model; init_prior_distribution = init_prior, d = d) + diff_model = DiffLatentModel(model, init_prior; d = d) @test diff_model.model == model @test diff_model.init_prior == filldist(init_prior, d) From b0a77316e233e740ceb119865133bfa664a57b47 Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 15 Mar 2024 11:32:03 +0000 Subject: [PATCH 56/57] fix final DiffLatent test --- .../src/EpiLatentModels/difflatentmodel.jl | 2 +- EpiAware/test/test_difflatentmodel.jl | 27 +++++++++---------- 2 files changed, 13 insertions(+), 16 deletions(-) diff --git a/EpiAware/src/EpiLatentModels/difflatentmodel.jl b/EpiAware/src/EpiLatentModels/difflatentmodel.jl index f72c569ba..f2e896265 100644 --- a/EpiAware/src/EpiLatentModels/difflatentmodel.jl +++ b/EpiAware/src/EpiLatentModels/difflatentmodel.jl @@ -172,7 +172,7 @@ unobserved latent process. Z_t ``` """ -@model function generate_latent(latent_model::DiffLatentModel, n) +@model function EpiAwareBase.generate_latent(latent_model::DiffLatentModel, n) d = latent_model.d @assert n>d "n must be longer than d" latent_init ~ latent_model.init_prior diff --git a/EpiAware/test/test_difflatentmodel.jl b/EpiAware/test/test_difflatentmodel.jl index 7b2b4a5bc..650900f43 100644 --- a/EpiAware/test/test_difflatentmodel.jl +++ b/EpiAware/test/test_difflatentmodel.jl @@ -26,17 +26,16 @@ end using DynamicPPL, Turing using Distributions using HypothesisTests: ExactOneSampleKSTest, pvalue - using EpiAware: _add_normals n = 100 d = 2 - model = EpiAware.RandomWalk(Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) + model = RandomWalk(Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) init_priors = [Normal(0.0, 1.0), Normal(1.0, 2.0)] - diff_model = EpiAware.DiffLatentModel(latentmodel = model, init_priors = init_priors) + diff_model = DiffLatentModel(model = model, init_priors = init_priors) - latent_model = EpiAware.generate_latent(diff_model, n) + latent_model = generate_latent(diff_model, n) fixed_model = fix( - latent_model, (σ²_RW = 0, rw_init = 0.0)) + latent_model, (σ_RW = 0, rw_init = 0.0)) n_samples = 2000 samples = sample(fixed_model, Prior(), n_samples) |> @@ -47,11 +46,11 @@ end #Because of the recursive d-times cumsum to undifference the process, #The distribution of the second day should be d lots of first day init distribution """ -Add two normal distributions `x` and `y` together. + Add two normal distributions `x` and `y` together. -# Returns -- `Normal`: The sum of `x` and `y` as a normal distribution. -""" + # Returns + - `Normal`: The sum of `x` and `y` as a normal distribution. + """ function _add_normals(x::Normal, y::Normal) return Normal(x.μ + y.μ, sqrt(x.σ^2 + y.σ^2)) end @@ -73,15 +72,13 @@ end n = 100 d = 2 - model = EpiAware.AR(EpiAware.default_ar_priors()[:damp_prior], - EpiAware.default_ar_priors()[:var_prior], - EpiAware.default_ar_priors()[:init_prior]) + model = AR() init_priors = [Normal(0.0, 1.0), Normal(1.0, 2.0)] - diff_model = EpiAware.DiffLatentModel(latentmodel = model, init_priors = init_priors) + diff_model = DiffLatentModel(model = model, init_priors = init_priors) - latent_model = EpiAware.generate_latent(diff_model, n) + latent_model = generate_latent(diff_model, n) fixed_model = fix(latent_model, - (latant_init = [0.0, 1.0], σ²_AR = 1.0, damp_AR = [0.8], ar_init = [0.0])) + (latent_init = [0.0, 1.0], σ_AR = 1.0, damp_AR = [0.8], ar_init = [0.0])) n_samples = 100 samples = sample(fixed_model, Prior(), n_samples) |> From 2f0efaee6af5c10715855a12ad27c41a271771b8 Mon Sep 17 00:00:00 2001 From: Sam Abbott Date: Fri, 15 Mar 2024 11:35:41 +0000 Subject: [PATCH 57/57] Update EpiAware.jl --- EpiAware/src/EpiAware.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/EpiAware/src/EpiAware.jl b/EpiAware/src/EpiAware.jl index fdf21fe48..378ae7fff 100644 --- a/EpiAware/src/EpiAware.jl +++ b/EpiAware/src/EpiAware.jl @@ -48,7 +48,7 @@ export spread_draws, scan, create_discrete_pmf include("EpiLatentModels/EpiLatentModels.jl") using .EpiLatentModels -export RandomWalk, AR +export RandomWalk, AR, DiffLatentModel include("EpiInfModels/EpiInfModels.jl") using .EpiInfModels