work through tests

CDCgov · Mar 12, 2024 · 130967b · 130967b
1 parent 98aa223
commit 130967b
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 31 deletions.
diff --git a/EpiAware/src/latentmodels/autoregressive.jl b/EpiAware/src/latentmodels/autoregressive.jl
@@ -6,7 +6,7 @@
 end
 
 @model function generate_latent(latent_model::AR, n)
-    p = length(damp_prior)
+    p = length(latent_model.damp_prior)
     ϵ_t ~ MvNormal(ones(n - p))
     σ_AR ~ latent_model.std_prior
     ar_init ~ latent_model.init_prior

diff --git a/EpiAware/test/test_autoregressive.jl b/EpiAware/test/test_autoregressive.jl
@@ -1,55 +1,48 @@
-@testitem "Testing default_ar_priors" begin
+@testitem "Testing AR constructor" begin
     using Distributions
 
+    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 = EpiAware.AR(damp_prior, std_prior, init_prior)
+
+    @test ar_process.damp_prior == damp_prior
+    @test ar_process.std_prior == std_prior
+    @test ar_process.init_prior == init_prior
+end
+
+@testitem "Test AR defaults" begin
+    using Distributions
+    ar = EpiAware.AR()
     @testset "damp_prior" begin
-        priors = EpiAware.default_ar_priors()
-        damp = rand(priors[:damp_prior][1])
+        damp = rand(ar.damp_prior)
         @test 0.0 <= damp <= 1.0
     end
 
-    @testset "var_prior" begin
-        priors = EpiAware.default_ar_priors()
-        var_AR = rand(priors[:var_prior])
-        @test var_AR >= 0.0
+    @testset "std_prior" begin
+        std_AR = rand(ar.std_prior)
+        @test std_AR >= 0.0
     end
 
     @testset "init_prior" begin
-        priors = EpiAware.default_ar_priors()
-        init_ar_value = rand(priors[:init_prior])
+        init_ar_value = rand(ar.init_prior)
         @test typeof(init_ar_value) == Float64
     end
 end
 
-@testitem "Testing AR constructor" begin
-    using Distributions
-
-    damp_prior = [truncated(Normal(0.0, 0.05), 0.0, 1)]
-    var_prior = truncated(Normal(0.0, 0.05), 0.0, Inf)
-    init_prior = Normal()
-    ar_process = EpiAware.AR(damp_prior, var_prior, init_prior)
-
-    @test ar_process.damp_prior == damp_prior
-    @test ar_process.var_prior == var_prior
-    @test ar_process.init_prior == init_prior
-    @test ar_process.p == 1
-end
-
 @testitem "Testing AR process against theoretical properties" begin
     using DynamicPPL, Turing
     using HypothesisTests: ExactOneSampleKSTest, pvalue
     using Distributions
 
-    ar_model = EpiAware.AR(EpiAware.default_ar_priors()[:damp_prior],
-        EpiAware.default_ar_priors()[:var_prior],
-        EpiAware.default_ar_priors()[:init_prior]
-    )
+    ar_model = EpiAware.AR()
     n = 1000
     damp = [0.1]
-    σ²_AR = 1.0
+    σ_AR = 1.0
     ar_init = [0.0]
 
     model = EpiAware.generate_latent(ar_model, n)
-    fixed_model = fix(model, (σ²_AR = σ²_AR, damp_AR = damp, ar_init = ar_init))
+    fixed_model = fix(model, (σ_AR = σ_AR, damp_AR = damp, ar_init = ar_init))
 
     n_samples = 100
     samples = sample(fixed_model, Prior(), n_samples) |>
@@ -58,7 +51,7 @@ end
     end
 
     theoretical_mean = 0.0
-    theoretical_var = σ²_AR / (1 - damp[1]^2)
+    theoretical_var = σ_AR^2 / (1 - damp[1]^2)
 
     @test isapprox(mean(samples), theoretical_mean, atol = 0.1)
     @test isapprox(var(samples), theoretical_var, atol = 0.2)