Faster istriu/istril/isdiag for AbstractToeplitz (#88)

* Faster istriu/istril/isdiag for AbstractToeplitz

* rectangular tests

Author: jishnub

src/ToeplitzMatrices.jl

@@ -9,13 +9,22 @@ import Base: AbstractMatrix
 import LinearAlgebra: Cholesky, Factorization
 import LinearAlgebra: ldiv!, factorize, lmul!, pinv, eigvals, eigvecs, eigen, Eigen, det
 import LinearAlgebra: cholesky!, cholesky, tril!, triu!, checksquare, rmul!, dot, mul!, tril, triu
+import LinearAlgebra: istriu, istril, isdiag
 import LinearAlgebra: UpperTriangular, LowerTriangular, Symmetric, Adjoint
 import AbstractFFTs: Plan, plan_fft!
 import StatsBase
 
 export AbstractToeplitz, Toeplitz, SymmetricToeplitz, Circulant, LowerTriangularToeplitz, UpperTriangularToeplitz, TriangularToeplitz, Hankel
 export durbin, trench, levinson
 
+@static if isdefined(Base, :require_one_based_indexing)
+    const require_one_based_indexing = Base.require_one_based_indexing
+else
+    function require_one_based_indexing(A...)
+        !Base.has_offset_axes(A...) || throw(ArgumentError("offset arrays are not supported but got an array with index other than 1"))
+    end
+end
+
 include("iterativeLinearSolvers.jl")
 
 # Abstract
@@ -33,6 +42,21 @@ convert(::Type{AbstractToeplitz{T}}, A::AbstractToeplitz) where T = AbstractToep
 isconcrete(A::AbstractToeplitz) = isconcretetype(typeof(A.vc)) && isconcretetype(typeof(A.vr))
 iszero(A::AbstractToeplitz) = iszero(A.vc) && iszero(A.vr)
 
+function istril(A::AbstractToeplitz, k::Integer=0)
+    vr, vc = _vr(A), _vc(A)
+    all(iszero, @view vr[max(1, k+2):end]) && all(iszero, @view vc[2:min(-k,end)])
+end
+
+function istriu(A::AbstractToeplitz, k::Integer=0)
+    vr, vc = _vr(A), _vc(A)
+    all(iszero, @view vc[max(1, -k+2):end]) && all(iszero, @view vr[2:min(k,end)])
+end
+
+function isdiag(A::AbstractToeplitz)
+    vr, vc = _vr(A), _vc(A)
+    all(iszero, @view vr[2:end]) && all(iszero, @view vc[2:end])
+end
+
 """
     ToeplitzFactorization
 

src/hankel.jl

@@ -5,6 +5,7 @@ struct Hankel{T, V<:AbstractVector{T}, S<:DimsInteger} <: AbstractMatrix{T}
 
     function Hankel{T,V,S}(v::V, (m,n)::DimsInteger) where {T, V<:AbstractVector{T}, S<:DimsInteger}
         (m < 0 || n < 0) && throw(ArgumentError("negative size: $s"))
+        require_one_based_indexing(v)
         length(v) ≥ m+n-1 || throw(ArgumentError("inconsistency between size and number of anti-diagonals"))
         new{T,V,S}(v, (m,n))
     end

src/special.jl

@@ -45,11 +45,16 @@ for TYPE in (:SymmetricToeplitz, :Circulant, :LowerTriangularToeplitz, :UpperTri
     @eval begin
         struct $TYPE{T, V<:AbstractVector{T}} <: AbstractToeplitzSingleVector{T}
             v::V
+            function $TYPE{T,V}(v::V) where {T,V<:AbstractVector{T}}
+                require_one_based_indexing(v)
+                new{T,V}(v)
+            end
         end
         function $TYPE{T}(v::AbstractVector) where T
             vT = convert(AbstractVector{T},v)
             $TYPE{T, typeof(vT)}(vT)
         end
+        $TYPE(v::V) where {T,V<:AbstractVector{T}} = $TYPE{T,V}(v)
 
         basetype(::Type{T}) where {T<:$TYPE} = $TYPE
 
@@ -111,8 +116,12 @@ function getproperty(A::UpperTriangularToeplitz, s::Symbol)
     end
 end
 
-_circulate(v::AbstractVector) = vcat(v[1],v[end:-1:2])
-_firstnonzero(v::AbstractVector) = vcat(v[1],zero(view(v,2:lastindex(v))))
+_circulate(v::AbstractVector) = reverse(v, 2)
+function _firstnonzero(v::AbstractVector)
+    w = zero(v)
+    w[1] = v[1]
+    w
+end
 
 # transpose
 transpose(A::SymmetricToeplitz) = A
@@ -242,3 +251,6 @@ function _trisame!(A::TriangularToeplitz, k::Integer)
 end
 tril!(A::LowerTriangularToeplitz, k::Integer) = _trisame!(A,k)
 triu!(A::UpperTriangularToeplitz, k::Integer) = _trisame!(A,-k)
+
+isdiag(A::Union{Circulant, LowerTriangularToeplitz, SymmetricToeplitz}) = all(iszero, @view _vc(A)[2:end])
+isdiag(A::UpperTriangularToeplitz) = all(iszero, @view _vr(A)[2:end])

src/toeplitz.jl

@@ -9,6 +9,7 @@ struct Toeplitz{T, VC<:AbstractVector{T}, VR<:AbstractVector{T}} <: AbstractToep
     vr::VR
 
     function Toeplitz{T, VC, VR}(vc::VC, vr::VR) where {T, VC<:AbstractVector{T}, VR<:AbstractVector{T}}
+        require_one_based_indexing(vr, vc)
         if first(vc) != first(vr)
             error("First element of the vectors must be the same")
         end

test/runtests.jl

@@ -347,6 +347,9 @@ end
         @test TA+TB == A+B
         @test TA-TB == A-B
 
+        @test all(k -> istril(TA, k) == istril(A, k), -5:5)
+        @test all(k -> istriu(TA, k) == istriu(A, k), -5:5)
+
         @test_throws ArgumentError reverse(TA,dims=3)
         if isa(TA,AbstractToeplitz)
             @test isa(reverse(TA),Hankel)
@@ -372,6 +375,28 @@ end
     end
     @test fill!(Toeplitz(zeros(2,2)),1) == ones(2,2)
 
+    @testset "istril/istriu/isdiag" begin
+        for (vc,vr) in (([1,2,0,0], [1,4,5,0]), ([0,0,0], [0,5,0]), ([3,0,0], [3,0,0]), ([0], [0]))
+            for T in (Toeplitz(vc, vr), Circulant(vr),
+                            SymmetricToeplitz(vc), LowerTriangularToeplitz(vc),
+                            UpperTriangularToeplitz(vr))
+                M = Matrix(T)
+                for k in -5:5, f in [istriu, istril]
+                    @test f(T, k) == f(M, k)
+                end
+                @test isdiag(T) == isdiag(M)
+            end
+        end
+
+        for (vr, vc) in (([1,2], [1,2,3,4]), ([1,2,3,4], [1,2]))
+            T = Toeplitz(vr, vc)
+            M = Matrix(T)
+            @testset for f in (istril, istriu)
+                @test all(k -> f(T,k) == f(M,k), -5:5)
+            end
+        end
+    end
+
     @testset "aliasing" begin
         v = [1,2,3]
         T = Toeplitz(v, v)