Merge pull request #208 from ReactiveBayes/fix-207

Add support for `Diagonal` covariance matrices
ReactiveBayes · Sep 16, 2024 · c0fa8ad · c0fa8ad
2 parents da9bb3a + 2934df2
commit c0fa8ad
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 4 deletions.
diff --git a/src/distributions/normal_family/mv_normal_mean_covariance.jl b/src/distributions/normal_family/mv_normal_mean_covariance.jl
@@ -103,8 +103,8 @@ end
 
 function BayesBase.prod(
     ::PreserveTypeProd{Distribution},
-    left::MvNormalMeanCovariance{T1},
-    right::MvNormalMeanCovariance{T2}
+    left::MvNormalMeanCovariance{T1, <:AbstractVector, <:Matrix},
+    right::MvNormalMeanCovariance{T2, <:AbstractVector, <:Matrix}
 ) where {T1 <: LinearAlgebra.BlasFloat, T2 <: LinearAlgebra.BlasFloat}
     xi, W = weightedmean_precision(left)
 

diff --git a/src/distributions/normal_family/mv_normal_mean_precision.jl b/src/distributions/normal_family/mv_normal_mean_precision.jl
@@ -93,8 +93,8 @@ end
 
 function BayesBase.prod(
     ::PreserveTypeProd{Distribution},
-    left::MvNormalMeanPrecision{T1},
-    right::MvNormalMeanPrecision{T2}
+    left::MvNormalMeanPrecision{T1, <:AbstractVector, <:Matrix},
+    right::MvNormalMeanPrecision{T2, <:AbstractVector, <:Matrix}
 ) where {T1 <: LinearAlgebra.BlasFloat, T2 <: LinearAlgebra.BlasFloat}
     W = precision(left) + precision(right)
 

diff --git a/test/distributions/normal_family/mv_normal_mean_covariance_tests.jl b/test/distributions/normal_family/mv_normal_mean_covariance_tests.jl
@@ -108,6 +108,13 @@ end
         Σ    = diagm([1.0, 2.0, 3.0])
         dist = MvNormalMeanCovariance(μ, Σ)
 
+        @test prod(strategy, dist, dist) ≈
+              MvNormalWeightedMeanPrecision([2.0, 2.0, 2.0], diagm([2.0, 1.0, 2 / 3]))
+
+        μ    = [1.0, 2.0, 3.0]
+        Σ    = Diagonal([1.0, 2.0, 3.0])
+        dist = MvNormalMeanCovariance(μ, Σ)
+
         @test prod(strategy, dist, dist) ≈
               MvNormalWeightedMeanPrecision([2.0, 2.0, 2.0], diagm([2.0, 1.0, 2 / 3]))
     end

diff --git a/test/distributions/normal_family/mv_normal_mean_precision_tests.jl b/test/distributions/normal_family/mv_normal_mean_precision_tests.jl
@@ -108,6 +108,13 @@ end
         Λ    = diagm(1 ./ [1.0, 2.0, 3.0])
         dist = MvNormalMeanPrecision(μ, Λ)
 
+        @test prod(strategy, dist, dist) ≈
+              MvNormalWeightedMeanPrecision([2.0, 2.0, 2.0], diagm([2.0, 1.0, 2 / 3]))
+
+        μ    = [1.0, 2.0, 3.0]
+        Λ    = Diagonal(1 ./ [1.0, 2.0, 3.0])
+        dist = MvNormalMeanPrecision(μ, Λ)
+
         @test prod(strategy, dist, dist) ≈
               MvNormalWeightedMeanPrecision([2.0, 2.0, 2.0], diagm([2.0, 1.0, 2 / 3]))
     end