Faster indexing for BandedGeneralizedEigenvectors (#364)

* Faster indexing for BandedGeneralizedEigenvectors

* merge getindex methods

Author: jishnub

src/symbanded/bandedeigen.jl

@@ -29,45 +29,46 @@ function eigvals!(A::Symmetric{T,<:BandedMatrix{T}}, B::Symmetric{T,<:BandedMatr
     eigvals!(A)
 end
 
-# V = G Q
-struct BandedEigenvectors{T} <: AbstractMatrix{T}
-    G::Vector{Givens{T}}
-    Q::Matrix{T}
-    z1::Vector{T}
-end
+abstract type AbstractBandedEigenvectors{T} <: AbstractMatrix{T} end
 
-size(B::BandedEigenvectors) = size(B.Q)
-getindex(B::BandedEigenvectors, i, j) = Matrix(B)[i,j]
-function _getindex_vec(B::BandedEigenvectors{T}, j) where {T}
-    z1 = B.z1
-    z2 = OneElement(one(T), j, size(B,2))
+getindex(B::AbstractBandedEigenvectors, i, j) = Matrix(B)[i,j]
+function _getindex_vec(B, j)
+    z1 = _get_scratch(B)
+    z2 = OneElement(one(eltype(B)), j, size(B,2))
     mul!(z1, B, z2)
 end
-function getindex(B::BandedEigenvectors, i::Int, j::Int)
+function getindex(B::AbstractBandedEigenvectors, i::Union{Int, Colon, AbstractVector{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})
+function getindex(B::AbstractBandedEigenvectors, ::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 = G Q
+struct BandedEigenvectors{T} <: AbstractBandedEigenvectors{T}
+    G::Vector{Givens{T}}
+    Q::Matrix{T}
+    z1::Vector{T} # scratch space, used in indexing
+end
+
+size(B::BandedEigenvectors) = size(B.Q)
+
+_get_scratch(B::BandedEigenvectors) = B.z1
+
 # V = S⁻¹ Q W
-struct BandedGeneralizedEigenvectors{T,M<:AbstractMatrix{T}} <: AbstractMatrix{T}
+struct BandedGeneralizedEigenvectors{T,M<:AbstractMatrix{T}} <: AbstractBandedEigenvectors{T}
     S::SplitCholesky{T,M}
     Q::Vector{Givens{T}}
     W::BandedEigenvectors{T}
 end
 
 size(B::BandedGeneralizedEigenvectors) = size(B.W)
-getindex(B::BandedGeneralizedEigenvectors, i, j) = Matrix(B)[i, j]
+_get_scratch(B::BandedGeneralizedEigenvectors) = _get_scratch(B.W)
 
 convert(::Type{Eigen{T, T, Matrix{T}, Vector{T}}}, F::Eigen{T, T, BandedEigenvectors{T}, Vector{T}}) where T = Eigen(F.values, Matrix(F.vectors))
 convert(::Type{GeneralizedEigen{T, T, Matrix{T}, Vector{T}}}, F::GeneralizedEigen{T, T, BandedGeneralizedEigenvectors{T}, Vector{T}}) where T = GeneralizedEigen(F.values, Matrix(F.vectors))

test/test_symbanded.jl

@@ -143,15 +143,25 @@ using Test
         @test V'AM*V ≈ Diagonal(Λ)
         @test V'Matrix(B)*V ≈ I
 
-        # misc mul/div tests
         VM = Matrix(V)
-        @test AM * V ≈ AM * VM
-        @test V * AM ≈ VM * AM
-        x = rand(T, size(V,2))
-        @test ldiv!(similar(x), V, x) ≈ ldiv!(similar(x), factorize(VM), x)
 
-        z = OneElement{T}(4, size(V,2))
-        @test V * z ≈ V * Vector(z)
+        @testset "indexing" begin
+            @test V[axes(V,1),1:3] ≈ VM[axes(V,1),1:3]
+            @test V[:,1:3] ≈ VM[:,1:3]
+            @test V[:,3] ≈ VM[:,3]
+            @test V[axes(V,1),3] ≈ VM[axes(V,1),3]
+            @test V[1,2] ≈ VM[1,2]
+        end
+
+        @testset "mul/div" begin
+            @test AM * V ≈ AM * VM
+            @test V * AM ≈ VM * AM
+            x = rand(T, size(V,2))
+            @test ldiv!(similar(x), V, x) ≈ ldiv!(similar(x), factorize(VM), x)
+
+            z = OneElement{T}(4, size(V,2))
+            @test V * z ≈ V * Vector(z)
+        end
     end
 
     @testset "eigen with mismatched parent bandwidths" begin