add findstructralnz for (bi/tri-)diagonal matrices

Author: huanglangwen

src/ArrayInterface.jl

@@ -1,6 +1,10 @@
 module ArrayInterface
 
 using Requires
+using LinearAlgebra
+using SparseArrays
+
+export findstructralnz
 
 function ismutable end
 
@@ -16,6 +20,69 @@ ismutable(x) = ismutable(typeof(x))
 ismutable(::Type{<:Array}) = true
 ismutable(::Type{<:Number}) = false
 
+"""
+    findstructralnz(x::AbstractArray)
+
+Return: (I,J) #indexable objects
+Find sparsity pattern of special matrices, similar to first two elements of findnz(::SparseCSCMatrix)
+"""
+function findstructralnz(x::Diagonal)
+  n=size(x,1)
+  (1:n,1:n)
+end
+
+abstract type MatrixIndex end
+
+struct BidiagonalIndex <: MatrixIndex
+  count::Int
+  isup::Bool
+end
+
+struct TridiagonalIndex <: MatrixIndex
+  count::Int
+  nsize::Int
+  isrow::Bool
+end
+
+Base.firstindex(ind::MatrixIndex)=1
+Base.lastindex(ind::MatrixIndex)=ind.count
+function Base.getindex(ind::BidiagonalIndex,i::Int)
+  1 <= i <= ind.count || throw(BoundsError(ind, i))
+  if ind.isup
+    ii=i+1
+  else
+    ii=i+1+1
+  end
+  convert(Int,floor(ii/2))
+end
+
+function Base.getindex(ind::TridiagonalIndex,i::Int)
+  1 <= i <= ind.count || throw(BoundsError(ind, i))
+  offsetu= ind.isrow ? 0 : 1
+  offsetl= ind.isrow ? 1 : 0
+  if 1 <= i <= ind.nsize
+    return i
+  elseif ind.nsize < i <= ind.nsize+ind.nsize-1
+    return i-ind.nsize+offsetu
+  else
+    return i-(ind.nsize+ind.nsize-1)+offsetl
+  end
+end
+
+function findstructralnz(x::Bidiagonal)
+  n=size(x,1)
+  isup= x.uplo=='U' ? true : false
+  rowind=BidiagonalIndex(n+n-1,isup)
+  colind=BidiagonalIndex(n+n-1,!isup)
+  (rowind,colind)
+end
+
+function findstructralnz(x::Union{Tridiagonal,SymTridiagonal})
+  n=size(x,1)
+  rowind=TridiagonalIndex(n+n-1+n-1,n,true)
+  colind=TridiagonalIndex(n+n-1+n-1,n,false)
+  (rowind,colind)
+end
 
 function __init__()
 

test/runtests.jl

@@ -3,5 +3,24 @@ using ArrayInterface, Test
 @test ArrayInterface.ismutable(rand(3))
 
 using StaticArrays
-ArrayInterface.ismutable(@SVector [1,2,3]) == false
-ArrayInterface.ismutable(@MVector [1,2,3]) == true
+@test ArrayInterface.ismutable(@SVector [1,2,3]) == false
+@test ArrayInterface.ismutable(@MVector [1,2,3]) == true
+
+using LinearAlgebra
+D=Diagonal([1,2,3,4])
+rowind,colind=findstructralnz(D)
+@test [D[rowind[i],colind[i]] for i in 1:4]==[1,2,3,4]
+
+Bu = Bidiagonal([1,2,3,4], [7,8,9], :U)
+rowind,colind=findstructralnz(Bu)
+@test [Bu[rowind[i],colind[i]] for i in 1:7]==[1,7,2,8,3,9,4]
+Bl = Bidiagonal([1,2,3,4], [7,8,9], :L)
+rowind,colind=findstructralnz(Bl)
+@test [Bl[rowind[i],colind[i]] for i in 1:7]==[1,7,2,8,3,9,4]
+
+Tri=Tridiagonal([1,2,3],[1,2,3,4],[4,5,6])
+rowind,colind=findstructralnz(Tri)
+@test [Tri[rowind[i],colind[i]] for i in 1:10]==[1,2,3,4,4,5,6,1,2,3]
+STri=SymTridiagonal([1,2,3,4],[5,6,7])
+rowind,colind=findstructralnz(STri)
+@test [STri[rowind[i],colind[i]] for i in 1:10]==[1,2,3,4,5,6,7,5,6,7]