add BandedMatrix

Project.toml

@@ -9,6 +9,7 @@ LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 VertexSafeGraphs = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f"
+BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 
 [compat]
 julia = "1"

src/SparseDiffTools.jl

@@ -1,6 +1,6 @@
 module SparseDiffTools
 
-using SparseArrays, LinearAlgebra, LightGraphs, VertexSafeGraphs
+using SparseArrays, LinearAlgebra, BandedMatrices, LightGraphs, VertexSafeGraphs
 using ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD
 
 export  contract_color,

src/coloring/high_level.jl

@@ -27,14 +27,24 @@ function matrix_colors(A::Union{Array,UpperTriangular,LowerTriangular})
     eachindex(1:size(A,2)) # Vector size matches number of rows
 end
 
+function _cycle(repetend,len)
+    repeat(repetend,div(len,length(repetend))+1)[1:len]
+end
+
 function matrix_colors(A::Diagonal)
     fill(1,size(A,2))
 end
 
 function matrix_colors(A::Bidiagonal)
-    repeat(1:2,div(size(A,2),2)+1)[1:size(A,2)]
+    _cycle(1:2,size(A,2))
 end
 
 function matrix_colors(A::Union{Tridiagonal,SymTridiagonal})
-    repeat(1:3,div(size(A,2),3)+1)[1:size(A,2)]
+    _cycle(1:3,size(A,2))
+end
+
+function matrix_colors(A::BandedMatrix)
+    u,l=bandwidths(A)
+    width=u+l+1
+    _cycle(1:width,size(A,2))
 end

test/test_specialmatrices.jl

@@ -1,16 +1,20 @@
 using SparseDiffTools
 using LinearAlgebra
+using BandedMatrices
 
 n=10
 dense=fill(1.,(n,n))
 uptri=UpperTriangular(dense)
 lotri=LowerTriangular(dense)
+
 diagonal=Diagonal(dense)
 bidiagonalU=Bidiagonal(dense,:U)
 bidiagonalL=Bidiagonal(dense,:L)
 tridiagonal=Tridiagonal(dense)
 symtridiagonal=SymTridiagonal(dense)
 
+banded=BandedMatrix(dense,(1,2))
+
 @test matrix_colors(dense)==1:n
 @test matrix_colors(uptri)==1:n
 @test matrix_colors(lotri)==1:n
@@ -19,4 +23,6 @@ symtridiagonal=SymTridiagonal(dense)
 @test matrix_colors(bidiagonalU)==[1,2,1,2,1,2,1,2,1,2]
 @test matrix_colors(bidiagonalL)==[1,2,1,2,1,2,1,2,1,2]
 @test matrix_colors(tridiagonal)==[1,2,3,1,2,3,1,2,3,1]
-@test matrix_colors(symtridiagonal)==[1,2,3,1,2,3,1,2,3,1]
+@test matrix_colors(symtridiagonal)==[1,2,3,1,2,3,1,2,3,1]
+
+@test matrix_colors(banded)==[1,2,3,4,1,2,3,4,1,2]