Propagate NaN and Inf to the unstored entries when scaling a sparse array by a Diagonal

src/linalg.jl

@@ -558,42 +558,136 @@ end
 
 ## diagonal scaling
 
+# An entry that is not stored scales to `op(0, d[j])` columnwise (`dim == 2`) or to
+# `op(d[i], 0)` rowwise (`dim == 1`), which is not zero for an infinite or NaN `d[j]`. Such a
+# column or row is stored in full before it is scaled, its new entries as stored zeros, so
+# that it scales as a dense one does. `_fullscaled` returns a mask of the columns or rows
+# that need it, or `nothing`. An eltype without a `zero` has only its stored values scaled.
+function _fullscaled(A, d::AbstractVector, op::F, dim::Int) where {F}
+    z = _implicit_zero(A)
+    z isa UnknownZero && return nothing
+    scaleszero(dk) = _iszero(dim == 1 ? op(dk, z) : op(z, dk))
+    all(scaleszero, d) && return nothing
+    nstored = _nstored(A, dim)
+    full = [!scaleszero(d[k]) && nstored[k] < size(A, 3 - dim) for k in axes(A, dim)]
+    return any(full) ? full : nothing
+end
+function _nstored(A::AbstractSparseMatrixCSC, dim::Int)
+    dim == 2 && return [length(nzrange(A, j)) for j in axes(A, 2)]
+    Arowval = getrowval(A)
+    counts = zeros(Int, size(A, 1))
+    for j in axes(A, 2), p in nzrange(A, j)
+        counts[Arowval[p]] += 1
+    end
+    return counts
+end
+function _nstored(x::AbstractSparseVector, dim::Int)
+    counts = zeros(Int, length(x))
+    counts[nonzeroinds(x)] .= 1
+    return counts
+end
+
+# `A` with the columns or rows that `full` marks stored in full
+function _storeinfull(A::AbstractSparseMatrixCSC, full::Vector{Bool}, dim::Int)
+    m, n = size(A)
+    Arowval, Anzval = getrowval(A), getnzval(A)
+    z = _densezero(A)
+    fullrows = dim == 1 ? findall(full) : collect(1:m)
+    instored = dim == 1 ? sum(j -> count(p -> full[Arowval[p]], nzrange(A, j)), axes(A, 2); init = 0) :
+                          sum(j -> full[j] ? length(nzrange(A, j)) : 0, axes(A, 2); init = 0)
+    nz = nnz(A) + widemul(count(full), size(A, 3 - dim)) - instored
+    _checknnzfits(indtype(A), nz)
+    colptr = Vector{indtype(A)}(undef, n + 1)
+    rowval = Vector{indtype(A)}(undef, nz)
+    nzval = Vector{promote_type(eltype(A), typeof(z))}(undef, nz)
+    colptr[1] = q = 1
+    for j in axes(A, 2)
+        rows = view(fullrows, 1:(dim == 1 || full[j] ? length(fullrows) : 0))
+        q = _mergestored!(rowval, nzval, q, Arowval, Anzval, nzrange(A, j), rows, z)
+        colptr[j + 1] = q
+    end
+    return SparseMatrixCSC(m, n, colptr, rowval, nzval)
+end
+function _storeinfull(x::AbstractSparseVector, full::Vector{Bool})
+    z = _densezero(x)
+    nz = nnz(x) + count(full)
+    nzind = Vector{indtype(x)}(undef, nz)
+    nzval = Vector{promote_type(eltype(x), typeof(z))}(undef, nz)
+    _mergestored!(nzind, nzval, 1, nonzeroinds(x), nonzeros(x), 1:nnz(x), findall(full), z)
+    return SparseVector(length(x), nzind, nzval)
+end
+# Write the indices `ind[ps]` and values `val[ps]`, merged with the sorted indices `inds` at
+# which they store `z` where `ind` has none, to `outind` and `outval` from `q` on, and return
+# the next `q`.
+function _mergestored!(outind, outval, q::Int, ind, val, ps, inds, z)
+    p, k = Int(first(ps)), 1
+    while p <= last(ps) || k <= length(inds)
+        i = p <= last(ps) ? Int(ind[p]) : typemax(Int)
+        r = k <= length(inds) ? inds[k] : typemax(Int)
+        outind[q] = min(i, r)
+        outval[q] = i <= r ? val[p] : z
+        i <= r && (p += 1)
+        r <= i && (k += 1)
+        q += 1
+    end
+    return q
+end
+
+# Scale `A`, with the columns or rows that `full` marks stored in full, into `C` with
+# `scale!`, which is `_scalecols!` or `_scalerows!`. A fixed `C` is written only if it stores
+# every entry of the result that is not zero.
+function _scaleinfull!(scale!::S, C, A, d, op::F, full, dim::Int) where {S,F}
+    B = _storeinfull(A, full, dim)
+    _is_fixed(C) && return _copyto_fixed!(C, scale!(_writablecopy(C), B, d, op))
+    return scale!(C, B, d, op)
+end
+
 # Scale the stored entries of `A` into `C`, which either is `A` or takes its pattern:
 # `C[i, j] = op(A[i, j], d[j])` columnwise, `C[i, j] = op(d[i], A[i, j])` rowwise. The
 # argument order of `op` is kept for non-commutative eltypes. The callers check that
-# `length(d)` matches the scaled dimension and that `size(C) == size(A)`.
+# `length(d)` matches the scaled dimension and that `size(C) == size(A)`. When a value might
+# not convert to `eltype(C)`, the values go to a new vector that `C` takes once all are written.
 function _scalecols!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, d::AbstractVector, op::F) where {F}
-    copyinds!(C, A)
-    Cnzval = getnzval(C)
+    full = _fullscaled(A, d, op, 2)
+    isnothing(full) || return _scaleinfull!(_scalecols!, C, A, d, op, full, 2)
     Anzval = getnzval(A)
-    resize!(Cnzval, length(Anzval))
+    direct = _storesafely(eltype(C), promote_op(op, eltype(A), eltype(d)))
+    Cnzval = direct ? _takepattern!(C, A) : Vector{eltype(C)}(undef, nnz(A))
     @inbounds for col in axes(A, 2)
         dcol = d[col]
         for p in nzrange(A, col)
             Cnzval[p] = op(Anzval[p], dcol)
         end
     end
+    direct || _storescaled!(C, A, Cnzval)
     return C
 end
 
 function _scalerows!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, d::AbstractVector, op::F) where {F}
-    copyinds!(C, A)
-    Cnzval = getnzval(C)
+    full = _fullscaled(A, d, op, 1)
+    isnothing(full) || return _scaleinfull!(_scalerows!, C, A, d, op, full, 1)
     Anzval = getnzval(A)
     Arowval = getrowval(A)
-    resize!(Cnzval, length(Anzval))
+    direct = _storesafely(eltype(C), promote_op(op, eltype(d), eltype(A)))
+    Cnzval = direct ? _takepattern!(C, A) : Vector{eltype(C)}(undef, nnz(A))
     @inbounds for col in axes(A, 2), p in nzrange(A, col)
         Cnzval[p] = op(d[Arowval[p]], Anzval[p])
     end
+    direct || _storescaled!(C, A, Cnzval)
     return C
 end
 
 function _scalerows!(x::AbstractSparseVector, d::AbstractVector, op::F) where {F}
+    full = _fullscaled(x, d, op, 1)
+    isnothing(full) || return copyto!(x, _scalerows!(_storeinfull(x, full), d, op))
     xnzval = nonzeros(x)
     xnzind = nonzeroinds(x)
+    direct = _storesafely(eltype(x), promote_op(op, eltype(d), eltype(x)))
+    vals = direct ? xnzval : similar(xnzval)
     @inbounds for p in eachindex(xnzind, xnzval)
-        xnzval[p] = op(d[xnzind[p]], xnzval[p])
+        vals[p] = op(d[xnzind[p]], xnzval[p])
     end
+    direct || copyto!(xnzval, vals)
     return x
 end
 
@@ -1353,6 +1447,19 @@ function copyinds!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC; copy_
     end
 end
 
+# Give `C` the pattern of `A`, and return the values of `C`, of the length of those of `A`.
+function _takepattern!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC)
+    copyinds!(C, A)
+    return resize!(getnzval(C), length(getnzval(A)))
+end
+# Store `nzval` in `C` with the pattern of `A`, converting it first so that a value that does
+# not convert to `eltype(C)` throws with `C` unchanged.
+function _storescaled!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, nzval::AbstractVector)
+    vals = convert(Vector{eltype(C)}, nzval)
+    copyto!(_takepattern!(C, A), vals)
+    return C
+end
+
 """
     rowcheck_index(A::AbstractSparseMatrixCSC, row::Integer, col::Integer)
 

src/matmul.jl

@@ -348,6 +348,8 @@ function *(A::Diagonal, b::AbstractSparseVector)
             DimensionMismatch(lazy"The dimension of the matrix A $(size(A)) and of the vector b $(length(b))")
         )
     end
+    full = _fullscaled(b, A.diag, *, 1)
+    isnothing(full) || return A * _storeinfull(b, full)
     T = promote_eltype(A, b)
     res = similar(b, T)
     nzind_b = nonzeroinds(b)
@@ -421,12 +423,16 @@ end
 function (*)(Da::Diagonal, A::SparseMatrixCSC, Db::Diagonal)
     (size(Da, 2) == size(A,1) && size(A,2) == size(Db,1)) ||
         throw(DimensionMismatch("incompatible sizes"))
+    da, db = map(parent, (Da, Db))
+    fullrows = _fullscaled(A, da, *, 1)
+    isnothing(fullrows) || return Da * _storeinfull(A, fullrows, 1) * Db
+    fullcols = _fullscaled(A, db, *, 2)
+    isnothing(fullcols) || return Da * _storeinfull(A, fullcols, 2) * Db
     T = promote_op(matprod, eltype(Da), promote_op(matprod, eltype(A), eltype(Db)))
     dest = similar(A, T)
     vals_dest = getnzval(dest)
     rows = getrowval(A)
     vals = getnzval(A)
-    da, db = map(parent, (Da, Db))
     @inbounds for col in axes(A,2)
         dbcol = db[col]
         for i in nzrange(A, col)
@@ -828,9 +834,10 @@ Store `A*D*α + C*β`, or `D*A*α + C*β`, in the sparse matrix `C`, and return
 columns, or the rows, of `A` scaled by the diagonal of `D`.
 
 Unlike a dense destination, `C` takes a new sparsity pattern: that of `A` when `β` is zero,
-and otherwise the union of the patterns of `A` and `C`. A fixed-pattern `C` keeps its
-pattern instead, and `mul!` throws, with `C` unchanged, if a nonzero of the result falls
-outside it.
+and otherwise the union of the patterns of `A` and `C`. A column, or row, that an infinite or
+NaN entry of `D` scales is stored in full, as its entries that `A` does not store become
+`NaN`. A fixed-pattern `C` keeps its pattern instead, and `mul!` throws, with `C` unchanged,
+if a nonzero of the result falls outside it.
 
 A sparse destination is supported for scaling only, by a `Diagonal` here or by a number
 with `mul!(C, A, b)`. Form other products of sparse matrices with `*`. With a dense
@@ -844,6 +851,8 @@ function mul!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, D::Diagona
     n == lb || throw(DimensionMismatch(lazy"A has size ($m, $n) but D has size ($lb, $lb)"))
     size(A)==size(C) || throw(DimensionMismatch(lazy"A has size ($m, $n), D has size ($lb, $lb), C has size $(size(C))"))
     iszero(alpha) && return _diagmul_scale!(C, beta)
+    full = _fullscaled(A, b, (a, bj) -> isone(alpha) ? a * bj : a * bj * alpha, 2)
+    isnothing(full) || return mul!(C, _storeinfull(A, full, 2), D, alpha, beta)
     R = _diagmul_eltype(C, promote_op(matprod, eltype(A), eltype(D)), alpha, beta)
     _storesafely(eltype(C), R) ||
         return _diagmul_store!(C, _spdiagmul!(_diagmul_widened(C, A, R, beta), A, D, alpha, beta))
@@ -951,6 +960,8 @@ function mul!(C::AbstractSparseMatrixCSC, D::Diagonal, A::AbstractSparseMatrixCS
     m == lb || throw(DimensionMismatch(lazy"D has size ($lb, $lb) but A has size ($m, $n)"))
     size(A)==size(C) || throw(DimensionMismatch(lazy"A has size ($m, $n), D has size ($lb, $lb), C has size $(size(C))"))
     iszero(alpha) && return _diagmul_scale!(C, beta)
+    full = _fullscaled(A, b, (bi, a) -> isone(alpha) ? bi * a : bi * a * alpha, 1)
+    isnothing(full) || return mul!(C, D, _storeinfull(A, full, 1), alpha, beta)
     R = _diagmul_eltype(C, promote_op(matprod, eltype(D), eltype(A)), alpha, beta)
     _storesafely(eltype(C), R) ||
         return _diagmul_store!(C, _diagspmul!(_diagmul_widened(C, A, R, beta), D, A, alpha, beta))
@@ -1022,18 +1033,18 @@ _lmulkeepszeros(b::Number, A) = (z = _implicit_zero(A); z isa UnknownZero || _is
 function mul!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, b::Number)
     size(A)==size(C) || throw(DimensionMismatch("A has size $(size(A)) but C has size $(size(C))"))
     (_is_fixed(C) || !_rmulkeepszeros(A, b)) && return broadcast!(*, C, A, b)
-    copyinds!(C, A)
-    resize!(getnzval(C), length(getnzval(A)))
-    mul!(getnzval(C), getnzval(A), b)
+    _storesafely(eltype(C), promote_op(*, eltype(A), typeof(b))) ||
+        return _storescaled!(C, A, view(getnzval(A), 1:nnz(A)) * b)
+    mul!(_takepattern!(C, A), getnzval(A), b)
     C
 end
 
 function mul!(C::AbstractSparseMatrixCSC, b::Number, A::AbstractSparseMatrixCSC)
     size(A)==size(C) || throw(DimensionMismatch("A has size $(size(A)) but C has size $(size(C))"))
     (_is_fixed(C) || !_lmulkeepszeros(b, A)) && return broadcast!(*, C, b, A)
-    copyinds!(C, A)
-    resize!(getnzval(C), length(getnzval(A)))
-    mul!(getnzval(C), b, getnzval(A))
+    _storesafely(eltype(C), promote_op(*, typeof(b), eltype(A))) ||
+        return _storescaled!(C, A, b * view(getnzval(A), 1:nnz(A)))
+    mul!(_takepattern!(C, A), b, getnzval(A))
     C
 end
 

test/linalg.jl

@@ -344,6 +344,15 @@ end
         C = spzeros(ComplexF64, 3, 3)
         @test_throws DimensionMismatch ldiv!(C, Dl, A)
         @test nnz(C) == 0
+        @static if COMPREHENSIVE
+        # a quotient that does not convert to the eltype of the destination throws with it unchanged
+        A0 = sparse([2 0; 0 3]); x0 = sparsevec([1, 2], [2, 3]); D2 = Diagonal([2, 2])
+        for (f!, X0) in ((C -> ldiv!(C, D2, A0), sparse([1 0; 0 0])), (B -> rdiv!(B, D2), A0), (y -> ldiv!(D2, y), x0))
+            X = copy(X0)
+            @test_throws InexactError f!(X)
+            @test mismatch(X, X0) === nothing
+        end
+        end
     end
 end
 

test/matmul.jl

@@ -448,6 +448,12 @@ end
     @test mismatch(Diagonal(b) * sA, Diagonal(b) * dA) === nothing
     @test mismatch(mul!(sC, Diagonal(b), sA), Diagonal(b) * dA) === nothing
     @test mismatch(lmul!(Diagonal(b), copy(sA)), Diagonal(b) * dA) === nothing
+    @static if COMPREHENSIVE
+    # an infinite entry of the Diagonal scales the entries that are not stored to NaN, as for dense
+    S = sparse([1.0 0; 0 1]); Dinf = Diagonal([Inf, 1])
+    @test mismatch(S * Dinf, Matrix(S) * Dinf) === nothing
+    @test mismatch(rmul!(copy(S), Dinf), Matrix(S) * Dinf) === nothing
+    end
 
     # adjoint/transpose of a sparse matrix with a Diagonal (issue #619)
     for (T, W) in ((ComplexF64, adjoint), (ComplexF64, transpose))
@@ -512,17 +518,18 @@ end
         @test mismatch(F * Dr, dA * Dr; approx=true) === nothing
         @test !_is_fixed(Dl * F) && !_is_fixed(F * Dr)
     end
-    # the kernels touch only the stored entries: exactly nnz(S) scalar multiplications,
-    # whereas the generic Diagonal kernel visits every element of the result
+    # the kernels touch only the stored entries: exactly nnz(S) scalar multiplications, and
+    # one per diagonal entry to find what it scales a zero to, whereas the generic Diagonal
+    # kernel visits every element of the result
     S = opcount_sparse(fixture(Float64, 20, 30))
     Dl = Diagonal(OpCount.(fixturedense(Float64, 30))); Dr = Diagonal(OpCount.(fixturedense(Float64, 20)))
     for W in (adjoint,)
-        @test mulcount(() -> W(S) * Dr) == nnz(S)
-        @test mulcount(() -> Dl * W(S)) == nnz(S)
+        @test mulcount(() -> W(S) * Dr) == nnz(S) + 20
+        @test mulcount(() -> Dl * W(S)) == nnz(S) + 30
         @static if COMPREHENSIVE
         C = similar(W(S))
-        @test mulcount(() -> mul!(C, W(S), Dr)) == nnz(S)
-        @test mulcount(() -> mul!(C, Dl, W(S))) == nnz(S)
+        @test mulcount(() -> mul!(C, W(S), Dr)) == nnz(S) + 20
+        @test mulcount(() -> mul!(C, Dl, W(S))) == nnz(S) + 30
         end
     end
     @static if COMPREHENSIVE
@@ -573,6 +580,15 @@ end
     @test mismatch(mul!(sC, sA, 0.5), 0.5 * dA) === nothing
     @test mismatch(lmul!(0.5, copy(sA)), 0.5 * dA) === nothing
     @test mul!(sC, 0.5, sA)   == mul!(sC, sA, 0.5)
+    @static if COMPREHENSIVE
+    # a product that does not convert to the eltype of `C` throws before `C` takes the pattern of `A`
+    A0 = sparse([1 0; 0 3]); C0 = sparse([2 0; 0 0])
+    for f! in (C -> mul!(C, A0, 0.5), C -> mul!(C, 0.5, A0))
+        C = copy(C0)
+        @test_throws InexactError f!(C)
+        @test mismatch(C, C0) === nothing
+    end
+    end
 
     @testset "inverse scaling with mul!" begin
         bi = inv.(b)