Merge pull request #527 from ReactiveBayes/approx-one-hot

use isapprox in the isonehot

Author: bvdmitri

src/helpers/algebra/common.jl

@@ -134,19 +134,20 @@ end
 """
     isonehot(vec::AbstractVector)
 
-Checks if the given vector `vec` is a one-hot vector, i.e., a vector with exactly one entry equal to `one(eltype(vec))` and all other entries equal to `zero(eltype(vec))`.
+Checks if the given vector `vec` is a one-hot vector, i.e., a vector with exactly one entry approximately equal to `one(eltype(vec))` and all other entries approximately equal to `zero(eltype(vec))`.
 
 Returns `true` if `vec` is one-hot, otherwise returns `false`.
 """
-function isonehot(vec::AbstractVector)
+function isonehot(vec::AbstractVector{T}) where {T}
     number_of_ones::Int = 0
+    atol = sqrt(eps(T))
     for e in vec
-        if isequal(e, one(e))
+        if isapprox(e, one(e); atol = atol)
             if number_of_ones > 1
                 return false
             end
             number_of_ones += 1
-        elseif !isequal(e, zero(e))
+        elseif !isapprox(e, zero(e); atol = atol)
             return false
         end
     end

src/rules/categorical/p.jl

@@ -4,21 +4,20 @@ using SpecialFunctions: logfactorial
 # https://github.com/JuliaStats/Distributions.jl/blob/master/src/univariate/discrete/categorical.jl#L27
 @rule Categorical(:p, Marginalisation) (q_out::Categorical,) = begin
     probs = probvec(q_out)
-    if !isonehot(probs)
-        throw(ArgumentError("q_out must be one-hot encoded. Got: $probs"))
-    end
     @logscale -logfactorial(length(probs))
     return Dirichlet(probs .+ one(eltype(probs)))
 end
 
 # https://github.com/ReactiveBayes/BayesBase.jl/blob/main/src/densities/pointmass.jl
 @rule Categorical(:p, Marginalisation) (q_out::PointMass{V},) where {T <: Real, V <: AbstractVector{T}} = begin
-    probs = probvec(q_out)
+    probs = mean(q_out)
     if !isonehot(probs)
         throw(ArgumentError("q_out must be one-hot encoded. Got: $probs"))
     end
     @logscale -logfactorial(length(probs))
     return Dirichlet(probs .+ one(eltype(probs)))
 end
 
-@rule Categorical(:p, Marginalisation) (q_out::Any,) = throw(ArgumentError("q_out is only defined for PointMass or Categorical over a one-hot vector. Got: $(typeof(q_out))"))
+@rule Categorical(:p, Marginalisation) (q_out::Any,) = throw(
+    ArgumentError("This rule is only defined for PointMass over a one-hot vector or a Categorical distribution. Got: $(typeof(q_out))")
+)

test/helpers/algebra/common_tests.jl

@@ -111,3 +111,19 @@ end
         end
     end
 end
+
+@testitem "isonehot approx" begin
+    import ReactiveMP: isonehot
+
+    @test isonehot([2.9999999999849994e-12, 2.9999999999849994e-12, 0.999999999994]) == true
+    @test isonehot([2.9999999999849994e-12, 0.999999999994, 2.9999999999849994e-12]) == true
+    @test isonehot([0.999999999994, 2.9999999999849994e-12, 2.9999999999849994e-12]) == true
+
+    @test isonehot([0.03, 0.03, 0.94]) == false
+    @test isonehot([0.03, 0.94, 0.03]) == false
+    @test isonehot([0.94, 0.03, 0.03]) == false
+
+    @test isonehot([0.0003, 0.0003, 0.9994]) == false
+    @test isonehot([0.0003, 0.9994, 0.0003]) == false
+    @test isonehot([0.9994, 0.0003, 0.0003]) == false
+end

test/rules/categorical/p_tests.jl

@@ -10,6 +10,7 @@
 
     @testset "Variational Message Passing: (q_out::Categorical)" begin
         @test_rules [check_type_promotion = false] Categorical(:p, Marginalisation) [(input = (q_out = Categorical([0.0, 1.0]),), output = Dirichlet([1.0, 2.0]))]
+        @test_rules [check_type_promotion = false] Categorical(:p, Marginalisation) [(input = (q_out = Categorical([0.5, 0.5]),), output = Dirichlet([1.5, 1.5]))]
     end
 
     # Test if the input q_out for the outgoing message towards p edge is not a scalar
@@ -20,13 +21,11 @@
     end
 
     # Test if the input q_out for the outgoing message towards p edge is a one-hot encoded vector
-    @testset "q_out =?= one-hot encoded vector" begin
+    @testset "q_out = PointMass one-hot encoded vector" begin
         # PointMass over a non-one-hot vector
-        @test_throws ArgumentError @call_rule Categorical(:p, Marginalisation) (q_out = PointMass([0.5, 0.5]),)
-        @test_throws ArgumentError @call_rule Categorical(:p, Marginalisation) (q_out = PointMass([1.0, 1.0]),)
-        # Categorical with non-one-hot probability vector
-        @test_throws ArgumentError @call_rule Categorical(:p, Marginalisation) (q_out = Categorical([0.3, 0.7]),)
+        @test_throws "q_out must be one-hot encoded" @call_rule Categorical(:p, Marginalisation) (q_out = PointMass([0.5, 0.5]),)
+        @test_throws "q_out must be one-hot encoded" @call_rule Categorical(:p, Marginalisation) (q_out = PointMass([1.0, 1.0]),)
         # Arbitrary non-distribution
-        @test_throws ArgumentError @call_rule Categorical(:p, Marginalisation) (q_out = Dirichlet([1.0, 2.0]),)
+        @test_throws "This rule is only defined for PointMass over a one-hot vector" @call_rule Categorical(:p, Marginalisation) (q_out = "Hello?",)
     end
 end