Accept abstract unit ranges as count levels

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

src/counts.jl

@@ -13,50 +13,65 @@ 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) ||
         throw(DimensionMismatch("x and wv must have the same length, got $(length(x)) and $(length(wv))"))
 
     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)
 
 """

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))