Merge 778c332 into 5441342

JuliaDiff · Nov 3, 2019 · cc0bcb0 · cc0bcb0
2 parents 5441342 + 778c332
commit cc0bcb0
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 5 deletions.
diff --git a/src/jacobians.jl b/src/jacobians.jl
@@ -154,7 +154,7 @@ function finite_difference_jacobian(
     end
     vecx = vec(x)
     vecx1 = vec(x1)
-    J = jac_prototype isa Nothing ? (sparsity isa Nothing ? false.*vecfx.*x' : zeros(eltype(x),size(sparsity))) : zero(jac_prototype)
+    J = jac_prototype isa Nothing ? (sparsity isa Nothing ? false.*vecfx.*vecx' : zeros(eltype(x),size(sparsity))) : zero(jac_prototype)
     nrows, ncols = size(J)
 
     if !(sparsity isa Nothing)
@@ -242,13 +242,25 @@ function finite_difference_jacobian!(J::AbstractMatrix,
     x::AbstractArray{<:Number},
     fdtype     :: Type{T1}=Val{:forward},
     returntype :: Type{T2}=eltype(x),
-    f_in       :: Union{T2,Nothing}=nothing;
+    f_in       :: Union{AbstractArray{<:T2},Nothing}=nothing;
     relstep=default_relstep(fdtype, eltype(x)),
     absstep=relstep,
     colorvec = 1:length(x),
-    sparsity = ArrayInterface.has_sparsestruct(J) ? J : nothing) where {T1,T2,T3}
-    cache = JacobianCache(x, fdtype, returntype)
-    finite_difference_jacobian!(J, f, x, cache, f_in; relstep=relstep, absstep=absstep, colorvec=colorvec, sparsity=sparsity)
+    sparsity = ArrayInterface.has_sparsestruct(J) ? J : nothing) where {T1,T2}
+    if f_in isa Nothing && fdtype == Val{:forward}
+        if size(J,1) == length(x)
+            fx = zero(x)
+        else
+            fx = zeros(returntype,size(J,1))
+        end
+        f(fx,x)
+        cache = JacobianCache(x, fx, fdtype, returntype)
+    elseif f_in isa Nothing
+        cache = JacobianCache(x, fdtype, returntype)
+    else
+        cache = JacobianCache(x, f_in, fdtype, returntype)
+    end
+    finite_difference_jacobian!(J, f, x, cache, cache.fx; relstep=relstep, absstep=absstep, colorvec=colorvec, sparsity=sparsity)
 end
 
 function finite_difference_jacobian!(

diff --git a/test/finitedifftests.jl b/test/finitedifftests.jl
@@ -391,6 +391,20 @@ end
     @test err_func(test_iipJac(J_ref, iipf, x, central_cache), J_ref) < 1e-8
     @test err_func(test_iipJac(J_ref, iipf, x, Val{:central}), J_ref) < 1e-8
 end
+
+# Non vector input
+x = rand(2,2)
+oopf(x) = x
+iipf(fx,x) = (fx.=x)
+J_ref = Matrix{Float64}(I,4,4)
+@time @testset "Jacobian for non-vector inputs" begin
+    @test err_func(DiffEqDiffTools.finite_difference_jacobian(oopf, x, Val{:forward}), J_ref) < 1e-8
+    @test err_func(DiffEqDiffTools.finite_difference_jacobian(oopf, x, Val{:central}), J_ref) < 1e-8
+    @test err_func(DiffEqDiffTools.finite_difference_jacobian(oopf, x, Val{:complex}), J_ref) < 1e-8
+    @test err_func(test_iipJac(J_ref, iipf, x, Val{:forward}, eltype(x), iipf(similar(x),x)), J_ref) < 1e-8
+    @test err_func(test_iipJac(J_ref, iipf, x, Val{:central}, eltype(x), iipf(similar(x),x)), J_ref) < 1e-8
+    @test err_func(test_iipJac(J_ref, iipf, x, Val{:complex}, eltype(x), iipf(similar(x),x)), J_ref) < 1e-8
+end
 # Hessian tests
 
 f(x) = sin(x[1]) + cos(x[2])