Merge pull request #226 from ReactiveBayes/logmean

Implement log mean for tensordirichlet

Author: bvdmitri

src/distributions/tensor_dirichlet.jl

@@ -75,6 +75,8 @@ getlogbasemeasure(::Type{TensorDirichlet}, conditioner) = (x) -> zero(Float64)
 getsufficientstatistics(::Type{TensorDirichlet}, conditioner) = (x -> vmap(log, x),)
 
 BayesBase.mean(dist::TensorDirichlet) = dist.a ./ dist.α0
+BayesBase.mean(::BroadcastFunction{typeof(log)}, dist::TensorDirichlet) = digamma.(dist.a) .- digamma.(dist.α0)
+
 function BayesBase.cov(dist::TensorDirichlet{T}) where {T}
     s = size(dist.a)
     news = (first(s), first(s), Base.tail(s)...)

test/distributions/tensor_dirichlet_test.jl

@@ -73,6 +73,28 @@ end
     end
 end
 
+@testitem "TensorDirichlet: logmean" begin
+    include("distributions_setuptests.jl")
+
+    for rank in (3, 5)
+        for d in (2, 5, 10)
+            for _ in 1:10
+                alpha = rand([d for _ in 1:rank]...)
+
+                distribution = TensorDirichlet(alpha)
+                mat_of_dir = Dirichlet.(eachslice(alpha, dims = Tuple(2:rank)))
+
+                temp = mean.(Base.Broadcast.BroadcastFunction(log), mat_of_dir)
+                mat_mean = similar(alpha)
+                for i in CartesianIndices(Base.tail(size(alpha)))
+                    mat_mean[:, i] = temp[i]
+                end
+                @test mean(Base.Broadcast.BroadcastFunction(log), distribution) ≈ mat_mean
+            end
+        end
+    end
+end
+
 @testitem "TensorDirichlet: std" begin
     include("distributions_setuptests.jl")