From d5d3daeda562aa0d3cc88fdc1067ce180f49d948 Mon Sep 17 00:00:00 2001
From: Kanav Gupta <kanav0610@gmail.com>
Date: Tue, 17 Nov 2020 20:12:36 +0530
Subject: [PATCH 1/5] [WIP] AD in Bisection

---
 src/scalar.jl    | 33 +++++++++++++++++++++++++++++++++
 test/runtests.jl | 17 +++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/src/scalar.jl b/src/scalar.jl
index 343543afd..4472d4ac2 100644
--- a/src/scalar.jl
+++ b/src/scalar.jl
@@ -19,6 +19,39 @@ function solve(prob::NonlinearProblem{<:Number}, ::NewtonRaphson, args...; xatol
   return NewtonSolution(x, MAXITERS_EXCEED)
 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...)
+  # f, x and p always satisfy
+  # f(x, p) = 0
+  # dx * f_x(x, p) + dp * f_p(x, p) = 0
+  # dx / dp = - f_p(x, p) / f_x(x, p)
+  f_p = (p) -> prob.f(sol.left, p)
+  f_x = (x) -> prob.f(x, ForwardDiff.value(prob.p))
+  d_p = ForwardDiff.derivative(f_p, ForwardDiff.value(prob.p))
+  d_x = ForwardDiff.derivative(f_x, sol.left)
+  partials = - d_p / d_x * ForwardDiff.partials(prob.p)
+  return BracketingSolution(ForwardDiff.Dual{T,V,P}(sol.left, partials), ForwardDiff.Dual{T,V,P}(sol.right, partials), sol.retcode)
+end
+
+# still WIP
+function solve(prob::NonlinearProblem{uType, iip, <:AbstractArray{<:ForwardDiff.Dual{T,V,P}, N}}, alg::Bisection, args...; kwargs...) where {uType, iip, T, V, P, N}
+  p_nodual = ForwardDiff.value.(prob.p)
+  prob_nodual = NonlinearProblem(prob.f, prob.u0, p_nodual; prob.kwargs...)
+  sol = solve(prob_nodual, alg, args...; kwargs...)
+  # f, x and p always satisfy
+  # f(x, p) = 0
+  # dx * f_x(x, p) + dp * f_p(x, p) = 0
+  # dx / dp = - f_p(x, p) / f_x(x, p)
+  f_p = (p) -> [ prob.f(sol.left, p) ]
+  f_x = (x) -> prob.f(x, p_nodual)
+  d_p = ForwardDiff.jacobian(f_p, p_nodual)
+  d_x = ForwardDiff.derivative(f_x, sol.left)
+  @. d_p =  - d_p / d_x
+  @show ForwardDiff.partials.(prob.p)
+  return ForwardDiff.Dual{T,V,P}(sol.left, d_p * ForwardDiff.partials.(prob.p))
+end
+
 function solve(prob::NonlinearProblem, ::Bisection, args...; maxiters = 1000, kwargs...)
   f = Base.Fix2(prob.f, prob.p)
   left, right = prob.u0
diff --git a/test/runtests.jl b/test/runtests.jl
index fc53f0904..74e12a2e0 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -69,6 +69,23 @@ for p in 1.1:0.1:100.0
     @test ForwardDiff.derivative(g, p) ≈ 1/(2*sqrt(p))
 end
 
+f, u0 = (u, p) -> p[1] * u * u - p[2], (1.0, 100.0)
+t = (p) -> [sqrt(p[2] / p[1])]
+g = function (p)
+    probN = NonlinearProblem{false}(f, u0, p)
+    sol = solve(probN, Bisection())
+    return [sol.left]
+end
+
+for p1 in 1.0:1.0:100.0
+    for p2 in 1.0:1.0:100.0
+        p = [p1, p2]
+        @show p
+        @test g(p) ≈ [sqrt(p[2] / p[1])]
+        @test ForwardDiff.jacobian(g, p) ≈ ForwardDiff.jacobian(t, p)
+    end
+end
+
 # Error Checks
 
 f, u0 = (u, p) -> u .* u .- 2.0, @SVector[1.0, 1.0]

From 08c9701cfca753738513fa6dfd99052e5f6e48c7 Mon Sep 17 00:00:00 2001
From: Yingbo Ma <mayingbo5@gmail.com>
Date: Fri, 20 Nov 2020 15:39:37 -0500
Subject: [PATCH 2/5] Fix scalar Newton AD

---
 src/scalar.jl    | 12 ++++++++++++
 src/utils.jl     |  2 +-
 test/runtests.jl |  4 ++--
 3 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/src/scalar.jl b/src/scalar.jl
index 4472d4ac2..3294fbe13 100644
--- a/src/scalar.jl
+++ b/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...)
diff --git a/src/utils.jl b/src/utils.jl
index c13635fd5..5e8e6f262 100644
--- a/src/utils.jl
+++ b/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
diff --git a/test/runtests.jl b/test/runtests.jl
index 74e12a2e0..440f02f67 100644
--- a/test/runtests.jl
+++ b/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)

From 45425432567a1bc368c18945d21dce20ec0460f0 Mon Sep 17 00:00:00 2001
From: Yingbo Ma <mayingbo5@gmail.com>
Date: Fri, 20 Nov 2020 16:28:17 -0500
Subject: [PATCH 3/5] Add scalar nonlinear solve AD

---
 src/NonlinearSolve.jl |  1 +
 src/scalar.jl         | 72 +++++++++++++++++++++----------------------
 src/types.jl          |  3 ++
 src/utils.jl          |  4 +++
 4 files changed, 44 insertions(+), 36 deletions(-)

diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl
index dbb6e0cb4..d7a894d5c 100644
--- a/src/NonlinearSolve.jl
+++ b/src/NonlinearSolve.jl
@@ -3,6 +3,7 @@ module NonlinearSolve
   using Reexport
   using UnPack: @unpack
   using FiniteDiff, ForwardDiff
+  using ForwardDiff: Dual
   using Setfield
   using StaticArrays
   using RecursiveArrayTools
diff --git a/src/scalar.jl b/src/scalar.jl
index 3294fbe13..ccdc61ab1 100644
--- a/src/scalar.jl
+++ b/src/scalar.jl
@@ -19,49 +19,49 @@ 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}
+function scalar_nlsolve_ad(prob, alg, args...; kwargs...)
   f = prob.f
-  p = ForwardDiff.value(prob.p)
-  u0 = ForwardDiff.value(prob.u0)
+  p = value(prob.p)
+  u0 = 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)
+
+  uu = getsolution(sol)
+  if p isa Number
+    f_p = ForwardDiff.derivative(Base.Fix1(f, uu), p)
+  else
+    f_p = ForwardDiff.gradient(Base.Fix1(f, uu), p)
+  end
+
+  f_x = ForwardDiff.derivative(Base.Fix2(f, p), uu)
+  pp = prob.p
+  sumfun = let f_x′ = -f_x
+    ((fp, p),) -> (fp / f_x′) * ForwardDiff.partials(p)
+  end
+  partials = sum(sumfun, zip(f_p, pp))
+  return sol, partials
 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...)
-  # f, x and p always satisfy
-  # f(x, p) = 0
-  # dx * f_x(x, p) + dp * f_p(x, p) = 0
-  # dx / dp = - f_p(x, p) / f_x(x, p)
-  f_p = (p) -> prob.f(sol.left, p)
-  f_x = (x) -> prob.f(x, ForwardDiff.value(prob.p))
-  d_p = ForwardDiff.derivative(f_p, ForwardDiff.value(prob.p))
-  d_x = ForwardDiff.derivative(f_x, sol.left)
-  partials = - d_p / d_x * ForwardDiff.partials(prob.p)
-  return BracketingSolution(ForwardDiff.Dual{T,V,P}(sol.left, partials), ForwardDiff.Dual{T,V,P}(sol.right, partials), sol.retcode)
+function solve(prob::NonlinearProblem{<:Number, iip, <:Dual{T,V,P}}, alg::NewtonRaphson, args...; kwargs...) where {iip, T, V, P}
+  sol, partials = scalar_nlsolve_ad(prob, alg, args...; kwargs...)
+  return NewtonSolution(Dual{T,V,P}(sol.u, partials), sol.retcode)
+end
+function solve(prob::NonlinearProblem{<:Number, iip, <:AbstractArray{<:Dual{T,V,P}}}, alg::NewtonRaphson, args...; kwargs...) where {iip, T, V, P}
+  sol, partials = scalar_nlsolve_ad(prob, alg, args...; kwargs...)
+  return NewtonSolution(Dual{T,V,P}(sol.u, partials), sol.retcode)
 end
 
-# still WIP
-function solve(prob::NonlinearProblem{uType, iip, <:AbstractArray{<:ForwardDiff.Dual{T,V,P}, N}}, alg::Bisection, args...; kwargs...) where {uType, iip, T, V, P, N}
-  p_nodual = ForwardDiff.value.(prob.p)
-  prob_nodual = NonlinearProblem(prob.f, prob.u0, p_nodual; prob.kwargs...)
-  sol = solve(prob_nodual, alg, args...; kwargs...)
-  # f, x and p always satisfy
-  # f(x, p) = 0
-  # dx * f_x(x, p) + dp * f_p(x, p) = 0
-  # dx / dp = - f_p(x, p) / f_x(x, p)
-  f_p = (p) -> [ prob.f(sol.left, p) ]
-  f_x = (x) -> prob.f(x, p_nodual)
-  d_p = ForwardDiff.jacobian(f_p, p_nodual)
-  d_x = ForwardDiff.derivative(f_x, sol.left)
-  @. d_p =  - d_p / d_x
-  @show ForwardDiff.partials.(prob.p)
-  return ForwardDiff.Dual{T,V,P}(sol.left, d_p * ForwardDiff.partials.(prob.p))
+# avoid ambiguities
+for Alg in [Bisection, Falsi]
+  @eval function solve(prob::NonlinearProblem{uType, iip, <:Dual{T,V,P}}, alg::$Alg, args...; kwargs...) where {uType, iip, T, V, P}
+    sol, partials = scalar_nlsolve_ad(prob, alg, args...; kwargs...)
+    return BracketingSolution(Dual{T,V,P}(sol.left, partials), Dual{T,V,P}(sol.right, partials), sol.retcode)
+  end
+  @eval function solve(prob::NonlinearProblem{uType, iip, <:AbstractArray{<:Dual{T,V,P}}}, alg::$Alg, args...; kwargs...) where {uType, iip, T, V, P}
+    sol, partials = scalar_nlsolve_ad(prob, alg, args...; kwargs...)
+    return BracketingSolution(Dual{T,V,P}(sol.left, partials), Dual{T,V,P}(sol.right, partials), sol.retcode)
+  end
 end
 
 function solve(prob::NonlinearProblem, ::Bisection, args...; maxiters = 1000, kwargs...)
diff --git a/src/types.jl b/src/types.jl
index c04012944..9f1b03ee9 100644
--- a/src/types.jl
+++ b/src/types.jl
@@ -78,3 +78,6 @@ function sync_residuals!(solver::BracketingImmutableSolver)
     @set! solver.fr = solver.f(solver.right, solver.p)
     solver
 end
+
+getsolution(sol::NewtonSolution) = sol.u
+getsolution(sol::BracketingSolution) = sol.left
diff --git a/src/utils.jl b/src/utils.jl
index 5e8e6f262..061ee3838 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -234,3 +234,7 @@ function value_derivative(f::F, x::R) where {F,R}
     out = f(ForwardDiff.Dual{T}(x, one(x)))
     ForwardDiff.value(out), ForwardDiff.extract_derivative(T, out)
 end
+
+value(x) = x
+value(x::Dual) = ForwardDiff.value(x)
+value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x)

From 5cacdb824945d996e8a46eff31af6489a910a534 Mon Sep 17 00:00:00 2001
From: Yingbo Ma <mayingbo5@gmail.com>
Date: Fri, 20 Nov 2020 16:28:37 -0500
Subject: [PATCH 4/5] Fix tests and add AD tests for NewtonRaphson and Falsi

---
 test/runtests.jl | 28 +++++++++++++++++-----------
 1 file changed, 17 insertions(+), 11 deletions(-)

diff --git a/test/runtests.jl b/test/runtests.jl
index 440f02f67..80699284d 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -71,20 +71,26 @@ end
 
 f, u0 = (u, p) -> p[1] * u * u - p[2], (1.0, 100.0)
 t = (p) -> [sqrt(p[2] / p[1])]
-g = function (p)
-    probN = NonlinearProblem{false}(f, u0, p)
-    sol = solve(probN, Bisection())
-    return [sol.left]
+p = [0.9, 50.0]
+for alg in [Bisection(), Falsi()]
+    global g, p
+    g = function (p)
+        probN = NonlinearProblem{false}(f, u0, p)
+        sol = solve(probN, Bisection())
+        return [sol.left]
+    end
+
+    @test g(p) ≈ [sqrt(p[2] / p[1])]
+    @test ForwardDiff.jacobian(g, p) ≈ ForwardDiff.jacobian(t, p)
 end
 
-for p1 in 1.0:1.0:100.0
-    for p2 in 1.0:1.0:100.0
-        p = [p1, p2]
-        @show p
-        @test g(p) ≈ [sqrt(p[2] / p[1])]
-        @test ForwardDiff.jacobian(g, p) ≈ ForwardDiff.jacobian(t, p)
-    end
+gnewton = function (p)
+    probN = NonlinearProblem{false}(f, 0.5, p)
+    sol = solve(probN, NewtonRaphson())
+    return [sol.u]
 end
+@test gnewton(p) ≈ [sqrt(p[2] / p[1])]
+@test ForwardDiff.jacobian(gnewton, p) ≈ ForwardDiff.jacobian(t, p)
 
 # Error Checks
 

From 3e21c45b784ff4b9d286be674837b35ecab3115d Mon Sep 17 00:00:00 2001
From: Yingbo Ma <mayingbo5@gmail.com>
Date: Sun, 22 Nov 2020 14:51:02 -0500
Subject: [PATCH 5/5] Add compat bound on julia

---
 Project.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/Project.toml b/Project.toml
index 6753ed3ee..6488517f2 100644
--- a/Project.toml
+++ b/Project.toml
@@ -20,6 +20,7 @@ Reexport = "0.2"
 Setfield = "0.7"
 StaticArrays = "0.11, 0.12"
 UnPack = "0.1, 1.0"
+julia = "1"
 
 [extras]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"