Repository: biaslab/ReactiveMP.jl (julia)

Issue:
Revise `mean(inv, ::PointMass)`
Currently, `mean(inv, ::PointMass)` calls `cholinv` if `PointMass` holds a matrix, but that is wrong in general case. We should neither assume that a matrix is squared not that it is positive-definite.

Pull request description:
This PR fixes #191 and adds `mean(::typeof(cholinv), distribution)` for cases where we know that matrix (or matrix-variate distribution) must be squared positive-definite. `mean(::typeof(inv), ::PointMass)` now uses `inv` by default.

Patch:
diff --git a/src/distributions/pointmass.jl b/src/distributions/pointmass.jl
index 0e7e751c8..86e51c453 100644
--- a/src/distributions/pointmass.jl
+++ b/src/distributions/pointmass.jl
@@ -61,6 +61,7 @@ Distributions.cov(distribution::PointMass{V}) where {T <: Real, V <: AbstractVec
 probvec(distribution::PointMass{V}) where {T <: Real, V <: AbstractVector{T}} = mean(distribution)
 
 mean(::typeof(inv), distribution::PointMass{V}) where {T <: Real, V <: AbstractVector{T}}       = error("mean of inverse of `::PointMass{ <: AbstractVector }` is not defined")
+mean(::typeof(cholinv), distribution::PointMass{V}) where {T <: Real, V <: AbstractVector{T}}   = error("mean of Cholesky inverse of `::PointMass{ <: AbstractVector }` is not defined")
 mean(::typeof(log), distribution::PointMass{V}) where {T <: Real, V <: AbstractVector{T}}       = log.(mean(distribution))
 mean(::typeof(clamplog), distribution::PointMass{V}) where {T <: Real, V <: AbstractVector{T}}  = clamplog.(mean(distribution))
 mean(::typeof(mirrorlog), distribution::PointMass{V}) where {T <: Real, V <: AbstractVector{T}} = error("mean of mirrorlog of `::PointMass{ <: AbstractVector }` is not defined")
@@ -90,7 +91,8 @@ Distributions.cov(distribution::PointMass{M}) where {T <: Real, M <: AbstractMat
 probvec(distribution::PointMass{M}) where {T <: Real, M <: AbstractMatrix{T}} =
     error("probvec(::PointMass{ <: AbstractMatrix }) is not defined")
 
-mean(::typeof(inv), distribution::PointMass{M}) where {T <: Real, M <: AbstractMatrix{T}}       = cholinv(mean(distribution))
+mean(::typeof(inv), distribution::PointMass{M}) where {T <: Real, M <: AbstractMatrix{T}}       = inv(mean(distribution))
+mean(::typeof(cholinv), distribution::PointMass{M}) where {T <: Real, M <: AbstractMatrix{T}}   = cholinv(mean(distribution))
 mean(::typeof(log), distribution::PointMass{M}) where {T <: Real, M <: AbstractMatrix{T}}       = log.(mean(distribution))
 mean(::typeof(clamplog), distribution::PointMass{M}) where {T <: Real, M <: AbstractMatrix{T}}  = clamplog.(mean(distribution))
 mean(::typeof(mirrorlog), distribution::PointMass{M}) where {T <: Real, M <: AbstractMatrix{T}} = error("mean of mirrorlog of `::PointMass{ <: AbstractMatrix }` is not defined")
diff --git a/src/distributions/wishart.jl b/src/distributions/wishart.jl
index d2c2ab655..b85a2c380 100644
--- a/src/distributions/wishart.jl
+++ b/src/distributions/wishart.jl
@@ -50,7 +50,11 @@ function Distributions.mean(::typeof(logdet), distribution::WishartDistributions
     return mapreduce(i -> digamma((ν + 1 - i) / 2), +, 1:d) + d * log(2) + logdet(S)
 end
 
-function Distributions.mean(::typeof(inv), distribution::WishartDistributionsFamily)
+function Distributions.mean(::typeof(inv), dist::WishartDistributionsFamily)
+    return mean(cholinv, dist)
+end
+
+function Distributions.mean(::typeof(cholinv), distribution::WishartDistributionsFamily)
     ν, S = params(distribution)
     return mean(InverseWishart(ν, cholinv(S)))
 end
diff --git a/src/distributions/wishart_inverse.jl b/src/distributions/wishart_inverse.jl
index e39e9ca34..9956a9af7 100644
--- a/src/distributions/wishart_inverse.jl
+++ b/src/distributions/wishart_inverse.jl
@@ -59,6 +59,10 @@ function Distributions.mean(::typeof(logdet), dist::InverseWishartDistributionsF
 end
 
 function Distributions.mean(::typeof(inv), dist::InverseWishartDistributionsFamily)
+    return mean(cholinv, dist)
+end
+
+function Distributions.mean(::typeof(cholinv), dist::InverseWishartDistributionsFamily)
     ν, S = params(dist)
     return mean(Wishart(ν, cholinv(S)))
 end
diff --git a/src/nodes/mv_normal_mean_covariance.jl b/src/nodes/mv_normal_mean_covariance.jl
index 27fe76bdf..475ca2f1c 100644
--- a/src/nodes/mv_normal_mean_covariance.jl
+++ b/src/nodes/mv_normal_mean_covariance.jl
@@ -12,7 +12,7 @@ conjugate_type(::Type{<:MvNormalMeanCovariance}, ::Type{Val{:Σ}})   = InverseWi
 
     m_mean, v_mean = mean_cov(q_μ)
     m_out, v_out   = mean_cov(q_out)
-    inv_m_Σ        = mean(inv, q_Σ)
+    inv_m_Σ        = mean(cholinv, q_Σ)
 
     result = zero(promote_type(eltype(q_out), eltype(q_μ), eltype(q_Σ)))
     result += mean(logdet, q_Σ)
@@ -32,7 +32,7 @@ end
     dim = div(ndims(q_out_μ), 2)
 
     m, V = mean_cov(q_out_μ)
-    inv_m_Σ = mean(inv, q_Σ)
+    inv_m_Σ = mean(cholinv, q_Σ)
 
     result = zero(promote_type(eltype(q_out_μ), eltype(q_Σ)))
     result += mean(logdet, q_Σ)
diff --git a/src/nodes/wishart.jl b/src/nodes/wishart.jl
index fa544dde6..d93730537 100644
--- a/src/nodes/wishart.jl
+++ b/src/nodes/wishart.jl
@@ -13,6 +13,6 @@ import Distributions: Wishart
     return (
         m_q_ν * (mean(logdet, q_S) + d * log(2)) -
         mean(logdet, q_out) * (m_q_ν - d - 1) +
-        tr(mean(inv, q_S) * mean(q_out)) + d * (d - 1) / 2 * logπ
+        tr(mean(cholinv, q_S) * mean(q_out)) + d * (d - 1) / 2 * logπ
     ) / 2 + mapreduce(i -> loggamma((m_q_ν + 1 - i) / 2), +, 1:d)
 end
diff --git a/src/nodes/wishart_inverse.jl b/src/nodes/wishart_inverse.jl
index 4ccaa5d75..ea888da27 100644
--- a/src/nodes/wishart_inverse.jl
+++ b/src/nodes/wishart_inverse.jl
@@ -14,6 +14,6 @@ import Distributions: InverseWishart
     return 0.5 * (
         m_q_ν * (-mean(logdet, q_S) + d * log(2)) +
         mean(logdet, q_out) * (m_q_ν + d + 1) +
-        tr(mean(q_S) * mean(inv, q_out)) + d * (d - 1) / 2 * logπ
+        tr(mean(q_S) * mean(cholinv, q_out)) + d * (d - 1) / 2 * logπ
     ) + mapreduce(i -> loggamma((m_q_ν + 1 - i) / 2), +, 1:d)
 end


Test patch:
diff --git a/test/distributions/test_pointmass.jl b/test/distributions/test_pointmass.jl
index e7cb128f9..661b01d88 100644
--- a/test/distributions/test_pointmass.jl
+++ b/test/distributions/test_pointmass.jl
@@ -104,6 +104,7 @@ import ReactiveMP: xtlog, mirrorlog
             @test @test_inferred(AbstractVector{T}, probvec(dist)) == vector
             @test @test_inferred(AbstractVector{T}, mean(log, dist)) == log.(vector)
             @test_throws ErrorException mean(inv, dist)
+            @test_throws ErrorException mean(cholinv, dist)
             @test_throws ErrorException mean(mirrorlog, dist)
             @test @test_inferred(AbstractVector{T}, mean(loggamma, dist)) == loggamma.(vector)
         end
@@ -154,7 +155,8 @@ import ReactiveMP: xtlog, mirrorlog
 
             @test_throws ErrorException probvec(dist)
             @test @test_inferred(AbstractMatrix{T}, mean(log, dist)) == log.(matrix)
-            @test @test_inferred(AbstractMatrix{T}, mean(inv, dist)) ≈ cholinv(matrix)
+            @test @test_inferred(AbstractMatrix{T}, mean(inv, dist)) ≈ inv(matrix)
+            @test @test_inferred(AbstractMatrix{T}, mean(cholinv, dist)) ≈ cholinv(matrix)
             @test_throws ErrorException mean(mirrorlog, dist)
             @test @test_inferred(AbstractMatrix{T}, mean(loggamma, dist)) == loggamma.(matrix)
         end
diff --git a/test/distributions/test_wishart.jl b/test/distributions/test_wishart.jl
index e3c7cf9b1..4fb74bd23 100644
--- a/test/distributions/test_wishart.jl
+++ b/test/distributions/test_wishart.jl
@@ -13,7 +13,7 @@ import ReactiveMP: WishartMessage
     # Wishart comes from Distributions.jl and most of the things should be covered there
     # Here we test some extra ReactiveMP.jl specific functionality
 
-    @testset "mean(:logdet)" begin
+    @testset "mean(::logdet)" begin
         @test mean(logdet, Wishart(3, [1.0 0.0; 0.0 1.0])) ≈ 0.845568670196936
         @test mean(
             logdet,
@@ -28,6 +28,14 @@ import ReactiveMP: WishartMessage
         ) ≈ -3.4633310802040693
     end
 
+    @testset "mean(::cholinv)" begin
+        L    = rand(2, 2)
+        S    = L * L' + diageye(2)
+        invS = cholinv(S)
+        @test mean(inv, Wishart(5, S)) ≈ mean(InverseWishart(5, invS))
+        @test mean(cholinv, Wishart(5, S)) ≈ mean(InverseWishart(5, invS))
+    end
+
     @testset "vague" begin
         @test_throws MethodError vague(Wishart)
 
diff --git a/test/distributions/test_wishart_inverse.jl b/test/distributions/test_wishart_inverse.jl
index bb9695a40..368d26864 100644
--- a/test/distributions/test_wishart_inverse.jl
+++ b/test/distributions/test_wishart_inverse.jl
@@ -102,10 +102,12 @@ import ReactiveMP: InverseWishartMessage
         ν, S = 2.0, [2.2658069783329573 -0.47934965873423374; -0.47934965873423374 1.4313564100863712]
         samples = rand(rng, InverseWishart(ν, S), Int(1e6))
         @test isapprox(mean(inv, InverseWishartMessage(ν, S)), mean(inv.(samples)), atol = 1e-2)
+        @test isapprox(mean(cholinv, InverseWishartMessage(ν, S)), mean(cholinv.(samples)), atol = 1e-2)
 
         ν, S = 4.0, diageye(3)
         samples = rand(rng, InverseWishart(ν, S), Int(1e6))
         @test isapprox(mean(inv, InverseWishartMessage(ν, S)), mean(inv.(samples)), atol = 1e-2)
+        @test isapprox(mean(cholinv, InverseWishartMessage(ν, S)), mean(cholinv.(samples)), atol = 1e-2)
     end
 
     @testset "prod" begin