From e335fe0756baa2a36f0ce16bd05abc6dcc461642 Mon Sep 17 00:00:00 2001 From: Jan Weidner Date: Fri, 14 Oct 2016 11:55:05 +0200 Subject: [PATCH] add scan, scan! - rename cumop! to scan! and export it - add specialized 1d method to scan! - add scan - add cummin!, cummax! - refactor cummin, cummax --- base/arraymath.jl | 2 +- base/deprecated.jl | 2 + base/exports.jl | 4 ++ base/multidimensional.jl | 117 ++++++++++++++++++++++----------------- doc/stdlib/arrays.rst | 32 +++++++++++ test/arrayops.jl | 31 +++++++++++ 6 files changed, 135 insertions(+), 53 deletions(-) diff --git a/base/arraymath.jl b/base/arraymath.jl index 26f60b7aed75af..e9b5e9130d383e 100644 --- a/base/arraymath.jl +++ b/base/arraymath.jl @@ -448,7 +448,7 @@ ctranspose{T}(x::AbstractVector{T}) = T[ ctranspose(v) for i=of_indices(x, OneTo # see discussion in #18364 ... we try not to widen type of the resulting array # from cumsum or cumprod, but in some cases (+, Bool) we may not have a choice. -rcum_promote_type{T<:Number}(op, ::Type{T}) = promote_op(op, T) +rcum_promote_type{T<:Number}(op, ::Type{T}) = promote_op(op, T, T) rcum_promote_type{T}(op, ::Type{T}) = T # handle sums of Vector{Bool} and similar. it would be nice to handle diff --git a/base/deprecated.jl b/base/deprecated.jl index 8fc573b6a110bf..b406e4ef3df998 100644 --- a/base/deprecated.jl +++ b/base/deprecated.jl @@ -144,6 +144,8 @@ end @deprecate chol(A::Number, ::Type{Val{:L}}) ctranspose(chol(A)) @deprecate chol(A::AbstractMatrix, ::Type{Val{:L}}) ctranspose(chol(A)) +@deprecate cumop! scan! + # Number updates # rem1 is inconsistent for x==0: The result should both have the same diff --git a/base/exports.jl b/base/exports.jl index e6d6dabbf8925a..3262b2c041ced0 100644 --- a/base/exports.jl +++ b/base/exports.jl @@ -498,11 +498,15 @@ export conj!, copy!, cummax, + cummax!, cummin, + cummin!, cumprod, cumprod!, cumsum, cumsum!, + scan, + scan!, cumsum_kbn, eachindex, extrema, diff --git a/base/multidimensional.jl b/base/multidimensional.jl index 69b6ef1a9887b4..cd4a8ae7f46e54 100644 --- a/base/multidimensional.jl +++ b/base/multidimensional.jl @@ -479,50 +479,11 @@ end end end -for (f, fmod, op) = ((:cummin, :_cummin!, :min), (:cummax, :_cummax!, :max)) - @eval function ($f)(v::AbstractVector) - n = length(v) - cur_val = v[1] - res = similar(v, n) - res[1] = cur_val - for i in 2:n - cur_val = ($op)(v[i], cur_val) - res[i] = cur_val - end - return res - end - - @eval function ($f)(A::AbstractArray, axis::Integer) - axis > 0 || throw(ArgumentError("axis must be a positive integer")) - res = similar(A) - axis > ndims(A) && return copy!(res, A) - inds = indices(A) - if isempty(inds[axis]) - return res - end - R1 = CartesianRange(inds[1:axis-1]) - R2 = CartesianRange(inds[axis+1:end]) - ($fmod)(res, A, R1, R2, axis) - end - - @eval @noinline function ($fmod)(res, A::AbstractArray, R1::CartesianRange, R2::CartesianRange, axis::Integer) - inds = indices(A, axis) - i1 = first(inds) - for I2 in R2 - for I1 in R1 - res[I1, i1, I2] = A[I1, i1, I2] - end - for i = i1+1:last(inds) - for I1 in R1 - res[I1, i, I2] = ($op)(A[I1, i, I2], res[I1, i-1, I2]) - end - end - end - res - end +cummax{T}(A::AbstractArray{T}, axis::Integer=1) = scan(max, A, axis) +cummax!(B, A, axis::Integer) = scan!(max, B, A, axis) - @eval ($f)(A::AbstractArray) = ($f)(A, 1) -end +cummin{T}(A::AbstractArray{T}, axis::Integer=1) = scan(min, A, axis) +cummin!(B, A, axis::Integer) = scan!(min, B, A, axis) """ cumsum(A, dim=1) @@ -548,8 +509,9 @@ julia> cumsum(a,2) 4 9 15 ``` """ -cumsum{T}(A::AbstractArray{T}, axis::Integer=1) = cumsum!(similar(A, Base.rcum_promote_type(+, T)), A, axis) -cumsum!(B, A::AbstractArray) = cumsum!(B, A, 1) +cumsum{T}(A::AbstractArray{T}, axis::Integer=1) = scan(+, A, axis) +cumsum!(B, A, axis::Integer) = scan!(+, B, A, axis) + """ cumprod(A, dim=1) @@ -574,13 +536,64 @@ julia> cumprod(a,2) 4 20 120 ``` """ -cumprod(A::AbstractArray, axis::Integer=1) = cumprod!(similar(A), A, axis) -cumprod!(B, A) = cumprod!(B, A, 1) +cumprod(A::AbstractArray, axis::Integer=1) = scan(*, A, axis) +cumprod!(B, A, axis::Integer) = scan!(*, B, A, axis) -cumsum!(B, A, axis::Integer) = cumop!(+, B, A, axis) -cumprod!(B, A, axis::Integer) = cumop!(*, B, A, axis) +""" + scan(op, A, dim=1) -function cumop!(op, B, A, axis::Integer) +Cumulative operation `op` along a dimension `dim` (defaults to 1). See also +[`scan!`](:func:`scan!`) to use a preallocated output array, both for performance and +to control the precision of the output (e.g. to avoid overflow). For common operations +there are specialized variants of scan, see: +[`cumsum`](:func:`cumsum`), [`cummin`](:func:`cummin`), [`cummax`](:func:`cummax`), [`cumprod`](:func:`cumprod`) + +```jldoctest +julia> scan(+, [1,2,3]) +3-element Array{Int64,1}: + 1 + 3 + 6 + +julia> scan(*, [1,2,3]) +3-element Array{Int64,1}: + 1 + 2 + 6 + +julia> scan(min, [1,2,-1]) +3-element Array{Int64,1}: + 1 + 1 + -1 +``` +""" +function scan(op, A, axis::Integer=1) + out = similar(A, Base.rcum_promote_type(op, eltype(A))) + scan!(op, out, A, axis) +end + +function scan!(op, out, v::AbstractVector, axis::Integer=1) + axis > 0 || throw(ArgumentError("axis must be a positive integer")) + size(out) == size(v) || throw(DimensionMismatch("shape of out must match v")) + isempty(v) && return out + axis > 1 && return copy!(out, v) + cur_val = v[1] + out[1] = cur_val + @inbounds for i in 2:length(v) + cur_val = op(v[i], cur_val) + out[i] = cur_val + end + return out +end + +""" + scan!(op, B, A, dim=1) + +Cumulative operation `op` on A along a dimension, storing the result in B. The dimension defaults to 1. +See also [`scan`](:func:`scan`). +""" +function scan!(op, B, A, axis::Integer=1) axis > 0 || throw(ArgumentError("axis must be a positive integer")) inds_t = indices(A) indices(B) == inds_t || throw(DimensionMismatch("shape of B must match A")) @@ -601,12 +614,12 @@ function cumop!(op, B, A, axis::Integer) else R1 = CartesianRange(indices(A)[1:axis-1]) # not type-stable R2 = CartesianRange(indices(A)[axis+1:end]) - _cumop!(op, B, A, R1, inds_t[axis], R2) # use function barrier + _scan!(op, B, A, R1, inds_t[axis], R2) # use function barrier end return B end -@noinline function _cumop!(op, B, A, R1, ind, R2) +@noinline function _scan!(op, B, A, R1, ind, R2) # Copy the initial element in each 1d vector along dimension `axis` i = first(ind) @inbounds for J in R2, I in R1 diff --git a/doc/stdlib/arrays.rst b/doc/stdlib/arrays.rst index 544d37754dadea..e043f70d49f5cd 100644 --- a/doc/stdlib/arrays.rst +++ b/doc/stdlib/arrays.rst @@ -1463,6 +1463,38 @@ Indexing, Assignment, and Concatenation Array functions --------------- +.. function:: scan(op, A, dim=1) + + .. Docstring generated from Julia source + + Cumulative operation ``op`` along a dimension ``dim`` (defaults to 1). See also :func:`scan!` to use a preallocated output array, both for performance and to control the precision of the output (e.g. to avoid overflow). For common operations there are specialized variants of scan, see: :func:`cumsum`\ , :func:`cummin`\ , :func:`cummax`\ , :func:`cumprod` + + .. doctest:: + + julia> scan(+, [1,2,3]) + 3-element Array{Int64,1}: + 1 + 3 + 6 + + julia> scan(*, [1,2,3]) + 3-element Array{Int64,1}: + 1 + 2 + 6 + + julia> scan(min, [1,2,-1]) + 3-element Array{Int64,1}: + 1 + 1 + -1 + +.. function:: scan!(op, B, A, dim=1) + + .. Docstring generated from Julia source + + Cumulative operation ``op`` on A along a dimension, storing the result in B. The dimension defaults to 1. See also :func:`scan`\ . + .. function:: cumprod(A, dim=1) .. Docstring generated from Julia source diff --git a/test/arrayops.jl b/test/arrayops.jl index 340173edc9acac..0dc14aa3dfb71a 100644 --- a/test/arrayops.jl +++ b/test/arrayops.jl @@ -1868,3 +1868,34 @@ end @test size(a) == size(b) end end + +@testset "scan, scan!" begin + + @test scan(+, [1,2,3]) == [1, 3, 6] + @test scan(min, [1 2; 3 4], 1) == [1 2; 1 2] + @test scan(max, [1 2; 3 0], 2) == [1 2; 3 3] + @test scan(+, Bool[]) == Int[] + @test scan(*, Bool[]) == Bool[] + @test scan(+, Float64[]) == Float64[] + + N = 5 + for arr in [rand(Float64, N), rand(Bool, N), rand(-2:2, N)] + for (op, cumop) in [(+, cumsum), (min, cummin), (max, cummax), (*, cumprod)] + @inferred scan(op, arr) + scan_arr = scan(op, arr) + @test scan_arr ≈ cumop(arr) + @test scan_arr[end] ≈ reduce(op, arr) + @test scan_arr[1] ≈ arr[1] + @test scan(op, arr, 10) ≈ arr + + if eltype(arr) in [Int, Float64] # eltype of out easy + out = similar(arr) + @test out ≉ scan_arr + @test scan!(op, out, arr) ≈ scan_arr + @test out ≈ scan_arr + end + + end + end + +end