From d7b7b3b510f5c7a4e1b35329f09d99c75c21ba62 Mon Sep 17 00:00:00 2001 From: Sam Abbott Date: Thu, 12 Dec 2024 11:15:20 +0000 Subject: [PATCH] Issue 408: ARMA and ARIMA models (#438) * draft MA method * draft MA methd * use IDD for epsilon in all models * add MA benchmark * Add docs and tests for IDD * make episilon_t a arg of the latent model constructors * improve MA correctness * fully import EpiAwareUtils * add a test for IDD * add tests for MA.jl * add doc tests and unit tests + start on helper fn * more updatres to AR appraoc * chase down partial arg changes * clean up AR * clean up and add arma and arima helpers * Contributions towards Arma/Arima models (#531) * Patch: Switch to fork of benchmarkCI (#520) * patch to fork of benchmarkCI * put fork version of BenchmarkCI in [sources] * swap order * add EpiAware [source] * fix path * rm benchmarkCI from project * Patch fix: add `Manifest.toml` to benchmarking (#524) * trigger * Update benchmark.yaml * Update benchmark.yaml * commit benchmark Manifest * try alternate approach * Update benchmark.yaml * Update EpiMethod.jl * Update benchmark.yaml * change baseline to origin/main * remove trigger * rm other trigger * Issue 465: Add an infection generating model for ODE problems (#510) * CompatHelper: bump compat for Turing to 0.35 for package EpiAware, (drop existing compat) (#516) * CompatHelper: bump compat for Turing to 0.35 for package EpiAware, (drop existing compat) * Update Project.toml * fix Project.toml --------- Co-authored-by: CompatHelper Julia Co-authored-by: Sam Abbott Co-authored-by: Samuel Brand <48288458+SamuelBrand1@users.noreply.github.com> Co-authored-by: Samuel Brand * CompatHelper: bump compat for DynamicPPL to 0.30 for package EpiAware, (drop existing compat) (#528) Co-authored-by: CompatHelper Julia * rename IDD -> IID * rename test file * Issue 529: Create null Latent model (#530) * Null Latent model * Null Latent model * fix doctest * fix generate_epiaware unit tests New usage of RW * fix turing method test underlying std of step size changed name * fix broadcast test Underlying std param changed name * fix HN unit test Default std prior had changed * fix AR unit tests --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: CompatHelper Julia Co-authored-by: Sam Abbott * revert define_ namming * clean out repeated utils from merge * fix MA tests * fix RW tests - feel made about RandomWalk vs AR naming * fix remaining unit tests that aren't doctests * update latent recovery test * try and fix doctests automatically * update all doctests to output nothing - this is awful * add doctests for arima and arma * fix doctest * clean up deps * update replication studies * add interface tests for combination functions and add benchmarks * add some basic theoretical properties tests * name change IDD -> IID benchmarks * moving all the constructors because this PR is too contained * catch missing using * update iid benchmark: * update extraction * remove old param namme from case study * get the dot * get the dot * fix initial guess point for MAP opt * Update index.jl * add a compile time branch for HN * add a compile time branch for HN * update test * add a new constructor to get old default behaviour * update docs * update docs - using the structs for priors is very brittle * reorder prior plots --------- Co-authored-by: Samuel Brand <48288458+SamuelBrand1@users.noreply.github.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: CompatHelper Julia Co-authored-by: Samuel Brand --- .../replications/chatzilena-2019/index.jl | 13 +- .../replications/mishra-2020/index.jl | 20 +-- EpiAware/src/EpiAwareUtils/HalfNormal.jl | 43 ++++--- .../src/EpiAwareUtils/SafeNegativeBinomial.jl | 30 +++-- EpiAware/src/EpiAwareUtils/SafePoisson.jl | 30 +++-- EpiAware/src/EpiAwareUtils/censored_pmf.jl | 55 ++------ .../src/EpiLatentModels/EpiLatentModels.jl | 11 +- .../src/EpiLatentModels/combinations/arima.jl | 42 ++++++ .../src/EpiLatentModels/combinations/arma.jl | 40 ++++++ .../manipulators/CombineLatentModels.jl | 10 +- .../manipulators/ConcatLatentModels.jl | 38 +++--- EpiAware/src/EpiLatentModels/models/AR.jl | 88 +++++++------ .../models/HierarchicalNormal.jl | 40 ++++-- EpiAware/src/EpiLatentModels/models/IID.jl | 53 ++++++++ .../src/EpiLatentModels/models/Intercept.jl | 26 +++- EpiAware/src/EpiLatentModels/models/MA.jl | 121 ++++++++++++++++++ .../src/EpiLatentModels/models/RandomWalk.jl | 101 +++++++-------- .../modifiers/DiffLatentModel.jl | 26 ++-- .../EpiObsModels/StackObservationModels.jl | 16 +-- .../src/EpiObsModels/modifiers/Aggregate.jl | 27 +++- .../src/EpiObsModels/modifiers/LatentDelay.jl | 14 +- .../modifiers/PrefixObservationModel.jl | 15 +++ .../modifiers/RecordExpectedObs.jl | 16 +++ .../modifiers/ascertainment/Ascertainment.jl | 32 ++--- .../test/EpiAwareUtils/generate_epiware.jl | 21 +-- EpiAware/test/EpiAwareUtils/turing-methods.jl | 7 +- .../EpiLatentModels/combinations/arima.jl | 97 ++++++++++++++ .../test/EpiLatentModels/combinations/arma.jl | 34 +++++ .../manipulators/CombineLatentModels.jl | 2 +- .../manipulators/broadcast/LatentModel.jl | 2 +- EpiAware/test/EpiLatentModels/models/AR.jl | 15 ++- .../models/HierarchicalNormal.jl | 5 +- EpiAware/test/EpiLatentModels/models/IID.jl | 57 +++++++++ EpiAware/test/EpiLatentModels/models/MA.jl | 84 ++++++++++++ .../test/EpiLatentModels/models/RandomWalk.jl | 23 ++-- .../modifiers/DiffLatentModel.jl | 5 +- .../EpiObsModels/modifiers/LatentDelay.jl | 14 +- .../bench/EpiLatentModels/EpiLatentModels.jl | 2 + .../EpiLatentModels/combinations/arima.jl | 5 + .../EpiLatentModels/combinations/arma.jl | 5 + benchmark/bench/EpiLatentModels/models/IID.jl | 5 + benchmark/bench/EpiLatentModels/models/MA.jl | 5 + 42 files changed, 977 insertions(+), 318 deletions(-) create mode 100644 EpiAware/src/EpiLatentModels/combinations/arima.jl create mode 100644 EpiAware/src/EpiLatentModels/combinations/arma.jl create mode 100644 EpiAware/src/EpiLatentModels/models/IID.jl create mode 100644 EpiAware/src/EpiLatentModels/models/MA.jl create mode 100644 EpiAware/test/EpiLatentModels/combinations/arima.jl create mode 100644 EpiAware/test/EpiLatentModels/combinations/arma.jl create mode 100644 EpiAware/test/EpiLatentModels/models/IID.jl create mode 100644 EpiAware/test/EpiLatentModels/models/MA.jl create mode 100644 benchmark/bench/EpiLatentModels/combinations/arima.jl create mode 100644 benchmark/bench/EpiLatentModels/combinations/arma.jl create mode 100644 benchmark/bench/EpiLatentModels/models/IID.jl create mode 100644 benchmark/bench/EpiLatentModels/models/MA.jl diff --git a/EpiAware/docs/src/showcase/replications/chatzilena-2019/index.jl b/EpiAware/docs/src/showcase/replications/chatzilena-2019/index.jl index 7ce473618..7363732cd 100644 --- a/EpiAware/docs/src/showcase/replications/chatzilena-2019/index.jl +++ b/EpiAware/docs/src/showcase/replications/chatzilena-2019/index.jl @@ -1,5 +1,5 @@ ### A Pluto.jl notebook ### -# v0.20.0 +# v0.20.3 using Markdown using InteractiveUtils @@ -428,8 +428,8 @@ We define the AR(1) process by matching means of `HalfNormal` prior distribution # ╔═╡ 71a26408-1c26-46cf-bc72-c6ba528dfadd ar = AR( damp_priors = [HalfNormal(mean(sampled_AR_damps))], - std_prior = HalfNormal(mean(sampled_AR_stds)), - init_priors = [Normal(0, 0.001)] + init_priors = [Normal(0, 0.001)], + ϵ_t = HierarchicalNormal(std_prior = HalfNormal(mean(sampled_AR_stds))) ) # ╔═╡ e1ffdaf6-ca2e-405d-8355-0d8848d005b0 @@ -578,9 +578,10 @@ rand(stochastic_mdl) initial_guess = [[mean(chn[:β]), mean(chn[:γ]), mean(chn[:S₀]), - mean(ar.std_prior), mean(ar.init_prior)[1], - mean(ar.damp_prior)[1]] + mean(ar.damp_prior)[1], + mean(ar.ϵ_t.std_prior) + ] zeros(13)] # ╔═╡ 685221ea-f268-4ddc-937f-e7620d065c28 @@ -611,7 +612,7 @@ chn2 = sample( describe(chn2) # ╔═╡ 37a016d8-8384-41c9-abdd-23e88b1f988d -pairplot(chn2[[:β, :γ, :S₀, Symbol(mdl_prefix * ".σ_AR"), +pairplot(chn2[[:β, :γ, :S₀, Symbol(mdl_prefix * ".std"), Symbol(mdl_prefix * ".ar_init[1]"), Symbol(mdl_prefix * ".damp_AR[1]")]]) # ╔═╡ 7df5d669-d3a2-4a66-83c3-f8618e39bec6 diff --git a/EpiAware/docs/src/showcase/replications/mishra-2020/index.jl b/EpiAware/docs/src/showcase/replications/mishra-2020/index.jl index 28b7a9ed0..1a7453b9b 100644 --- a/EpiAware/docs/src/showcase/replications/mishra-2020/index.jl +++ b/EpiAware/docs/src/showcase/replications/mishra-2020/index.jl @@ -1,5 +1,5 @@ ### A Pluto.jl notebook ### -# v0.20.0 +# v0.20.3 using Markdown using InteractiveUtils @@ -114,10 +114,10 @@ In _Mishra et al_ the standard deviation of the _stationary distribution_ of $Z_ # ╔═╡ c88bbbd6-0101-4c04-97c9-c5887ef23999 ar = AR( - damp_priors = reverse([truncated(Normal(0.8, 0.05), 0, 1), - truncated(Normal(0.1, 0.05), 0, 1)]), - std_prior = HalfNormal(0.5), - init_priors = [Normal(-1.0, 0.1), Normal(-1.0, 0.5)] + damp_priors = [truncated(Normal(0.1, 0.05), 0, 1), + truncated(Normal(0.8, 0.05), 0, 1)], + init_priors = [Normal(-1.0, 0.1), Normal(-1.0, 0.5)], + ϵ_t = HierarchicalNormal(std_prior = HalfNormal(0.5)) ) # ╔═╡ 31ee2757-0409-45df-b193-60c552797a3d @@ -561,11 +561,11 @@ let sub_chn = inference_results.samples[inference_results.samples.name_map.parameters[[1:5; end]]] fig = pairplot(sub_chn) - lines!(fig[1, 1], ar.std_prior, label = "Prior") - lines!(fig[2, 2], ar.init_prior.v[1], label = "Prior") - lines!(fig[3, 3], ar.init_prior.v[2], label = "Prior") - lines!(fig[4, 4], ar.damp_prior.v[1], label = "Prior") - lines!(fig[5, 5], ar.damp_prior.v[2], label = "Prior") + lines!(fig[1, 1], ar.init_prior.v[1], label = "Prior") + lines!(fig[2, 2], ar.init_prior.v[2], label = "Prior") + lines!(fig[3, 3], ar.damp_prior.v[1], label = "Prior") + lines!(fig[4, 4], ar.damp_prior.v[2], label = "Prior") + lines!(fig[5, 5], ar.ϵ_t.std_prior, label = "Prior") lines!(fig[6, 6], epi.initialisation_prior, label = "Prior") fig diff --git a/EpiAware/src/EpiAwareUtils/HalfNormal.jl b/EpiAware/src/EpiAwareUtils/HalfNormal.jl index 9d8c78da5..347ea40b7 100644 --- a/EpiAware/src/EpiAwareUtils/HalfNormal.jl +++ b/EpiAware/src/EpiAwareUtils/HalfNormal.jl @@ -11,49 +11,62 @@ Create a half-normal prior distribution with the specified mean. # Examples: -```jldoctest HalfNormal +```jldoctest HalfNormal; output = false using EpiAware, Distributions hn = HalfNormal(1.0) +nothing + # output -EpiAware.EpiAwareUtils.HalfNormal{Float64}(μ=1.0) + ``` -# filter out all the values that are less than 0 -```jldoctest HalfNormal; filter = r\"\b\d+(\.\d+)?\b\" => \"*\" +```jldoctest HalfNormal; output = false rand(hn) +nothing + # output -0.4508533245229199 + ``` -```jldoctest HalfNormal +```jldoctest HalfNormal; output = false cdf(hn, 2) +nothing + # output -0.8894596502772643 + ``` -```jldoctest HalfNormal +```jldoctest HalfNormal; output = false quantile(hn, 0.5) +nothing + # output -0.8453475393951495 + ``` -```jldoctest HalfNormal +```jldoctest HalfNormal; output = false logpdf(hn, 2) +nothing + # output --3.1111166111445083 + ``` -```jldoctest HalfNormal +```jldoctest HalfNormal; output = false mean(hn) +nothing + # output -1.0 + ``` -```jldoctest HalfNormal +```jldoctest HalfNormal; output = false var(hn) +nothing + # output -0.5707963267948966 + ``` " struct HalfNormal{T <: Real} <: ContinuousUnivariateDistribution diff --git a/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl b/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl index e57714b80..f141215c6 100644 --- a/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl +++ b/EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl @@ -27,7 +27,7 @@ parameterisation is useful for specifying the distribution in a way that is easi # Examples: -```jldoctest SafeNegativeBinomial +```jldoctest SafeNegativeBinomial; output = false using EpiAware, Distributions bigμ = exp(48.0) #Large value of μ @@ -37,32 +37,42 @@ bigμ = exp(48.0) #Large value of μ p = bigμ / σ² r = bigμ * p / (1 - p) d = SafeNegativeBinomial(r, p) +nothing + # output -EpiAware.EpiAwareUtils.SafeNegativeBinomial{Float64}(r=20.0, p=2.85032816548187e-20) + ``` -```jldoctest SafeNegativeBinomial +```jldoctest SafeNegativeBinomial; output = false cdf(d, 100) +nothing + # output -0.0 + ``` -```jldoctest SafeNegativeBinomial +```jldoctest SafeNegativeBinomial; output = false logpdf(d, 100) +nothing + # output --850.1397180331871 + ``` -```jldoctest SafeNegativeBinomial +```jldoctest SafeNegativeBinomial; output = false mean(d) +nothing + # output -7.016735912097631e20 + ``` -```jldoctest SafeNegativeBinomial +```jldoctest SafeNegativeBinomial; output = false var(d) +nothing + # output -2.4617291430060293e40 + ``` " struct SafeNegativeBinomial{T <: Real} <: SafeDiscreteUnivariateDistribution diff --git a/EpiAware/src/EpiAwareUtils/SafePoisson.jl b/EpiAware/src/EpiAwareUtils/SafePoisson.jl index c85243ebc..3d0ba2061 100644 --- a/EpiAware/src/EpiAwareUtils/SafePoisson.jl +++ b/EpiAware/src/EpiAwareUtils/SafePoisson.jl @@ -12,37 +12,47 @@ when the mean is too large. # Examples: -```jldoctest SafePoisson +```jldoctest SafePoisson; output = false using EpiAware, Distributions bigλ = exp(48.0) #Large value of λ d = SafePoisson(bigλ) +nothing + # output -EpiAware.EpiAwareUtils.SafePoisson{Float64}(λ=7.016735912097631e20) + ``` -```jldoctest SafePoisson +```jldoctest SafePoisson; output = false cdf(d, 2) +nothing + # output -0.0 + ``` -```jldoctest SafePoisson +```jldoctest SafePoisson; output = false logpdf(d, 100) +nothing + # output --7.016735912097631e20 + ``` -```jldoctest SafePoisson +```jldoctest SafePoisson; output = false mean(d) +nothing + # output -7.016735912097631e20 + ``` -```jldoctest SafePoisson +```jldoctest SafePoisson; output = false var(d) +nothing + # output -7.016735912097631e20 + ``` " struct SafePoisson{T <: Real} <: SafeDiscreteUnivariateDistribution diff --git a/EpiAware/src/EpiAwareUtils/censored_pmf.jl b/EpiAware/src/EpiAwareUtils/censored_pmf.jl index e12d44279..c5a7abce7 100644 --- a/EpiAware/src/EpiAwareUtils/censored_pmf.jl +++ b/EpiAware/src/EpiAwareUtils/censored_pmf.jl @@ -23,27 +23,17 @@ Raises: # Examples -```jldoctest filter +```jldoctest; output = false using Distributions using EpiAware.EpiAwareUtils dist = Exponential(1.0) -censored_pmf(dist, Val(:single_censored); D = 10) |> - p -> round.(p, digits=3) +censored_pmf(dist, Val(:single_censored); D = 10) +nothing # output -10-element Vector{Float64}: - 0.393 - 0.383 - 0.141 - 0.052 - 0.019 - 0.007 - 0.003 - 0.001 - 0.0 - 0.0 + ``` " function censored_pmf(dist::Distribution, @@ -122,28 +112,17 @@ to nearest multiple of `Δd`. # Examples -```jldoctest filter +```jldoctest filter; output = false using Distributions using EpiAware.EpiAwareUtils dist = Exponential(1.0) -censored_cdf(dist; D = 10) |> - p -> round.(p, digits=3) +censored_cdf(dist; D = 10) +nothing # output -11-element Vector{Float64}: - 0.0 - 0.368 - 0.767 - 0.914 - 0.969 - 0.988 - 0.996 - 0.998 - 0.999 - 1.0 - 1.0 + ``` " function censored_cdf(dist::Distribution; Δd = 1.0, D = nothing, upper = 0.999) @@ -184,27 +163,17 @@ to nearest multiple of `Δd`. # Examples -```jldoctest filter +```jldoctest filter; output = false using Distributions using EpiAware.EpiAwareUtils dist = Exponential(1.0) -censored_pmf(dist; D = 10) |> - p -> round.(p, digits=3) +censored_pmf(dist; D = 10) +nothing # output -10-element Vector{Float64}: - 0.368 - 0.4 - 0.147 - 0.054 - 0.02 - 0.007 - 0.003 - 0.001 - 0.0 - 0.0 + ``` " function censored_pmf(dist::Distribution; Δd = 1.0, D = nothing, upper = 0.99) diff --git a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl index 1dabab681..e5b34e726 100644 --- a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl +++ b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl @@ -5,7 +5,7 @@ module EpiLatentModels using ..EpiAwareBase -using ..EpiAwareUtils: HalfNormal, prefix_submodel, accumulate_scan +using ..EpiAwareUtils using LogExpFunctions: softmax @@ -15,7 +15,7 @@ using Turing, Distributions, DocStringExtensions, LinearAlgebra, SparseArrays, OrdinaryDiffEq #Export models -export FixedIntercept, Intercept, RandomWalk, AR, HierarchicalNormal, Null +export Null, FixedIntercept, Intercept, IID, RandomWalk, AR, MA, HierarchicalNormal #Export ODE definitions export SIRParams, SEIRParams @@ -32,13 +32,20 @@ export broadcast_rule, broadcast_dayofweek, broadcast_weekly, equal_dimensions # Export tools for modifying latent models export DiffLatentModel, TransformLatentModel, PrefixLatentModel, RecordExpectedLatent +# Export combinations of models and modifiers +export arma, arima + include("docstrings.jl") include("utils.jl") include("models/Intercept.jl") +include("models/IID.jl") include("models/RandomWalk.jl") include("models/AR.jl") +include("models/MA.jl") include("models/HierarchicalNormal.jl") include("models/Null.jl") +include("combinations/arma.jl") +include("combinations/arima.jl") include("odemodels/SIRParams.jl") include("odemodels/SEIRParams.jl") include("modifiers/DiffLatentModel.jl") diff --git a/EpiAware/src/EpiLatentModels/combinations/arima.jl b/EpiAware/src/EpiLatentModels/combinations/arima.jl new file mode 100644 index 000000000..409bfa497 --- /dev/null +++ b/EpiAware/src/EpiLatentModels/combinations/arima.jl @@ -0,0 +1,42 @@ +@doc raw""" +Define an ARIMA model by wrapping `define_arma` and applying differencing via `DiffLatentModel`. + +# Arguments +- `ar_init`: Prior distribution for AR initial conditions. + A vector of distributions. +- `diff_init`: Prior distribution for differencing initial conditions. + A vector of distributions. +- `θ`: Prior distribution for MA coefficients. + A vector of distributions. +- `damp`: Prior distribution for AR damping coefficients. + A vector of distributions. +- `ϵ_t`: Distribution of the error term. + Default is `HierarchicalNormal()`. + +# Returns +An ARIMA model consisting of AR and MA components with differencing applied. + +# Example + +```jldoctest ARIMA; output = false +using EpiAware, Distributions + +ARIMA = arima() +arima_model = generate_latent(ARIMA, 10) +arima_model() +nothing +# output + +``` +""" +function arima(; + ar_init = [Normal()], + diff_init = [Normal()], + damp = [truncated(Normal(0.0, 0.05), 0, 1)], + θ = [truncated(Normal(0.0, 0.05), -1, 1)], + ϵ_t = HierarchicalNormal() +) + arma_model = arma(; init = ar_init, damp = damp, θ = θ, ϵ_t = ϵ_t) + arima_model = DiffLatentModel(; model = arma_model, init_priors = diff_init) + return arima_model +end diff --git a/EpiAware/src/EpiLatentModels/combinations/arma.jl b/EpiAware/src/EpiLatentModels/combinations/arma.jl new file mode 100644 index 000000000..f212a2ef8 --- /dev/null +++ b/EpiAware/src/EpiLatentModels/combinations/arma.jl @@ -0,0 +1,40 @@ +@doc raw""" +Define an ARMA model using AR and MA components. + +# Arguments +- `init`: Prior distribution for AR initial conditions. + A vector of distributions. +- `θ`: Prior distribution for MA coefficients. + A vector of distributions. +- `damp`: Prior distribution for AR damping coefficients. + A vector of distributions. +- `ϵ_t`: Distribution of the error term. + Default is `HierarchicalNormal()`. + +# Returns +An AR model with an MA model as its error term, effectively creating an ARMA model. + +# Example + +```jldoctest ARMA; output = false +using EpiAware, Distributions + +ARMA = arma(; + θ = [truncated(Normal(0.0, 0.02), -1, 1)], + damp = [truncated(Normal(0.0, 0.02), 0, 1)] +) +arma_model = generate_latent(ARMA, 10) +arma_model() +nothing +# output +``` +""" +function arma(; + init = [Normal()], + damp = [truncated(Normal(0.0, 0.05), 0, 1)], + θ = [truncated(Normal(0.0, 0.05), -1, 1)], + ϵ_t = HierarchicalNormal()) + ma = MA(; θ_priors = θ, ϵ_t = ϵ_t) + ar = AR(; damp_priors = damp, init_priors = init, ϵ_t = ma) + return ar +end diff --git a/EpiAware/src/EpiLatentModels/manipulators/CombineLatentModels.jl b/EpiAware/src/EpiLatentModels/manipulators/CombineLatentModels.jl index d19347d13..fc401b139 100644 --- a/EpiAware/src/EpiLatentModels/manipulators/CombineLatentModels.jl +++ b/EpiAware/src/EpiLatentModels/manipulators/CombineLatentModels.jl @@ -37,12 +37,12 @@ latent_model() return new{AbstractVector{<:AbstractTuringLatentModel}, AbstractVector{<:String}}( prefix_models, prefixes) end +end - function CombineLatentModels(models::M) where { - M <: AbstractVector{<:AbstractTuringLatentModel}} - prefixes = "Combine." .* string.(1:length(models)) - return CombineLatentModels(models, prefixes) - end +function CombineLatentModels(models::M) where { + M <: AbstractVector{<:AbstractTuringLatentModel}} + prefixes = "Combine." .* string.(1:length(models)) + return CombineLatentModels(models, prefixes) end @doc raw" diff --git a/EpiAware/src/EpiLatentModels/manipulators/ConcatLatentModels.jl b/EpiAware/src/EpiLatentModels/manipulators/ConcatLatentModels.jl index 79c306b6c..d922e9894 100644 --- a/EpiAware/src/EpiLatentModels/manipulators/ConcatLatentModels.jl +++ b/EpiAware/src/EpiLatentModels/manipulators/ConcatLatentModels.jl @@ -52,29 +52,29 @@ struct ConcatLatentModels{ AbstractVector{<:String}}( prefix_models, no_models, dimension_adaptor, prefixes) end +end - function ConcatLatentModels(models::M, dimension_adaptor::Function; - prefixes = nothing) where { - M <: AbstractVector{<:AbstractTuringLatentModel}} - no_models = length(models) - if isnothing(prefixes) - prefixes = "Concat." .* string.(1:no_models) - end - return ConcatLatentModels(models, no_models, dimension_adaptor, prefixes) +function ConcatLatentModels(models::M, dimension_adaptor::Function; + prefixes = nothing) where { + M <: AbstractVector{<:AbstractTuringLatentModel}} + no_models = length(models) + if isnothing(prefixes) + prefixes = "Concat." .* string.(1:no_models) end + return ConcatLatentModels(models, no_models, dimension_adaptor, prefixes) +end - function ConcatLatentModels(models::M; - dimension_adaptor::Function = equal_dimensions, - prefixes = nothing) where { - M <: AbstractVector{<:AbstractTuringLatentModel}} - return ConcatLatentModels(models, dimension_adaptor; prefixes = prefixes) - end +function ConcatLatentModels(models::M; + dimension_adaptor::Function = equal_dimensions, + prefixes = nothing) where { + M <: AbstractVector{<:AbstractTuringLatentModel}} + return ConcatLatentModels(models, dimension_adaptor; prefixes = prefixes) +end - function ConcatLatentModels(; models::M, - dimension_adaptor::Function = equal_dimensions, prefixes = nothing) where { - M <: AbstractVector{<:AbstractTuringLatentModel}} - return ConcatLatentModels(models, dimension_adaptor; prefixes = prefixes) - end +function ConcatLatentModels(; models::M, + dimension_adaptor::Function = equal_dimensions, prefixes = nothing) where { + M <: AbstractVector{<:AbstractTuringLatentModel}} + return ConcatLatentModels(models, dimension_adaptor; prefixes = prefixes) end @doc raw" diff --git a/EpiAware/src/EpiLatentModels/models/AR.jl b/EpiAware/src/EpiLatentModels/models/AR.jl index ad17556d9..6d1a82370 100644 --- a/EpiAware/src/EpiLatentModels/models/AR.jl +++ b/EpiAware/src/EpiLatentModels/models/AR.jl @@ -2,61 +2,70 @@ 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_prior::Sampleable, init_prior::Sampleable; p::Int = 1, ϵ_t::AbstractTuringLatentModel = HierarchicalNormal())`: Constructs an AR model with the specified prior distributions for damping coefficients 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. +- `AR(; damp_priors::Vector{D} = [truncated(Normal(0.0, 0.05), 0, 1)], init_priors::Vector{I} = [Normal()], ϵ_t::AbstractTuringLatentModel = HierarchicalNormal()) where {D <: Sampleable, I <: Sampleable}`: Constructs an AR model with the specified prior distributions for damping coefficients and initial conditions. The order of the AR model is determined by the length of the `damp_priors` vector. # Examples -```julia -using Distributions -using EpiAware +```jldoctest AutoRegressive; output = false +using Distributions, Turing, EpiAware ar = AR() -ar_model = generate_latent(ar, 10) -rand(ar_model) +ar +nothing +# output +``` + +```jldoctest AutoRegressive; output = false +mdl = generate_latent(ar, 10) +mdl() +nothing +# output +``` + +```jldoctest AutoRegressive; output = false +rand(mdl) +nothing +# output ``` " -struct AR{D <: Sampleable, S <: Sampleable, I <: Sampleable, P <: Int} <: - AbstractTuringLatentModel +struct AR{D <: Sampleable, I <: Sampleable, + P <: Int, E <: AbstractTuringLatentModel} <: AbstractTuringLatentModel "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 + "Prior distribution for the error term." + ϵ_t::E - function AR(; damp_priors::Vector{D} = [truncated(Normal(0.0, 0.05), 0, 1)], - std_prior::Distribution = HalfNormal(0.1), - 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) + function AR(damp_prior::Sampleable, init_prior::Sampleable, + p::Int, ϵ_t::AbstractTuringLatentModel) @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 - ) + @assert p==length(damp_prior)==length(init_prior) "p must be equal to the length of damp_prior and init_prior" + new{typeof(damp_prior), typeof(init_prior), typeof(p), typeof(ϵ_t)}( + damp_prior, init_prior, p, ϵ_t) end end +function AR(damp_prior::Sampleable, init_prior::Sampleable; p::Int = 1, + ϵ_t::AbstractTuringLatentModel = HierarchicalNormal()) + damp_priors = fill(damp_prior, p) + init_priors = fill(init_prior, p) + return AR(; damp_priors = damp_priors, init_priors = init_priors, ϵ_t = ϵ_t) +end + +function AR(; damp_priors::Vector{D} = [truncated(Normal(0.0, 0.05), 0, 1)], + init_priors::Vector{I} = [Normal()], + ϵ_t::AbstractTuringLatentModel = HierarchicalNormal()) where { + D <: Sampleable, I <: Sampleable} + p = length(damp_priors) + damp_prior = _expand_dist(damp_priors) + init_prior = _expand_dist(init_priors) + return AR(damp_prior, init_prior, p, ϵ_t) +end + @doc raw" Generate a latent AR series. @@ -76,12 +85,11 @@ Generate a latent AR series. p = latent_model.p @assert n>p "n must be longer than order of the autoregressive process" - σ_AR ~ latent_model.std_prior ar_init ~ latent_model.init_prior damp_AR ~ latent_model.damp_prior - ϵ_t ~ filldist(Normal(), n - p) + @submodel ϵ_t = generate_latent(latent_model.ϵ_t, n - p) - ar = accumulate_scan(ARStep(damp_AR), ar_init, σ_AR * ϵ_t) + ar = accumulate_scan(ARStep(damp_AR), ar_init, ϵ_t) return ar end diff --git a/EpiAware/src/EpiLatentModels/models/HierarchicalNormal.jl b/EpiAware/src/EpiLatentModels/models/HierarchicalNormal.jl index 7f77fd47e..9ddb9a6f2 100644 --- a/EpiAware/src/EpiLatentModels/models/HierarchicalNormal.jl +++ b/EpiAware/src/EpiLatentModels/models/HierarchicalNormal.jl @@ -4,20 +4,40 @@ The `HierarchicalNormal` struct represents a non-centered hierarchical normal di ## Constructors - `HierarchicalNormal(mean, std_prior)`: Constructs a `HierarchicalNormal` instance with the specified mean and standard deviation prior. -- `HierarchicalNormal(; mean = 0.0, std_prior = truncated(Normal(0,1), 0, Inf))`: Constructs a `HierarchicalNormal` instance with the specified mean and standard deviation prior using named arguments and with default values. +- `HierarchicalNormal(; mean = 0.0, std_prior = truncated(Normal(0,0.1), 0, Inf))`: Constructs a `HierarchicalNormal` instance with the specified mean and standard deviation prior using named arguments and with default values. +- `HierarchicalNormal(std_prior)`: Constructs a `HierarchicalNormal` instance with the specified standard deviation prior. +- `HierarchicalNormal(mean, std_prior)`: Constructs a `HierarchicalNormal` instance with the specified mean and standard deviation prior. ## Examples -```julia -using Distributions, EpiAware -hnorm = HierarchicalNormal(0.0, truncated(Normal(0, 1), 0, Inf)) -hnorm_model = generate_latent(hnorm, 10) -hnorm_model() +```jldoctest HierarchicalNormal; output = false +using Distributions, Turing, EpiAware +hn = HierarchicalNormal() + +mdl = generate_latent(hn, 10) +mdl() + +rand(mdl) +nothing +# output ``` " -@kwdef struct HierarchicalNormal{R <: Real, D <: Sampleable} <: AbstractTuringLatentModel +@kwdef struct HierarchicalNormal{R <: Real, D <: Sampleable, M <: Bool} <: + AbstractTuringLatentModel + "Mean of the normal distribution." mean::R = 0.0 - std_prior::D = truncated(Normal(0, 1), 0, Inf) + "Prior distribution for the standard deviation." + std_prior::D = truncated(Normal(0, 0.1), 0, Inf) + "Flag to indicate if mean should be added (false when mean = 0)" + add_mean::M = mean != 0 +end + +function HierarchicalNormal(std_prior::Distribution) + return HierarchicalNormal(; std_prior = std_prior) +end + +function HierarchicalNormal(mean::Real, std_prior::Distribution) + return HierarchicalNormal(mean, std_prior, mean != 0) end @doc raw" @@ -34,8 +54,8 @@ Generate latent variables from the hierarchical normal distribution. " @model function EpiAwareBase.generate_latent(obs_model::HierarchicalNormal, n) std ~ obs_model.std_prior - ϵ_t ~ filldist(Normal(), n) + @submodel ϵ_t = generate_latent(IID(Normal()), n) - η_t = obs_model.mean .+ std .* ϵ_t + η_t = obs_model.add_mean ? obs_model.mean .+ std * ϵ_t : std * ϵ_t return η_t end diff --git a/EpiAware/src/EpiLatentModels/models/IID.jl b/EpiAware/src/EpiLatentModels/models/IID.jl new file mode 100644 index 000000000..b41195b88 --- /dev/null +++ b/EpiAware/src/EpiLatentModels/models/IID.jl @@ -0,0 +1,53 @@ +@doc raw" +Model latent process ``\epsilon_t`` as independent and identically distributed random variables. + +## Mathematical specification + +The IID process ``\epsilon_t`` is specified as a sequence of independent and identically distributed random variables, + +```math +\epsilon_t \sim \text{Prior}, \quad t = 1, 2, \ldots +``` + +where Prior is the specified distribution. + +## Constructors + +- `IID(prior::Distribution = Normal(0, 1))`: Create an IID model with the specified prior distribution. + +## Examples + +```jldoctest IID; filter = r\"\b\d+(\.\d+)?\b\" => \"*\" +using EpiAware, Distributions + +model = IID(Normal(0, 1)) +# output +IID{Normal{Float64}}(Distributions.Normal{Float64}(μ=0.0, σ=1.0)) + +``` + +```jldoctest IID; output = false +idd = generate_latent(model, 10) +idd() +nothing +# output +``` +" +@kwdef struct IID{D <: Sampleable} <: AbstractTuringLatentModel + ϵ_t::D = Normal(0, 1) +end + +@doc raw" + Generate latent variables from the IID (Independent and Identically Distributed) model. + +# Arguments +- `model::IID`: The IID model. +- `n`: Number of latent variables to generate. + +# Returns +- `ϵ_t`: Generated latent variables. +" +@model function EpiAwareBase.generate_latent(model::IID, n) + ϵ_t ~ filldist(model.ϵ_t, n) + return ϵ_t +end diff --git a/EpiAware/src/EpiLatentModels/models/Intercept.jl b/EpiAware/src/EpiLatentModels/models/Intercept.jl index 506ff86cd..ea0569959 100644 --- a/EpiAware/src/EpiLatentModels/models/Intercept.jl +++ b/EpiAware/src/EpiLatentModels/models/Intercept.jl @@ -9,12 +9,26 @@ broadcasts a single intercept value to a length `n` latent process. ## Examples -```julia +```jldoctest Intercept using Distributions, Turing, EpiAware int = Intercept(Normal(0, 1)) -int_model = generate_latent(int, 10) -rand(int_model) -int_model() +int +# output +Intercept{Normal{Float64}}(Distributions.Normal{Float64}(μ=0.0, σ=1.0)) +``` + +```jldoctest Intercept; output = false +mdl = generate_latent(int, 10) +mdl() +nothing +# output + +``` + +```jldoctest Intercept; output = false +rand(mdl) +nothing +# output ``` " @kwdef struct Intercept{D <: Sampleable} <: AbstractTuringIntercept @@ -49,11 +63,13 @@ A variant of the `Intercept` struct that represents a fixed intercept value for # Examples -```julia +```jldoctest FixedIntercept; output = false using EpiAware fi = FixedIntercept(2.0) fi_model = generate_latent(fi, 10) fi_model() +nothing +# output ``` " @kwdef struct FixedIntercept{F <: Real} <: AbstractTuringIntercept diff --git a/EpiAware/src/EpiLatentModels/models/MA.jl b/EpiAware/src/EpiLatentModels/models/MA.jl new file mode 100644 index 000000000..0ba395bd6 --- /dev/null +++ b/EpiAware/src/EpiLatentModels/models/MA.jl @@ -0,0 +1,121 @@ +@doc raw" +The moving average (MA) model struct. + +# Constructors +- `MA(θ::Distribution, σ::Distribution; q::Int = 1, ϵ::AbstractTuringLatentModel = IID(Normal()))`: Constructs an MA model with the specified prior distributions. + +- `MA(; θ::Vector{C} = [truncated(Normal(0.0, 0.05), -1, 1)], ϵ::AbstractTuringLatentModel = HierarchicalNormal) where {C <: Distribution}`: Constructs an MA model with the specified prior distributions. + +- `MA(θ::Distribution, q::Int, ϵ_t::AbstractTuringLatentModel)`: Constructs an MA model with the specified prior distributions and order. + +# Parameters +- `θ`: Prior distribution for the MA coefficients. For MA(q), this should be a vector of q distributions or a multivariate distribution of dimension q. +- `q`: Order of the MA model, i.e., the number of lagged error terms. +- `ϵ_t`: Distribution of the error term, typically standard normal. Defaults to `IID(Normal())`. + +# Examples + +```jldoctest MA; output = false +using Distributions, Turing, EpiAware +ma = MA() +ma +nothing +# output + +``` + +```jldoctest MA; output = false +mdl = generate_latent(ma, 10) +mdl() +nothing +# output +``` + +```jldoctest MA; output = false +rand(mdl) +nothing +# output +``` +" + +struct MA{C <: Sampleable, Q <: Int, E <: AbstractTuringLatentModel} <: + AbstractTuringLatentModel + "Prior distribution for the MA coefficients. For MA(q), this should be a vector of q distributions or a multivariate distribution of dimension q" + θ::C + "Order of the MA model, i.e., the number of lagged error terms." + q::Q + "Prior distribution for the error term." + ϵ_t::E + + function MA(θ::Distribution, q::Int, ϵ_t::AbstractTuringLatentModel) + @assert q>0 "q must be greater than 0" + @assert q==length(θ) "q must be equal to the length of θ" + new{typeof(θ), typeof(q), typeof(ϵ_t)}(θ, q, ϵ_t) + end +end + +function MA(θ::Distribution; + q::Int = 1, ϵ_t::AbstractTuringLatentModel = HierarchicalNormal()) + θ_priors = fill(θ, q) + return MA(; θ_priors = θ_priors, ϵ_t = ϵ_t) +end + +function MA(; θ_priors::Vector{C} = [truncated(Normal(0.0, 0.05), -1, 1)], + ϵ_t::AbstractTuringLatentModel = HierarchicalNormal()) where {C <: + Distribution} + q = length(θ_priors) + θ = _expand_dist(θ_priors) + return MA(θ, q, ϵ_t) +end + +@doc raw" +Generate a latent MA series. + +# Arguments + +- `latent_model::MA`: The MA model. +- `n::Int`: The length of the MA series. + +# Returns +- `ma::Vector{Float64}`: The generated MA series. + +# Notes +- `n` must be longer than the order of the moving average process. +" +@model function EpiAwareBase.generate_latent(latent_model::MA, n) + q = latent_model.q + @assert n>q "n must be longer than order of the moving average process" + θ ~ latent_model.θ + @submodel ϵ_t = generate_latent(latent_model.ϵ_t, n) + + ma = accumulate_scan( + MAStep(θ), + (; val = 0, state = ϵ_t[1:q]), ϵ_t[(q + 1):end]) + + return ma +end + +@doc raw" +The moving average (MA) step function struct +" +struct MAStep{C <: AbstractVector{<:Real}} <: AbstractAccumulationStep + θ::C +end + +@doc raw" +The moving average (MA) step function for use with `accumulate_scan`. +" +function (ma::MAStep)(state, ϵ) + new_val = ϵ + dot(ma.θ, state.state) + new_state = vcat(ϵ, state.state[1:(end - 1)]) + return (; val = new_val, state = new_state) +end + +@doc raw" +The MA step function method for get_state. +" +function EpiAwareUtils.get_state(acc_step::MAStep, initial_state, state) + init_vals = initial_state.state + new_vals = state .|> x -> x.val + return vcat(init_vals, new_vals) +end diff --git a/EpiAware/src/EpiLatentModels/models/RandomWalk.jl b/EpiAware/src/EpiLatentModels/models/RandomWalk.jl index bce69031a..95fe49506 100644 --- a/EpiAware/src/EpiLatentModels/models/RandomWalk.jl +++ b/EpiAware/src/EpiLatentModels/models/RandomWalk.jl @@ -13,81 +13,78 @@ Z_t = Z_0 + \sigma \sum_{i = 1}^t \epsilon_t Constructing a random walk requires specifying: - An `init_prior` as a prior for ``Z_0``. Default is `Normal()`. - A `std_prior` for ``\sigma``. The default is HalfNormal with a mean of 0.25. +- An `ϵ_t` prior for the white noise sequence. The default is `IID(Normal())`. ## Constructors -- `RandomWalk(; init_prior, std_prior)` +- `RandomWalk(init_prior::Sampleable, ϵ_t::AbstractTuringLatentModel)`: Constructs a random walk model with the specified prior distributions for the initial condition and white noise sequence. +- `RandomWalk(; init_prior::Sampleable = Normal(), ϵ_t::AbstractTuringLatentModel = HierarchicalNormal())`: Constructs a random walk model with the specified prior distributions for the initial condition and white noise sequence. -## Example usage with `generate_latent` +## Example usage -`generate_latent` can be used to construct a `Turing` model for the random walk ``Z_t``. - -First, we construct a `RandomWalk` struct with priors, - -```julia +```jldoctest RandomWalk; output = false using Distributions, Turing, EpiAware - -# Create a RandomWalk model -rw = RandomWalk(init_prior = Normal(2., 1.), - std_prior = HalfNormal(0.1)) +rw = RandomWalk() +rw +nothing +# output ``` -Then, we can use `generate_latent` to construct a Turing model for a 10 step random walk. - -```julia -# Construct a Turing model -rw_model = generate_latent(rw, 10) +```jldoctest RandomWalk; output = false +mdl = generate_latent(rw, 10) +mdl() +nothing +# output ``` -Now we can use the `Turing` PPL API to sample underlying parameters and generate the -unobserved infections. +```jldoctest RandomWalk; output = false +rand(mdl) +nothing +# output -```julia -#Sample random parameters from prior -θ = rand(rw_model) -#Get random walk sample path as a generated quantities from the model -Z_t, _ = generated_quantities(rw_model, θ) ``` " -@kwdef struct RandomWalk{D <: Sampleable, S <: Sampleable} <: AbstractTuringLatentModel +@kwdef struct RandomWalk{ + D <: Sampleable, E <: AbstractTuringLatentModel} <: + AbstractTuringLatentModel init_prior::D = Normal() - std_prior::S = HalfNormal(0.25) + ϵ_t::E = HierarchicalNormal() end @doc raw" -Implement the `generate_latent` function for the `RandomWalk` model. +Generate a latent RW series using accumulate_scan. -## Example usage of `generate_latent` with `RandomWalk` type of latent process model +# Arguments -```julia -using Distributions, Turing, EpiAware +- `latent_model::RandomWalk`: The RandomWalk model. +- `n::Int`: The length of the RW series. -# Create a RandomWalk model -rw = RandomWalk(init_prior = Normal(2., 1.), - std_prior = HalfNormal(0.1)) -``` +# Returns +- `rw::Vector{Float64}`: The generated RW series. -Then, we can use `generate_latent` to construct a Turing model for a 10 step random walk. - -```julia -# Construct a Turing model -rw_model = generate_latent(rw, 10) -``` - -Now we can use the `Turing` PPL API to sample underlying parameters and generate the -unobserved infections. - -```julia -#Sample random parameters from prior -θ = rand(rw_model) -#Get random walk sample path as a generated quantities from the model -Z_t, _ = generated_quantities(rw_model, θ) -``` +# Notes +- `n` must be greater than 0. " @model function EpiAwareBase.generate_latent(latent_model::RandomWalk, n) - σ_RW ~ latent_model.std_prior + @assert n>0 "n must be greater than 0" + rw_init ~ latent_model.init_prior - ϵ_t ~ filldist(Normal(), n - 1) - rw = rw_init .+ vcat(0.0, σ_RW .* cumsum(ϵ_t)) + @submodel ϵ_t = generate_latent(latent_model.ϵ_t, n - 1) + + rw = accumulate_scan(RWStep(), rw_init, ϵ_t) + return rw end + +@doc raw" +The random walk (RW) step function struct +" +struct RWStep <: AbstractAccumulationStep end + +@doc raw" +The random walk (RW) step function for use with `accumulate_scan`. +" +function (rw::RWStep)(state, ϵ) + new_val = state + ϵ + return new_val +end diff --git a/EpiAware/src/EpiLatentModels/modifiers/DiffLatentModel.jl b/EpiAware/src/EpiLatentModels/modifiers/DiffLatentModel.jl index 85b69cd8f..241b41fc1 100644 --- a/EpiAware/src/EpiLatentModels/modifiers/DiffLatentModel.jl +++ b/EpiAware/src/EpiLatentModels/modifiers/DiffLatentModel.jl @@ -86,19 +86,6 @@ struct DiffLatentModel{M <: AbstractTuringLatentModel, P <: Distribution} <: "Number of times differenced." d::Int - function DiffLatentModel( - model::AbstractTuringLatentModel, init_prior::Distribution; d::Int) - init_priors = fill(init_prior, d) - return DiffLatentModel(; model = model, init_priors = init_priors) - end - - function DiffLatentModel(; model::AbstractTuringLatentModel, - init_priors::Vector{D} where {D <: Distribution} = [Normal()]) - d = length(init_priors) - init_prior = _expand_dist(init_priors) - return DiffLatentModel(model, init_prior, d) - end - function DiffLatentModel( model::AbstractTuringLatentModel, init_prior::Distribution, d::Int) @assert d>0 "d must be greater than 0" @@ -107,6 +94,19 @@ struct DiffLatentModel{M <: AbstractTuringLatentModel, P <: Distribution} <: end end +function DiffLatentModel( + model::AbstractTuringLatentModel, init_prior::Distribution; d::Int) + init_priors = fill(init_prior, d) + return DiffLatentModel(; model = model, init_priors = init_priors) +end + +function DiffLatentModel(; model::AbstractTuringLatentModel, + init_priors::Vector{D} where {D <: Distribution} = [Normal()]) + d = length(init_priors) + init_prior = _expand_dist(init_priors) + return DiffLatentModel(model, init_prior, d) +end + """ Generate a `Turing` model for `n`-step latent process ``Z_t`` using a differenced latent model defined by `latent_model`. diff --git a/EpiAware/src/EpiObsModels/StackObservationModels.jl b/EpiAware/src/EpiObsModels/StackObservationModels.jl index 778de5e9e..763d4b8a5 100644 --- a/EpiAware/src/EpiObsModels/StackObservationModels.jl +++ b/EpiAware/src/EpiObsModels/StackObservationModels.jl @@ -53,15 +53,15 @@ deaths_y_t new{AbstractVector{<:AbstractTuringObservationModel}, typeof(model_names)}( wrapped_models, model_names) end +end - function StackObservationModels(models::NamedTuple{ - names, T}) where {names, T <: Tuple{Vararg{AbstractTuringObservationModel}}} - model_names = models |> - keys .|> - string |> - collect - return StackObservationModels(collect(values(models)), model_names) - end +function StackObservationModels(models::NamedTuple{ + names, T}) where {names, T <: Tuple{Vararg{AbstractTuringObservationModel}}} + model_names = models |> + keys .|> + string |> + collect + return StackObservationModels(collect(values(models)), model_names) end @doc raw" diff --git a/EpiAware/src/EpiObsModels/modifiers/Aggregate.jl b/EpiAware/src/EpiObsModels/modifiers/Aggregate.jl index b38c7f1cb..4826423f8 100644 --- a/EpiAware/src/EpiObsModels/modifiers/Aggregate.jl +++ b/EpiAware/src/EpiObsModels/modifiers/Aggregate.jl @@ -34,12 +34,33 @@ struct Aggregate{M <: AbstractTuringObservationModel, new{typeof(model), typeof(aggregation), typeof(present)}( model, aggregation, present) end +end - function Aggregate(; model, aggregation) - return Aggregate(model, aggregation) - end +function Aggregate(; model, aggregation) + return Aggregate(model, aggregation) end +@doc raw" +Generate observations using an aggregation model. + +# Arguments +- `ag::Aggregate`: The aggregation model. +- `y_t`: The current state of the observations. If missing, a vector of missing values is created. +- `Y_t`: The expected observations. + +# Returns +- A vector of observations where entries are aggregated according to the aggregation model's + specification. Only positions marked as present in the aggregation model contain non-zero + values. + +# Details +The function: +1. Broadcasts the aggregation vector to match the length of observations +2. Broadcasts the present vector to match the length of observations +3. For each present position, sums the expected observations over the aggregation window +4. Generates observations for the aggregated values using the underlying model +5. Returns a vector with observations only in the present positions +" @model function EpiAwareBase.generate_observations(ag::Aggregate, y_t, Y_t) if ismissing(y_t) y_t = Vector{Missing}(missing, length(Y_t)) diff --git a/EpiAware/src/EpiObsModels/modifiers/LatentDelay.jl b/EpiAware/src/EpiObsModels/modifiers/LatentDelay.jl index 4459131b6..275bea1e1 100644 --- a/EpiAware/src/EpiObsModels/modifiers/LatentDelay.jl +++ b/EpiAware/src/EpiObsModels/modifiers/LatentDelay.jl @@ -32,13 +32,6 @@ struct LatentDelay{M <: AbstractTuringObservationModel, T <: AbstractVector{<:Re model::M rev_pmf::T - function LatentDelay(model::M, distribution::C; D = nothing, - Δd = 1.0) where { - M <: AbstractTuringObservationModel, C <: ContinuousDistribution} - pmf = censored_pmf(distribution; Δd = Δd, D = D) - return LatentDelay(model, pmf) - end - function LatentDelay(model::M, pmf::T) where {M <: AbstractTuringObservationModel, T <: AbstractVector{<:Real}} @assert all(pmf .>= 0) "Delay interval must be non-negative" @@ -48,6 +41,13 @@ struct LatentDelay{M <: AbstractTuringObservationModel, T <: AbstractVector{<:Re end end +function LatentDelay(model::M, distribution::C; D = nothing, + Δd = 1.0) where { + M <: AbstractTuringObservationModel, C <: ContinuousDistribution} + pmf = censored_pmf(distribution; Δd = Δd, D = D) + return LatentDelay(model, pmf) +end + @doc raw" Generates observations based on the `LatentDelay` observation model. diff --git a/EpiAware/src/EpiObsModels/modifiers/PrefixObservationModel.jl b/EpiAware/src/EpiObsModels/modifiers/PrefixObservationModel.jl index 5be17d99c..fbb50c139 100644 --- a/EpiAware/src/EpiObsModels/modifiers/PrefixObservationModel.jl +++ b/EpiAware/src/EpiObsModels/modifiers/PrefixObservationModel.jl @@ -21,6 +21,21 @@ prefix::P end +@doc raw" +Generate observations using a prefixed observation model. + +# Arguments +- `observation_model::PrefixObservationModel`: The prefixed observation model +- `y_t`: The current observations +- `Y_t`: The expected observations + +# Returns +The observations generated by the underlying model with prefixed parameter names + +# Details +This function wraps the underlying observation model's `generate_observations` function using +`prefix_submodel` to prefix all parameter names with the specified prefix. +" @model function EpiAwareBase.generate_observations( observation_model::PrefixObservationModel, y_t, Y_t) @submodel submodel = prefix_submodel( diff --git a/EpiAware/src/EpiObsModels/modifiers/RecordExpectedObs.jl b/EpiAware/src/EpiObsModels/modifiers/RecordExpectedObs.jl index 02f3ccb9c..34c3aff92 100644 --- a/EpiAware/src/EpiObsModels/modifiers/RecordExpectedObs.jl +++ b/EpiAware/src/EpiObsModels/modifiers/RecordExpectedObs.jl @@ -22,6 +22,22 @@ struct RecordExpectedObs{M <: AbstractTuringObservationModel} <: model::M end +@doc raw" +Generate observations using a model that records the expected observations. + +# Arguments +- `model::RecordExpectedObs`: The recording model. +- `y_t`: The current state of the observations. If missing, a vector of missing values is created. +- `Y_t`: The expected observations. + +# Returns +- The observations generated by the underlying model. Additionally records `Y_t` as `exp_y_t` + using Turing's `:=` syntax. + +# Details +This function wraps the underlying observation model's `generate_observations` function and +additionally records the expected observations `Y_t` as `exp_y_t` in the model. +" @model function EpiAwareBase.generate_observations(model::RecordExpectedObs, y_t, Y_t) exp_y_t := Y_t @submodel y_t = generate_observations(model.model, y_t, Y_t) diff --git a/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl b/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl index fd74d8a4f..63b0a8a56 100644 --- a/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl +++ b/EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl @@ -39,24 +39,24 @@ struct Ascertainment{ return new{M, AbstractTuringLatentModel, F, P}( model, prefix_model, transform, latent_prefix) end +end - function Ascertainment(model::M, - latent_model::T; - transform::F = (x, y) -> xexpy.(x, y), - latent_prefix::P = "Ascertainment") where { - M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, - F <: Function, P <: String} - return Ascertainment(model, latent_model, transform, latent_prefix) - end +function Ascertainment(model::M, + latent_model::T; + transform::F = (x, y) -> xexpy.(x, y), + latent_prefix::P = "Ascertainment") where { + M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, + F <: Function, P <: String} + return Ascertainment(model, latent_model, transform, latent_prefix) +end - function Ascertainment(; model::M, - latent_model::T, - transform::F = (x, y) -> xexpy.(x, y), - latent_prefix::P = "Ascertainment") where { - M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, - F <: Function, P <: String} - return Ascertainment(model, latent_model, transform, latent_prefix) - end +function Ascertainment(; model::M, + latent_model::T, + transform::F = (x, y) -> xexpy.(x, y), + latent_prefix::P = "Ascertainment") where { + M <: AbstractTuringObservationModel, T <: AbstractTuringLatentModel, + F <: Function, P <: String} + return Ascertainment(model, latent_model, transform, latent_prefix) end @doc raw" diff --git a/EpiAware/test/EpiAwareUtils/generate_epiware.jl b/EpiAware/test/EpiAwareUtils/generate_epiware.jl index 0cbeb0bf5..b2a7ac579 100644 --- a/EpiAware/test/EpiAwareUtils/generate_epiware.jl +++ b/EpiAware/test/EpiAwareUtils/generate_epiware.jl @@ -8,9 +8,8 @@ #Define the epi_model epi_model = DirectInfections(data, Normal()) - #Define the latent process model - rwp = RandomWalk(Normal(0.0, 1.0), - truncated(Normal(0.0, 0.05), 0.0, Inf)) + #Define the latent process model using defaults + rwp = RandomWalk() #Define the observation model delay_distribution = Gamma(2.0, 5 / 2) @@ -38,7 +37,7 @@ end @testitem "`generate_epiaware` with Exp growth rate and RW latent process runs" begin using Distributions, Turing, DynamicPPL # Define test inputs - y_t = missing# rand(1:10, 365) # Data will be generated from the model + y_t = missing # Data will be generated from the model data = EpiData([0.2, 0.3, 0.5], exp) #Define the epi_model @@ -46,9 +45,10 @@ end #Define the latent process model r_3 = log(2) / 3.0 - rwp = RandomWalk( - truncated(Normal(0.0, r_3 / 3), -r_3, r_3), # 3 day doubling time at 3 sigmas in prior - truncated(Normal(0.0, 0.01), 0.0, 0.1)) + init_prior = truncated(Normal(0.0, r_3 / 3), -r_3, r_3)# initial 3 day doubling time at 3 sigmas in prior + step_size_dist = HierarchicalNormal(std_prior = truncated(Normal(0.0, 0.01), 0.0, 0.1)) + rwp = RandomWalk(init_prior = init_prior, + ϵ_t = step_size_dist) #Define the observation model - no delay model time_horizon = 5 @@ -83,9 +83,10 @@ end #Define the latent process model r_3 = log(2) / 3.0 - rwp = RandomWalk( - truncated(Normal(0.0, r_3 / 3), -r_3, r_3), # 3 day doubling time at 3 sigmas in prior - truncated(Normal(0.0, 0.01), 0.0, 0.1)) + init_prior = truncated(Normal(0.0, r_3 / 3), -r_3, r_3)# initial 3 day doubling time at 3 sigmas in prior + step_size_dist = HierarchicalNormal(std_prior = truncated(Normal(0.0, 0.01), 0.0, 0.1)) + rwp = RandomWalk(init_prior = init_prior, + ϵ_t = step_size_dist) #Define the observation model - no delay model time_horizon = 5 diff --git a/EpiAware/test/EpiAwareUtils/turing-methods.jl b/EpiAware/test/EpiAwareUtils/turing-methods.jl index 47e77c9e2..b175f64ce 100644 --- a/EpiAware/test/EpiAwareUtils/turing-methods.jl +++ b/EpiAware/test/EpiAwareUtils/turing-methods.jl @@ -67,11 +67,8 @@ end @testitem "Turing method for generate_epiaware with two latent processes" begin using Distributions, Turing - # Latent model - 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) + # Latent model - default AR model + ar_process = AR() # Used again in obs model diff --git a/EpiAware/test/EpiLatentModels/combinations/arima.jl b/EpiAware/test/EpiLatentModels/combinations/arima.jl new file mode 100644 index 000000000..8fea33737 --- /dev/null +++ b/EpiAware/test/EpiLatentModels/combinations/arima.jl @@ -0,0 +1,97 @@ +@testitem "Testing ARIMA constructor" begin + using Distributions, Turing + + # Test default constructor + arima_model = arima() + @test arima_model isa DiffLatentModel + @test arima_model.model isa AR + @test arima_model.model.ϵ_t isa MA + @test length(arima_model.model.damp_prior) == 1 + @test length(arima_model.model.init_prior) == 1 + @test length(arima_model.model.ϵ_t.θ) == 1 + @test arima_model.model.damp_prior == + filldist(truncated(Normal(0.0, 0.05), 0, 1), 1) + @test arima_model.model.ϵ_t.θ == + filldist(truncated(Normal(0.0, 0.05), -1, 1), 1) + + # Test with custom parameters + ar_init_prior = Normal(1.0, 0.5) + diff_init_prior = Normal(0.0, 0.3) + damp_prior = truncated(Normal(0.0, 0.04), 0, 1) + θ_prior = truncated(Normal(0.0, 0.06), -1, 1) + + custom_arima = arima(; + ar_init = [ar_init_prior, ar_init_prior], + diff_init = [diff_init_prior, diff_init_prior], + damp = [damp_prior, damp_prior], + θ = [θ_prior, θ_prior], + ϵ_t = HierarchicalNormal() + ) + + @test custom_arima isa DiffLatentModel + @test custom_arima.model isa AR + @test custom_arima.model.ϵ_t isa MA + @test length(custom_arima.model.damp_prior) == 2 + @test length(custom_arima.model.init_prior) == 2 + @test length(custom_arima.model.ϵ_t.θ) == 2 + @test custom_arima.model.damp_prior == filldist(damp_prior, 2) + @test custom_arima.model.init_prior == filldist(ar_init_prior, 2) + @test custom_arima.model.ϵ_t.θ == filldist(θ_prior, 2) +end + +@testitem "Testing ARIMA process against theoretical properties" begin + using DynamicPPL, Turing + using HypothesisTests: ExactOneSampleKSTest, pvalue + using Distributions + using Statistics + + # Set up simple ARIMA model + arima_model = arima() + n = 1000 + damp = [0.1] + σ_AR = 1.0 + ar_init = [0.0] + diff_init = [0.0] + θ = [0.2] # Add MA component + + # Generate and fix model parameters + model = generate_latent(arima_model, n) + fixed_model = fix(model, + ( + std = σ_AR, + damp_AR = damp, + ar_init = ar_init, + diff_init = diff_init, + θ = θ + )) + + # Generate samples + n_samples = 100 + samples = sample(fixed_model, Prior(), n_samples; progress = false) |> + chn -> mapreduce(vcat, generated_quantities(fixed_model, chn)) do gen + gen + end + + # Compare with pure AR with differencing + ar_base = AR() + ar_model = DiffLatentModel(; model = ar_base, init_priors = [Normal()]) + ar_fixed = fix( + generate_latent(ar_model, n), + (std = σ_AR, damp_AR = damp, ar_init = ar_init, diff_init = diff_init) + ) + + ar_samples = sample(ar_fixed, Prior(), n_samples; progress = false) |> + chn -> mapreduce(vcat, generated_quantities(ar_fixed, chn)) do gen + gen + end + + # Test that ARIMA produces different distribution than pure differenced AR + # This tests that the MA component has an effect + ks_test = ExactOneSampleKSTest(samples, fit(Normal, ar_samples)) + @test pvalue(ks_test) < 1e-6 + + # Test for stationarity of differences + diff_samples = diff(samples) + @test isapprox(mean(diff_samples), 0.0, atol = 0.1) + @test std(diff_samples) > 0 +end diff --git a/EpiAware/test/EpiLatentModels/combinations/arma.jl b/EpiAware/test/EpiLatentModels/combinations/arma.jl new file mode 100644 index 000000000..19704332d --- /dev/null +++ b/EpiAware/test/EpiLatentModels/combinations/arma.jl @@ -0,0 +1,34 @@ +@testitem "Testing ARMA constructor" begin + using Distributions, Turing + + # Test default constructor + arma_model = arma() + @test arma_model isa AR + @test arma_model.ϵ_t isa MA + @test length(arma_model.damp_prior) == 1 + @test length(arma_model.init_prior) == 1 + @test length(arma_model.ϵ_t.θ) == 1 + @test arma_model.damp_prior == filldist(truncated(Normal(0.0, 0.05), 0, 1), 1) + @test arma_model.ϵ_t.θ == filldist(truncated(Normal(0.0, 0.05), -1, 1), 1) + + # Test with custom parameters + damp_prior = truncated(Normal(0.0, 0.04), 0, 1) + θ_prior = truncated(Normal(0.0, 0.06), 0, 1) + init_prior = Normal(1.0, 0.5) + + custom_arma = arma(; + init = [init_prior, init_prior], + damp = [damp_prior, damp_prior], + θ = [θ_prior, θ_prior], + ϵ_t = HierarchicalNormal() + ) + + @test custom_arma isa AR + @test custom_arma.ϵ_t isa MA + @test length(custom_arma.damp_prior) == 2 + @test length(custom_arma.init_prior) == 2 + @test length(custom_arma.ϵ_t.θ) == 2 + @test custom_arma.damp_prior == filldist(damp_prior, 2) + @test custom_arma.init_prior == filldist(init_prior, 2) + @test custom_arma.ϵ_t.θ == filldist(θ_prior, 2) +end diff --git a/EpiAware/test/EpiLatentModels/manipulators/CombineLatentModels.jl b/EpiAware/test/EpiLatentModels/manipulators/CombineLatentModels.jl index 881baa599..ec03cd829 100644 --- a/EpiAware/test/EpiLatentModels/manipulators/CombineLatentModels.jl +++ b/EpiAware/test/EpiLatentModels/manipulators/CombineLatentModels.jl @@ -84,7 +84,7 @@ end ns_regression_mdl = no_slope_linear_regression(y) | (var"Combine.2.damp_AR" = [0.0], var"Combine.2.σ_AR" = 1.0, var"Combine.2.ϵ_t" = zeros(n - 1), var"Combine.2.ar_init" = [0.0]) - chain = sample(ns_regression_mdl, NUTS(), 5000, progress = false) + chain = sample(ns_regression_mdl, NUTS(), 5000; progress = false) # Theoretical posterior distribution for intercept # if \beta ~ int.intercept_prior = N(\mu_0, \sigma_0) and \sigma^2 = 1 for diff --git a/EpiAware/test/EpiLatentModels/manipulators/broadcast/LatentModel.jl b/EpiAware/test/EpiLatentModels/manipulators/broadcast/LatentModel.jl index 9843bf140..3afdb70c7 100644 --- a/EpiAware/test/EpiLatentModels/manipulators/broadcast/LatentModel.jl +++ b/EpiAware/test/EpiLatentModels/manipulators/broadcast/LatentModel.jl @@ -28,7 +28,7 @@ end @test length(rand_model.ϵ_t) == 2 fix_model = fix( broadcasted_model, - (σ_RW = 2.0, rw_init = 1.0, ϵ_t = [1, 2]) + (std = 2.0, rw_init = 1.0, ϵ_t = [1, 2]) ) out = fix_model() @test out == vcat(fill(1.0, 5), fill(3.0, 5), fill(7.0, 5)) diff --git a/EpiAware/test/EpiLatentModels/models/AR.jl b/EpiAware/test/EpiLatentModels/models/AR.jl index 0e99b7b25..e3e98358c 100644 --- a/EpiAware/test/EpiLatentModels/models/AR.jl +++ b/EpiAware/test/EpiLatentModels/models/AR.jl @@ -3,11 +3,12 @@ damp_prior = truncated(Normal(0.0, 0.05), 0.0, 1) std_prior = truncated(Normal(0.0, 0.05), 0.0, Inf) + ϵ_t = HierarchicalNormal(; std_prior) init_prior = Normal() - ar_process = AR(damp_prior, std_prior, init_prior) + ar_process = AR(damp_prior, init_prior; ϵ_t) @test ar_process.damp_prior == filldist(damp_prior, 1) - @test ar_process.std_prior == std_prior + @test ar_process.ϵ_t.std_prior == std_prior @test ar_process.init_prior == filldist(init_prior, 1) end @@ -20,7 +21,7 @@ end end @testset "std_prior" begin - std_AR = rand(ar.std_prior) + std_AR = rand(ar.ϵ_t.std_prior) @test std_AR >= 0.0 end @@ -32,10 +33,12 @@ end @testitem "Test AR(2)" begin using Distributions + std_prior = truncated(Normal(0.0, 0.05), 0.0, Inf) + ϵ_t = HierarchicalNormal(; std_prior) 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), + ϵ_t = ϵ_t, init_priors = [Normal(), Normal()] ) @testset "damp_prior" begin @@ -46,7 +49,7 @@ end end @testset "std_prior" begin - std_AR = rand(ar.std_prior) + std_AR = rand(ar.ϵ_t.std_prior) @test std_AR >= 0.0 end @@ -70,7 +73,7 @@ end ar_init = [0.0] model = generate_latent(ar_model, n) - fixed_model = fix(model, (σ_AR = σ_AR, damp_AR = damp, ar_init = ar_init)) + fixed_model = fix(model, (std = σ_AR, damp_AR = damp, ar_init = ar_init)) n_samples = 100 samples = sample(fixed_model, Prior(), n_samples; progress = false) |> diff --git a/EpiAware/test/EpiLatentModels/models/HierarchicalNormal.jl b/EpiAware/test/EpiLatentModels/models/HierarchicalNormal.jl index 5fd843f79..4f6287447 100644 --- a/EpiAware/test/EpiLatentModels/models/HierarchicalNormal.jl +++ b/EpiAware/test/EpiLatentModels/models/HierarchicalNormal.jl @@ -8,9 +8,10 @@ int_def = HierarchicalNormal() @test typeof(int_def) <: AbstractTuringLatentModel @test int_def.mean == 0.0 - @test int_def.std_prior == truncated(Normal(0, 1), 0, Inf) + @test int_def.std_prior == truncated(Normal(0, 0.1), 0, Inf) - @test int == HierarchicalNormal(mean = 0.1, std_prior = truncated(Normal(0, 2), 0, Inf)) + @test int == HierarchicalNormal( + mean = 0.1, std_prior = truncated(Normal(0, 2), 0, Inf)) end @testitem "HierarchicalNormal generate_latent" begin diff --git a/EpiAware/test/EpiLatentModels/models/IID.jl b/EpiAware/test/EpiLatentModels/models/IID.jl new file mode 100644 index 000000000..478e0385e --- /dev/null +++ b/EpiAware/test/EpiLatentModels/models/IID.jl @@ -0,0 +1,57 @@ +@testitem "Testing IID constructor" begin + using Distributions, Turing + + normal_prior = Normal(0.0, 1.0) + idd_process = IID(normal_prior) + + @test idd_process.ϵ_t == normal_prior +end + +@testitem "Test IID defaults" begin + using Distributions + idd = IID() + @test idd.ϵ_t == Normal(0, 1) +end + +@testitem "Test IID with different distributions" begin + using Distributions + + @testset "Uniform distribution" begin + idd = IID(Uniform(0, 1)) + sample = rand(idd.ϵ_t) + @test 0 <= sample <= 1 + end + + @testset "Exponential distribution" begin + idd = IID(Exponential(1)) + sample = rand(idd.ϵ_t) + @test sample >= 0 + end +end + +@testitem "Testing IID process against theoretical properties" begin + using DynamicPPL, Turing + using HypothesisTests: ExactOneSampleKSTest, pvalue + using Distributions + + idd_model = IID(Normal(2, 3)) + n = 1000 + + model = generate_latent(idd_model, n) + + n_samples = 100 + samples = sample(model, Prior(), n_samples; progress = false) |> + chn -> mapreduce(vcat, generated_quantities(model, chn)) do gen + gen + end + + theoretical_mean = 2.0 + theoretical_var = 3.0^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/EpiLatentModels/models/MA.jl b/EpiAware/test/EpiLatentModels/models/MA.jl new file mode 100644 index 000000000..d52d4017d --- /dev/null +++ b/EpiAware/test/EpiLatentModels/models/MA.jl @@ -0,0 +1,84 @@ +@testitem "Testing MA constructor" begin + using Distributions, Turing + + θ_prior = truncated(Normal(0.0, 0.05), -1, 1) + ma_process = MA(θ_prior; q = 1, ϵ_t = HierarchicalNormal()) + + @test ma_process.q == 1 + @test ma_process.ϵ_t isa HierarchicalNormal + @test length(ma_process.θ) == 1 +end + +@testitem "Test MA(2)" begin + using Distributions, Turing + θ_prior = truncated(Normal(0.0, 0.05), -1, 1) + ma = MA(; + θ_priors = [θ_prior, θ_prior], + ϵ_t = HierarchicalNormal() + ) + @test ma.q == 2 + @test ma.ϵ_t isa HierarchicalNormal + @test length(ma.θ) == 2 +end + +@testitem "Testing MA process against theoretical properties" begin + using DynamicPPL, Turing + using HypothesisTests: ExactOneSampleKSTest, pvalue + using Distributions + + # Test MA(1) process + θ = [0.1] + @testset "MA(1) with θ = $θ" begin + ma_model = MA(; θ_priors = [truncated(Normal(0.0, 0.05), -1, 1)], + ϵ_t = IID(Normal())) + n = 1000 + model = generate_latent(ma_model, n) + fixed_model = fix(model, (θ = θ,)) + + n_samples = 100 + samples = sample(fixed_model, Prior(), n_samples; progress = false) |> + chn -> mapreduce(vcat, generated_quantities(fixed_model, chn)) do gen + gen + end + + # For MA(1), mean should be 0 + @test isapprox(mean(samples), 0.0, atol = 0.1) + + # For MA(1) with standard normal errors, variance is 1 + θ^2 + theoretical_var = 1 + sum(θ .^ 2) + @test isapprox(var(samples), theoretical_var, atol = 0.2) + + # Test distribution is approximately normal + ks_test = ExactOneSampleKSTest( + samples, Normal(0.0, sqrt(theoretical_var))) + @test pvalue(ks_test) > 1e-6 + end + + # Test MA(2) process + θ = [0.3, 0.2] + @testset "MA(2) with θ = $θ" begin + ma_model = MA(; θ_priors = fill(truncated(Normal(0.0, 0.05), -1, 1), 2), + ϵ_t = IID(Normal())) + n = 1000 + model = generate_latent(ma_model, n) + fixed_model = fix(model, (θ = θ,)) + + n_samples = 100 + samples = sample(fixed_model, Prior(), n_samples; progress = false) |> + chn -> mapreduce(vcat, generated_quantities(fixed_model, chn)) do gen + gen + end + + # For MA(2), mean should be 0 + @test isapprox(mean(samples), 0.0, atol = 0.1) + + # For MA(2) with standard normal errors, variance is 1 + θ_1^2 + θ_2^2 + theoretical_var = 1 + sum(θ .^ 2) + @test isapprox(var(samples), theoretical_var, atol = 0.2) + + # Test distribution is approximately normal + ks_test = ExactOneSampleKSTest( + samples, Normal(0.0, sqrt(theoretical_var))) + @test pvalue(ks_test) > 1e-6 + end +end diff --git a/EpiAware/test/EpiLatentModels/models/RandomWalk.jl b/EpiAware/test/EpiLatentModels/models/RandomWalk.jl index d2efa4a3b..7d7aac220 100644 --- a/EpiAware/test/EpiLatentModels/models/RandomWalk.jl +++ b/EpiAware/test/EpiLatentModels/models/RandomWalk.jl @@ -3,9 +3,9 @@ using HypothesisTests: ExactOneSampleKSTest, pvalue n = 5 - rw_process = RandomWalk(Normal(0.0, 1.0), HalfNormal(0.05)) + rw_process = RandomWalk(; ϵ_t = IID(Normal())) model = generate_latent(rw_process, n) - fixed_model = fix(model, (σ_RW = 1.0, init_rw_value = 0.0)) #Fixing the standard deviation of the random walk process + fixed_model = fix(model, (rw_init = 0.0,)) n_samples = 1000 samples_day_5 = sample(fixed_model, Prior(), n_samples; progress = false) |> chn -> mapreduce(vcat, generated_quantities(fixed_model, chn)) do gen @@ -17,26 +17,25 @@ end @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_prior" begin priors = RandomWalk() init_value = rand(priors.init_prior) @test typeof(init_value) == Float64 end + + @testset "ϵ_t" begin + priors = RandomWalk() + @test priors.ϵ_t isa HierarchicalNormal + end end @testitem "Testing RandomWalk constructor" begin - using Distributions: Normal, truncated + using Distributions: Normal init_prior = Normal(0.0, 1.0) - std_prior = HalfNormal(0.05) - rw_process = RandomWalk(init_prior, std_prior) + ϵ_t = HierarchicalNormal() + rw_process = RandomWalk(init_prior, ϵ_t) @test rw_process.init_prior == init_prior - @test rw_process.std_prior == std_prior + @test rw_process.ϵ_t == ϵ_t end @testitem "Testing RandomWalk parameter recovery: Negative Binomial errors on log rw" begin diff --git a/EpiAware/test/EpiLatentModels/modifiers/DiffLatentModel.jl b/EpiAware/test/EpiLatentModels/modifiers/DiffLatentModel.jl index 86ea0112e..b2403c0b0 100644 --- a/EpiAware/test/EpiLatentModels/modifiers/DiffLatentModel.jl +++ b/EpiAware/test/EpiLatentModels/modifiers/DiffLatentModel.jl @@ -31,7 +31,10 @@ end n = 100 d = 2 - model = RandomWalk(Normal(0.0, 1.0), truncated(Normal(0.0, 0.05), 0.0, Inf)) + model = RandomWalk( + init_prior = Normal(0.0, 1.0), + ϵ_t = HierarchicalNormal(truncated(Normal(0.0, 0.05), 0.0, Inf)) + ) init_priors = [Normal(0.0, 1.0), Normal(1.0, 2.0)] diff_model = DiffLatentModel(model = model, init_priors = init_priors) diff --git a/EpiAware/test/EpiObsModels/modifiers/LatentDelay.jl b/EpiAware/test/EpiObsModels/modifiers/LatentDelay.jl index 34ccc1721..7fe379695 100644 --- a/EpiAware/test/EpiObsModels/modifiers/LatentDelay.jl +++ b/EpiAware/test/EpiObsModels/modifiers/LatentDelay.jl @@ -196,18 +196,20 @@ end return (; init_prior, std_prior) end end - function set_latent_process(epimodel, latentprocess_type) init_prior, std_prior = set_init_and_std_prior(epimodel) if latentprocess_type == RandomWalk - return RandomWalk(init_prior, std_prior) + return RandomWalk(; init_prior = init_prior, + ϵ_t = HierarchicalNormal(; std_prior = std_prior)) elseif latentprocess_type == AR - return AR(damp_priors = [Beta(2, 8; check_args = false)], - std_prior = std_prior, init_priors = [init_prior]) + return AR(; damp_priors = [Beta(2, 8; check_args = false)], + init_priors = [init_prior], + ϵ_t = HierarchicalNormal(; std_prior = std_prior)) elseif latentprocess_type == DiffLatentModel return DiffLatentModel( - AR(damp_priors = [Beta(2, 8; check_args = false)], - std_prior = std_prior, init_priors = [Normal(0.0, 0.25)]), + AR(; damp_priors = [Beta(2, 8; check_args = false)], + init_priors = [Normal(0.0, 0.25)], + ϵ_t = HierarchicalNormal(; std_prior = std_prior)), init_prior; d = 1) end end diff --git a/benchmark/bench/EpiLatentModels/EpiLatentModels.jl b/benchmark/bench/EpiLatentModels/EpiLatentModels.jl index dd0e5e869..b90fb0285 100644 --- a/benchmark/bench/EpiLatentModels/EpiLatentModels.jl +++ b/benchmark/bench/EpiLatentModels/EpiLatentModels.jl @@ -4,7 +4,9 @@ using BenchmarkTools, TuringBenchmarking, EpiAware, DynamicPPL suite = BenchmarkGroup() include("../../make_epiaware_suite.jl") +include("models/IID.jl") include("models/AR.jl") +include("models/MA.jl") include("models/RandomWalk.jl") include("models/Intercept.jl") include("models/FixedIntercept.jl") diff --git a/benchmark/bench/EpiLatentModels/combinations/arima.jl b/benchmark/bench/EpiLatentModels/combinations/arima.jl new file mode 100644 index 000000000..628e95308 --- /dev/null +++ b/benchmark/bench/EpiLatentModels/combinations/arima.jl @@ -0,0 +1,5 @@ +let + model = arima() + mdl = generate_latent(model, 10) + suite["arima"] = make_epiaware_suite(mdl) +end diff --git a/benchmark/bench/EpiLatentModels/combinations/arma.jl b/benchmark/bench/EpiLatentModels/combinations/arma.jl new file mode 100644 index 000000000..d953c22b5 --- /dev/null +++ b/benchmark/bench/EpiLatentModels/combinations/arma.jl @@ -0,0 +1,5 @@ +let + model = arma() + mdl = generate_latent(model, 10) + suite["arma"] = make_epiaware_suite(mdl) +end diff --git a/benchmark/bench/EpiLatentModels/models/IID.jl b/benchmark/bench/EpiLatentModels/models/IID.jl new file mode 100644 index 000000000..ce05cac9d --- /dev/null +++ b/benchmark/bench/EpiLatentModels/models/IID.jl @@ -0,0 +1,5 @@ +let + model = IID() + mdl = generate_latent(model, 10) + suite["IID"] = make_epiaware_suite(mdl) +end diff --git a/benchmark/bench/EpiLatentModels/models/MA.jl b/benchmark/bench/EpiLatentModels/models/MA.jl new file mode 100644 index 000000000..43f00dd7a --- /dev/null +++ b/benchmark/bench/EpiLatentModels/models/MA.jl @@ -0,0 +1,5 @@ +let + latent = MA() + mdl = generate_latent(latent, 10) + suite["MA"] = make_epiaware_suite(mdl) +end