Enable ECE with vector of scalars (#109)

bors r+

Co-authored-by: David Widmann <[email protected]>

Author: devmotion

Project.toml

@@ -1,7 +1,7 @@
 name = "CalibrationErrors"
 uuid = "33913031-fe46-5864-950f-100836f47845"
 authors = ["David Widmann <[email protected]>"]
-version = "0.5.20"
+version = "0.5.21"
 
 [deps]
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"

src/CalibrationErrors.jl

@@ -5,8 +5,8 @@ using Reexport
 using DataStructures
 @reexport using Distances
 @reexport using KernelFunctions
-using StatsBase
-using UnPack
+using StatsBase: StatsBase
+using UnPack: @unpack
 
 using LinearAlgebra
 using Statistics

src/binning/generic.jl

@@ -1,16 +1,23 @@
 abstract type AbstractBinningAlgorithm end
 
-mutable struct Bin{T<:Real}
+mutable struct Bin{T}
     """Number of samples."""
     nsamples::Int
     """Mean of predictions."""
-    mean_predictions::Vector{T}
+    mean_predictions::T
     """Proportions of targets."""
-    proportions_targets::Vector{T}
+    proportions_targets::T
 
     function Bin{T}(
-        nsamples::Int, mean_predictions::Vector{T}, proportions_targets::Vector{T}
-    ) where {T}
+        nsamples::Int, mean_predictions::T, proportions_targets::T
+    ) where {T<:Real}
+        nsamples ≥ 0 || throw(ArgumentError("the number of samples must be non-negative"))
+        return new{T}(nsamples, mean_predictions, proportions_targets)
+    end
+
+    function Bin{T}(
+        nsamples::Int, mean_predictions::T, proportions_targets::T
+    ) where {T<:AbstractVector{<:Real}}
         nsamples ≥ 0 || throw(ArgumentError("the number of samples must be non-negative"))
         nclasses = length(mean_predictions)
         nclasses > 1 || throw(ArgumentError("the number of classes must be greater than 1"))
@@ -24,9 +31,7 @@ mutable struct Bin{T<:Real}
     end
 end
 
-function Bin(
-    nsamples::Int, mean_predictions::Vector{T}, proportions_targets::Vector{T}
-) where {T<:Real}
+function Bin(nsamples::Int, mean_predictions::T, proportions_targets::T) where {T}
     return Bin{T}(nsamples, mean_predictions, proportions_targets)
 end
 
@@ -35,6 +40,15 @@ end
 
 Create bin of `predictions` and corresponding `targets`.
 """
+function Bin(predictions::AbstractVector{<:Real}, targets::AbstractVector{Bool})
+    # compute mean of predictions
+    mean_predictions = mean(predictions)
+
+    # compute proportion of targets
+    proportions_targets = mean(targets)
+
+    return Bin(length(predictions), mean_predictions, proportions_targets)
+end
 function Bin(
     predictions::AbstractVector{<:AbstractVector{<:Real}},
     targets::AbstractVector{<:Integer},
@@ -44,16 +58,26 @@ function Bin(
 
     # compute proportion of targets
     nclasses = length(predictions[1])
-    proportions_targets = proportions(targets, nclasses)
+    proportions_targets = StatsBase.proportions(targets, nclasses)
 
     return Bin(length(predictions), mean_predictions, proportions_targets)
 end
 
 """
     Bin(prediction, target)
 
-Create bin of a signle `prediction` and corresponding `target`.
+Create bin of a single `prediction` and corresponding `target`.
 """
+function Bin(prediction::Real, target::Bool)
+    # compute mean of predictions
+    mean_predictions = prediction / 1
+
+    # compute proportion of targets
+    proportions_targets = target / 1
+
+    return Bin(1, mean_predictions, proportions_targets)
+end
+
 function Bin(prediction::AbstractVector{<:Real}, target::Integer)
     # compute mean of predictions
     mean_predictions = prediction ./ 1
@@ -73,6 +97,22 @@ end
 Update running statistics of the `bin` by integrating one additional pair of `prediction`s
 and `target`.
 """
+function adddata!(bin::Bin, prediction::Real, target::Bool)
+    @unpack mean_predictions, proportions_targets = bin
+
+    # update number of samples
+    nsamples = (bin.nsamples += 1)
+
+    # update mean of predictions
+    mean_predictions += (prediction - mean_predictions) / nsamples
+    bin.mean_predictions = mean_predictions
+
+    # update proportions of targets
+    proportions_targets += (target - proportions_targets) / nsamples
+    bin.proportions_targets = proportions_targets
+
+    return nothing
+end
 function adddata!(bin::Bin, prediction::AbstractVector{<:Real}, target::Integer)
     @unpack mean_predictions, proportions_targets = bin
 

src/binning/medianvariance.jl

@@ -25,9 +25,9 @@ MedianVarianceBinning(minsize::Int=10) = MedianVarianceBinning(minsize, typemax(
 
 function perform(
     alg::MedianVarianceBinning,
-    predictions::AbstractVector{<:AbstractVector{T}},
+    predictions::AbstractVector{<:AbstractVector{<:Real}},
     targets::AbstractVector{<:Integer},
-) where {T<:Real}
+)
     @unpack minsize, maxbins = alg
 
     # check if binning is not possible
@@ -50,8 +50,7 @@ function perform(
             (idxs_predictions, argmax_var_predictions) => max_var_predictions,
             Base.Order.Reverse,
         )
-        S = typeof(zero(T) / 1)
-        bins = Vector{Bin{S}}(undef, 0)
+        bins = Vector{typeof(Bin(predictions, targets))}(undef, 0)
 
         nbins = 1
         while nbins < maxbins && !isempty(queue)

src/binning/uniform.jl

@@ -16,9 +16,41 @@ end
 
 function perform(
     binning::UniformBinning,
-    predictions::AbstractVector{<:AbstractVector{T}},
+    predictions::AbstractVector{<:Real},
+    targets::AbstractVector{Bool},
+)
+    @unpack nbins = binning
+
+    # create dictionary of bins
+    T = eltype(float(zero(eltype(predictions))))
+    bins = Dict{Int,Bin{T}}()
+
+    # reserve some memory (very rough guess)
+    nsamples = length(predictions)
+    sizehint!(bins, min(nbins, nsamples))
+
+    # for all other samples
+    @inbounds for (prediction, target) in zip(predictions, targets)
+        # compute index of bin
+        index = binindex(prediction, nbins)
+
+        # create new bin or update existing one
+        bin = get(bins, index, nothing)
+        if bin === nothing
+            bins[index] = Bin(prediction, target)
+        else
+            adddata!(bin, prediction, target)
+        end
+    end
+
+    return values(bins)
+end
+
+function perform(
+    binning::UniformBinning,
+    predictions::AbstractVector{<:AbstractVector{<:Real}},
     targets::AbstractVector{<:Integer},
-) where {T<:Real}
+)
     return _perform(binning, predictions, targets, Val(length(predictions[1])))
 end
 
@@ -36,8 +68,9 @@ function _perform(
 
     # reserve some memory (very rough guess)
     nsamples = length(predictions)
-    sizehint!(bins, nsamples)
-    sizehint!(binindices, nsamples)
+    guess = min(nbins, nsamples)
+    sizehint!(bins, guess)
+    sizehint!(binindices, guess)
 
     # for all other samples
     @inbounds for i in 2:nsamples

src/ece.jl

@@ -31,14 +31,11 @@ In particular, distance measures of the package
 ECE(binning::AbstractBinningAlgorithm) = ECE(binning, TotalVariation())
 
 # estimate ECE
-function (ece::ECE)(
-    predictions::AbstractVector{<:AbstractVector{<:Real}},
-    targets::AbstractVector{<:Integer},
-)
+function (ece::ECE)(predictions::AbstractVector, targets::AbstractVector)
     @unpack binning, distance = ece
 
     # check number of samples
-    nsamples = check_nsamples(predictions, targets)
+    check_nsamples(predictions, targets)
 
     # bin predictions and labels
     bins = perform(binning, predictions, targets)
@@ -50,17 +47,20 @@ function (ece::ECE)(
 
     # evaluate the distance in the first bin
     @inbounds begin
-        bin = bins[1]
+        bin, state = iterate(bins) # there is always at least one bin
         x = distance(bin.mean_predictions, bin.proportions_targets)
 
         # initialize the estimate
         estimate = x / 1
 
         # for all other bins
         n = bin.nsamples
-        for i in 2:nbins
+        while true
+            bin_state = iterate(bins, state)
+            bin_state === nothing && break
+
             # evaluate the distance
-            bin = bins[i]
+            bin, state = bin_state
             x = distance(bin.mean_predictions, bin.proportions_targets)
 
             # update the estimate

test/binning/generic.jl

@@ -1,4 +1,28 @@
 @testset "generic.jl" begin
+    @testset "Simple example" begin
+        # sample predictions and outcomes
+        nsamples = 1_000
+        predictions = rand(nsamples)
+        outcomes = rand(Bool, nsamples)
+
+        # create bin with all predictions and outcomes
+        bin = CalibrationErrors.Bin(predictions, outcomes)
+
+        # check statistics
+        @test bin.nsamples == nsamples
+        @test bin.mean_predictions ≈ mean(predictions)
+        @test bin.proportions_targets == mean(outcomes)
+
+        # compare with adding data
+        bin2 = CalibrationErrors.Bin(predictions[1], outcomes[1])
+        for i in 2:nsamples
+            CalibrationErrors.adddata!(bin2, predictions[i], outcomes[i])
+        end
+        @test bin2.nsamples == bin.nsamples
+        @test bin2.mean_predictions ≈ bin.mean_predictions
+        @test bin2.proportions_targets ≈ bin.proportions_targets
+    end
+
     @testset "Simple example ($nclasses classes)" for nclasses in (2, 10, 100)
         # sample predictions and targets
         nsamples = 1_000

test/binning/uniform.jl

@@ -22,6 +22,36 @@
         @test_throws ArgumentError CalibrationErrors.binindex([1.5, 0.5], 10, Val(2))
     end
 
+    @testset "Basic tests" begin
+        # sample predictions and targets
+        nsamples = 1_000
+        predictions = rand(nsamples)
+        targets = rand(Bool, nsamples)
+
+        for nbins in (1, 10, 100, 500, 1_000)
+            # bin data in bins of uniform width
+            bins = @inferred(
+                CalibrationErrors.perform(UniformBinning(nbins), predictions, targets)
+            )
+
+            # check all bins
+            for bin in bins
+                # compute index of bin from average prediction
+                idx = CalibrationErrors.binindex(bin.mean_predictions, nbins)
+
+                # compute indices of all predictions in the same bin
+                idxs = filter(
+                    i -> idx == CalibrationErrors.binindex(predictions[i], nbins),
+                    1:nsamples,
+                )
+
+                @test bin.nsamples == length(idxs)
+                @test bin.mean_predictions ≈ mean(predictions[idxs])
+                @test bin.proportions_targets ≈ mean(targets[idxs])
+            end
+        end
+    end
+
     @testset "Basic tests ($nclasses classes)" for nclasses in (2, 10, 100)
         # sample predictions and targets
         nsamples = 1_000

test/ece.jl

@@ -1,6 +1,8 @@
 @testset "ece.jl" begin
     @testset "Trivial tests" begin
         ece = ECE(UniformBinning(10))
+
+        # categorical distributions
         for predictions in ([[0, 1], [1, 0]], ColVecs([0 1; 1 0]), RowVecs([0 1; 1 0]))
             @test iszero(@inferred(ece(predictions, [2, 1])))
         end
@@ -11,12 +13,19 @@
         )
             @test iszero(@inferred(ece(predictions, [2, 2, 1, 1])))
         end
+
+        # probabilities
+        for predictions in ([0, 1], [0.0, 1.0])
+            @test iszero(@inferred(ece(predictions, [false, true])))
+        end
+        @test iszero(@inferred(ece([0, 0.5, 0.5, 1], [false, false, true, true])))
     end
 
     @testset "Uniform binning: Basic properties" begin
         ece = ECE(UniformBinning(10))
         estimates = Vector{Float64}(undef, 1_000)
 
+        # categorical distributions
         for nclasses in (2, 10, 100)
             dist = Dirichlet(nclasses, 1.0)
 
@@ -32,6 +41,18 @@
 
             @test all(x -> zero(x) < x < one(x), estimates)
         end
+
+        # probabilities
+        predictions = Vector{Float64}(undef, 20)
+        targets = Vector{Bool}(undef, 20)
+        for i in 1:length(estimates)
+            rand!(predictions)
+            map!(targets, predictions) do p
+                rand() < p
+            end
+            estimates[i] = ece(predictions, targets)
+        end
+        @test all(x -> zero(x) < x < one(x), estimates)
     end
 
     @testset "Median variance binning: Basic properties" begin