add diag (#86)

* add `diag`

* Update README.md

* Update Project.toml

* Update README.md

* support diag(A,k)

* fix

* fix

* add tests

* fix and add tests

* add test and fix rectangular cases

* Update Project.toml

---------

Co-authored-by: Jishnu Bhattacharya <jishnub.github@gmail.com>

README.md

@@ -65,6 +65,10 @@ The `reverse` can transform between Hankel and Toeplitz. It is used to achieve f
 ## Implemented interface
 
 ### Operations
+
+<details>
+  <summary>Full list</summary>
+
 - ✓ implemented
 - ✗ error
 - _ fall back to `Matrix`
@@ -97,10 +101,13 @@ The `reverse` can transform between Hankel and Toeplitz. It is used to achieve f
 |istril|||||||
 |iszero|✓|✓|✓|✓|✓||
 |isone|||||||
+|diag|✓|✓|✓|✓|✓|✓|
 |copyto!|✓|✓|✓|✓|✓|✓|
 |reverse|✓|✓|✓|✓|✓|✓|
 |broadcast|||||||
 |broadcast!|||||||
+
+</details>
 
 Note that scalar multiplication, `conj`, `+` and `-` could be removed once `broadcast` is implemented.
 

src/ToeplitzMatrices.jl

@@ -7,7 +7,7 @@ import Base: ==, +, -, *, \
 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: cholesky!, cholesky, tril!, triu!, checksquare, rmul!, dot, mul!, tril, triu, diag
 import LinearAlgebra: istriu, istril, isdiag
 import LinearAlgebra: UpperTriangular, LowerTriangular, Symmetric, Adjoint
 import LinearAlgebra: issymmetric, ishermitian
@@ -47,6 +47,15 @@ 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 diag(A::AbstractToeplitz, k::Integer=0)
+    if k >= size(A, 2) || -k >= size(A, 1)
+        Fill(zero(eltype(A)), 0)
+    elseif k >= 0
+        Fill(A.vr[k + 1], diaglenpos(size(A)..., k))
+    else
+        Fill(A.vc[-k + 1], diaglenneg(size(A)..., k))
+    end
+end
 
 function istril(A::AbstractToeplitz, k::Integer=0)
     vr, vc = _vr(A), _vc(A)
@@ -88,6 +97,10 @@ Return real-valued part of `x` if `T` is a type of a real number, and `x` otherw
 maybereal(::Type, x) = x
 maybereal(::Type{<:Real}, x) = real(x)
 
+# length of LinearAlgebra.diagind, for positive and negative k respectively
+@inline diaglenpos(m::Integer, n::Integer, k::Integer=0) = min(m, n-k)
+@inline diaglenneg(m::Integer, n::Integer, k::Integer=0) = min(m+k, n)
+
 include("directLinearSolvers.jl")
 
 if !isdefined(Base, :get_extension)

src/hankel.jl

@@ -63,6 +63,8 @@ convert(::Type{Hankel{T}}, A::Hankel) where {T<:Number} = Hankel{T}(convert(Abst
 
 # Size
 size(H::Hankel) = H.size
+diag(H::Hankel, k::Integer=0) = H.v[hankeldiagind(size(H)..., k)]
+@inline hankeldiagind(m::Integer, n::Integer, k::Integer=0) = k ≥ 0 ? range(1 + k, step = 2, length = diaglenpos(m, n, k)) : range(1 - k, step = 2, length = diaglenneg(m, n, k))
 
 # Retrieve an entry by two indices
 Base.@propagate_inbounds function getindex(A::Hankel, i::Integer, j::Integer)

test/runtests.jl

@@ -1,6 +1,6 @@
 using Pkg
 
-using ToeplitzMatrices, Test, LinearAlgebra, Aqua, Random
+using ToeplitzMatrices, Test, LinearAlgebra, Aqua, FillArrays, Random
 import StatsBase
 using FillArrays
 using FFTW: fft
@@ -377,6 +377,22 @@ end
         T=copy(TA)
     end
     @test fill!(Toeplitz(zeros(2,2)),1) == ones(2,2)
+
+	@testset "diag" begin
+		H = Hankel(1:11, 4, 8)
+		@test diag(H) ≡ 1:2:7
+		@test diag(H, 1) ≡ 2:2:8
+		@test diag(H, -1) ≡ 2:2:6
+		@test diag(H, 5) ≡ 6:2:10
+		@test diag(H, 100) == diag(H, -100) == []
+
+		T = Toeplitz(1:4, 1:8)
+		@test diag(T) ≡ Fill(1, 4)
+		@test diag(T, 1) ≡ Fill(2, 4)
+		@test diag(T, -1) ≡ Fill(2, 3)
+		@test diag(T, 5) ≡ Fill(6, 3)
+		@test diag(T, 100) == diag(T, -100) == []
+	end
 
     @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]))