add findstructralnz support for BlockBandedMatrix

Author: huanglangwen

src/ArrayInterface.jl

@@ -111,17 +111,85 @@ end
 
 struct BandedMatrixIndex <: MatrixIndex
   count::Int
-  nsize::Int
-  bandinds::Array{Int}
-  bandsizes::Array{Int}
+  rowsize::Int
+  colsize::Int
+  bandinds::Array{Int,1}
+  bandsizes::Array{Int,1}
   isrow::Bool
 end
 
-function BandedMatrixIndex(nsize,lowerbandwidth,upperbandwidth,isrow)
+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=[nsize-abs(band) for band in bandinds]
-  BandedMatrixIndex(sum(bandsizes),nsize,bandinds,bandsizes,isrow)
+  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[2])
+  sort!(sortedinds,by=x->x[1],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
 
 Base.firstindex(ind::MatrixIndex)=1
@@ -159,14 +227,29 @@ function Base.getindex(ind::BandedMatrixIndex,i::Int)
     p+=1
   end
   bandind=ind.bandinds[p]
-  startfromone=ind.isrow & (bandind>0)
+  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 findstructralnz(x::Bidiagonal)
   n=size(x,1)
   isup= x.uplo=='U' ? true : false
@@ -259,9 +342,9 @@ function __init__()
   @require BandedMatrices="aae01518-5342-5314-be14-df237901396f" begin
     function findstructralnz(x::BandedMatrices.BandedMatrix)
       l,u=BandedMatrices.bandwidths(x)
-      nsize=size(x,1)
-      rowind=BandedMatrixIndex(nsize,l,u,true)
-      colind=BandedMatrixIndex(nsize,l,u,false)
+      rowsize,colsize=size(x)
+      rowind=BandedMatrixIndex(rowsize,colsize,l,u,true)
+      colind=BandedMatrixIndex(rowsize,colsize,l,u,false)
       (rowind,colind)
     end
 
@@ -277,9 +360,19 @@ function __init__()
 
   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
+
+    has_sparsestruct(::Type{<:BlockBandedMatrices.BlockBandedMatrix}) = 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
 

test/runtests.jl

@@ -45,4 +45,20 @@ 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]
+@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]