From 6613377e3165d0cfa947e5b32ed170961c360c59 Mon Sep 17 00:00:00 2001 From: Sacha Verweij Date: Sat, 5 Nov 2016 19:15:03 -0700 Subject: [PATCH] Make broadcast over a single SparseMatrixCSC (and possibly broadcast scalars) more generic. --- base/sparse/sparsematrix.jl | 142 +++++++++++++++--------------------- test/sparse/sparse.jl | 13 +--- 2 files changed, 64 insertions(+), 91 deletions(-) diff --git a/base/sparse/sparsematrix.jl b/base/sparse/sparsematrix.jl index f1bfb662f567f..fc210ef173336 100644 --- a/base/sparse/sparsematrix.jl +++ b/base/sparse/sparsematrix.jl @@ -1390,31 +1390,19 @@ function one{T}(S::SparseMatrixCSC{T}) speye(T, m) end -## Unary arithmetic and boolean operators -""" -Helper macro for the unary broadcast definitions below. Takes parent method `fp` and a set -of desired child methods `fcs`, and builds an expression defining each of the child methods -such that `broadcast(::typeof(fc), A::SparseMatrixCSC) = fp(fc, A)`. -""" -macro _enumerate_childmethods(fp, fcs...) - fcexps = Expr(:block) - for fc in fcs - push!(fcexps.args, :( $(esc(:broadcast))(::typeof($(esc(fc))), A::SparseMatrixCSC) = $(esc(fp))($(esc(fc)), A) ) ) - end - return fcexps -end +## Broadcast operations involving a single sparse matrix and possibly broadcast scalars -# Operations that map zeros to zeros and may map nonzeros to zeros, yielding a sparse matrix -""" -Takes unary function `f` that maps zeros to zeros and may map nonzeros to zeros, and returns -a new `SparseMatrixCSC{TiA,TvB}` `B` generated by applying `f` to each nonzero entry in -`A` and retaining only the resulting nonzeros. -""" -function _broadcast_unary_nz2z_z2z_T{TvA,TiA,TvB}(f::Function, A::SparseMatrixCSC{TvA,TiA}, ::Type{TvB}) - Bcolptr = Array{TiA}(A.n + 1) - Browval = Array{TiA}(nnz(A)) - Bnzval = Array{TvB}(nnz(A)) +function broadcast{Tf}(f::Tf, A::SparseMatrixCSC) + fofzero = f(zero(eltype(A))) + fpreszero = fofzero == zero(fofzero) + return fpreszero ? _broadcast_zeropres(f, A) : _broadcast_notzeropres(f, A) +end +"Returns a `SparseMatrixCSC` storing only the nonzero entries of `broadcast(f, Matrix(A))`." +function _broadcast_zeropres{Tf}(f::Tf, A::SparseMatrixCSC) + Bcolptr = similar(A.colptr, A.n + 1) + Browval = similar(A.rowval, nnz(A)) + Bnzval = similar(A.nzval, promote_eltype_op(f, A), nnz(A)) Bk = 1 @inbounds for j in 1:A.n Bcolptr[j] = Bk @@ -1432,69 +1420,59 @@ function _broadcast_unary_nz2z_z2z_T{TvA,TiA,TvB}(f::Function, A::SparseMatrixCS resize!(Bnzval, Bk - 1) return SparseMatrixCSC(A.m, A.n, Bcolptr, Browval, Bnzval) end -function _broadcast_unary_nz2z_z2z{Tv}(f::Function, A::SparseMatrixCSC{Tv}) - _broadcast_unary_nz2z_z2z_T(f, A, Tv) -end -@_enumerate_childmethods(_broadcast_unary_nz2z_z2z, - sin, sinh, sind, asin, asinh, asind, - tan, tanh, tand, atan, atanh, atand, - sinpi, cosc, ceil, floor, trunc, round) -broadcast(::typeof(real), A::SparseMatrixCSC) = copy(A) -broadcast{Tv,Ti}(::typeof(imag), A::SparseMatrixCSC{Tv,Ti}) = spzeros(Tv, Ti, A.m, A.n) -broadcast{TTv}(::typeof(real), A::SparseMatrixCSC{Complex{TTv}}) = _broadcast_unary_nz2z_z2z_T(real, A, TTv) -broadcast{TTv}(::typeof(imag), A::SparseMatrixCSC{Complex{TTv}}) = _broadcast_unary_nz2z_z2z_T(imag, A, TTv) -ceil{To}(::Type{To}, A::SparseMatrixCSC) = _broadcast_unary_nz2z_z2z_T(ceil, A, To) -floor{To}(::Type{To}, A::SparseMatrixCSC) = _broadcast_unary_nz2z_z2z_T(floor, A, To) -trunc{To}(::Type{To}, A::SparseMatrixCSC) = _broadcast_unary_nz2z_z2z_T(trunc, A, To) -round{To}(::Type{To}, A::SparseMatrixCSC) = _broadcast_unary_nz2z_z2z_T(round, A, To) - -# Operations that map zeros to zeros and map nonzeros to nonzeros, yielding a sparse matrix -""" -Takes unary function `f` that maps zeros to zeros and nonzeros to nonzeros, and returns a -new `SparseMatrixCSC{TiA,TvB}` `B` generated by applying `f` to each nonzero entry in `A`. -""" -function _broadcast_unary_nz2nz_z2z{TvA,TiA,Tf<:Function}(f::Tf, A::SparseMatrixCSC{TvA,TiA}) - Bcolptr = Vector{TiA}(A.n + 1) - Browval = Vector{TiA}(nnz(A)) - copy!(Bcolptr, 1, A.colptr, 1, A.n + 1) - copy!(Browval, 1, A.rowval, 1, nnz(A)) - Bnzval = broadcast(f, A.nzval) - resize!(Bnzval, nnz(A)) - return SparseMatrixCSC(A.m, A.n, Bcolptr, Browval, Bnzval) -end -@_enumerate_childmethods(_broadcast_unary_nz2nz_z2z, - log1p, expm1, abs, abs2, conj) -function conj!(A::SparseMatrixCSC) - @inbounds @simd for k in 1:nnz(A) - A.nzval[k] = conj(A.nzval[k]) - end - return A -end - -# Operations that map both zeros and nonzeros to zeros, yielding a dense matrix """ -Takes unary function `f` that maps both zeros and nonzeros to nonzeros, constructs a new -`Matrix{TvB}` `B`, populates all entries of `B` with the result of a single `f(one(zero(Tv)))` -call, and then, for each stored entry in `A`, calls `f` on the entry's value and stores the -result in the corresponding location in `B`. +Returns a (dense) `SparseMatrixCSC` with `f(zero(eltype(A)))` stored in place of each +unstored entry in `A` and `f(A[i,j])` stored in place of each stored entry `A[i,j]` in `A`. """ -function _broadcast_unary_nz2nz_z2nz{Tv}(f::Function, A::SparseMatrixCSC{Tv}) - B = fill(f(zero(Tv)), size(A)) - @inbounds for j in 1:A.n - for k in nzrange(A, j) - i = A.rowval[k] - x = A.nzval[k] - B[i,j] = f(x) - end +function _broadcast_notzeropres{Tf}(f::Tf, A::SparseMatrixCSC) + nnzB = A.m * A.n + # Build structure + Bcolptr = similar(A.colptr, A.n + 1) + copy!(Bcolptr, 1:A.m:(nnzB + 1)) + Browval = similar(A.rowval, nnzB) + for k in 1:A.m:(nnzB - A.m + 1) + copy!(Browval, k, 1:A.m) + end + # Populate values + Bnzval = fill(f(zero(eltype(A))), nnzB) + @inbounds for (j, jo) in zip(1:A.n, 0:A.m:(nnzB - 1)), k in nzrange(A, j) + Bnzval[jo + A.rowval[k]] = f(A.nzval[k]) end - return B + # NOTE: Combining the fill call into the loop above to avoid multiple sweeps over / + # nonsequential access of Bnzval does not appear to improve performance + return SparseMatrixCSC(A.m, A.n, Bcolptr, Browval, Bnzval) end -@_enumerate_childmethods(_broadcast_unary_nz2nz_z2nz, - log, log2, log10, exp, exp2, exp10, sinc, cospi, - cos, cosh, cosd, acos, acosd, - cot, coth, cotd, acot, acotd, - sec, sech, secd, asech, - csc, csch, cscd, acsch) + +# Cover common broadcast operations involving a single sparse matrix and one or more +# broadcast scalars. +# +# TODO: The minimal snippet below is not satisfying: A better solution would achieve +# the same for (1) all broadcast scalar types (Base.Broadcast.containertype(x) == Any?) and +# (2) any combination (number, order, type mixture) of broadcast scalars. One might achieve +# (2) via a generated function. But how one might simultaneously and gracefully achieve (1) +# eludes me. Thoughts much appreciated. +# +broadcast{Tf}(f::Tf, x::Union{Number,Bool}, A::SparseMatrixCSC) = broadcast(y -> f(x, y), A) +broadcast{Tf}(f::Tf, A::SparseMatrixCSC, y::Union{Number,Bool}) = broadcast(x -> f(x, y), A) +# NOTE: The following two method definitions work around #19096. These definitions should +# be folded into the two preceding definitions on resolution of #19096. +broadcast{Tf,T}(f::Tf, ::Type{T}, A::SparseMatrixCSC) = broadcast(y -> f(T, y), A) +broadcast{Tf,T}(f::Tf, A::SparseMatrixCSC, ::Type{T}) = broadcast(x -> f(x, T), A) + +# TODO: The following definitions should be deprecated. +ceil{To}(::Type{To}, A::SparseMatrixCSC) = ceil.(To, A) +floor{To}(::Type{To}, A::SparseMatrixCSC) = floor.(To, A) +trunc{To}(::Type{To}, A::SparseMatrixCSC) = trunc.(To, A) +round{To}(::Type{To}, A::SparseMatrixCSC) = round.(To, A) + +# TODO: Nix these two specializations, or keep them around? +broadcast{Tf<:typeof(real)}(::Tf, A::SparseMatrixCSC) = copy(A) +broadcast{Tf<:typeof(imag),Tv,Ti}(::Tf, A::SparseMatrixCSC{Tv,Ti}) = spzeros(Tv, Ti, A.m, A.n) +# NOTE: Avoiding ambiguities with the general method (broadcast{Tf}(f::Tf, A::SparseMatrixCSC)) +# necessitates the less-than-graceful Tf<:typeof(fn) type parameters in the above definitions + +# TODO: More appropriate location? +conj!(A::SparseMatrixCSC) = (broadcast!(conj, A.nzval); A) ## Broadcasting kernels specialized for returning a SparseMatrixCSC diff --git a/test/sparse/sparse.jl b/test/sparse/sparse.jl index 72aef1cbe1ef0..9b65bead02622 100644 --- a/test/sparse/sparse.jl +++ b/test/sparse/sparse.jl @@ -499,7 +499,6 @@ end @test R == real.(S) @test I == imag.(S) @test real.(sparse(R)) == R - @test nnz(imag.(sparse(R))) == 0 @test abs.(S) == abs.(D) @test abs2.(S) == abs2.(D) end @@ -1627,14 +1626,10 @@ let A = sparse(UInt32[1,2,3], UInt32[1,2,3], [1.0,2.0,3.0]) @test A[1,1:3] == A[1,:] == [1,0,0] end -# Check that `broadcast` methods specialized for unary operations over -# `SparseMatrixCSC`s are called. (Issue #18705.) -let - A = spdiagm(1.0:5.0) - @test isa(sin.(A), SparseMatrixCSC) # representative for _unary_nz2z_z2z class - @test isa(abs.(A), SparseMatrixCSC) # representative for _unary_nz2nz_z2z class - @test isa(exp.(A), Array) # representative for _unary_nz2nz_z2nz class -end +# Check that `broadcast` methods specialized for unary operations over `SparseMatrixCSC`s +# are called. (Issue #18705.) EDIT: #19239 unified broadcast over a single sparse matrix, +# eliminating the former operation classes. +@test isa(sin.(spdiagm(1.0:5.0)), SparseMatrixCSC) # 19225 let X = sparse([1 -1; -1 1])