Loop over diagind in diag for banded matrices

Author: jishnub

src/bidiag.jl

@@ -405,14 +405,15 @@ end
 function diag(M::Bidiagonal, n::Integer=0)
     # every branch call similar(..., ::Int) to make sure the
     # same vector type is returned independent of n
-    v = similar(M.dv, max(0, length(M.dv)-abs(n)))
+    dinds = diagind(M, n, IndexStyle(M))
+    v = similar(M.dv, length(dinds))
     if n == 0
         copyto!(v, M.dv)
     elseif (n == 1 && M.uplo == 'U') ||  (n == -1 && M.uplo == 'L')
         copyto!(v, M.ev)
     elseif -size(M,1) <= n <= size(M,1)
-        for i in eachindex(v)
-            v[i] = M[BandIndex(n,i)]
+        for i in eachindex(v, dinds)
+            @inbounds v[i] = M[BandIndex(n,i)]
         end
     end
     return v

src/diagonal.jl

@@ -836,12 +836,13 @@ permutedims(D::Diagonal, perm) = (Base.checkdims_perm(axes(D), axes(D), perm); D
 function diag(D::Diagonal, k::Integer=0)
     # every branch call similar(..., ::Int) to make sure the
     # same vector type is returned independent of k
-    v = similar(D.diag, max(0, length(D.diag)-abs(k)))
+    dinds = diagind(D, k, IndexStyle(D))
+    v = similar(D.diag, length(dinds))
     if k == 0
         copyto!(v, D.diag)
     else
-        for i in eachindex(v)
-            v[i] = D[BandIndex(k, i)]
+        for i in eachindex(v, dinds)
+            @inbounds v[i] = D[dinds[i]]
         end
     end
     return v

src/tridiag.jl

@@ -191,16 +191,17 @@ _eviter_transposed(M::SymTridiagonal) = (transpose(x) for x in _evview(M))
 function diag(M::SymTridiagonal, n::Integer=0)
     # every branch call similar(..., ::Int) to make sure the
     # same vector type is returned independent of n
-    v = similar(M.dv, max(0, length(M.dv)-abs(n)))
+    dinds = diagind(M, n, IndexStyle(M))
+    v = similar(M.dv, length(dinds))
     if n == 0
         return copyto!(v, _diagiter(M))
     elseif n == 1
         return copyto!(v, _evview(M))
     elseif n == -1
         return copyto!(v, _eviter_transposed(M))
     else
-        for i in eachindex(v)
-            v[i] = M[BandIndex(n,i)]
+        for i in eachindex(v, dinds)
+            v[i] = M[dinds[i]]
         end
     end
     return v
@@ -662,15 +663,16 @@ issymmetric(S::Tridiagonal) = all(issymmetric, S.d) && all(Iterators.map((x, y)
 function diag(M::Tridiagonal, n::Integer=0)
     # every branch call similar(..., ::Int) to make sure the
     # same vector type is returned independent of n
-    v = similar(M.d, max(0, length(M.d)-abs(n)))
+    dinds = diagind(M, n, IndexStyle(M))
+    v = similar(M.d, length(dinds))
     if n == 0
         copyto!(v, M.d)
     elseif n == -1
         copyto!(v, M.dl)
     elseif n == 1
         copyto!(v, M.du)
     elseif abs(n) <= size(M,1)
-        for i in eachindex(v)
+        for i in eachindex(v, dinds)
             v[i] = M[BandIndex(n,i)]
         end
     end