fix further tests

CDCgov · Jul 10, 2024 · d405992 · d405992
1 parent 4e02a56
commit d405992
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 6 deletions.
diff --git a/EpiAware/test/EpiLatentModels/models/RandomWalk.jl b/EpiAware/test/EpiLatentModels/models/RandomWalk.jl
@@ -9,7 +9,7 @@
     n_samples = 1000
     samples_day_5 = sample(fixed_model, Prior(), n_samples; progress = false) |>
                     chn -> mapreduce(vcat, generated_quantities(fixed_model, chn)) do gen
-        gen[1][5] #Extracting day 5 samples
+        gen[5] #Extracting day 5 samples
     end
     #Check that the samples are drawn from the correct distribution which is Normal(mean = 0, var = 5)
     ks_test_pval = ExactOneSampleKSTest(samples_day_5, Normal(0.0, sqrt(5))) |> pvalue

diff --git a/EpiAware/test/EpiLatentModels/modifiers/DiffLatentModel.jl b/EpiAware/test/EpiLatentModels/modifiers/DiffLatentModel.jl
@@ -42,7 +42,7 @@ end
     n_samples = 2000
     samples = sample(fixed_model, Prior(), n_samples; progress = false) |>
               chn -> mapreduce(hcat, generated_quantities(fixed_model, chn)) do gen
-        gen[1]
+        gen
     end
 
     #Because of the recursive d-times cumsum to undifference the process,
@@ -86,7 +86,7 @@ end
     n_samples = 100
     samples = sample(fixed_model, Prior(), n_samples; progress = false) |>
               chn -> mapreduce(hcat, generated_quantities(fixed_model, chn)) do gen
-        gen[1]
+        gen
     end
 
     @test size(samples) == (n, n_samples)

diff --git a/EpiAware/test/EpiLatentModels/modifiers/TransformLatentModel.jl b/EpiAware/test/EpiLatentModels/modifiers/TransformLatentModel.jl
@@ -8,10 +8,11 @@
 end
 
 @testitem "TransformLatentModel generate_latent method" begin
-    using Turing, Distributions
+    using Turing, Distributions, DynamicPPL
 
     trans = TransformLatentModel(Intercept(Normal(2, 0.2)), x -> x .|> exp)
     trans_model = generate_latent(trans, 5)
-    returns = trans_model()
-    @test returns[1][1] == exp(returns[2].latent_aux.intercept)
+    fix_model = fix(trans_model, (intercept = 2.0))
+    returns = fix_model()
+    @test returns[1] == exp(2)
 end