Make fillband! public and specialize for banded matrices

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 @@
                 ),
                 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

test/bidiag.jl

@@ -1246,4 +1246,56 @@ end
     end
 end
 
+@testset "fillband!" begin
+    @testset "uplo = :U" begin
+        B = Bidiagonal(zeros(4), zeros(3), :U)
+        LinearAlgebra.fillband!(B, 2, 1, 1)
+        @test all(==(2), diagview(B,1))
+        LinearAlgebra.fillband!(B, 3, 0, 0)
+        @test all(==(3), diagview(B,0))
+        @test all(==(2), diagview(B,1))
+        LinearAlgebra.fillband!(B, 4, 0, 1)
+        @test all(==(4), diagview(B,0))
+        @test all(==(4), diagview(B,1))
+        @test_throws ArgumentError LinearAlgebra.fillband!(B, 3, -1, 0)
+
+        LinearAlgebra.fillstored!(B, 1)
+        LinearAlgebra.fillband!(B, 0, -3, 3)
+        @test iszero(B)
+        LinearAlgebra.fillband!(B, 0, -10, 10)
+        @test iszero(B)
+        LinearAlgebra.fillstored!(B, 1)
+        B2 = copy(B)
+        LinearAlgebra.fillband!(B, 0, -1, -3)
+        @test B == B2
+        LinearAlgebra.fillband!(B, 0, 10, 10)
+        @test B == B2
+    end
+
+    @testset "uplo = :L" begin
+        B = Bidiagonal(zeros(4), zeros(3), :L)
+        LinearAlgebra.fillband!(B, 2, -1, -1)
+        @test all(==(2), diagview(B,-1))
+        LinearAlgebra.fillband!(B, 3, 0, 0)
+        @test all(==(3), diagview(B,0))
+        @test all(==(2), diagview(B,-1))
+        LinearAlgebra.fillband!(B, 4, -1, 0)
+        @test all(==(4), diagview(B,0))
+        @test all(==(4), diagview(B,-1))
+        @test_throws ArgumentError LinearAlgebra.fillband!(B, 3, 0, 1)
+
+        LinearAlgebra.fillstored!(B, 1)
+        LinearAlgebra.fillband!(B, 0, -3, 3)
+        @test iszero(B)
+        LinearAlgebra.fillband!(B, 0, -10, 10)
+        @test iszero(B)
+        LinearAlgebra.fillstored!(B, 1)
+        B2 = copy(B)
+        LinearAlgebra.fillband!(B, 0, -1, -3)
+        @test B == B2
+        LinearAlgebra.fillband!(B, 0, 10, 10)
+        @test B == B2
+    end
+end
+
 end # module TestBidiagonal

test/diagonal.jl

@@ -1542,4 +1542,26 @@ end
     end
 end
 
+@testset "fillband!" begin
+    D = Diagonal(zeros(4))
+    LinearAlgebra.fillband!(D, 2, 0, 0)
+    @test all(==(2), diagview(D,0))
+    @test all(==(0), diagview(D,-1))
+    @test_throws ArgumentError LinearAlgebra.fillband!(D, 3, -2, 2)
+
+    LinearAlgebra.fillstored!(D, 1)
+    LinearAlgebra.fillband!(D, 0, -3, 3)
+    @test iszero(D)
+    LinearAlgebra.fillstored!(D, 1)
+    LinearAlgebra.fillband!(D, 0, -10, 10)
+    @test iszero(D)
+
+    LinearAlgebra.fillstored!(D, 1)
+    D2 = copy(D)
+    LinearAlgebra.fillband!(D, 0, -1, -3)
+    @test D == D2
+    LinearAlgebra.fillband!(D, 0, 10, 10)
+    @test D == D2
+end
+
 end # module TestDiagonal

test/hessenberg.jl

@@ -300,4 +300,24 @@ end
     @test_throws DimensionMismatch hessenberg(zeros(0,0)).Q * ones(1, 2)
 end
 
+@testset "fillband" begin
+    U = UpperHessenberg(zeros(4,4))
+    @test_throws ArgumentError LinearAlgebra.fillband!(U, 1, -2, 1)
+    @test iszero(U)
+
+    LinearAlgebra.fillband!(U, 10, -1, 2)
+    @test all(==(10), diagview(U,-1))
+    @test all(==(10), diagview(U,2))
+    @test all(==(0), diagview(U,3))
+
+    LinearAlgebra.fillband!(U, 0, -5, 5)
+    @test iszero(U)
+
+    U2 = copy(U)
+    LinearAlgebra.fillband!(U, -10, 1, -2)
+    @test U == U2
+    LinearAlgebra.fillband!(U, -10, 10, 10)
+    @test U == U2
+end
+
 end # module TestHessenberg

test/symmetric.jl

@@ -1257,4 +1257,24 @@ end
     end
 end
 
+@testset "fillband!" begin
+    A = zeros(4,4)
+    @testset for T in (Symmetric, Hermitian), uplo in (:U, :L)
+        A .= 0
+        S = T(A, uplo)
+        LinearAlgebra.fillband!(S, 2, -2, 2)
+        @test all(all(==(2), diagview(S, k)) for k in -2:2)
+        @test iszero(diagview(S, -3))
+        @test iszero(diagview(S, 3))
+        LinearAlgebra.fillband!(S, 4, -1, 1)
+        @test all(all(==(4), diagview(S, k)) for k in -1:1)
+        @test all(==(2), diagview(S, -2))
+        @test all(==(2), diagview(S, 2))
+    end
+    msg = "cannot fill Hermitian matrix with a non-hermitian value"
+    @test_throws msg LinearAlgebra.fillband!(Hermitian(A), 2im, -3, 3)
+    msg = "lower and upper bands must be equal in magnitude and opposite in sign, got l=0, u=1"
+    @test_throws msg LinearAlgebra.fillband!(Symmetric(A), 2, 0, 1)
+end
+
 end # module TestSymmetric