make counting more robust to input datatype

src/counts.jl

@@ -23,17 +23,15 @@ array `r`. If a weighting vector `wv` is specified, the sum of weights is used
 rather than the raw counts.
 """
 function addcounts!(r::AbstractArray, x::IntegerArray, levels::IntUnitRange)
-    # add counts of integers from x to r
+    # add counts of integers from x that fall within levels to r
 
-    k = length(levels)
-    length(r) == k || throw(DimensionMismatch())
+    @boundscheck checkbounds(r, axes(levels)...)
 
-    m0 = levels[1]
-    m1 = levels[end]
+    m0 = first(levels)
+    m1 = last(levels)
     b = m0 - 1
 
-    @inbounds for i in 1 : length(x)
-        xi = x[i]
+    @inbounds for xi in x
         if m0 <= xi <= m1
             r[xi - b] += 1
         end
@@ -42,14 +40,15 @@ function addcounts!(r::AbstractArray, x::IntegerArray, levels::IntUnitRange)
 end
 
 function addcounts!(r::AbstractArray, x::IntegerArray, levels::IntUnitRange, wv::AbstractWeights)
-    k = length(levels)
-    length(r) == k || throw(DimensionMismatch())
+    # add wv weighted counts of integers from x that fall within levels to r
+
+    @boundscheck checkbounds(r, axes(levels)...)
 
-    m0 = levels[1]
-    m1 = levels[end]
+    m0 = first(levels)
+    m1 = last(levels)
     b = m0 - 1
 
-    @inbounds for i in 1 : length(x)
+    @inbounds for i in eachindex(x, wv)
         xi = x[i]
         if m0 <= xi <= m1
             r[xi - b] += wv[i]
@@ -112,24 +111,21 @@ proportions(x::IntegerArray, wv::AbstractWeights) = proportions(x, span(x), wv)
 function addcounts!(r::AbstractArray, x::IntegerArray, y::IntegerArray, levels::NTuple{2,IntUnitRange})
     # add counts of integers from x to r
 
-    n = length(x)
-    length(y) == n || throw(DimensionMismatch())
 
     xlevels, ylevels = levels
 
-    kx = length(xlevels)
-    ky = length(ylevels)
-    size(r) == (kx, ky) || throw(DimensionMismatch())
+    @boundscheck checkbounds(r, axes(xlevels, 1), axes(ylevels, 1))
 
-    mx0 = xlevels[1]
-    mx1 = xlevels[end]
-    my0 = ylevels[1]
-    my1 = ylevels[end]
+
+    mx0 = first(xlevels)
+    mx1 = last(xlevels)
+    my0 = first(ylevels)
+    my1 = last(ylevels)
 
     bx = mx0 - 1
     by = my0 - 1
 
-    for i = 1:n
+    for i in eachindex(x, y)
         xi = x[i]
         yi = y[i]
         if (mx0 <= xi <= mx1) && (my0 <= yi <= my1)
@@ -143,24 +139,20 @@ function addcounts!(r::AbstractArray, x::IntegerArray, y::IntegerArray,
                     levels::NTuple{2,IntUnitRange}, wv::AbstractWeights)
     # add counts of integers from x to r
 
-    n = length(x)
-    length(y) == length(wv) == n || throw(DimensionMismatch())
 
     xlevels, ylevels = levels
 
-    kx = length(xlevels)
-    ky = length(ylevels)
-    size(r) == (kx, ky) || throw(DimensionMismatch())
+    @boundscheck checkbounds(r, axes(xlevels, 1), axes(ylevels, 1))
 
-    mx0 = xlevels[1]
-    mx1 = xlevels[end]
-    my0 = ylevels[1]
-    my1 = ylevels[end]
+    mx0 = first(xlevels)
+    mx1 = last(xlevels)
+    my0 = first(ylevels)
+    my1 = last(ylevels)
 
     bx = mx0 - 1
     by = my0 - 1
 
-    for i = 1:n
+    for i in eachindex(x, y, wv)
         xi = x[i]
         yi = y[i]
         if (mx0 <= xi <= mx1) && (my0 <= yi <= my1)
@@ -284,9 +276,9 @@ function addcounts_dict!(cm::Dict{T}, x) where T
 end
 
 # If the bits type is of small size i.e. it can have up to 65536 distinct values
-# then it is always better to apply a counting-sort like reduce algorithm for 
+# then it is always better to apply a counting-sort like reduce algorithm for
 # faster results and less memory usage. However we still wish to enable others
-# to write generic algorithms, therefore the methods below still accept the 
+# to write generic algorithms, therefore the methods below still accept the
 # `alg` argument but it is ignored.
 function _addcounts!(::Type{Bool}, cm::Dict{Bool}, x::AbstractArray{Bool}; alg = :ignored)
     sumx = sum(x)
@@ -337,23 +329,22 @@ radixsort_safe(::Type{T}) where T = T<:BaseRadixSortSafeTypes
 
 function _addcounts_radix_sort_loop!(cm::Dict{T}, sx::AbstractArray{T}) where T
     isempty(sx) && return cm
-    last_sx = sx[1]
-    tmpcount = get(cm, last_sx, 0) + 1
+    last_sx = first(sx)
+    start_i = firstindex(sx)
 
     # now the data is sorted: can just run through and accumulate values before
     # adding into the Dict
-    @inbounds for i in 2:length(sx)
+    @inbounds for i in start_i+1:lastindex(sx)
         sxi = sx[i]
-        if last_sx == sxi
-            tmpcount += 1
-        else
-            cm[last_sx] = tmpcount
+        if last_sx != sxi
+            cm[last_sx] = get(cm, last_sx, 0) + i - start_i
             last_sx = sxi
-            tmpcount = get(cm, last_sx, 0) + 1
+            start_i = i
         end
     end
 
-    cm[sx[end]] = tmpcount
+    last_sx = last(sx)
+    cm[last_sx] = get(cm, last_sx, 0) + firstindex(sx) + 1 - start_i
 
     return cm
 end
@@ -369,17 +360,16 @@ function addcounts_radixsort!(cm::Dict{T}, x::AbstractArray{T}) where T
 end
 
 # fall-back for `x` an iterator
-function addcounts_radixsort!(cm::Dict{T}, x) where T 
+function addcounts_radixsort!(cm::Dict{T}, x) where T
     sx = sort!(collect(x), alg = RadixSort)
     return _addcounts_radix_sort_loop!(cm, sx)
 end
 
 function addcounts!(cm::Dict{T}, x::AbstractArray{T}, wv::AbstractVector{W}) where {T,W<:Real}
-    n = length(x)
-    length(wv) == n || throw(DimensionMismatch())
+
     z = zero(W)
 
-    for i = 1 : n
+    for i in eachindex(x, wv)
         @inbounds xi = x[i]
         @inbounds wi = wv[i]
         cm[xi] = get(cm, xi, z) + wi
@@ -418,5 +408,5 @@ countmap(x::AbstractArray{T}, wv::AbstractVector{W}) where {T,W<:Real} = addcoun
 Return a dictionary mapping each unique value in `x` to its
 proportion in `x`.
 """
-proportionmap(x::AbstractArray) = _normalize_countmap(countmap(x), length(x))
+proportionmap(x::AbstractArray; alg = :auto) = _normalize_countmap(countmap(x; alg = alg), length(x))
 proportionmap(x::AbstractArray, wv::AbstractWeights) = _normalize_countmap(countmap(x, wv), sum(wv))

src/weights.jl

@@ -21,6 +21,7 @@ length(wv::AbstractWeights) = length(wv.values)
 sum(wv::AbstractWeights) = wv.sum
 isempty(wv::AbstractWeights) = isempty(wv.values)
 size(wv::AbstractWeights) = size(wv.values)
+Base.axes(wv::AbstractWeights) = Base.axes(wv.values)
 
 Base.convert(::Type{Vector}, wv::AbstractWeights) = convert(Vector, wv.values)
 
@@ -310,7 +311,7 @@ julia> uweights(3)
  1
  1
  1
- 
+
 julia> uweights(Float64, 3)
 3-element UnitWeights{Float64}:
  1.0

test/counts.jl

@@ -1,5 +1,6 @@
 using StatsBase
 using Test
+using OffsetArrays
 
 n = 5000
 
@@ -104,7 +105,7 @@ cm_missing = countmap(skipmissing(xx))
 @test cm_missing isa Dict{Int, Int}
 @test cm_missing == cm
 
-cm_any_itr = countmap((i for i in xx)) 
+cm_any_itr = countmap((i for i in xx))
 @test cm_any_itr isa Dict{Any,Int} # no knowledge about type
 @test cm_missing == cm
 
@@ -164,3 +165,24 @@ end
     X = view([1,1,1,2,2], 1:5)
     @test countmap(X) == countmap(copy(X))
 end
+
+@testset "offset arrays" begin
+    x = rand(1:5, n)
+    w = rand(n)
+    xw = weights(w)
+    y = OffsetArray(x, n÷2)
+    yw = weights(OffsetArray(w, n÷2))
+    z = OffsetArray(x, -2n)
+    zw = weights(OffsetArray(w, -2n))
+
+    # proportions calls counts which calls addcounts!
+    @test proportions(x) == proportions(y) == proportions(z)
+    @test proportions(x, xw) == proportions(y, yw) == proportions(z, zw)
+    @test proportionmap(x) == proportionmap(y) == proportionmap(z)
+    @test proportionmap(x, xw) == proportionmap(y, yw) == proportionmap(z, zw)
+    @test (proportionmap(x) == proportionmap(x; alg = :dict) == proportionmap(x; alg = :radixsort)
+        == proportionmap(y) == proportionmap(y; alg = :dict) == proportionmap(y; alg = :radixsort)
+        == proportionmap(z) == proportionmap(z; alg = :dict) == proportionmap(z; alg = :radixsort))
+    @test proportionmap(x, xw) == proportionmap(y, yw) == proportionmap(z, zw)
+    # countmap and proportionmap only support the :dict algorithm for weighted sums.
+end