Skip to content

Commit

Permalink
Issue 408: ARMA and ARIMA models (#438)
Browse files Browse the repository at this point in the history
* 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 <compathelper_noreply@julialang.org>
Co-authored-by: Sam Abbott <azw1@cdc.gov>
Co-authored-by: Samuel Brand <48288458+SamuelBrand1@users.noreply.github.com>
Co-authored-by: Samuel Brand <sam055@mac.com>

* CompatHelper: bump compat for DynamicPPL to 0.30 for package EpiAware, (drop existing compat) (#528)

Co-authored-by: CompatHelper Julia <compathelper_noreply@julialang.org>

* 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 <compathelper_noreply@julialang.org>
Co-authored-by: Sam Abbott <azw1@cdc.gov>

* 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 <compathelper_noreply@julialang.org>
Co-authored-by: Samuel Brand <sam055@mac.com>
5 people authored Dec 12, 2024
1 parent 8dc444f commit d7b7b3b
Showing 42 changed files with 977 additions and 318 deletions.
13 changes: 7 additions & 6 deletions EpiAware/docs/src/showcase/replications/chatzilena-2019/index.jl
Original file line number Diff line number Diff line change
@@ -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
20 changes: 10 additions & 10 deletions EpiAware/docs/src/showcase/replications/mishra-2020/index.jl
Original file line number Diff line number Diff line change
@@ -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
43 changes: 28 additions & 15 deletions EpiAware/src/EpiAwareUtils/HalfNormal.jl
Original file line number Diff line number Diff line change
@@ -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
30 changes: 20 additions & 10 deletions EpiAware/src/EpiAwareUtils/SafeNegativeBinomial.jl
Original file line number Diff line number Diff line change
@@ -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
30 changes: 20 additions & 10 deletions EpiAware/src/EpiAwareUtils/SafePoisson.jl
Original file line number Diff line number Diff line change
@@ -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
55 changes: 12 additions & 43 deletions EpiAware/src/EpiAwareUtils/censored_pmf.jl
Original file line number Diff line number Diff line change
@@ -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)
11 changes: 9 additions & 2 deletions EpiAware/src/EpiLatentModels/EpiLatentModels.jl
Original file line number Diff line number Diff line change
@@ -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")
42 changes: 42 additions & 0 deletions EpiAware/src/EpiLatentModels/combinations/arima.jl
Original file line number Diff line number Diff line change
@@ -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
40 changes: 40 additions & 0 deletions EpiAware/src/EpiLatentModels/combinations/arma.jl
Original file line number Diff line number Diff line change
@@ -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
10 changes: 5 additions & 5 deletions EpiAware/src/EpiLatentModels/manipulators/CombineLatentModels.jl
Original file line number Diff line number Diff line change
@@ -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"
38 changes: 19 additions & 19 deletions EpiAware/src/EpiLatentModels/manipulators/ConcatLatentModels.jl
Original file line number Diff line number Diff line change
@@ -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"
88 changes: 48 additions & 40 deletions EpiAware/src/EpiLatentModels/models/AR.jl
Original file line number Diff line number Diff line change
@@ -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
40 changes: 30 additions & 10 deletions EpiAware/src/EpiLatentModels/models/HierarchicalNormal.jl
Original file line number Diff line number Diff line change
@@ -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
53 changes: 53 additions & 0 deletions EpiAware/src/EpiLatentModels/models/IID.jl
Original file line number Diff line number Diff line change
@@ -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
26 changes: 21 additions & 5 deletions EpiAware/src/EpiLatentModels/models/Intercept.jl
Original file line number Diff line number Diff line change
@@ -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
121 changes: 121 additions & 0 deletions EpiAware/src/EpiLatentModels/models/MA.jl
Original file line number Diff line number Diff line change
@@ -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
101 changes: 49 additions & 52 deletions EpiAware/src/EpiLatentModels/models/RandomWalk.jl
Original file line number Diff line number Diff line change
@@ -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
26 changes: 13 additions & 13 deletions EpiAware/src/EpiLatentModels/modifiers/DiffLatentModel.jl
Original file line number Diff line number Diff line change
@@ -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`.
16 changes: 8 additions & 8 deletions EpiAware/src/EpiObsModels/StackObservationModels.jl
Original file line number Diff line number Diff line change
@@ -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"
27 changes: 24 additions & 3 deletions EpiAware/src/EpiObsModels/modifiers/Aggregate.jl
Original file line number Diff line number Diff line change
@@ -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))
14 changes: 7 additions & 7 deletions EpiAware/src/EpiObsModels/modifiers/LatentDelay.jl
Original file line number Diff line number Diff line change
@@ -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.
15 changes: 15 additions & 0 deletions EpiAware/src/EpiObsModels/modifiers/PrefixObservationModel.jl
Original file line number Diff line number Diff line change
@@ -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(
16 changes: 16 additions & 0 deletions EpiAware/src/EpiObsModels/modifiers/RecordExpectedObs.jl
Original file line number Diff line number Diff line change
@@ -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)
32 changes: 16 additions & 16 deletions EpiAware/src/EpiObsModels/modifiers/ascertainment/Ascertainment.jl
Original file line number Diff line number Diff line change
@@ -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"
21 changes: 11 additions & 10 deletions EpiAware/test/EpiAwareUtils/generate_epiware.jl
Original file line number Diff line number Diff line change
@@ -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,17 +37,18 @@ 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
epi_model = ExpGrowthRate(data, Normal())

#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
7 changes: 2 additions & 5 deletions EpiAware/test/EpiAwareUtils/turing-methods.jl
Original file line number Diff line number Diff line change
@@ -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

97 changes: 97 additions & 0 deletions EpiAware/test/EpiLatentModels/combinations/arima.jl
Original file line number Diff line number Diff line change
@@ -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
34 changes: 34 additions & 0 deletions EpiAware/test/EpiLatentModels/combinations/arma.jl
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
@@ -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))
15 changes: 9 additions & 6 deletions EpiAware/test/EpiLatentModels/models/AR.jl
Original file line number Diff line number Diff line change
@@ -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) |>
5 changes: 3 additions & 2 deletions EpiAware/test/EpiLatentModels/models/HierarchicalNormal.jl
Original file line number Diff line number Diff line change
@@ -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
57 changes: 57 additions & 0 deletions EpiAware/test/EpiLatentModels/models/IID.jl
Original file line number Diff line number Diff line change
@@ -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
84 changes: 84 additions & 0 deletions EpiAware/test/EpiLatentModels/models/MA.jl
Original file line number Diff line number Diff line change
@@ -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
23 changes: 11 additions & 12 deletions EpiAware/test/EpiLatentModels/models/RandomWalk.jl
Original file line number Diff line number Diff line change
@@ -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
5 changes: 4 additions & 1 deletion EpiAware/test/EpiLatentModels/modifiers/DiffLatentModel.jl
Original file line number Diff line number Diff line change
@@ -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)

14 changes: 8 additions & 6 deletions EpiAware/test/EpiObsModels/modifiers/LatentDelay.jl
Original file line number Diff line number Diff line change
@@ -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
2 changes: 2 additions & 0 deletions benchmark/bench/EpiLatentModels/EpiLatentModels.jl
Original file line number Diff line number Diff line change
@@ -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")
5 changes: 5 additions & 0 deletions benchmark/bench/EpiLatentModels/combinations/arima.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
let
model = arima()
mdl = generate_latent(model, 10)
suite["arima"] = make_epiaware_suite(mdl)
end
5 changes: 5 additions & 0 deletions benchmark/bench/EpiLatentModels/combinations/arma.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
let
model = arma()
mdl = generate_latent(model, 10)
suite["arma"] = make_epiaware_suite(mdl)
end
5 changes: 5 additions & 0 deletions benchmark/bench/EpiLatentModels/models/IID.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
let
model = IID()
mdl = generate_latent(model, 10)
suite["IID"] = make_epiaware_suite(mdl)
end
5 changes: 5 additions & 0 deletions benchmark/bench/EpiLatentModels/models/MA.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
let
latent = MA()
mdl = generate_latent(latent, 10)
suite["MA"] = make_epiaware_suite(mdl)
end

0 comments on commit d7b7b3b

Please sign in to comment.