diff --git a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl index 802df6957..50a4fd251 100644 --- a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl +++ b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl @@ -28,6 +28,9 @@ 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 define_arma, define_arima + include("docstrings.jl") include("models/Intercept.jl") include("models/IDD.jl") @@ -44,6 +47,8 @@ include("manipulators/ConcatLatentModels.jl") include("manipulators/broadcast/LatentModel.jl") include("manipulators/broadcast/rules.jl") include("manipulators/broadcast/helpers.jl") +include("combinations/define_arma.jl") +include("combinations/define_arima.jl") include("utils.jl") end diff --git a/EpiAware/src/EpiLatentModels/combinations/arma.jl b/EpiAware/src/EpiLatentModels/combinations/arma.jl deleted file mode 100644 index a43489d94..000000000 --- a/EpiAware/src/EpiLatentModels/combinations/arma.jl +++ /dev/null @@ -1,3 +0,0 @@ -function arma() - AR(; ϵ_t = MA()) -end diff --git a/EpiAware/src/EpiLatentModels/combinations/define_arima.jl b/EpiAware/src/EpiLatentModels/combinations/define_arima.jl new file mode 100644 index 000000000..8af9b2a32 --- /dev/null +++ b/EpiAware/src/EpiLatentModels/combinations/define_arima.jl @@ -0,0 +1,44 @@ +@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. +- `d_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 + +```julia +using EpiAware, Distributions + +ARIMA = arima( + ar_init = [Normal(0.0, 1.0)], + d_init = [Normal()], + θ = [truncated(Normal(0.0, 0.02), -1, 1)], + damp = [truncated(Normal(0.0, 0.02), 0, 1)] +) +arma_model = generate_latent(ARIMA, 10) +arma_model() +``` +""" +function arima(; + ar_init = [Normal()], + d_init = [Normal()], + damp = [truncated(Normal(0.0, 0.05), 0, 1)], + θ = [truncated(Normal(0.0, 0.05), -1, 1)], + ϵ_t = HierarchicalNormal() +) + arma = define_arma(; init = ar_init, damp = damp, θ = θ, ϵ_t = ϵ_t) + arima_model = DiffLatentModel(; model = arma, init_priors = d_init) + return arima_model +end diff --git a/EpiAware/src/EpiLatentModels/combinations/define_arma.jl b/EpiAware/src/EpiLatentModels/combinations/define_arma.jl new file mode 100644 index 000000000..4fabac7b1 --- /dev/null +++ b/EpiAware/src/EpiLatentModels/combinations/define_arma.jl @@ -0,0 +1,38 @@ +@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 + +```@example +using EpiAware, Distributions + +ARMA = define_arma(; + θ = [truncated(Normal(0.0, 0.02), -1, 1)], + damp = [truncated(Normal(0.0, 0.02), 0, 1)] +) +arma = generate_latent(ARMA, 10) +arma() +``` +""" +function define_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/models/AR.jl b/EpiAware/src/EpiLatentModels/models/AR.jl index 15c00796c..8ab9516a0 100644 --- a/EpiAware/src/EpiLatentModels/models/AR.jl +++ b/EpiAware/src/EpiLatentModels/models/AR.jl @@ -27,7 +27,7 @@ rand(mdl) # output ``` " -struct AR{D <: Sampleable, S <: Sampleable, I <: Sampleable, +struct AR{D <: Sampleable, I <: Sampleable, P <: Int, E <: AbstractTuringLatentModel} <: AbstractTuringLatentModel "Prior distribution for the damping coefficients." damp_prior::D @@ -60,8 +60,7 @@ struct AR{D <: Sampleable, S <: Sampleable, I <: Sampleable, @assert p>0 "p must be greater than 0" @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 - ) + damp_prior, init_prior, p, ϵ_t) end end @@ -84,12 +83,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 @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/RandomWalk.jl b/EpiAware/src/EpiLatentModels/models/RandomWalk.jl index 455d84a56..2592eb44e 100644 --- a/EpiAware/src/EpiLatentModels/models/RandomWalk.jl +++ b/EpiAware/src/EpiLatentModels/models/RandomWalk.jl @@ -17,7 +17,8 @@ Constructing a random walk requires specifying: ## Constructors -- `RandomWalk(; init_prior, std_prior, ϵ_t)` +- `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 @@ -26,7 +27,7 @@ using Distributions, Turing, EpiAware rw = RandomWalk() rw # output -RandomWalk{Normal{Float64}, HalfNormal{Float64}, IDD{Normal{Float64}}}(init_prior=Normal{Float64}(μ=0.0, σ=1.0), std_prior=HalfNormal{Float64}(σ=0.25), ϵ_t=IDD{Normal{Float64}}(ϵ_t=Normal{Float64}(μ=0.0, σ=1.0))) +RandomWalk{Normal{Float64}, HierarchicalNormal{Float64}}(init_prior=Normal{Float64}(μ=0.0, σ=1.0), ϵ_t=HierarchicalNormal{Float64}(mean=0.0, std_prior=Truncated{Normal{Float64}, Continuous, Float64}(a=0.0, b=Inf, x=Normal{Float64}(μ=0.0, σ=0.1)))) ``` ```jldoctest RandomWalk; filter=r\"\b\d+(\.\d+)?\b\" => \"*\" @@ -41,51 +42,46 @@ rand(mdl) ``` " @kwdef struct RandomWalk{ - D <: Sampleable, S <: Sampleable, E <: AbstractTuringLatentModel} <: + D <: Sampleable, E <: AbstractTuringLatentModel} <: AbstractTuringLatentModel init_prior::D = Normal() - std_prior::S = HalfNormal(0.25) - ϵ_t::E = IDD(Normal()) -end - -function RandomWalk(init_prior::D, std_prior::S) where {D <: Sampleable, S <: Sampleable} - return RandomWalk(; init_prior = init_prior, std_prior = std_prior) + ϵ_t::E = HierarchicalNormal() end @doc raw" -Implement the `generate_latent` function for the `RandomWalk` model. - -## Example usage of `generate_latent` with `RandomWalk` type of latent process model - -```julia -using Distributions, Turing, EpiAware +Generate a latent RW series using accumulate_scan. -# Create a RandomWalk model -rw = RandomWalk(init_prior = Normal(2., 1.), - std_prior = HalfNormal(0.1)) -``` - -Then, we can use `generate_latent` to construct a Turing model for a 10 step random walk. +# Arguments -```julia -# Construct a Turing model -rw_model = generate_latent(rw, 10) -``` +- `latent_model::RandomWalk`: The RandomWalk model. +- `n::Int`: The length of the RW series. -Now we can use the `Turing` PPL API to sample underlying parameters and generate the -unobserved infections. +# Returns +- `rw::Vector{Float64}`: The generated RW series. -```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 @submodel ϵ_t = generate_latent(latent_model.ϵ_t, n - 1) - rw = rw_init .+ vcat(0.0, σ_RW .* cumsum(ϵ_t)) + + 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