diff --git a/EpiAware/test/EpiLatentModels/combinations/arima.jl b/EpiAware/test/EpiLatentModels/combinations/arima.jl new file mode 100644 index 000000000..c78fc4b4f --- /dev/null +++ b/EpiAware/test/EpiLatentModels/combinations/arima.jl @@ -0,0 +1,40 @@ +@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 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/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