Fix scalar Newton AD

Author: YingboMa

src/scalar.jl

@@ -19,6 +19,18 @@ function solve(prob::NonlinearProblem{<:Number}, ::NewtonRaphson, args...; xatol
   return NewtonSolution(x, MAXITERS_EXCEED)
 end
 
+function solve(prob::NonlinearProblem{<:Number, iip, <:ForwardDiff.Dual{T,V,P}}, alg::NewtonRaphson, args...; kwargs...) where {uType, iip, T, V, P}
+  f = prob.f
+  p = ForwardDiff.value(prob.p)
+  u0 = ForwardDiff.value(prob.u0)
+  newprob = NonlinearProblem(f, u0, p; prob.kwargs...)
+  sol = solve(newprob, alg, args...; kwargs...)
+  f_p = ForwardDiff.derivative(Base.Fix1(f, sol.u), p)
+  f_x = ForwardDiff.derivative(Base.Fix2(f, p), sol.u)
+  partials = (-f_p / f_x) * ForwardDiff.partials(prob.p)
+  return NewtonSolution(ForwardDiff.Dual{T,V,P}(sol.u, partials), sol.retcode)
+end
+
 function solve(prob::NonlinearProblem{uType, iip, <:ForwardDiff.Dual{T,V,P}}, alg::Bisection, args...; kwargs...) where {uType, iip, T, V, P}
   prob_nodual = NonlinearProblem(prob.f, prob.u0, ForwardDiff.value(prob.p); prob.kwargs...)
   sol = solve(prob_nodual, alg, args...; kwargs...)

src/utils.jl

@@ -216,7 +216,7 @@ Move `x` one floating point towards x0.
 function prevfloat_tdir(x::T, x0::T, x1::T)::T where {T}
   x1 > x0 ? prevfloat(x) : nextfloat(x)
 end
-  
+
 function nextfloat_tdir(x::T, x0::T, x1::T)::T where {T}
   x1 > x0 ? nextfloat(x) : prevfloat(x)
 end

test/runtests.jl

@@ -57,12 +57,12 @@ end
 f, u0 = (u, p) -> u * u - p, 1.0
 
 g = function (p)
-    probN = NonlinearProblem{false}(f, u0, p)
+    probN = NonlinearProblem{false}(f, oftype(p, u0), p)
     sol = solve(probN, NewtonRaphson())
     return sol.u
 end
 
-@test_broken ForwardDiff.derivative(g, 1.0) ≈ 0.5
+@test ForwardDiff.derivative(g, 1.0) ≈ 0.5
 
 for p in 1.1:0.1:100.0
     @test g(p) ≈ sqrt(p)