clean up and add arma and arima helpers

Author: seabbs

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

EpiAware/src/EpiLatentModels/combinations/arma.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

EpiAware/src/EpiLatentModels/combinations/define_arima.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

EpiAware/src/EpiLatentModels/combinations/define_arma.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

EpiAware/src/EpiLatentModels/models/AR.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