consistent band indexing for structured matrices (#376)

* consistent band indexing for structured matrices

* Add test for Eye

* bump FillArrays compat bound

Author: jishnub

Project.toml

@@ -13,7 +13,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Aqua = "0.6"
 ArrayLayouts = "1"
 Documenter = "0.27"
-FillArrays = "1.0.1"
+FillArrays = "1.3"
 PrecompileTools = "1"
 julia = "1.6"
 

src/interfaceimpl.jl

@@ -67,27 +67,6 @@ function rot180(A::AbstractBandedMatrix)
     _BandedMatrix(bandeddata(A)[end:-1:1,end:-1:1], m, u+sh,l-sh)
 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)))
-end
-
-function getindex(D::Tridiagonal{T,V}, b::Band) where {T,V}
-    b.i == -1 && return copy(D.dl)
-    iszero(b.i) && return copy(D.d)
-    b.i == 1 && return copy(D.du)
-    convert(V, Zeros{T}(size(D,1)-abs(b.i)))
-end
-
-function getindex(D::SymTridiagonal{T,V}, b::Band) where {T,V}
-    iszero(b.i) && return copy(D.dv)
-    abs(b.i) == 1 && return copy(D.ev)
-    convert(V, Zeros{T}(size(D,1)-abs(b.i)))
-end
-
-function getindex(D::Bidiagonal{T,V}, b::Band) where {T,V}
-    iszero(b.i) && return copy(D.dv)
-    D.uplo == 'L' && b.i == -1 && return copy(D.ev)
-    D.uplo == 'U' && b.i == 1 && return copy(D.ev)
-    convert(V, Zeros{T}(size(D,1)-abs(b.i)))
+for MT in (:Diagonal, :SymTridiagonal, :Tridiagonal, :Bidiagonal)
+    @eval getindex(D::$MT, b::Band) = diag(D, b.i)
 end

test/test_interface.jl

@@ -58,6 +58,11 @@ LinearAlgebra.fill!(A::PseudoBandedMatrix, v) = fill!(A.data,v)
         @test B * Eye(5) == B
         @test muladd!(2.0, Eye(5), B, 0.0, zeros(5,5)) == 2B
         @test muladd!(2.0, B, Eye(5), 0.0, zeros(5,5)) == 2B
+
+        E = Eye(4)
+        @test (@inferred E[band(0)]) == Ones(4)
+        @test (@inferred E[band(1)]) == Zeros(3)
+        @test (@inferred E[band(-1)]) == Zeros(3)
     end
 
     @testset "Diagonal" begin
@@ -79,9 +84,9 @@ 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 B[band(1)] ≡ B[band(-1)] ≡ Fill(0,4)
-        @test B[band(2)] ≡ B[band(-2)] ≡ Fill(0,3)
+        @test (@inferred B[band(0)]) == Fill(1,5)
+        @test (@inferred B[band(1)]) == B[band(-1)] == Fill(0,4)
+        @test (@inferred B[band(2)]) == B[band(-2)] == Fill(0,3)
     end
 
     @testset "SymTridiagonal" begin
@@ -93,32 +98,32 @@ LinearAlgebra.fill!(A::PseudoBandedMatrix, v) = fill!(A.data,v)
         @test A[1,1] == 2
 
         B = SymTridiagonal(Fill(1,5), Fill(2,4))
-        @test 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)
+        @test (@inferred B[band(0)]) == Fill(1,5)
+        @test (@inferred B[band(1)]) == B[band(-1)] == Fill(2,4)
+        @test (@inferred B[band(2)]) == B[band(-2)] == Fill(0,3)
     end
 
     @testset "Tridiagonal" begin
         B = Tridiagonal(Fill(1,4), Fill(2,5), Fill(3,4))
-        @test 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)
+        @test (@inferred B[band(0)]) == Fill(2,5)
+        @test (@inferred B[band(1)]) == Fill(3,4)
+        @test (@inferred B[band(-1)]) == Fill(1,4)
+        @test B[band(2)] == B[band(-2)] == Fill(0,3)
     end
 
     @testset "Bidiagonal" begin
         L = Bidiagonal(Fill(2,5), Fill(1,4), :L)
-        @test 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 (@inferred L[band(0)]) == Fill(2,5)
+        @test (@inferred L[band(1)]) == Fill(0,4)
+        @test (@inferred L[band(-1)]) == Fill(1,4)
+        @test (@inferred 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 U[band(1)] ≡ Fill(1,4)
-        @test U[band(-1)] ≡ Fill(0,4)
-        @test U[band(2)] ≡ U[band(-2)] ≡ Fill(0,3)
+        @test (@inferred U[band(0)]) == Fill(2,5)
+        @test (@inferred U[band(1)]) == Fill(1,4)
+        @test (@inferred U[band(-1)]) == Fill(0,4)
+        @test (@inferred U[band(2)]) == U[band(-2)] == Fill(0,3)
         @test BandedMatrix(U) == U
     end