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/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 9d22427d2..71b317641 100644 --- a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl +++ b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl @@ -8,8 +8,10 @@ using ..EpiAwareBase using Turing, Distributions, DocStringExtensions -export RandomWalk, default_rw_priors +export RandomWalk, AR +include("docstrings.jl") include("randomwalk.jl") - +include("autoregressive.jl") +include("utils.jl") end diff --git a/EpiAware/src/EpiLatentModels/autoregressive.jl b/EpiAware/src/EpiLatentModels/autoregressive.jl new file mode 100644 index 000000000..82ecd1c63 --- /dev/null +++ b/EpiAware/src/EpiLatentModels/autoregressive.jl @@ -0,0 +1,95 @@ +@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 + "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) + 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) + end + + 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} + p = length(damp_priors) + damp_prior = _expand_dist(damp_priors) + 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" +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 EpiAwareBase.generate_latent(latent_model::AR, n) + p = latent_model.p + ϵ_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/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/EpiLatentModels/randomwalk.jl b/EpiAware/src/EpiLatentModels/randomwalk.jl index 5f47bd4b1..03911d717 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 = Normal() + std_prior::S = truncated(Normal(0.0, 0.05), 0.0, Inf) 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/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/test/test_autoregressive.jl b/EpiAware/test/test_autoregressive.jl new file mode 100644 index 000000000..7ef7c97fc --- /dev/null +++ b/EpiAware/test/test_autoregressive.jl @@ -0,0 +1,90 @@ +@testitem "Testing AR constructor" begin + 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 = AR(damp_prior, std_prior, init_prior) + + @test ar_process.damp_prior == filldist(damp_prior, 1) + @test ar_process.std_prior == std_prior + @test ar_process.init_prior == filldist(init_prior, 1) +end + +@testitem "Test AR defaults" begin + using Distributions + ar = AR() + @testset "damp_prior" begin + damp = rand(ar.damp_prior) + @test 0.0 <= damp[1] <= 1.0 + 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) + @test typeof(init_ar_value[1]) == Float64 + 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 = AR() + n = 1000 + damp = [0.1] + σ_AR = 1.0 + ar_init = [0.0] + + model = 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^2 / (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 75% rename from EpiAware/test/test_latent-models.jl rename to EpiAware/test/test_randomwalk.jl index e34f90dff..ab9591311 100644 --- a/EpiAware/test/test_latent-models.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