Make fillband! public and specialize for banded matrices (#1345)

In parallel with `fillstored!`, this allows one to fill a structured
matrix. In addition, this is a useful function in structured matrix
multiplication operations where only certain bands need to be populated.
It would be good to have this public, so that sparse matrix types may
add methods.

Author: jishnub

src/LinearAlgebra.jl

@@ -180,7 +180,8 @@ public AbstractTriangular,
         symmetric_type,
         zeroslike,
         matprod_dest,
-        fillstored!
+        fillstored!,
+        fillband!
 
 const BlasFloat = Union{Float64,Float32,ComplexF64,ComplexF32}
 const BlasReal = Union{Float64,Float32}

src/adjtrans.jl

@@ -576,3 +576,8 @@ diagview(A::Adjoint, k::Integer = 0) = _vecadjoint(diagview(parent(A), -k))
 # triu and tril
 triu!(A::AdjOrTransAbsMat, k::Integer = 0) = wrapperop(A)(tril!(parent(A), -k))
 tril!(A::AdjOrTransAbsMat, k::Integer = 0) = wrapperop(A)(triu!(parent(A), -k))
+
+function fillband!(A::AdjOrTrans, v, k1, k2)
+    fillband!(parent(A), wrapperop(A)(v), -k2, -k1)
+    return A
+end

src/bidiag.jl

@@ -1543,3 +1543,30 @@ function Base._sum(A::Bidiagonal, dims::Integer)
     end
     res
 end
+
+function fillband!(B::Bidiagonal, x, l, u)
+    if l > u
+        return B
+    end
+    if ((B.uplo == 'U' && (l < 0 || u > 1)) ||
+            (B.uplo == 'L' && (l < -1 || u > 0))) && !iszero(x)
+        throw_fillband_error(l, u, x)
+    else
+        if B.uplo == 'U'
+            if l <= 1 <= u
+                fill!(B.ev, x)
+            end
+            if l <= 0 <= u
+                fill!(B.dv, x)
+            end
+        else # B.uplo == 'L'
+            if l <= 0 <= u
+                fill!(B.dv, x)
+            end
+            if l <= -1 <= u
+                fill!(B.ev, x)
+            end
+        end
+    end
+    return B
+end

src/dense.jl

@@ -205,6 +205,23 @@ tril(M::Matrix, k::Integer) = tril!(copy(M), k)
     fillband!(A::AbstractMatrix, x, l, u)
 
 Fill the band between diagonals `l` and `u` with the value `x`.
+
+# Examples
+```jldoctest
+julia> A = zeros(4,4)
+4×4 Matrix{Float64}:
+ 0.0  0.0  0.0  0.0
+ 0.0  0.0  0.0  0.0
+ 0.0  0.0  0.0  0.0
+ 0.0  0.0  0.0  0.0
+
+julia> LinearAlgebra.fillband!(A, 2, 0, 1)
+4×4 Matrix{Float64}:
+ 2.0  2.0  0.0  0.0
+ 0.0  2.0  2.0  0.0
+ 0.0  0.0  2.0  2.0
+ 0.0  0.0  0.0  2.0
+```
 """
 function fillband!(A::AbstractMatrix{T}, x, l, u) where T
     require_one_based_indexing(A)

src/diagonal.jl

@@ -1216,3 +1216,18 @@ end
 
 uppertriangular(D::Diagonal) = D
 lowertriangular(D::Diagonal) = D
+
+throw_fillband_error(l, u, x) = throw(ArgumentError(lazy"cannot set bands $l:$u to a nonzero value ($x)"))
+
+function fillband!(D::Diagonal, x, l, u)
+    if l > u
+        return D
+    end
+    if (l < 0 || u > 0) && !iszero(x)
+        throw_fillband_error(l, u, x)
+    end
+    if l <= 0 <= u
+        fill!(D.diag, x)
+    end
+    return D
+end

src/hessenberg.jl

@@ -124,6 +124,17 @@ lmul!(x::Number, H::UpperHessenberg) = (lmul!(x, H.data); H)
 
 fillstored!(H::UpperHessenberg, x) = (fillband!(H.data, x, -1, size(H,2)-1); H)
 
+function fillband!(H::UpperHessenberg, x, l, u)
+    if l > u
+        return H
+    end
+    if l < -1 && !iszero(x)
+        throw_fillband_error(l, u, x)
+    end
+    fillband!(H.data, x, l, u)
+    return H
+end
+
 +(A::UpperHessenberg, B::UpperHessenberg) = UpperHessenberg(A.data+B.data)
 -(A::UpperHessenberg, B::UpperHessenberg) = UpperHessenberg(A.data-B.data)
 

src/symmetric.jl

@@ -414,6 +414,19 @@ function fillstored!(A::HermOrSym{T}, x) where T
     return A
 end
 
+function fillband!(A::HermOrSym, x, l, u)
+    if isa(A, Hermitian)
+        ishermitian(x) || throw(ArgumentError("cannot fill Hermitian matrix with a non-hermitian value"))
+    elseif isa(A, Symmetric)
+        issymmetric(x) || throw(ArgumentError("cannot fill Symmetric matrix with an asymmetric value"))
+    end
+    l == -u || throw(ArgumentError(lazy"lower and upper bands must be equal in magnitude and opposite in sign, got l=$(l), u=$(u)"))
+    lp = A.uplo == 'U' ? 0 : l
+    up = A.uplo == 'U' ? u : 0
+    applytri(A -> fillband!(A, x, lp, up), A)
+    return A
+end
+
 Base.isreal(A::HermOrSym{<:Real}) = true
 function Base.isreal(A::HermOrSym)
     n = size(A, 1)

src/triangular.jl

@@ -931,6 +931,36 @@ fillstored!(A::UnitLowerTriangular, x) = (fillband!(A.data, x, 1-size(A,1), -1);
 fillstored!(A::UpperTriangular, x)     = (fillband!(A.data, x, 0, size(A,2)-1); A)
 fillstored!(A::UnitUpperTriangular, x) = (fillband!(A.data, x, 1, size(A,2)-1); A)
 
+function fillband!(A::LowerOrUnitLowerTriangular, x, l, u)
+    if l > u
+        return A
+    end
+    if u > 0 && !iszero(x)
+        throw_fillband_error(l, u, x)
+    end
+    isunit = A isa UnitLowerTriangular
+    if isunit && u >= 0 && x != oneunit(x)
+        throw(ArgumentError(lazy"cannot set the diagonal band to a non-unit value ($x)"))
+    end
+    fillband!(A.data, x, l, min(u, -isunit))
+    return A
+end
+
+function fillband!(A::UpperOrUnitUpperTriangular, x, l, u)
+    if l > u
+        return A
+    end
+    if l < 0 && !iszero(x)
+        throw_fillband_error(l, u, x)
+    end
+    isunit = A isa UnitUpperTriangular
+    if isunit && l <= 0 && x != oneunit(x)
+        throw(ArgumentError(lazy"cannot set the diagonal band to a non-unit value ($x)"))
+    end
+    fillband!(A.data, x, max(l, isunit), u)
+    return A
+end
+
 # Binary operations
 # use broadcasting if the parents are strided, where we loop only over the triangular part
 function +(A::UpperTriangular, B::UpperTriangular)

src/tridiag.jl

@@ -1189,3 +1189,44 @@ function _opnorm1Inf(A::SymTridiagonal, p::Real)
                 ),
                 normfirst, normend)
 end
+
+function fillband!(T::Tridiagonal, x, l, u)
+    if l > u
+        return T
+    end
+    if (l < -1 || u > 1) && !iszero(x)
+        throw_fillband_error(l, u, x)
+    else
+        if l <= -1 <= u
+            fill!(T.dl, x)
+        end
+        if l <= 0 <= u
+            fill!(T.d, x)
+        end
+        if l <= 1 <= u
+            fill!(T.du, x)
+        end
+    end
+    return T
+end
+
+function fillband!(T::SymTridiagonal, x, l, u)
+    if l > u
+        return T
+    end
+    if (l <= 1 <= u) != (l <= -1 <= u)
+        throw(ArgumentError(lazy"cannot set only one off-diagonal band of a SymTridiagonal"))
+    elseif (l < -1 || u > 1) && !iszero(x)
+        throw_fillband_error(l, u, x)
+    elseif l <= 0 <= u && !issymmetric(x)
+        throw(ArgumentError(lazy"cannot set entries in the diagonal band of a SymTridiagonal to an asymmetric value $x"))
+    else
+        if l <= 0 <= u
+            fill!(T.dv, x)
+        end
+        if l <= 1 <= u
+            fill!(T.ev, x)
+        end
+    end
+    return T
+end

test/adjtrans.jl

@@ -809,4 +809,16 @@ end
     end
 end
 
+@testset "fillband!" begin
+    for A in (rand(4, 4), rand(ComplexF64,4,4))
+        B = similar(A)
+        for op in (adjoint, transpose), k in -3:3
+            B .= op(A)
+            LinearAlgebra.fillband!(op(A), 1, k, k)
+            LinearAlgebra.fillband!(B, 1, k, k)
+            @test op(A) == B
+        end
+    end
+end
+
 end # module TestAdjointTranspose