sum_kbn(fn, A) (issue 3, in part)

README.md

@@ -14,14 +14,14 @@ These functions were formerly part of Julia's Base library.
 ```julia
 julia> using KahanSummation
 
-julia> vec = [1.0, 1.0e16, -1.0e16, -0.5];
+julia> avec = [1.0, 1.0e16, -1.0e16, -0.5];
 
-julia> sum(vec), sum_kbn(vec)
+julia> sum(avec), sum_kbn(avec)
 (-0.5, 0.5)
 
-julia> vec = [1.0, 1.0e16, 1.0, -1.0e16];
+julia> avec = [1.0, 1.0e16, 1.0, -1.0e16];
 
-julia> cumsum_kbn(vec) .- cumsum(vec)
+julia> cumsum_kbn(avec) .- cumsum(avec)
 4-element Array{Float64,1}:
  0.0
  0.0

src/KahanSummation.jl

@@ -64,6 +64,64 @@ function cumsum_kbn(A::AbstractArray{T}, axis::Integer) where T<:AbstractFloat
     return B + C
 end
 
+if VERSION >= v"0.7.0-DEV.3000" # TODO: More specific bound
+    if isdefined(Base, :sum_kbn) # Deprecated
+        import Base: sum_kbn, cumsum_kbn
+    else
+        export sum_kbn, cumsum_kbn
+    end
+end
+
+if isdefined(Base, Symbol("@default_eltype"))
+    using Base: @default_eltype
+else
+    macro default_eltype(itr)
+        quote
+            Core.Inference.return_type(first, Tuple{$(esc(itr))})
+        end
+    end
+end
+
+if isdefined(Base, :promote_sys_size_add)
+    using Base: promote_sys_size_add
+else
+    promote_sys_size_add(x::T) where {T} = Base.r_promote(+, zero(T)::T)
+end
+
+"""
+    cumsum_kbn(A, dim::Integer)
+Cumulative sum along a dimension, using the Kahan-Babuska-Neumaier compensated summation
+algorithm for additional accuracy.
+"""
+function cumsum_kbn(A::AbstractArray{T}, axis::Integer) where T<:AbstractFloat
+    dimsA = size(A)
+    ndimsA = ndims(A)
+    axis_size = dimsA[axis]
+    axis_stride = 1
+    for i = 1:(axis-1)
+        axis_stride *= size(A, i)
+    end
+    axis_size <= 1 && return A
+    B = similar(A)
+    C = similar(A)
+    for i = 1:length(A)
+        if div(i-1, axis_stride) % axis_size == 0
+            B[i] = A[i]
+            C[i] = zero(T)
+        else
+            s = B[i-axis_stride]
+            Ai = A[i]
+            B[i] = t = s + Ai
+            if abs(s) >= abs(Ai)
+                C[i] = C[i-axis_stride] + ((s-t) + Ai)
+            else
+                C[i] = C[i-axis_stride] + ((Ai-t) + s)
+            end
+        end
+    end
+    return B + C
+end
+
 function cumsum_kbn(v::AbstractVector{T}) where T<:AbstractFloat
     r = similar(v)
     isempty(v) && return r
@@ -86,31 +144,34 @@ function cumsum_kbn(v::AbstractVector{T}) where T<:AbstractFloat
 end
 
 """
-    sum_kbn(A)
+    sum_kbn([f,] A)
 
 Return the sum of all elements of `A`, using the Kahan-Babuska-Neumaier compensated
-summation algorithm for additional accuracy.
+summation algorithm for additional accuracy.  When a function `f` is supplied, the
+result of calling f on each element of `A` is summed.
 """
-function sum_kbn(A)
-    T = @default_eltype(typeof(A))
+function sum_kbn(f::Function, A)
+    T = Base.return_types(f, (@default_eltype(typeof(A)),))[1]
     c = promote_sys_size_add(zero(T)::T)
     i = start(A)
     if done(A, i)
         return c
     end
     Ai, i = next(A, i)
-    s = Ai - c
+    s = f(Ai) - c
     while !(done(A, i))
         Ai, i = next(A, i)
-        t = s + Ai
-        if abs(s) >= abs(Ai)
-            c -= ((s-t) + Ai)
+        fAi = f(Ai); t = s + fAi
+        if abs(s) >= abs(fAi)
+            c -= ((s-t) + fAi)
         else
-            c -= ((Ai-t) + s)
+            c -= ((fAi-t) + s)
         end
         s = t
     end
     s - c
 end
 
+sum_kbn(A) = sum_kbn(identity, A)
+
 end # module

test/runtests.jl

@@ -31,6 +31,12 @@ import KahanSummation: sum_kbn, cumsum_kbn
     @test isequal(c[1,:], cv2)
     @test isequal(c[3,:], cv)
     @test isequal(c[:,4], [2.0,2.0,2.0,2.0]*1000)
+
+    amat = reshape(collect(1.0:16.0),(4,4))
+
+    @test cumsum_kbn(amat, 2)[3,:] == [3.0, 10.0, 21.0, 36.0]
+    @test diag(cumsum_kbn(amat, 1)) == [1.0, 11.0, 30.0, 58.0]
+
 end
 
 @testset "sum_kbn" begin
@@ -44,5 +50,25 @@ end
     @test sum_kbn([7 8 9]) === sum_kbn([8 9 7])
     @test sum_kbn(i for i=1:1:10) === sum_kbn(i for i=10:-1:1)
     @test sum_kbn([-0.0]) === -0.0
-    @test sum_kbn([-0.0,-0.0]) === -0.0
+
+    @test sum_kbn(identity, [1,1e100,1,-1e100]) === 2.0
+    @test sum_kbn(identity, Float64[]) === 0.0
+    @test sum_kbn(identity, i for i=1.0:1.0:10.0) === 55.0
+    @test sum_kbn(identity, i for i=1:1:10) === 55
+    @test sum_kbn(identity, [1 2 3]) === 6
+    @test sum_kbn(identity, [2+im 3-im]) === 5+0im
+    @test sum_kbn(identity, [1+im 2+3im]) === 3+4im
+    @test sum_kbn(identity, [7 8 9]) === sum_kbn([8 9 7])
+    @test sum_kbn(identity, i for i=1:1:10) === sum_kbn(i for i=10:-1:1)
+    @test sum_kbn(identity, [-0.0]) === -0.0
+
+    @test sum_kbn([-0.0]) === -0.0
+    @test sum_kbn(x->x*x, [1.0]) === 1.0
+    @test sum_kbn(x->x*x, [2.0]) === 4.0
+    @test sum_kbn(x->x+1, Float64[]) === 0.0
+    @test sum_kbn(x->x+1, [6 7 8]) === sum_kbn([8 9 7])
+    @test sum_kbn(sqrt, i for i=1.0:1.0:10.0) === 22.4682781862041
+    @test sum_kbn(x->x*im,[1+im 2+3im]) === -4+3im
+    @test sum_kbn(sqrt, [i for i in 1.0:1.0:10.0]) === sum_kbn(sqrt, [i for i in 10.0:-1.0:1.0])
+    @test sum_kbn([1.0, 1.0e16, 1.0, -1.0e16]) === sum_kbn(x->x+1, [1.0, 1.0e16, 1.0, -1.0e16] - 1.0)
 end