Merge pull request #21 from JuliaDiffEq/high_level

ChrisRackauckas · web-flow · commit f1ed12892147 · 2019-06-21T09:37:57.000-04:00
add high level interface
diff --git a/LICENSE b/LICENSE
@@ -1,4 +1,4 @@
-Copyright (c) 2019 Pankaj Mishra <pankajmishra1511@gmail.com>
+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
diff --git a/Project.toml b/Project.toml
@@ -1,21 +1,21 @@
 name = "SparseDiffTools"
 uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
-authors = ["Pankaj Mishra <pankajmishra1511@gmail.com>"]
+authors = ["Pankaj Mishra <pankajmishra1511@gmail.com>","Chris Rackauckas <contact@chrisrackauckas.com>"]
 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"]
diff --git a/src/SparseDiffTools.jl b/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
diff --git a/src/coloring/bsc_algo.jl b/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
diff --git a/src/coloring/contraction_coloring.jl b/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)
diff --git a/src/coloring/greedy_d1_coloring.jl b/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
diff --git a/src/coloring/high_level.jl b/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
diff --git a/src/coloring/matrix2graph.jl b/src/coloring/matrix2graph.jl
@@ -1,7 +1,3 @@
-using SparseArrays
-using LightGraphs
-using VertexSafeGraphs
-
 """
         matrix2graph(SparseMatrix)
 
diff --git a/src/differentiation/compute_jacobian_ad.jl b/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
diff --git a/test/test_contraction.jl b/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)
diff --git a/test/test_greedy_d1.jl b/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)
diff --git a/test/test_integration.jl b/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

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-Copyright (c) 2019 Pankaj Mishra <[email protected]>`
	`1`	`+Copyright (c) 2019 JuliaDiffEq`
`2`	`2`
`3`	`3`	`Permission is hereby granted, free of charge, to any person obtaining a copy`
`4`	`4`	`of this software and associated documentation files (the "Software"), to deal`
Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,3 @@`
`1`		`-using ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD`
`2`		`-`
`3`	`1`	`struct ForwardColorJacCache{T,T2,T3,T4,T5}`
`4`	`2`	`t::T`
`5`	`3`	`fx::T2`