Make broadcast over a single SparseMatrixCSC (and possibly broadcast scalars) more generic. (#19239)

Author: Sacha0

base/sparse/sparsematrix.jl

@@ -1398,31 +1398,19 @@ end
 
 sparse(S::UniformScaling, m::Integer, n::Integer=m) = speye_scaled(S.λ, m, n)
 
-## 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, fofzero, 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, Base.Broadcast.promote_eltype_op(f, A), nnz(A))
     Bk = 1
     @inbounds for j in 1:A.n
         Bcolptr[j] = Bk
@@ -1440,69 +1428,51 @@ 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 `fillvalue` 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, fillvalue, 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(fillvalue, 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.
+#
+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: More appropriate location?
+conj!(A::SparseMatrixCSC) = (broadcast!(conj, A.nzval, A.nzval); A)
 
 
 ## Broadcasting kernels specialized for returning a SparseMatrixCSC

test/sparse/sparse.jl

@@ -1627,14 +1627,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])