diff --git a/docs/src/counts.md b/docs/src/counts.md index e4498a1d7..2765c8248 100644 --- a/docs/src/counts.md +++ b/docs/src/counts.md @@ -7,7 +7,7 @@ The package provides functions to count the occurrences of distinct values. ```@docs counts proportions -addcounts!(r::AbstractArray, x::AbstractArray{<:Integer}, levels::UnitRange{<:Integer}) +addcounts!(r::AbstractArray, x::AbstractArray{<:Integer}, levels::AbstractUnitRange{<:Integer}) ``` ## Counting over arbitrary distinct values diff --git a/src/counts.jl b/src/counts.jl index 43fa1d495..6cf39dc72 100644 --- a/src/counts.jl +++ b/src/counts.jl @@ -13,33 +13,47 @@ else end #### functions for counting a single list of integers (1D) + +_count_modulo(x::Union{Base.BitInteger,BigInt}) = x % UInt +_count_modulo(::Integer) = nothing + +@inline _count_index(x::Union{Base.BitInteger,BigInt}, m0::Integer, u0::UInt) = + ((x % UInt) - u0) % Int + 1 +@inline _count_index(x::Integer, m0::Integer, u0) = + Int(x - (m0 + zero(Int))) + 1 + """ - addcounts!(r, x, levels::UnitRange{<:Integer}, [wv::AbstractWeights]) + addcounts!(r, x, levels::AbstractUnitRange{<:Integer}, [wv::AbstractWeights]) Add the number of occurrences in `x` of each value in `levels` to an existing -array `r`. For each `xi ∈ x`, if `xi == levels[j]`, then we increment `r[j]`. +array `r`. `r[j]` counts the `j`th value in `levels`, independently of the +indices of `levels`. If a weighting vector `wv` is specified, the sum of weights is used rather than the raw counts. """ -function addcounts!(r::AbstractArray, x::AbstractArray{<:Integer}, levels::UnitRange{<:Integer}) +function addcounts!(r::AbstractArray, x::AbstractArray{<:Integer}, levels::AbstractUnitRange{<:Integer}) # add counts of integers from x that fall within levels to r - checkbounds(r, axes(levels)...) + checkbounds(r, 1:length(levels)) m0 = first(levels) m1 = last(levels) - b = m0 - firstindex(levels) # firstindex(levels) == 1 because levels::UnitRange{<:Integer} + # `m0 <= xi <= m1` implies `0 <= xi - m0 < length(levels) <= length(r) <= typemax(Int)` + # (by the `checkbounds` call above), so the offset is computed exactly in modular + # `UInt` arithmetic for built-in integers, without per-element `BigInt` allocations. + # Other Integer types retain ordinary arithmetic rather than requiring modular conversion. + u0 = _count_modulo(m0) for xi in x if m0 <= xi <= m1 - r[xi - b] += 1 + r[_count_index(xi, m0, u0)] += 1 end end return r end -function addcounts!(r::AbstractArray, x::AbstractArray{<:Integer}, levels::UnitRange{<:Integer}, wv::AbstractWeights) +function addcounts!(r::AbstractArray, x::AbstractArray{<:Integer}, levels::AbstractUnitRange{<:Integer}, wv::AbstractWeights) # add wv weighted counts of integers from x that fall within levels to r length(x) == length(wv) || @@ -47,16 +61,17 @@ function addcounts!(r::AbstractArray, x::AbstractArray{<:Integer}, levels::UnitR xv = vec(x) # discard shape because weights() discards shape - checkbounds(r, axes(levels)...) + checkbounds(r, 1:length(levels)) m0 = first(levels) m1 = last(levels) - b = m0 - 1 + # See the unweighted method: the offset `xi - m0` is exact in modular `UInt` arithmetic. + u0 = _count_modulo(m0) for i in eachindex(xv, wv) xi = xv[i] if m0 <= xi <= m1 - r[xi - b] += wv[i] + r[_count_index(xi, m0, u0)] += wv[i] end end return r @@ -65,7 +80,7 @@ end """ counts(x, [wv::AbstractWeights]) - counts(x, levels::UnitRange{<:Integer}, [wv::AbstractWeights]) + counts(x, levels::AbstractUnitRange{<:Integer}, [wv::AbstractWeights]) counts(x, k::Integer, [wv::AbstractWeights]) Count the number of times each value in `x` occurs. If `levels` is provided, only values @@ -80,9 +95,9 @@ The output is a vector of length `length(levels)`. """ function counts end -counts(x::AbstractArray{<:Integer}, levels::UnitRange{<:Integer}) = +counts(x::AbstractArray{<:Integer}, levels::AbstractUnitRange{<:Integer}) = addcounts!(zeros(Int, length(levels)), x, levels) -counts(x::AbstractArray{<:Integer}, levels::UnitRange{<:Integer}, wv::AbstractWeights) = +counts(x::AbstractArray{<:Integer}, levels::AbstractUnitRange{<:Integer}, wv::AbstractWeights) = addcounts!(zeros(eltype(wv), length(levels)), x, levels, wv) counts(x::AbstractArray{<:Integer}, k::Integer) = counts(x, 1:k) counts(x::AbstractArray{<:Integer}, k::Integer, wv::AbstractWeights) = counts(x, 1:k, wv) @@ -99,8 +114,8 @@ Equivalent to `counts(x, levels) / length(x)`. If a vector of weights `wv` is provided, the proportion of weights is computed rather than the proportion of raw counts. """ -proportions(x::AbstractArray{<:Integer}, levels::UnitRange{<:Integer}) = counts(x, levels) / length(x) -proportions(x::AbstractArray{<:Integer}, levels::UnitRange{<:Integer}, wv::AbstractWeights) = +proportions(x::AbstractArray{<:Integer}, levels::AbstractUnitRange{<:Integer}) = counts(x, levels) / length(x) +proportions(x::AbstractArray{<:Integer}, levels::AbstractUnitRange{<:Integer}, wv::AbstractWeights) = counts(x, levels, wv) / sum(wv) """ diff --git a/test/counts.jl b/test/counts.jl index cb16abe3f..3e1683093 100644 --- a/test/counts.jl +++ b/test/counts.jl @@ -4,6 +4,80 @@ using OffsetArrays n = 5000 +@testset "abstract unit ranges as count levels" begin + x = [1, 1, 2, 3] + w = weights([1.0, 2.0, 3.0, 4.0]) + @test counts(x, Base.OneTo(4)) == [2, 1, 1, 0] + @test counts(x, Base.OneTo(5)) == [2, 1, 1, 0, 0] + @test isempty(counts(x, Base.OneTo(0))) + + # Bin positions depend on the values, even for a range with offset axes. + levels = Base.IdentityUnitRange(2:5) + @test counts(x, levels) == [1, 1, 0, 0] + @test addcounts!(ones(Int, 4), x, levels) == [2, 2, 1, 1] + @test counts(x, levels, w) == [3.0, 4.0, 0.0, 0.0] + @test proportions(x, levels) == [0.25, 0.25, 0.0, 0.0] + @test proportions(x, levels, w) ≈ [0.3, 0.4, 0.0, 0.0] + @test addcounts!(ones(4), x, levels, w) == [4.0, 5.0, 1.0, 1.0] + + # `r` must have linear indices `1:length(levels)`. + @test_throws BoundsError addcounts!(zeros(Int, 3), x, levels) + @test_throws BoundsError addcounts!(zeros(3), x, levels, w) + @test_throws BoundsError addcounts!(OffsetVector(zeros(Int, 4), 2:5), x, levels) + @test_throws BoundsError addcounts!(OffsetVector(zeros(4), 2:5), x, levels, w) + + # Examples from #908. + @test counts(x, Base.OneTo(4), w) == [3.0, 3.0, 4.0, 0.0] + @test proportions(x, Base.OneTo(4)) == [0.5, 0.25, 0.25, 0.0] + @test proportions(x, Base.OneTo(4), w) ≈ [0.3, 0.3, 0.4, 0.0] +end + +@testset "integer count bin arithmetic" begin + # The offset must not overflow during subtraction of small integers. + levels = Int8(-100):Int8(100) + x = Int8[-101, -100, 100, 101] + w = weights([1.0, 2.0, 3.0, 4.0]) + expected = zeros(Int, 201) + expected[[1, 201]] .= 1 + expected_w = zeros(201) + expected_w[[1, 201]] = [2.0, 3.0] + @test counts(x, levels) == expected + @test counts(x, levels, w) == expected_w + + lo = typemin(Int) + limits = Base.IdentityUnitRange(lo:lo+1) + wide = BigInt[lo, big(lo) + 1] + @test counts(wide, limits) == [1, 1] + @test counts(wide, limits, weights([2.0, 3.0])) == [2.0, 3.0] + + for T in (UInt64, UInt128) + zero_limits = zero(T):one(T) + @test counts(BigInt[0, 1], zero_limits) == [1, 1] + @test counts(BigInt[0, 1], zero_limits, weights([2.0, 3.0])) == [2.0, 3.0] + end + +end + +struct PromotingCountInteger <: Integer + value::BigInt +end +Base.convert(::Type{BigInt}, x::PromotingCountInteger) = x.value +Base.BigInt(x::PromotingCountInteger) = x.value +Base.convert(::Type{UInt}, x::PromotingCountInteger) = UInt(x.value) +Base.promote_rule(::Type{PromotingCountInteger}, ::Type{T}) where {T<:Integer} = BigInt + +@testset "integer count bin arithmetic with generic promotion" begin + # Ordinary arithmetic works without a specialized modular conversion. + x = PromotingCountInteger.(BigInt[-1, 1]) + @test counts(x, -1:1) == [1, 0, 1] + @test counts(x, -1:1, weights([2.0, 3.0])) == [2.0, 0.0, 3.0] + + lo = big(10)^50 + wide = PromotingCountInteger.([lo, lo + 1]) + @test counts(wide, lo:lo+1) == [1, 1] + @test counts(wide, lo:lo+1, weights([2.0, 3.0])) == [2.0, 3.0] +end + @testset "1D integer counts" begin x = rand(1:5, n) w = weights(rand(n))