Band indexing for structured matrices with Fill bands (#378)

* Band indexing for structured matrices with Fill bands

* bump version to v0.17.31

* Test for infinite matrices

Project.toml

@@ -1,6 +1,6 @@
 name = "BandedMatrices"
 uuid = "aae01518-5342-5314-be14-df237901396f"
-version = "0.17.30"
+version = "0.17.31"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -13,7 +13,8 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Aqua = "0.6"
 ArrayLayouts = "1"
 Documenter = "0.27"
-FillArrays = "1.0.1"
+FillArrays = "1.3"
+InfiniteArrays = "0.12"
 PrecompileTools = "1"
 julia = "1.6"
 
@@ -22,8 +23,9 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 GenericLinearAlgebra = "14197337-ba66-59df-a3e3-ca00e7dcff7a"
+InfiniteArrays = "4858937d-0d70-526a-a4dd-2d5cb5dd786c"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Base64", "Documenter", "GenericLinearAlgebra", "Random", "Test"]
+test = ["Aqua", "Base64", "Documenter", "GenericLinearAlgebra", "InfiniteArrays", "Random", "Test"]

src/interfaceimpl.jl

@@ -67,6 +67,10 @@ function rot180(A::AbstractBandedMatrix)
     _BandedMatrix(bandeddata(A)[end:-1:1,end:-1:1], m, u+sh,l-sh)
 end
 
+for MT in (:Diagonal, :SymTridiagonal, :Tridiagonal, :Bidiagonal)
+    @eval getindex(D::$MT{T,<:AbstractFill{T,1}}, b::Band) where {T<:Number} = diag(D, b.i)
+end
+
 function getindex(D::Diagonal{T,V}, b::Band) where {T,V}
     iszero(b.i) && return copy(D.diag)
     convert(V, Zeros{T}(size(D,1)-abs(b.i)))

test/test_interface.jl

@@ -2,6 +2,7 @@ using BandedMatrices, LinearAlgebra, ArrayLayouts, FillArrays, Test, Base64
 import BandedMatrices: banded_mul!, isbanded, AbstractBandedLayout, BandedStyle,
                         rowsupport, colsupport, _BandedMatrix, BandedColumns, bandeddata
 import ArrayLayouts: OnesLayout, UnknownLayout
+using InfiniteArrays
 
 struct PseudoBandedMatrix{T} <: AbstractMatrix{T}
     data::Array{T}
@@ -79,43 +80,80 @@ LinearAlgebra.fill!(A::PseudoBandedMatrix, v) = fill!(A.data,v)
         @test A[band(0)] == [2; ones(4)]
 
         B = Diagonal(Fill(1,5))
-        @test B[band(0)] ≡ Fill(1,5)
+        @test @inferred(B[band(0)]) ≡ Fill(1,5)
         @test B[band(1)] ≡ B[band(-1)] ≡ Fill(0,4)
         @test B[band(2)] ≡ B[band(-2)] ≡ Fill(0,3)
+
+        B = Diagonal(Ones(5))
+        @test @inferred(B[band(0)]) ≡ Fill(1.0,5)
+        @test B[band(1)] ≡ B[band(-1)] ≡ Fill(0.0,4)
+        @test B[band(2)] ≡ B[band(-2)] ≡ Fill(0.0,3)
+
+        B = Diagonal(1:∞)
+        @test @inferred(B[Band(0)]) == 1:∞
     end
 
     @testset "SymTridiagonal" begin
         A = SymTridiagonal([1,2,3],[4,5])
+        @test @inferred(A[Band(0)]) == [1,2,3]
+        @test A[Band(1)] == A[Band(-1)] == [4,5]
+        @test A[Band(2)] == A[Band(-2)] == [0]
+        @test A[Band(3)] == A[Band(-3)] == Int[]
         @test isbanded(A)
         @test bandwidths(A) == (1,1)
         @test BandedMatrices.inbands_getindex(A, 1,1) == 1
         BandedMatrices.inbands_setindex!(A, 2, 1,1)
         @test A[1,1] == 2
 
+        S = SymTridiagonal(1:∞, 1:∞)
+        @test @inferred(S[Band(0)]) == S[Band(1)] == S[Band(-1)] == 1:∞
+
         B = SymTridiagonal(Fill(1,5), Fill(2,4))
-        @test B[band(0)] ≡ Fill(1,5)
+        @test @inferred(B[band(0)]) ≡ Fill(1,5)
         @test B[band(1)] ≡ B[band(-1)] ≡ Fill(2,4)
         @test B[band(2)] ≡ B[band(-2)] ≡ Fill(0,3)
+
+        B = SymTridiagonal(Ones(5), Ones(4))
+        @test @inferred(B[band(0)]) ≡ Fill(1.0,5)
+        @test B[band(1)] ≡ B[band(-1)] ≡ Fill(1.0,4)
+        @test B[band(2)] ≡ B[band(-2)] ≡ Fill(0.0,3)
     end
 
     @testset "Tridiagonal" begin
+        B = Tridiagonal([1:3;], [1:4;], [1:3;])
+        @test @inferred(B[Band(0)]) == 1:4
+        @test B[Band(1)] == B[Band(-1)] == 1:3
+        @test B[Band(2)] == B[Band(-2)] == [0,0]
+        @test B[Band(5)] == B[Band(-5)] == Int[]
+
+        T = Tridiagonal(1:∞, 1:∞, 1:∞)
+        @test @inferred(T[Band(0)]) == T[Band(1)] == T[Band(-1)] == 1:∞
+
         B = Tridiagonal(Fill(1,4), Fill(2,5), Fill(3,4))
-        @test B[band(0)] ≡ Fill(2,5)
+        @test @inferred(B[band(0)]) ≡ Fill(2,5)
         @test B[band(1)] ≡ Fill(3,4)
         @test B[band(-1)] ≡ Fill(1,4)
         @test B[band(2)] ≡ B[band(-2)] ≡ Fill(0,3)
     end
 
     @testset "Bidiagonal" begin
+        L = Bidiagonal([1:5;], [1:4;], :L)
+        @test @inferred(L[Band(0)]) == 1:5
+        @test L[Band(-1)] == 1:4
+        @test L[Band(1)] == zeros(Int,4)
+
+        L = Bidiagonal(1:∞, 1:∞, :L)
+        @test @inferred(L[Band(0)]) == L[Band(-1)] == 1:∞
+
         L = Bidiagonal(Fill(2,5), Fill(1,4), :L)
-        @test L[band(0)] ≡ Fill(2,5)
+        @test @inferred(L[band(0)]) ≡ Fill(2,5)
         @test L[band(1)] ≡ Fill(0,4)
         @test L[band(-1)] ≡ Fill(1,4)
         @test L[band(2)] ≡ L[band(-2)] ≡ Fill(0,3)
         @test BandedMatrix(L) == L
 
         U = Bidiagonal(Fill(2,5), Fill(1,4), :U)
-        @test U[band(0)] ≡ Fill(2,5)
+        @test @inferred(U[band(0)]) ≡ Fill(2,5)
         @test U[band(1)] ≡ Fill(1,4)
         @test U[band(-1)] ≡ Fill(0,4)
         @test U[band(2)] ≡ U[band(-2)] ≡ Fill(0,3)