Generalize istriu/istril to accept a band index

Author: jishnub

src/sparsematrix.jl

@@ -4141,7 +4141,7 @@ function is_hermsym(A::AbstractSparseMatrixCSC, check::Function)
     return true
 end
 
-function istriu(A::AbstractSparseMatrixCSC)
+function istriu(A::AbstractSparseMatrixCSC, k::Integer=0)
     m, n = size(A)
     colptr = getcolptr(A)
     rowval = rowvals(A)
@@ -4150,7 +4150,8 @@ function istriu(A::AbstractSparseMatrixCSC)
     for col = 1:min(n, m-1)
         l1 = colptr[col+1]-1
         for i = 0 : (l1 - colptr[col])
-            if rowval[l1-i] <= col
+            if rowval[l1-i] <= col - k
+                # rows preceeding the index would also lie above the band
                 break
             end
             if _isnotzero(nzval[l1-i])
@@ -4161,15 +4162,16 @@ function istriu(A::AbstractSparseMatrixCSC)
     return true
 end
 
-function istril(A::AbstractSparseMatrixCSC)
+function istril(A::AbstractSparseMatrixCSC, k::Integer=0)
     m, n = size(A)
     colptr = getcolptr(A)
     rowval = rowvals(A)
     nzval  = nonzeros(A)
 
     for col = 2:n
         for i = colptr[col] : (colptr[col+1]-1)
-            if rowval[i] >= col
+            if rowval[i] >= col - k
+                # subsequent rows would also lie below the band
                 break
             end
             if _isnotzero(nzval[i])

test/sparsematrix_ops.jl

@@ -626,4 +626,17 @@ end
     @test_throws ArgumentError copytrito!(M, S, 'M')
 end
 
+@testset "istriu/istril" begin
+    for T in Any[Tridiagonal(1:3, 1:4, 1:3),
+                    Bidiagonal(1:4, 1:3, :U), Bidiagonal(1:4, 1:3, :L),
+                    Diagonal(1:4),
+                    diagm(-2=>1:2, 2=>1:2)]
+        S = sparse(T)
+        for k in -5:5
+            @test istriu(S, k) == istriu(T, k)
+            @test istril(S, k) == istril(T, k)
+        end
+    end
+end
+
 end # module