Update MvNormal constructor with FillArrays.Fill (#272)

* Update `MvNormal` constructor with `FillArrays.Fill` * fix missed MvNormal depreciation
CDCgov · Jun 11, 2024 · d498669 · d498669
1 parent 9a57509
commit d498669
Show file tree

Hide file tree

Showing 6 changed files with 12 additions and 6 deletions.
diff --git a/EpiAware/Project.toml b/EpiAware/Project.toml
@@ -11,6 +11,7 @@ DataFramesMeta = "1313f7d8-7da2-5740-9ea0-a2ca25f37964"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
@@ -31,6 +32,7 @@ DataFramesMeta = "0.15"
 Distributions = "0.25"
 DocStringExtensions = "0.9"
 DynamicPPL = "0.27"
+FillArrays = "1.11"
 LinearAlgebra = ">= 1.9"
 LogExpFunctions = "0.3"
 MCMCChains = "6.0"

diff --git a/EpiAware/src/EpiLatentModels/AR.jl b/EpiAware/src/EpiLatentModels/AR.jl
@@ -75,12 +75,12 @@ Generate a latent AR series.
 "
 @model function EpiAwareBase.generate_latent(latent_model::AR, n)
     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
-    ϵ_t ~ MvNormal(n - p, one(eltype(σ_AR)))
-
-    @assert n>p "n must be longer than order of the autoregressive process"
+    ϵ_t ~ MvNormal(Diagonal(Fill(one(eltype(σ_AR)), n - p)))
 
     # Initialize the AR series with the initial values
     ar = Vector{eltype(ϵ_t)}(undef, n)

diff --git a/EpiAware/src/EpiLatentModels/EpiLatentModels.jl b/EpiAware/src/EpiLatentModels/EpiLatentModels.jl
@@ -9,6 +9,8 @@ using ..EpiAwareUtils: HalfNormal
 
 using LogExpFunctions: softmax
 
+using FillArrays: Fill
+
 using Turing, Distributions, DocStringExtensions, LinearAlgebra
 
 #Export models

diff --git a/EpiAware/src/EpiLatentModels/HierarchicalNormal.jl b/EpiAware/src/EpiLatentModels/HierarchicalNormal.jl
@@ -35,7 +35,8 @@ Generate latent variables from the hierarchical normal distribution.
 "
 @model function EpiAwareBase.generate_latent(obs_model::HierarchicalNormal, n)
     std ~ obs_model.std_prior
-    ϵ_t ~ MvNormal(n, one(eltype(std)))
+    ϵ_t ~ MvNormal(Diagonal(Fill(one(eltype(std)), n)))
+
     η_t = obs_model.mean .+ std .* ϵ_t
     return η_t, (; std = std)
 end
diff --git a/EpiAware/src/EpiLatentModels/RandomWalk.jl b/EpiAware/src/EpiLatentModels/RandomWalk.jl
@@ -87,7 +87,7 @@ Z_t, _ = generated_quantities(rw_model, θ)
 @model function EpiAwareBase.generate_latent(latent_model::RandomWalk, n)
     σ_RW ~ latent_model.std_prior
     rw_init ~ latent_model.init_prior
-    ϵ_t ~ MvNormal(n, one(eltype(σ_RW)))
+    ϵ_t ~ MvNormal(Diagonal(Fill(one(eltype(σ_RW)), n)))
     rw = Vector{eltype(σ_RW)}(undef, n)
 
     rw[1] = rw_init + σ_RW * ϵ_t[1]

diff --git a/EpiAware/test/EpiLatentModels/CombineLatentModels.jl b/EpiAware/test/EpiLatentModels/CombineLatentModels.jl
@@ -34,6 +34,7 @@ end
     using Turing
     using Distributions: Normal
     using HypothesisTests: ExactOneSampleKSTest, pvalue
+    using LinearAlgebra: Diagonal
 
     int = Intercept(Normal(0, 1))
     ar = AR()
@@ -56,7 +57,7 @@ end
     # Fit no-slope linear regression as a model test
     @model function no_slope_linear_regression(y)
         @submodel y_pred, θ = generate_latent(comb, n)
-        y ~ MvNormal(y_pred, ones(n))
+        y ~ MvNormal(y_pred, Diagonal(ones(n)))
     end
 
     ns_regression_mdl = no_slope_linear_regression(y) |