Automatic Differentiation with coloring

Project.toml

@@ -4,6 +4,7 @@ authors = ["Pankaj Mishra <[email protected]>"]
 version = "0.1.0"
 
 [deps]
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 VertexSafeGraphs = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f"

src/SparseDiffTools.jl

@@ -2,11 +2,12 @@ module SparseDiffTools
 
 export  contract_color,
         greedy_d1,
-        matrix2graph
+        matrix2graph,
+        compute_jacobian!
 
 include("contraction_coloring.jl")
 include("greedy_d1_coloring.jl")
 include("matrix2graph.jl")
-
+include("compute_jacobian_ad.jl")
 
 end # module

src/compute_jacobian_ad.jl

@@ -0,0 +1,40 @@
+using ForwardDiff: Dual, jacobian, partials
+
+function compute_jacobian!(J::AbstractMatrix{<:Number},
+                f,
+                x::AbstractArray{<:Number};
+                color=1:length(x))
+
+    partials_array = Array{Float64}(undef, length(x), maximum(color))
+    for color_i in 1:maximum(color)
+        for i in 1:length(x)
+            if color[i]==color_i
+                partials_array[i,color_i] = 1
+            else
+                partials_array[i,color_i] = 0
+            end
+        end
+    end
+
+    p = Tuple.(eachrow(partials_array))
+    t = zeros(Dual{Nothing, Float64, maximum(color)},0)
+    for i in 1:length(x)
+       push!(t, Dual(x[i], p[i]))
+    end
+
+    fx = similar(t)
+    f(fx, t)
+
+    rows_index, cols_index, val = findnz(J)
+    for color_i in 1:maximum(color)
+        du = partials.(fx,color_i)
+        for i in 1:length(cols_index)
+            if color[cols_index[i]]==color_i
+                J[rows_index[i],cols_index[i]] = du[rows_index[i]]
+            end
+        end
+    end
+
+    J = Array(J)
+
+end

test/runtests.jl

@@ -5,4 +5,5 @@ using Test
 @testset "Exact coloring via contraction" begin include("test_contraction.jl") end
 @testset "Greedy distance-1 coloring" begin include("test_greedy_d1.jl") end
 @testset "Matrix to graph conversion" begin include("test_matrix2graph.jl") end
+@testset "AD using color vector" begin include("test_ad.jl") end
 @testset "Integration test" begin include("test_integration.jl") end

test/test_ad.jl

@@ -0,0 +1,37 @@
+using SparseDiffTools
+using SparseArrays
+using ForwardDiff: Dual, jacobian
+
+fcalls = 0
+function f(dx,x)
+  global fcalls += 1
+  for i in 2:length(x)-1
+    dx[i] = x[i-1] - 2x[i] + x[i+1]
+  end
+  dx[1] = -2x[1] + x[2]
+  dx[end] = x[end-1] - 2x[end]
+  nothing
+end
+
+function second_derivative_stencil(N)
+  A = zeros(N,N)
+  for i in 1:N, j in 1:N
+      (j-i==-1 || j-i==1) && (A[i,j]=1)
+      j-i==0 && (A[i,j]=-2)
+  end
+  A
+end
+
+x = rand(30)
+dx = rand(30)
+
+J = jacobian(f, dx, x)
+@test J ≈ second_derivative_stencil(30)
+_J = sparse(J)
+@test fcalls == 3
+
+fcalls = 0
+_J1 = similar(_J)
+compute_jacobian!(_J1, f, x, color = repeat(1:3,10))
+@test _J1 ≈ J
+@test fcalls == 1