Don't create a reshaped array when compute a Jacobian of a StaticVector

valued function of size one.

Author: andreasnoack

Project.toml

@@ -19,8 +19,8 @@ julia = "1.2"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = ["Test", "BlockBandedMatrices", "BandedMatrices", "Pkg", "SafeTestsets"]

src/FiniteDiff.jl

@@ -8,6 +8,10 @@ import Base: resize!
 _vec(x) = vec(x)
 _vec(x::Number) = x
 
+_mat(x::AbstractMatrix) = x
+_mat(x::StaticVector)   = reshape(x, (axes(x, 1), SOneTo(1)))
+_mat(x::AbstractVector) = reshape(x, (axes(x, 1),  OneTo(1)))
+
 include("iteration_utils.jl")
 include("epsilons.jl")
 include("derivatives.jl")

src/jacobians.jl

@@ -121,7 +121,7 @@ end
 
 function _make_Ji(::AbstractArray, xtype, dx, color_i, nrows, ncols)
     Ji = mapreduce(i -> i==color_i ? dx : zero(dx), hcat, 1:ncols)
-    size(Ji)!=(nrows, ncols) ? reshape(Ji,(nrows,ncols)) : Ji #branch when size(dx) == (1,) => size(Ji) == (1,) while size(J) == (1,1)
+    size(Ji) != (nrows, ncols) ? reshape(Ji, (nrows, ncols)) : Ji #branch when size(dx) == (1,) => size(Ji) == (1,) while size(J) == (1,1)
 end
 
 function finite_difference_jacobian(f, x,
@@ -186,17 +186,15 @@ function finite_difference_jacobian(
             x_save = ArrayInterface.allowed_getindex(vecx, i)
             epsilon = compute_epsilon(Val{:forward}, x_save, relstep, absstep, dir)
             _vecx1 = Base.setindex(vecx, x_save+epsilon, i)
-            _x1 = reshape(_vecx1,size(x))
+            _x1 = reshape(_vecx1, axes(x))
             vecfx1 = _vec(f(_x1))
             dx = (vecfx1-vecfx) / epsilon
             return dx
         end
 
         if jac_prototype isa Nothing && sparsity isa Nothing
             J = mapreduce(calculate_Ji_forward, hcat, 1:maximum(colorvec))
-            if maximum(colorvec) == 1
-                J = reshape(J, 1, 1)
-            end
+            J = _mat(J)
         else
             @inbounds for color_i ∈ 1:maximum(colorvec)
                 if sparsity isa Nothing
@@ -206,7 +204,7 @@ function finite_difference_jacobian(
                     tmp = norm(vecx .* (colorvec .== color_i))
                     epsilon = compute_epsilon(Val{:forward}, sqrt(tmp), relstep, absstep, dir)
                     _vecx = @. vecx + epsilon * (colorvec == color_i)
-                    _x = reshape(_vecx,size(x))
+                    _x = reshape(_vecx, axes(x))
                     vecfx1 = _vec(f(_x))
                     dx = (vecfx1-vecfx)/epsilon
                     Ji = _make_Ji(J,rows_index,cols_index,dx,colorvec,color_i,nrows,ncols)
@@ -222,8 +220,8 @@ function finite_difference_jacobian(
             epsilon = compute_epsilon(Val{:forward}, x1_save, relstep, absstep, dir)
             _vecx1 = Base.setindex(vecx1,x1_save+epsilon,i)
             _vecx = Base.setindex(vecx,x_save-epsilon,i)
-            _x1 = reshape(_vecx1,size(x))
-            _x = reshape(_vecx,size(x))
+            _x1 = reshape(_vecx1, axes(x))
+            _x = reshape(_vecx, axes(x))
             vecfx1 = _vec(f(_x1))
             vecfx = _vec(f(_x))
             dx = (vecfx1-vecfx)/(2epsilon)
@@ -232,9 +230,7 @@ function finite_difference_jacobian(
 
         if jac_prototype isa Nothing && sparsity isa Nothing
             J = mapreduce(calculate_Ji_central, hcat, 1:maximum(colorvec))
-            if maximum(colorvec) == 1
-                J = reshape(J, 1, 1)
-            end
+            J = _mat(J)
         else
             @inbounds for color_i ∈ 1:maximum(colorvec)
                 if sparsity isa Nothing
@@ -245,8 +241,8 @@ function finite_difference_jacobian(
                     epsilon = compute_epsilon(Val{:forward}, sqrt(tmp), relstep, absstep, dir)
                     _vecx1 = @. vecx1 + epsilon * (colorvec == color_i)
                     _vecx = @. vecx - epsilon * (colorvec == color_i)
-                    _x1 = reshape(_vecx1,size(x))
-                    _x = reshape(_vecx, size(x))
+                    _x1 = reshape(_vecx1, axes(x))
+                    _x = reshape(_vecx, axes(x))
                     vecfx1 = _vec(f(_x1))
                     vecfx = _vec(f(_x))
                     dx = (vecfx1-vecfx)/(2epsilon)
@@ -257,29 +253,27 @@ function finite_difference_jacobian(
         end
     elseif fdtype == Val{:complex} && returntype <: Real
         epsilon = eps(eltype(x))
-        
+
         function calculate_Ji_complex(i)
             x_save = ArrayInterface.allowed_getindex(vecx,i)
             _vecx = Base.setindex(complex.(vecx),x_save+im*epsilon,i)
-            _x = reshape(_vecx,size(x))
+            _x = reshape(_vecx, axes(x))
             vecfx = _vec(f(_x))
             dx = imag(vecfx)/epsilon
             return dx
         end
-        
+
         if jac_prototype isa Nothing && sparsity isa Nothing
             J = mapreduce(calculate_Ji_complex, hcat, 1:maximum(colorvec))
-            if maximum(colorvec) == 1
-                J = reshape(J, 1, 1)
-            end
+            J = _mat(J)
         else
             @inbounds for color_i ∈ 1:maximum(colorvec)
                 if sparsity isa Nothing
                     dx = calculate_Ji_complex(color_i)
                     J = J + _make_Ji(J, eltype(x), dx, color_i, nrows, ncols)
                 else
                     _vecx = @. vecx + im * epsilon * (colorvec == color_i)
-                    _x = reshape(_vecx,size(x))
+                    _x = reshape(_vecx, axes(x))
                     vecfx = _vec(f(_x))
                     dx = imag(vecfx)/epsilon
                     Ji = _make_Ji(J,rows_index,cols_index,dx,colorvec,color_i,nrows,ncols)
@@ -332,7 +326,7 @@ function finite_difference_jacobian!(
     dir = true) where {T1,T2,T3,cType,sType,fdtype,returntype}
 
     m, n = size(J)
-    _color = reshape(colorvec,size(x)...)
+    _color = reshape(colorvec, axes(x)...)
 
     x1, fx, fx1 = cache.x1, cache.fx, cache.fx1
     copyto!(x1, x)

test/out_of_place_tests.jl

@@ -16,28 +16,45 @@ function second_derivative_stencil(N)
 end
 
 x = @SVector ones(30)
-J = FiniteDiff.finite_difference_jacobian(f,x, Val{:forward}, eltype(x))
+J = FiniteDiff.finite_difference_jacobian(f, x, Val{:forward}, eltype(x))
 @test J ≈ second_derivative_stencil(30)
 
-J = FiniteDiff.finite_difference_jacobian(f,x, Val{:central}, eltype(x))
+J = FiniteDiff.finite_difference_jacobian(f, x, Val{:central}, eltype(x))
 @test J ≈ second_derivative_stencil(30)
 
-J = FiniteDiff.finite_difference_jacobian(f,x, Val{:complex}, eltype(x))
+J = FiniteDiff.finite_difference_jacobian(f, x, Val{:complex}, eltype(x))
 @test J ≈ second_derivative_stencil(30)
 
 spJ = sparse(second_derivative_stencil(30))
-J = FiniteDiff.finite_difference_jacobian(f,x, Val{:forward}, eltype(x),jac_prototype=spJ)
+J = FiniteDiff.finite_difference_jacobian(f, x, Val{:forward}, eltype(x), jac_prototype=spJ)
 @test J ≈ second_derivative_stencil(30)
 @test typeof(J) == typeof(spJ)
-J = FiniteDiff.finite_difference_jacobian(f,x, Val{:forward}, eltype(x),colorvec=repeat(1:3,10),sparsity=spJ,jac_prototype=spJ)
+J = FiniteDiff.finite_difference_jacobian(f, x, Val{:forward}, eltype(x),
+    colorvec=SVector{30}(repeat(1:3, 10)), sparsity=spJ, jac_prototype=spJ)
 @test J ≈ second_derivative_stencil(30)
 @test typeof(J) == typeof(spJ)
+
 #1x1 SVector test
 x = SVector{1}([1.])
 f(x) = x
 J = FiniteDiff.finite_difference_jacobian(f, x, Val{:forward}, eltype(x))
-@test J ≈ SMatrix{1,1}([1.])
+@test J[1, 1] ≈ 1.0
+@test J isa SMatrix{1,1}
+J = FiniteDiff.finite_difference_jacobian(f, x, Val{:central}, eltype(x))
+@test J[1, 1] ≈ 1.0
+@test J isa SMatrix{1,1}
+J = FiniteDiff.finite_difference_jacobian(f, x, Val{:complex}, eltype(x))
+@test J[1, 1] ≈ 1.0
+@test J isa SMatrix{1,1}
+
+x = SVector{1}([1.])
+f(x) = vcat(x, x)
+J = FiniteDiff.finite_difference_jacobian(f, x, Val{:forward}, eltype(x))
+@test J ≈ fill(1.0, 2, 1)
+@test J isa SMatrix{2,1}
 J = FiniteDiff.finite_difference_jacobian(f, x, Val{:central}, eltype(x))
-@test J ≈ SMatrix{1,1}([1.])
+@test J ≈ fill(1.0, 2, 1)
+@test J isa SMatrix{2,1}
 J = FiniteDiff.finite_difference_jacobian(f, x, Val{:complex}, eltype(x))
-@test J ≈ SMatrix{1,1}([1.])
+@test J ≈ fill(1.0, 2, 1)
+@test J isa SMatrix{2,1}