CVI Support for the Categorical Distribution

src/distributions/categorical.jl

@@ -1,4 +1,5 @@
 export Categorical
+export CategoricalNaturalParameters
 
 import Distributions: Categorical, probs
 
@@ -20,3 +21,46 @@ probvec(dist::Categorical) = probs(dist)
 function compute_logscale(new_dist::Categorical, left_dist::Categorical, right_dist::Categorical)
     return log(dot(probvec(left_dist), probvec(right_dist)))
 end
+
+struct CategoricalNaturalParameters{T <: Real, M <: AbstractArray{T}} <: NaturalParameters
+    η::M
+end
+
+Base.convert(::Type{CategoricalNaturalParameters}, vec::AbstractVector) = convert(CategoricalNaturalParameters{eltype(vec)}, vec)
+
+Base.convert(::Type{CategoricalNaturalParameters{T}}, vec::AbstractVector) where {T} = CategoricalNaturalParameters(convert(AbstractVector{T}, vec))
+
+as_naturalparams(::Type{T}, args...) where {T <: CategoricalNaturalParameters} = convert(CategoricalNaturalParameters, args...)
+
+function Base.convert(::Type{Distribution}, params::CategoricalNaturalParameters)
+    params = [params.η..., 0]
+    return Categorical(softmax(params))
+end
+
+function naturalparams(dist::Categorical)
+    p = probvec(dist)
+    η = log.(p / p[end])[1:(end - 1)]
+    return CategoricalNaturalParameters(η)
+end
+
+function Base.vec(params::CategoricalNaturalParameters)
+    return params.η
+end
+
+function Base.:(==)(left::CategoricalNaturalParameters, right::CategoricalNaturalParameters)
+    return left.η == right.η
+end
+
+function isproper(::CategoricalNaturalParameters)
+    return true
+end
+
+function lognormalizer(params::CategoricalNaturalParameters)
+    return log(sum(exp.(params.η)) + 1)
+end
+
+logpdf(params::CategoricalNaturalParameters, x) = logpdf(convert(Distribution, params), x)
+
+function Base.:-(left::CategoricalNaturalParameters, right::CategoricalNaturalParameters)
+    return CategoricalNaturalParameters(left.η - right.η)
+end

test/approximations/test_cvi.jl

@@ -8,7 +8,7 @@ using Distributions
 using Zygote
 using Optimisers
 using DiffResults
-import StatsFuns: logistic
+import StatsFuns: logistic, softmax
 
 import ReactiveMP: naturalparams, NaturalParameters, AbstractContinuousGenericLogPdf
 
@@ -143,6 +143,20 @@ end
         end
     end
 
+    @testset "Categorical x Categorical" begin
+        rng = StableRNG(42)
+
+        method = CVI(StableRNG(42), 1, 1000, Optimisers.Descent(0.007), ForwardDiffGrad(), 10, Val(true), true)
+
+        c1 = Categorical(softmax(rand(rng, 3)))
+        c2 = Categorical(softmax(rand(rng, 3)))
+
+        c_analytical = prod(ProdAnalytical(), c1, c2)
+        c_cvi = prod(method, c1, c2)
+
+        @test probvec(c_analytical) ≈ probvec(c_cvi) atol = 1e-1
+    end
+
     @testset "cvi `prod` tests (n_gradpoints = 60)" begin
         rng = StableRNG(42)
 

test/distributions/test_categorical.jl

@@ -4,6 +4,7 @@ using Test
 using ReactiveMP
 using Distributions
 using Random
+using StatsFuns
 
 @testset "Categorical" begin
 
@@ -35,6 +36,42 @@ using Random
         @test probvec(Categorical([1 / 3, 1 / 3, 1 / 3])) == [1 / 3, 1 / 3, 1 / 3]
         @test probvec(Categorical([0.8, 0.1, 0.1])) == [0.8, 0.1, 0.1]
     end
+
+    @testset "CategoricalNaturalParameters" begin
+        @testset "Constructor" begin
+            for i in 1:10
+                @test convert(Distribution, CategoricalNaturalParameters([0 for _ in 1:(i - 1)])) ≈ Categorical([1 / i for _ in 1:i])
+                @test convert(CategoricalNaturalParameters, [0 for _ in 1:i]) == CategoricalNaturalParameters([0 for _ in 1:i])
+                @test convert(CategoricalNaturalParameters{Float64}, [0 for _ in 1:i]) == CategoricalNaturalParameters([0 for _ in 1:i])
+                @test as_naturalparams(CategoricalNaturalParameters, [0 for _ in 1:i]) == CategoricalNaturalParameters([0 for _ in 1:i])
+                @test naturalparams(Categorical([1 / i for _ in 1:i])) == CategoricalNaturalParameters([0 for _ in 1:(i - 1)])
+            end
+        end
+
+        @testset "logpdf" begin
+            for i in 1:10
+                distribution = Categorical(softmax([rand() for _ in 1:i]))
+                cat_np = naturalparams(distribution)
+                for j in 1:i
+                    @test logpdf(distribution, j) ≈ logpdf(cat_np, j)
+                end
+            end
+        end
+
+        @testset "lognormalizer" begin
+            for i in 1:10
+                distribution = Categorical(softmax([rand() for _ in 1:i]))
+                cat_np = naturalparams(distribution)
+                @test lognormalizer(cat_np) ≈ log(sum(exp.(cat_np.η)) + 1)
+            end
+        end
+
+        @testset "isproper" begin
+            for i in 1:10
+                @test isproper(CategoricalNaturalParameters([rand() for _ in 1:i])) === true
+            end
+        end
+    end
 end
 
 end