Band indexing for adj/trans

Author: jishnub

src/adjtrans.jl

@@ -352,6 +352,12 @@ IndexStyle(::Type{<:AdjOrTransAbsVec}) = IndexLinear()
 @propagate_inbounds getindex(v::AdjOrTransAbsVec, ::Colon, is::AbstractArray{Int}) = wrapperop(v)(v.parent[is])
 @propagate_inbounds getindex(v::AdjOrTransAbsVec, ::Colon, ::Colon) = wrapperop(v)(v.parent[:])
 
+# band indexing
+@propagate_inbounds function getindex(A::AdjOrTransAbsMat{T}, b::BandIndex) where {T}
+    require_one_based_indexing(A)
+    wrapperop(A)(A.parent[BandIndex(-b.band, b.index)])::T
+end
+
 # conversion of underlying storage
 convert(::Type{Adjoint{T,S}}, A::Adjoint) where {T,S} = Adjoint{T,S}(convert(S, A.parent))::Adjoint{T,S}
 convert(::Type{Transpose{T,S}}, A::Transpose) where {T,S} = Transpose{T,S}(convert(S, A.parent))::Transpose{T,S}

test/adjtrans.jl

@@ -749,4 +749,36 @@ end
     end
 end
 
+@testset "band indexing" begin
+    n = 3
+    A = UnitUpperTriangular(reshape(1:n^2, n, n))
+    @testset "every index" begin
+        Aadj = A'
+        for k in -(n-1):n-1
+            di = diagind(Aadj, k, IndexStyle(Aadj))
+            for (i,d) in enumerate(di)
+                @test Aadj[LinearAlgebra.BandIndex(k,i)] == Aadj[d]
+            end
+        end
+    end
+    @testset "inference for structured matrices" begin
+        v = @inferred ((A,i) -> Val(A'[LinearAlgebra.BandIndex(0,i)]))(A,1)
+        @test v == Val(1)
+    end
+    @testset "non-square matrix" begin
+        r = reshape(1:6, 2, 3)
+        for d in (r, r*im)
+            @test d'[LinearAlgebra.BandIndex(1,1)] == adjoint(d[2,1])
+            @test d'[LinearAlgebra.BandIndex(-1,2)] == adjoint(d[2,3])
+            @test transpose(d)[LinearAlgebra.BandIndex(1,1)] == transpose(d[2,1])
+            @test transpose(d)[LinearAlgebra.BandIndex(-1,2)] == transpose(d[2,3])
+        end
+    end
+    @testset "block matrix" begin
+        B = reshape([[1 2; 3 4]*i for i in 1:4], 2, 2)
+        @test B'[LinearAlgebra.BandIndex(1,1)] == adjoint(B[2,1])
+        @test transpose(B)[LinearAlgebra.BandIndex(1,1)] == transpose(B[2,1])
+    end
+end
+
 end # module TestAdjointTranspose