Skip to content

Commit

Permalink
Make broadcast over a single SparseMatrixCSC (and possibly broadcast …
Browse files Browse the repository at this point in the history 
…scalars) more generic.
  • Loading branch information
@Sacha0
Sacha0 committed Nov 12, 2016
1 parent ce6a0c3 commit 1bba3b3
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 90 deletions.
134 changes: 52 additions & 82 deletions base/sparse/sparsematrix.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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, 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
Expand All @@ -1432,69 +1420,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
Expand Down
12 changes: 4 additions & 8 deletions test/sparse/sparse.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down

0 comments on commit 1bba3b3

Please sign in to comment.