Revert non-square and add downstream test

Author: ChrisRackauckas

src/jacobians.jl

@@ -373,8 +373,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
-                idx = findfirst(isequal(color_i), _color)
-                tmp = norm(x1[idx])
+                @. fx1 = x1 * (_color == color_i)
+                tmp = norm(fx1)
                 epsilon = compute_epsilon(Val(:forward), sqrt(tmp), relstep, absstep, dir)
                 @. x1 = x1 + epsilon * (_color == color_i)
                 f(fx1, x1)
@@ -416,8 +416,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
-                idx = findfirst(isequal(color_i), _color)
-                tmp = norm(x1[idx])
+                @. fx1 = x1 * (_color == color_i)
+                tmp = norm(fx1)
                 epsilon = compute_epsilon(Val(:central), sqrt(tmp), relstep, absstep, dir)
                 @. x1 = x1 + epsilon * (_color == color_i)
                 @. x  = x  - epsilon * (_color == color_i)

test/downstream/Project.toml

@@ -1,3 +1,4 @@
 [deps]
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"

test/downstream/ordinarydiffeq_tridiagonal_solve.jl

@@ -0,0 +1,27 @@
+using OrdinaryDiffEq, ForwardDiff, Test
+
+const nknots = 10
+const h = 1.0/(nknots+1)
+x = range(0, step=h, length=nknots)
+u0 = sin.(π*x)
+
+@inline function f(du,u,p,t)
+  du .= zero(eltype(u))
+  u₃ = @view u[3:end]
+  u₂ = @view u[2:end-1]
+  u₁ = @view u[1:end-2]
+  @. du[2:end-1] = p[1]*((u₃ - 2*u₂ + u₁)/(h^2.0))
+  nothing
+end
+
+p_true = [0.42]
+jac_proto = Tridiagonal(similar(u0,nknots-1), similar(u0), similar(u0, nknots-1))
+prob = ODEProblem(ODEFunction(f,jac_prototype=jac_proto), u0, (0.0,1.0), p_true)
+sol_true = solve(prob, Rodas4P(), saveat=0.1)
+
+function loss(p)
+  _prob = remake(prob, p=p)
+  sol = solve(_prob, Rodas4P(autodiff=false), saveat=0.1)
+  sum((sol .- sol_true).^2)
+end
+@test ForwardDiff.gradient(loss, [1.0])[1] ≈ 0.6662949361011025

test/runtests.jl

@@ -23,6 +23,7 @@ if GROUP == "All" || GROUP == "Downstream"
   activate_downstream_env()
   @time @safetestset "ODEs" begin
     import OrdinaryDiffEq
+    @time @safetestset "OrdinaryDiffEq Tridiagonal" begin include("ordinarydiffeq_tridiagonal_solve.jl") end
     include(joinpath(dirname(pathof(OrdinaryDiffEq)), "..", "test/interface/sparsediff_tests.jl"))
   end
 end