Merge pull request #133 from dingraha/master

Make sparse non-square Jacobians work

Author: ChrisRackauckas

Project.toml

@@ -1,6 +1,6 @@
 name = "FiniteDiff"
 uuid = "6a86dc24-6348-571c-b903-95158fe2bd41"
-version = "2.11.0"
+version = "2.11.1"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/jacobians.jl

@@ -1,5 +1,6 @@
 mutable struct JacobianCache{CacheType1,CacheType2,CacheType3,ColorType,SparsityType,fdtype,returntype}
     x1  :: CacheType1
+    x2  :: CacheType1
     fx  :: CacheType2
     fx1 :: CacheType3
     colorvec :: ColorType
@@ -96,7 +97,8 @@ function JacobianCache(
         @assert eltype(fx1) == T2
         _fx = fx
     end
-    JacobianCache{typeof(_x1),typeof(_fx),typeof(fx1),typeof(colorvec),typeof(sparsity),fdtype,returntype}(_x1,_fx,fx1,colorvec,sparsity)
+    _x2 = similar(_x1)
+    JacobianCache{typeof(_x1),typeof(_fx),typeof(fx1),typeof(colorvec),typeof(sparsity),fdtype,returntype}(_x1,_x2,_fx,fx1,colorvec,sparsity)
 end
 
 function _make_Ji(::SparseMatrixCSC, rows_index,cols_index,dx,colorvec,color_i,nrows,ncols)
@@ -334,7 +336,7 @@ function finite_difference_jacobian!(
     m, n = size(J)
     _color = reshape(colorvec, axes(x)...)
 
-    x1, fx, fx1 = cache.x1, cache.fx, cache.fx1
+    x1, x2, fx, fx1 = cache.x1, cache.x2, cache.fx, cache.fx1
     copyto!(x1, x)
     vfx = _vec(fx)
 
@@ -377,8 +379,8 @@ function finite_difference_jacobian!(
                 # Now return x1 back to its original value
                 ArrayInterface.allowed_setindex!(x1, x1_save, color_i)
             else # Perturb along the colorvec vector
-                @. fx1 = x1 * (_color == color_i)
-                tmp = norm(fx1)
+                @. x2 = x1 * (_color == color_i)
+                tmp = norm(x2)
                 epsilon = compute_epsilon(Val(:forward), sqrt(tmp), relstep, absstep, dir)
                 @. x1 = x1 + epsilon * (_color == color_i)
                 f(fx1, x1)
@@ -420,8 +422,8 @@ function finite_difference_jacobian!(
                 @. J[:,color_i] = (vfx1 - vfx) / 2epsilon
                 ArrayInterface.allowed_setindex!(x1, x_save, color_i)
             else # Perturb along the colorvec vector
-                @. fx1 = x1 * (_color == color_i)
-                tmp = norm(fx1)
+                @. x2 = x1 * (_color == color_i)
+                tmp = norm(x2)
                 epsilon = compute_epsilon(Val(:central), sqrt(tmp), relstep, absstep, dir)
                 @. x1 = x1 + epsilon * (_color == color_i)
                 @. x  = x  - epsilon * (_color == color_i)

test/coloring_tests.jl

@@ -117,3 +117,43 @@ Jbb = BlockBandedMatrix(similar(Jsparse),fill(100, 100), fill(100, 100),(1,1));
 colorsbb = ArrayInterface.matrix_colors(Jbb)
 FiniteDiff.finite_difference_jacobian!(Jbb, f, x, colorvec=colorsbb)
 @test Jbb ≈ Jsparse
+
+
+# Non-square Jacobian test.
+# The Jacobian of f_nonsquare! has size (n, 2*n).
+function f_nonsquare!(y, x)
+    global fcalls += 1
+    @assert length(x) == 2*length(y)
+    n = length(x) ÷ 2
+    x1 = @view x[1:n]
+    x2 = @view x[n+1:end]
+
+    @. y = (x1 .- 3).^2 .+ x1.*x2 .+ (x2 .+ 4).^2 .- 3
+    return nothing
+end
+
+n = 4
+x0 = vcat(ones(n).*(1:n) .+ 0.5, ones(n).*(1:n) .+ 1.5)
+y0 = zeros(n)
+rows = vcat([i for i in 1:n], [i for i in 1:n])
+cols = vcat([i for i in 1:n], [i+n for i in 1:n])
+sparsity = sparse(rows, cols, ones(length(rows)))
+colorvec = vcat(fill(1, n), fill(2, n))
+
+J_nonsquare1 = zeros(size(sparsity))
+FiniteDiff.finite_difference_jacobian!(J_nonsquare1, f_nonsquare!, x0)
+
+J_nonsquare2 = similar(sparsity)
+for method in [Val(:forward), Val(:central), Val(:complex)]
+    cache = FiniteDiff.JacobianCache(copy(x0), copy(y0), copy(y0), method; sparsity, colorvec)
+    global fcalls = 0
+    FiniteDiff.finite_difference_jacobian!(J_nonsquare2, f_nonsquare!, x0, cache)
+    if method == Val(:central)
+        @test fcalls == 2*maximum(colorvec)
+    elseif method == Val(:complex)
+        @test fcalls == maximum(colorvec)
+    else
+        @test fcalls == maximum(colorvec) + 1
+    end
+    @test isapprox(J_nonsquare2, J_nonsquare1; rtol=1e-6)
+end