add high level interface

LICENSE

@@ -1,4 +1,4 @@
-Copyright (c) 2019 Pankaj Mishra <[email protected]>
+Copyright (c) 2019 JuliaDiffEq
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

Project.toml

@@ -1,21 +1,21 @@
 name = "SparseDiffTools"
 uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
-authors = ["Pankaj Mishra <[email protected]>"]
+authors = ["Pankaj Mishra <[email protected]>","Chris Rackauckas <[email protected]>"]
 version = "0.1.0"
 
 [deps]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 VertexSafeGraphs = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f"
 
 [compat]
 julia = "1"
 
 [extras]
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 DiffEqDiffTools = "01453d9d-ee7c-5054-8395-0335cb756afa"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "LinearAlgebra", "DiffEqDiffTools"]
+test = ["Test", "DiffEqDiffTools"]

src/SparseDiffTools.jl

@@ -1,14 +1,19 @@
 module SparseDiffTools
 
+using SparseArrays, LinearAlgebra, LightGraphs, VertexSafeGraphs
+using ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD
+
 export  contract_color,
         greedy_d1,
         matrix2graph,
+        matrix_colors,
         forwarddiff_color_jacobian!,
         ForwardColorJacCache
 
+include("coloring/high_level.jl")
 include("coloring/contraction_coloring.jl")
 include("coloring/greedy_d1_coloring.jl")
-include("differentiation/matrix2graph.jl")
+include("coloring/matrix2graph.jl")
 include("differentiation/compute_jacobian_ad.jl")
 
 end # module

src/coloring/bsc_algo.jl

@@ -1,7 +1,9 @@
-#Backtracking Sequential Coloring algorithm
-using VertexSafeGraphs
+"""
+    BSCColor
 
-function bsc_color(G::VSafeGraph)
+    Backtracking Sequential Coloring algorithm
+"""
+function color_graph(G::VSafeGraph,::BSCColor)
     V = nv(G)
     F = zeros(Int64, V)
     freeColors = [Vector{Int64}() for _ in 1:V] #set of free colors for each vertex

src/coloring/contraction_coloring.jl

@@ -1,13 +1,10 @@
-using VertexSafeGraphs
-using LightGraphs
-
 """
-    colorGraph(g)
+    ColorContraction
 
 Find a coloring of the graph g such that no two vertices connected
 by an edge have the same color.
 """
-function contract_color(G::VSafeGraph)
+function color_graph(G::VSafeGraph,::ContractionColor)
 
     colornumber = 0
     V = nv(G)

src/coloring/greedy_d1_coloring.jl

@@ -1,16 +1,13 @@
-using VertexSafeGraphs
-
 """
-
-        greedy_d1(G)
+        GreedyD1 Coloring
 
 Find a coloring of a given input graph such that
 no two vertices connected by an edge have the same
 color using greedy approach. The number of colors
 used may be equal or greater than the chromatic
 number χ(G) of the graph.
 """
-function greedy_d1(G::VSafeGraph)
+function color_graph(G::VSafeGraph,alg::GreedyD1Color)
     V = nv(G)
     result = zeros(Int64, V)
     result[1] = 1

src/coloring/high_level.jl

@@ -0,0 +1,26 @@
+abstract type ColoringAlgorithm end
+struct GreedyD1Color <: ColoringAlgorithm end
+struct BSCColor <: ColoringAlgorithm end
+struct ContractionColor <: ColoringAlgorithm end
+
+"""
+    matrix_colors(A,alg::ColoringAlgorithm = GreedyD1Color())
+
+    Returns the color vector for the matrix A using the chosen coloring
+    algorithm. If a known analytical solution exists, that is used instead.
+    The coloring defaults to a greedy distance-1 coloring.
+
+"""
+function matrix_colors(A::AbstractMatrix,alg::ColoringAlgorithm = GreedyD1Color())
+    _A = A isa SparseMatrixCSC ? A : sparse(A) # Avoid the copy
+    A_graph = matrix2graph(_A)
+    color_graph(A_graph,alg)
+end
+
+function matrix_colors(A::Array)
+    eachindex(size(A,2)) # Vector size matches number of rows
+end
+
+function matrix_colors(A::Tridiagonal)
+    repeat(1:3,div(size(A,2),3)+1)[1:size(A,2)]
+end

src/differentiation/matrix2graph.jl → src/coloring/matrix2graph.jl

@@ -1,7 +1,3 @@
-using SparseArrays
-using LightGraphs
-using VertexSafeGraphs
-
 """
         matrix2graph(SparseMatrix)
 

src/differentiation/compute_jacobian_ad.jl

@@ -1,5 +1,3 @@
-using ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD
-
 struct ForwardColorJacCache{T,T2,T3,T4,T5}
     t::T
     fx::T2

test/test_contraction.jl

@@ -23,7 +23,7 @@ end
 
 for i in 1:20
     g = test_graphs[i]
-    out_colors = contract_color(g)
+    out_colors = SparseDiffTools.color_graph(g,SparseDiffTools.ContractionColor())
     for v = 1:nv(g)
         color = out_colors[v]
         for j in inneighbors(g, v)

test/test_greedy_d1.jl

@@ -23,7 +23,7 @@ end
 
 for i in 1:20
     g = test_graphs[i]
-    out_colors = greedy_d1(g)
+    out_colors = SparseDiffTools.color_graph(g,SparseDiffTools.GreedyD1Color())
     for v = 1:nv(g)
         color = out_colors[v]
         for j in inneighbors(g, v)

test/test_integration.jl

@@ -1,4 +1,4 @@
-using SparseDiffTools
+using SparseDiffTools, SparseArrays
 using LinearAlgebra
 using DiffEqDiffTools
 using Test
@@ -30,11 +30,9 @@ function generate_sparsity_pattern(N::Integer)
     return Tridiagonal(dl,d,du)
 end
 
-sparsity_pattern = sparse(generate_sparsity_pattern(30))
-
-_graph = matrix2graph(sparsity_pattern)
-coloring_vector = greedy_d1(_graph)
-@test coloring_vector == repeat(1:3,10)
+true_jac = sparse(generate_sparsity_pattern(30))
+colors = matrix_colors(true_jac)
+@test colors == repeat(1:3,10)
 
 #Jacobian computed without coloring vector
 J = DiffEqDiffTools.finite_difference_jacobian(f, rand(30))
@@ -43,7 +41,7 @@ J = DiffEqDiffTools.finite_difference_jacobian(f, rand(30))
 
 #Jacobian computed with coloring vectors
 fcalls = 0
-_J = sparsity_pattern
-DiffEqDiffTools.finite_difference_jacobian!(_J, f, rand(30), color = coloring_vector)
+_J = 200 .* true_jac
+DiffEqDiffTools.finite_difference_jacobian!(_J, f, rand(30), color = colors)
 @test fcalls == 4
 @test _J ≈ J