Support UCME with probabilities and boolean targets

Author: devmotion

src/ucme.jl

@@ -88,10 +88,13 @@ function unsafe_ucme_eval(
     testp::AbstractVector{<:Real},
     testy::Integer,
 )
-    a = kernel((p, y), (testp, testy))
-    b = sum(p[z] * kernel((p, z), (testp, testy)) for z in 1:length(p))
-
-    return a - b
+    res = sum(((z == y) - pz) * kernel((p, z), (testp, testy)) for (z, pz) in enumerate(p))
+    return res
+end
+function unsafe_ucme_eval(kernel::Kernel, p::Real, y::Bool, testp::Real, testy::Bool)
+    noty = !y
+    return (y - p) * kernel((p, y), (testp, testy)) +
+           (noty - p) * kernel((p, noty), (testp, testy))
 end
 
 function unsafe_ucme_eval(kernel::KernelTensorProduct, p, y, testp, testy)
@@ -109,6 +112,30 @@ function unsafe_ucme_eval(
     κpredictions, κtargets = kernel.kernels
     return unsafe_ucme_eval_targets(κtargets, p, y, testp, testy) * κpredictions(p, testp)
 end
+function unsafe_ucme_eval(
+    kernel::KernelTensorProduct, p::Real, y::Bool, testp::Real, testy::Bool
+)
+    κpredictions, κtargets = kernel.kernels
+    return unsafe_ucme_eval_targets(κtargets, p, y, testp, testy) * κpredictions(p, testp)
+end
+
+function unsafe_ucme_eval_targets(
+    kernel::Kernel,
+    p::AbstractVector{<:Real},
+    y::Integer,
+    testp::AbstractVector{<:Real},
+    testy::Integer,
+)
+    res = sum(((z == y) - pz) * kernel(z, testy) for (z, pz) in enumerate(p))
+    return res
+end
+function unsafe_ucme_eval_targets(
+    kernel::Kernel, p::Real, y::Bool, testp::Real, testy::Bool
+)
+    noty = !y
+    res = (y - p) * kernel(y, testy) + (noty - p) * kernel(noty, testy)
+    return res
+end
 
 function unsafe_ucme_eval_targets(
     κtargets::WhiteKernel,
@@ -120,3 +147,8 @@ function unsafe_ucme_eval_targets(
     @inbounds res = (y == testy) - p[testy]
     return res
 end
+function unsafe_ucme_eval_targets(
+    kernel::WhiteKernel, p::Real, y::Bool, testp::Real, testy::Bool
+)
+    return (2 * testy - 1) * (y - p)
+end

test/runtests.jl

@@ -8,7 +8,7 @@ using Random
 using Statistics
 using Test
 
-using CalibrationErrors: unsafe_skce_eval
+using CalibrationErrors: unsafe_skce_eval, unsafe_ucme_eval
 
 Random.seed!(1234)
 

test/skce/generic.jl

@@ -25,9 +25,9 @@
 
         # corresponding values and kernel for full categorical distribution
         pfull = [p, 1 - p]
-        yint = y ? 1 : 2
+        yint = 2 - y
         p̃full = [p̃, 1 - p̃]
-        ỹint = ỹ ? 1 : 2
+        ỹint = 2 - ỹ
         kernelfull = SqExponentialKernel() ∘ ScaleTransform(scale / sqrt(2))
 
         @test unsafe_skce_eval(kernelfull ⊗ WhiteKernel(), pfull, yint, p̃full, ỹint) ≈ val

test/ucme.jl

@@ -1,44 +1,145 @@
 @testset "ucme.jl" begin
-    @testset "UCME: Two-dimensional example" begin
-        # three test locations
+    @testset "UCME: Binary examples" begin
+        # categorical distributions
         ucme = UCME(
             SqExponentialKernel() ⊗ WhiteKernel(),
             [[1.0, 0], [0.5, 0.5], [0.0, 1]],
             [1, 1, 2],
         )
-
-        # two predictions
         @test iszero(@inferred(ucme([[1, 0], [0, 1]], [1, 2])))
         @test @inferred(ucme([[1, 0], [0, 1]], [1, 1])) ≈ (exp(-2) + exp(-0.5) + 1) / 12
         @test @inferred(ucme([[1, 0], [0, 1]], [2, 1])) ≈ (1 - exp(-1))^2 / 6
         @test @inferred(ucme([[1, 0], [0, 1]], [2, 2])) ≈ (exp(-2) + exp(-0.5) + 1) / 12
+
+        # probabilities
+        ucme = UCME(
+            (SqExponentialKernel() ∘ ScaleTransform(sqrt(2))) ⊗ WhiteKernel(),
+            [1.0, 0.5, 0.0],
+            [true, true, false],
+        )
+        @test iszero(@inferred(ucme([1, 0], [true, false])))
+        @test @inferred(ucme([1, 0], [true, true])) ≈ (exp(-2) + exp(-0.5) + 1) / 12
+        @test @inferred(ucme([1, 0], [false, true])) ≈ (1 - exp(-1))^2 / 6
+        @test @inferred(ucme([1, 0], [false, false])) ≈ (exp(-2) + exp(-0.5) + 1) / 12
     end
 
     @testset "UCME: Basic properties" begin
         estimates = Vector{Float64}(undef, 1_000)
 
-        for ntest in (1, 5, 10), nclasses in (2, 10, 100)
-            dist = Dirichlet(nclasses, 1.0)
+        for ntest in (1, 5, 10)
+            # categorical distributions
+            for nclasses in (2, 10, 100)
+                dist = Dirichlet(nclasses, 1.0)
+
+                testpredictions = [rand(dist) for _ in 1:ntest]
+                testtargets = rand(1:nclasses, ntest)
+                ucme = UCME(
+                    (ExponentialKernel() ∘ ScaleTransform(0.1)) ⊗ WhiteKernel(),
+                    testpredictions,
+                    testtargets,
+                )
+
+                predictions = [Vector{Float64}(undef, nclasses) for _ in 1:20]
+                targets = Vector{Int}(undef, 20)
 
-            testpredictions = [rand(dist) for _ in 1:ntest]
-            testtargets = rand(1:nclasses, ntest)
+                for i in 1:length(estimates)
+                    rand!.(Ref(dist), predictions)
+                    targets .= rand.(Categorical.(predictions))
+
+                    estimates[i] = ucme(predictions, targets)
+                end
+
+                @test all(x > zero(x) for x in estimates)
+            end
+
+            # probabilities
+            testpredictions = rand(ntest)
+            testtargets = rand(Bool, ntest)
             ucme = UCME(
                 (ExponentialKernel() ∘ ScaleTransform(0.1)) ⊗ WhiteKernel(),
                 testpredictions,
                 testtargets,
             )
 
-            predictions = [Vector{Float64}(undef, nclasses) for _ in 1:20]
-            targets = Vector{Int}(undef, 20)
+            predictions = Vector{Float64}(undef, 20)
+            targets = Vector{Bool}(undef, 20)
 
             for i in 1:length(estimates)
-                rand!.(Ref(dist), predictions)
-                targets .= rand.(Categorical.(predictions))
+                rand!(predictions)
+                map!(targets, predictions) do p
+                    return rand() < p
+                end
 
                 estimates[i] = ucme(predictions, targets)
             end
 
             @test all(x > zero(x) for x in estimates)
         end
     end
+
+    # alternative implementation of white kernel
+    struct WhiteKernel2 <: Kernel end
+    (::WhiteKernel2)(x, y) = x == y
+
+    # alternative implementation TensorProductKernel
+    struct TensorProduct2{K1<:Kernel,K2<:Kernel} <: Kernel
+        kernel1::K1
+        kernel2::K2
+    end
+    function (kernel::TensorProduct2)((x1, x2), (y1, y2))
+        return kernel.kernel1(x1, y1) * kernel.kernel2(x2, y2)
+    end
+
+    @testset "binary classification" begin
+        # probabilities and boolean targets
+        p, testp = rand(2)
+        y, testy = rand(Bool, 2)
+        scale = rand()
+        kernel = SqExponentialKernel() ∘ ScaleTransform(scale)
+        val = unsafe_ucme_eval(kernel ⊗ WhiteKernel(), p, y, testp, testy)
+        @test unsafe_ucme_eval(kernel ⊗ WhiteKernel2(), p, y, testp, testy) ≈ val
+        @test unsafe_ucme_eval(TensorProduct2(kernel, WhiteKernel()), p, y, testp, testy) ≈
+              val
+        @test unsafe_ucme_eval(TensorProduct2(kernel, WhiteKernel2()), p, y, testp, testy) ≈
+              val
+
+        # corresponding values and kernel for full categorical distribution
+        pfull = [p, 1 - p]
+        yint = 2 - y
+        testpfull = [testp, 1 - testp]
+        testyint = 2 - testy
+        kernelfull = SqExponentialKernel() ∘ ScaleTransform(scale / sqrt(2))
+
+        @test unsafe_ucme_eval(
+            kernelfull ⊗ WhiteKernel(), pfull, yint, testpfull, testyint
+        ) ≈ val
+        @test unsafe_ucme_eval(
+            kernelfull ⊗ WhiteKernel2(), pfull, yint, testpfull, testyint
+        ) ≈ val
+        @test unsafe_ucme_eval(
+            TensorProduct2(kernelfull, WhiteKernel()), pfull, yint, testpfull, testyint
+        ) ≈ val
+        @test unsafe_ucme_eval(
+            TensorProduct2(kernelfull, WhiteKernel2()), pfull, yint, testpfull, testyint
+        ) ≈ val
+    end
+
+    @testset "multi-class classification" begin
+        n = 10
+        p = rand(n)
+        p ./= sum(p)
+        y = rand(1:n)
+        testp = rand(n)
+        testp ./= sum(testp)
+        testy = rand(1:n)
+
+        kernel = SqExponentialKernel() ∘ ScaleTransform(rand())
+        val = unsafe_ucme_eval(kernel ⊗ WhiteKernel(), p, y, testp, testy)
+
+        @test unsafe_ucme_eval(kernel ⊗ WhiteKernel2(), p, y, testp, testy) ≈ val
+        @test unsafe_ucme_eval(TensorProduct2(kernel, WhiteKernel()), p, y, testp, testy) ≈
+              val
+        @test unsafe_ucme_eval(TensorProduct2(kernel, WhiteKernel2()), p, y, testp, testy) ≈
+              val
+    end
 end