Merge pull request #14 from huanglangwen/bbbmatrix

add findstructralnz support for BandedMatrix, BlockBandedMatrix and BandedBlockBandedMatrix

Author: ChrisRackauckas

Project.toml

@@ -14,6 +14,8 @@ julia = "1"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 
 [targets]
-test = ["Test", "LabelledArrays", "StaticArrays"]
+test = ["Test", "LabelledArrays", "StaticArrays","BandedMatrices","BlockBandedMatrices"]

src/ArrayInterface.jl

@@ -104,11 +104,134 @@ struct BidiagonalIndex <: MatrixIndex
 end
 
 struct TridiagonalIndex <: MatrixIndex
-  count::Int
+  count::Int#count==nsize+nsize-1+nsize-1
   nsize::Int
   isrow::Bool
 end
 
+struct BandedMatrixIndex <: MatrixIndex
+  count::Int
+  rowsize::Int
+  colsize::Int
+  bandinds::Array{Int,1}
+  bandsizes::Array{Int,1}
+  isrow::Bool
+end
+
+function _bandsize(bandind,rowsize,colsize)
+  -(rowsize-1) <= bandind <= colsize-1 || throw(ErrorException("Invalid Bandind"))
+  if (bandind*(colsize-rowsize)>0) & (abs(bandind)<=abs(colsize-rowsize))
+    return min(rowsize,colsize)
+  elseif bandind*(colsize-rowsize)<=0
+    return min(rowsize,colsize)-abs(bandind)
+  else
+    return min(rowsize,colsize)-abs(bandind)+abs(colsize-rowsize)
+  end
+end
+
+function BandedMatrixIndex(rowsize,colsize,lowerbandwidth,upperbandwidth,isrow)
+  upperbandwidth>-lowerbandwidth || throw(ErrorException("Invalid Bandwidths"))
+  bandinds=upperbandwidth:-1:-lowerbandwidth
+  bandsizes=[_bandsize(band,rowsize,colsize) for band in bandinds]
+  BandedMatrixIndex(sum(bandsizes),rowsize,colsize,bandinds,bandsizes,isrow)
+end
+
+struct BlockBandedMatrixIndex <: MatrixIndex
+  count::Int
+  refinds::Array{Int,1}
+  refcoords::Array{Int,1}#storing col or row inds at ref points
+  isrow::Bool
+end
+
+function BlockBandedMatrixIndex(nrowblock,ncolblock,rowsizes,colsizes,l,u)
+  blockrowind=BandedMatrixIndex(nrowblock,ncolblock,l,u,true)
+  blockcolind=BandedMatrixIndex(nrowblock,ncolblock,l,u,false)
+  sortedinds=sort([(blockrowind[i],blockcolind[i]) for i in 1:length(blockrowind)],by=x->x[1])
+  sort!(sortedinds,by=x->x[2],alg=InsertionSort)#stable sort keeps the second index in order
+  refinds=Array{Int,1}()
+  refrowcoords=Array{Int,1}()
+  refcolcoords=Array{Int,1}()
+  rowheights=cumsum(pushfirst!(copy(rowsizes),1))
+  blockheight=0
+  blockrow=1
+  blockcol=1
+  currenti=1
+  lastrowind=sortedinds[1][1]-1
+  lastcolind=sortedinds[1][2]
+  for ind in sortedinds
+    rowind,colind=ind
+    if colind==lastcolind
+      if rowind>lastrowind
+        blockheight+=rowsizes[rowind]
+      end
+    else
+      for j in blockcol:blockcol+colsizes[lastcolind]-1
+        push!(refinds,currenti)
+        push!(refrowcoords,blockrow)
+        push!(refcolcoords,j)
+        currenti+=blockheight
+      end
+      blockcol+=colsizes[lastcolind]
+      blockrow=rowheights[rowind]
+      blockheight=rowsizes[rowind]
+    end
+    lastcolind=colind
+    lastrowind=rowind
+  end
+  for j in blockcol:blockcol+colsizes[lastcolind]-1
+    push!(refinds,currenti)
+    push!(refrowcoords,blockrow)
+    push!(refcolcoords,j)
+    currenti+=blockheight
+  end
+  push!(refinds,currenti)#guard
+  push!(refrowcoords,-1)
+  push!(refcolcoords,-1)
+  rowindobj=BlockBandedMatrixIndex(currenti-1,refinds,refrowcoords,true)
+  colindobj=BlockBandedMatrixIndex(currenti-1,refinds,refcolcoords,false)
+  rowindobj,colindobj
+end
+
+struct BandedBlockBandedMatrixIndex <: MatrixIndex
+  count::Int
+  refinds::Array{Int,1}
+  refcoords::Array{Int,1}#storing col or row inds at ref points
+  reflocalinds::Array{BandedMatrixIndex,1}
+  isrow::Bool
+end
+
+function BandedBlockBandedMatrixIndex(nrowblock,ncolblock,rowsizes,colsizes,l,u,lambda,mu)
+  blockrowind=BandedMatrixIndex(nrowblock,ncolblock,l,u,true)
+  blockcolind=BandedMatrixIndex(nrowblock,ncolblock,l,u,false)
+  sortedinds=sort([(blockrowind[i],blockcolind[i]) for i in 1:length(blockrowind)],by=x->x[1])
+  sort!(sortedinds,by=x->x[2],alg=InsertionSort)#stable sort keeps the second index in order
+  rowheights=cumsum(pushfirst!(copy(rowsizes),1))
+  colwidths=cumsum(pushfirst!(copy(colsizes),1))
+  currenti=1
+  refinds=Array{Int,1}()
+  refrowcoords=Array{Int,1}()
+  refcolcoords=Array{Int,1}()
+  reflocalrowinds=Array{BandedMatrixIndex,1}()
+  reflocalcolinds=Array{BandedMatrixIndex,1}()
+  for ind in sortedinds
+    rowind,colind=ind
+    localrowind=BandedMatrixIndex(rowsizes[rowind],colsizes[colind],lambda,mu,true)
+    localcolind=BandedMatrixIndex(rowsizes[rowind],colsizes[colind],lambda,mu,false)
+    push!(refinds,currenti)
+    push!(refrowcoords,rowheights[rowind])
+    push!(refcolcoords,colwidths[colind])
+    push!(reflocalrowinds,localrowind)
+    push!(reflocalcolinds,localcolind)
+    currenti+=localrowind.count
+  end
+  push!(refinds,currenti)
+  push!(refrowcoords,-1)
+  push!(refcolcoords,-1)
+  rowindobj=BandedBlockBandedMatrixIndex(currenti-1,refinds,refrowcoords,reflocalrowinds,true)
+  colindobj=BandedBlockBandedMatrixIndex(currenti-1,refinds,refcolcoords,reflocalcolinds,false)
+  rowindobj,colindobj
+end
+
 Base.firstindex(ind::MatrixIndex)=1
 Base.lastindex(ind::MatrixIndex)=ind.count
 Base.length(ind::MatrixIndex)=ind.count
@@ -135,6 +258,49 @@ function Base.getindex(ind::TridiagonalIndex,i::Int)
   end
 end
 
+function Base.getindex(ind::BandedMatrixIndex,i::Int)
+  1 <= i <= ind.count || throw(BoundsError(ind, i))
+  _i=i
+  p=1
+  while _i-ind.bandsizes[p]>0
+    _i-=ind.bandsizes[p]
+    p+=1
+  end
+  bandind=ind.bandinds[p]
+  startfromone=!xor(ind.isrow,(bandind>0))
+  if startfromone
+    return _i
+  else
+    return _i+abs(bandind)
+  end
+end
+
+function Base.getindex(ind::BlockBandedMatrixIndex,i::Int)
+  1 <= i <= ind.count || throw(BoundsError(ind, i))
+  p=1
+  while i-ind.refinds[p]>=0
+    p+=1
+  end
+  p-=1
+  _i=i-ind.refinds[p]
+  if ind.isrow
+    return ind.refcoords[p]+_i
+  else
+    return ind.refcoords[p]
+  end
+end
+
+function Base.getindex(ind::BandedBlockBandedMatrixIndex,i::Int)
+  1 <= i <= ind.count || throw(BoundsError(ind, i))
+  p=1
+  while i-ind.refinds[p]>=0
+    p+=1
+  end
+  p-=1
+  _i=i-ind.refinds[p]+1
+  ind.reflocalinds[p][_i]+ind.refcoords[p]-1
+end
+
 function findstructralnz(x::Bidiagonal)
   n=size(x,1)
   isup= x.uplo=='U' ? true : false
@@ -225,28 +391,58 @@ function __init__()
   end
 
   @require BandedMatrices="aae01518-5342-5314-be14-df237901396f" begin
+    function findstructralnz(x::BandedMatrices.BandedMatrix)
+      l,u=BandedMatrices.bandwidths(x)
+      rowsize,colsize=size(x)
+      rowind=BandedMatrixIndex(rowsize,colsize,l,u,true)
+      colind=BandedMatrixIndex(rowsize,colsize,l,u,false)
+      (rowind,colind)
+    end
+
+    has_sparsestruct(::Type{<:BandedMatrices.BandedMatrix}) = true
     is_structured(::Type{<:BandedMatrices.BandedMatrix}) = true
     fast_matrix_colors(::Type{<:BandedMatrices.BandedMatrix}) = true
 
     function matrix_colors(A::BandedMatrices.BandedMatrix)
-        u,l=bandwidths(A)
+        l,u=BandedMatrices.bandwidths(A)
         width=u+l+1
         _cycle(1:width,size(A,2))
     end
 
   end
 
-  @require BlockBandedMatrices="aae01518-5342-5314-be14-df237901396f" begin
-    is_structured(::Type{<:BandedMatrices.BlockBandedMatrix}) = true
-    is_structured(::Type{<:BandedMatrices.BandedBlockBandedMatrix}) = true
+  @require BlockBandedMatrices="ffab5731-97b5-5995-9138-79e8c1846df0" begin
+    function findstructralnz(x::BlockBandedMatrices.BlockBandedMatrix)
+      l,u=BlockBandedMatrices.blockbandwidths(x)
+      nrowblock=BlockBandedMatrices.nblocks(x,1)
+      ncolblock=BlockBandedMatrices.nblocks(x,2)
+      rowsizes=[BlockBandedMatrices.blocksize(x,(i,1))[1] for i in 1:nrowblock]
+      colsizes=[BlockBandedMatrices.blocksize(x,(1,i))[2] for i in 1:ncolblock]
+      BlockBandedMatrixIndex(nrowblock,ncolblock,rowsizes,colsizes,l,u)
+    end
+
+    function findstructralnz(x::BlockBandedMatrices.BandedBlockBandedMatrix)
+      l,u=BlockBandedMatrices.blockbandwidths(x)
+      lambda,mu=BlockBandedMatrices.subblockbandwidths(x)
+      nrowblock=BlockBandedMatrices.nblocks(x,1)
+      ncolblock=BlockBandedMatrices.nblocks(x,2)
+      rowsizes=[BlockBandedMatrices.blocksize(x,(i,1))[1] for i in 1:nrowblock]
+      colsizes=[BlockBandedMatrices.blocksize(x,(1,i))[2] for i in 1:ncolblock]
+      BandedBlockBandedMatrixIndex(nrowblock,ncolblock,rowsizes,colsizes,l,u,lambda,mu)
+    end
+
+    has_sparsestruct(::Type{<:BlockBandedMatrices.BlockBandedMatrix}) = true
+    has_sparsestruct(::Type{<:BlockBandedMatrices.BandedBlockBandedMatrix}) = true
+    is_structured(::Type{<:BlockBandedMatrices.BlockBandedMatrix}) = true
+    is_structured(::Type{<:BlockBandedMatrices.BandedBlockBandedMatrix}) = true
     fast_matrix_colors(::Type{<:BlockBandedMatrices.BlockBandedMatrix}) = true
     fast_matrix_colors(::Type{<:BlockBandedMatrices.BandedBlockBandedMatrix}) = true
 
     function matrix_colors(A::BlockBandedMatrices.BlockBandedMatrix)
-        l,u=blockbandwidths(A)
+        l,u=BlockBandedMatrices.blockbandwidths(A)
         blockwidth=l+u+1
-        nblock=nblocks(A,2)
-        cols=[blocksize(A,(1,i))[2] for i in 1:nblock]
+        nblock=BlockBandedMatrices.nblocks(A,2)
+        cols=[BlockBandedMatrices.blocksize(A,(1,i))[2] for i in 1:nblock]
         blockcolors=_cycle(1:blockwidth,nblock)
         #the reserved number of colors of a block is the maximum length of columns of blocks with the same block color
         ncolors=[maximum(cols[i:blockwidth:nblock]) for i in 1:blockwidth]
@@ -257,12 +453,12 @@ function __init__()
     end
 
     function matrix_colors(A::BlockBandedMatrices.BandedBlockBandedMatrix)
-        l,u=blockbandwidths(A)
-        lambda,mu=subblockbandwidths(A)
+        l,u=BlockBandedMatrices.blockbandwidths(A)
+        lambda,mu=BlockBandedMatrices.subblockbandwidths(A)
         blockwidth=l+u+1
         subblockwidth=lambda+mu+1
-        nblock=nblocks(A,2)
-        cols=[blocksize(A,(1,i))[2] for i in 1:nblock]
+        nblock=BlockBandedMatrices.nblocks(A,2)
+        cols=[BlockBandedMatrices.blocksize(A,(1,i))[2] for i in 1:nblock]
         blockcolors=_cycle(1:blockwidth,nblock)
         #the reserved number of colors of a block is the min of subblockwidth and the largest length of columns of blocks with the same block color
         ncolors=[min(subblockwidth,maximum(cols[i:blockwidth:nblock])) for i in 1:min(blockwidth,nblock)]

test/runtests.jl

@@ -37,3 +37,38 @@ Sp=sparse([1,2,3],[1,2,3],[1,2,3])
 @test has_sparsestruct(Sp)
 rowind,colind=findstructralnz(Sp)
 @test [Tri[rowind[i],colind[i]] for i in 1:length(rowind)]==[1,2,3]
+
+using BandedMatrices
+
+B=BandedMatrix(Ones(5,5), (-1,2))
+B[band(1)].=[1,2,3,4]
+B[band(2)].=[5,6,7]
+@test has_sparsestruct(B)
+rowind,colind=findstructralnz(B)
+@test [B[rowind[i],colind[i]] for i in 1:length(rowind)]==[5,6,7,1,2,3,4]
+B=BandedMatrix(Ones(4,6), (-1,2))
+B[band(1)].=[1,2,3,4]
+B[band(2)].=[5,6,7,8]
+rowind,colind=findstructralnz(B)
+@test [B[rowind[i],colind[i]] for i in 1:length(rowind)]==[5,6,7,8,1,2,3,4]
+
+using BlockBandedMatrices
+BB=BlockBandedMatrix(Ones(10,10),([1,2,3,4],[4,3,2,1]),(1,0))
+BB[Block(1,1)].=[1 2 3 4]
+BB[Block(2,1)].=[5 6 7 8;9 10 11 12]
+rowind,colind=findstructralnz(BB)
+@test [BB[rowind[i],colind[i]] for i in 1:length(rowind)]==
+    [1,5,9,2,6,10,3,7,11,4,8,12,
+    1,1,1,1,1,1,1,1,1,1,1,1,1,1,
+    1,1,1,1,1,1,1,1,1,1,1,1,1,1,
+    1,1,1,1,1]
+
+dense=collect(Ones(8,8))
+for i in 1:8
+    dense[:,i].=[1,2,3,4,5,6,7,8]
+end
+BBB=BandedBlockBandedMatrix(dense, ([4, 4] ,[4, 4]), (1, 1), (1, 1))
+rowind,colind=findstructralnz(BBB)
+@test [BBB[rowind[i],colind[i]] for i in 1:length(rowind)]==
+    [1,2,3,1,2,3,4,2,3,4,5,6,7,5,6,7,8,6,7,8,
+     1,2,3,1,2,3,4,2,3,4,5,6,7,5,6,7,8,6,7,8]