Faster getindex for BandedEigenvectors (#354)

* faster getindex for BandedEigenvectors with two caches

* revert to using Array

* Add getindex test

* restrict getindex to Int

* getindex with colon

Author: jishnub

src/symbanded/bandedeigen.jl

@@ -2,10 +2,33 @@
 struct BandedEigenvectors{T} <: AbstractMatrix{T}
     G::Vector{Givens{T}}
     Q::Matrix{T}
+    z1::Vector{T}
+    z2::Vector{T}
 end
 
 size(B::BandedEigenvectors) = size(B.Q)
-getindex(B::BandedEigenvectors, i, j) = Matrix(B)[i, j]
+getindex(B::BandedEigenvectors, i, j) = Matrix(B)[i,j]
+function _getindex_vec(B::BandedEigenvectors{T}, j) where {T}
+    z1, z2 = B.z1, B.z2
+    z2 .= zero(T)
+    z2[j] = oneunit(T)
+    mul!(z1, B, z2)
+end
+function getindex(B::BandedEigenvectors, i::Int, j::Int)
+    z = _getindex_vec(B, j)
+    z[i]
+end
+function getindex(B::BandedEigenvectors, ::Colon, j::Int)
+    z = _getindex_vec(B, j)
+    copy(z)
+end
+function getindex(B::BandedEigenvectors, ::Colon, jr::AbstractVector{<:Int})
+    M = similar(B, size(B,1), length(jr))
+    for (ind, j) in enumerate(jr)
+        M[:, ind] = _getindex_vec(B, j)
+    end
+    return M
+end
 
 # V = S⁻¹ Q W
 struct BandedGeneralizedEigenvectors{T,M<:AbstractMatrix{T}} <: AbstractMatrix{T}
@@ -39,7 +62,7 @@ function eigen!(A::Symmetric{T,<:BandedMatrix{T}}) where T <: Real
     AB = symbandeddata(A)
     sbtrd!('V', A.uplo, N, KD, AB, D, E, G, WORK)
     Λ, Q = eigen(SymTridiagonal(D, E))
-    Eigen(Λ, BandedEigenvectors(G, Q))
+    Eigen(Λ, BandedEigenvectors(G, Q, similar(Q, size(Q,1)), similar(Q, size(Q,2))))
 end
 
 function eigen!(A::Symmetric{T,<:BandedMatrix{T}}, B::Symmetric{T,<:BandedMatrix{T}}) where T <: Real

test/test_symbanded.jl

@@ -73,6 +73,11 @@ import BandedMatrices: MemoryLayout, SymmetricLayout, HermitianLayout, BandedCol
     FD = convert(Eigen{Float64, Float64, Matrix{Float64}, Vector{Float64}}, F)
     @test FD.vectors'Matrix(A)*FD.vectors ≈ Diagonal(F.values)
 
+    MQ = Matrix(Q)
+    @test Q[axes(Q,1),1:3] ≈ MQ[axes(Q,1),1:3]
+    @test Q[:,1:3] ≈ MQ[:,1:3]
+    @test Q[:,3] ≈ MQ[:,3]
+    @test Q[1,2] ≈ MQ[1,2]
 
     function An(::Type{T}, N::Int) where {T}
         A = Symmetric(BandedMatrix(Zeros{T}(N,N), (0, 2)))