Fix Limited Memory Broyden

Author: avik-pal

README.md

@@ -50,3 +50,5 @@ For more details on the bracketing methods, refer to the [Tutorials](https://doc
   - `Broyden` and `Klement` have been renamed to `SimpleBroyden` and `SimpleKlement` to
     avoid conflicts with `NonlinearSolve.jl`'s `GeneralBroyden` and `GeneralKlement`, which
     will be renamed to `Broyden` and `Klement` in the future.
+  - `LBroyden` has been renamed to `SimpleLimitedMemoryBroyden` to make it consistent with
+    `NonlinearSolve.jl`'s `LimitedMemoryBroyden`.

src/SimpleNonlinearSolve.jl

@@ -13,7 +13,7 @@ import PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidat
     import ForwardDiff: Dual
     import MaybeInplace: @bb, setindex_trait, CanSetindex, CannotSetindex
     import SciMLBase: AbstractNonlinearAlgorithm, build_solution, isinplace
-    import StaticArraysCore: StaticArray, SVector, SMatrix, SArray, MArray
+    import StaticArraysCore: StaticArray, SVector, SMatrix, SArray, MArray, MMatrix, Size
 end
 
 @reexport using ADTypes, SciMLBase
@@ -24,16 +24,16 @@ abstract type AbstractNewtonAlgorithm <: AbstractSimpleNonlinearSolveAlgorithm e
 
 include("utils.jl")
 
-# Nonlinear Solvera
+## Nonlinear Solvers
 include("nlsolve/raphson.jl")
 include("nlsolve/broyden.jl")
-# include("nlsolve/lbroyden.jl")
+include("nlsolve/lbroyden.jl")
 include("nlsolve/klement.jl")
 include("nlsolve/trustRegion.jl")
 # include("nlsolve/halley.jl")
 # include("nlsolve/dfsane.jl")
 
-# Interval Nonlinear Solvers
+## Interval Nonlinear Solvers
 include("bracketing/bisection.jl")
 include("bracketing/falsi.jl")
 include("bracketing/ridder.jl")
@@ -42,7 +42,7 @@ include("bracketing/alefeld.jl")
 include("bracketing/itp.jl")
 
 # AD
-# include("ad.jl")
+include("ad.jl")
 
 ## Default algorithm
 
@@ -58,34 +58,22 @@ end
 
 @setup_workload begin
     for T in (Float32, Float64)
-        prob_no_brack = NonlinearProblem{false}((u, p) -> u .* u .- p, T(0.1), T(2))
-        algs = [SimpleNewtonRaphson(), SimpleBroyden(), SimpleKlement(),
-            SimpleTrustRegion()]
-
-        @compile_workload begin
-            for alg in algs
-                solve(prob_no_brack, alg, abstol = T(1e-2))
-            end
-        end
+        prob_no_brack_scalar = NonlinearProblem{false}((u, p) -> u .* u .- p, T(0.1), T(2))
+        prob_no_brack_iip = NonlinearProblem{true}((du, u, p) -> du .= u .* u .- p,
+            T.([1.0, 1.0, 1.0]), T(2))
+        prob_no_brack_oop = NonlinearProblem{false}((u, p) -> u .* u .- p,
+            T.([1.0, 1.0, 1.0]), T(2))
 
-        prob_no_brack = NonlinearProblem{true}((du, u, p) -> du .= u .* u .- p,
-            T.([1.0, 1.0]), T(2))
+        algs = [SimpleNewtonRaphson(), SimpleBroyden(), SimpleKlement(),
+            SimpleTrustRegion(), SimpleLimitedMemoryBroyden(; threshold = 2)]
 
         @compile_workload begin
             for alg in algs
-                solve(prob_no_brack, alg, abstol = T(1e-2))
-            end
-        end
-
-        #=
-        for alg in (SimpleNewtonRaphson,)
-            for u0 in ([1., 1.], StaticArraysCore.SA[1.0, 1.0])
-                u0 = T.(.1)
-                probN = NonlinearProblem{false}((u,p) -> u .* u .- p, u0, T(2))
-                solve(probN, alg(), tol = T(1e-2))
+                solve(prob_no_brack_scalar, alg, abstol = T(1e-2))
+                solve(prob_no_brack_iip, alg, abstol = T(1e-2))
+                solve(prob_no_brack_oop, alg, abstol = T(1e-2))
             end
         end
-        =#
 
         prob_brack = IntervalNonlinearProblem{false}((u, p) -> u * u - p,
             T.((0.0, 2.0)), T(2))
@@ -98,9 +86,9 @@ end
     end
 end
 
-export SimpleBroyden,
-    SimpleGaussNewton, SimpleKlement, SimpleNewtonRaphson, SimpleTrustRegion
-# SimpleDFSane, SimpleHalley, LBroyden
+export SimpleBroyden, SimpleGaussNewton, SimpleKlement, SimpleLimitedMemoryBroyden,
+    SimpleNewtonRaphson, SimpleTrustRegion
+# SimpleDFSane, SimpleHalley
 export Alefeld, Bisection, Brent, Falsi, ITP, Ridder
 
 end # module

src/ad.jl

@@ -1,7 +1,7 @@
 function scalar_nlsolve_ad(prob, alg, args...; kwargs...)
     f = prob.f
     p = value(prob.p)
-
+    u0 = value(prob.u0)
     if prob isa IntervalNonlinearProblem
         tspan = value(prob.tspan)
         newprob = IntervalNonlinearProblem(f, tspan, p; prob.kwargs...)
@@ -13,66 +13,57 @@ function scalar_nlsolve_ad(prob, alg, args...; kwargs...)
     sol = solve(newprob, alg, args...; kwargs...)
 
     uu = sol.u
-    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_p = scalar_nlsolve_∂f_∂p(f, uu, p)
+    f_x = scalar_nlsolve_∂f_∂u(f, uu, p)
+
+    z_arr = -inv(f_x) * f_p
 
-    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)
+    sumfun = ((z, p),) -> map(zᵢ -> zᵢ * ForwardDiff.partials(p), z)
+    if uu isa Number
+        partials = sum(sumfun, zip(z_arr, pp))
+    elseif p isa Number
+        partials = sumfun((z_arr, pp))
+    else
+        partials = sum(sumfun, zip(eachcol(z_arr), pp))
     end
-    partials = sum(sumfun, zip(f_p, pp))
+
     return sol, partials
 end
 
-function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number, StaticArraysCore.SVector},
-            iip,
-            <:Dual{T, V, P}},
-        alg::AbstractSimpleNonlinearSolveAlgorithm,
-        args...; kwargs...) where {iip, T, V, P}
+function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number, SVector, <:AbstractArray},
+            false, <:Dual{T, V, P}}, alg::AbstractSimpleNonlinearSolveAlgorithm, args...;
+        kwargs...) where {T, V, P}
     sol, partials = scalar_nlsolve_ad(prob, alg, args...; kwargs...)
-    return SciMLBase.build_solution(prob, alg, Dual{T, V, P}(sol.u, partials), sol.resid;
-        retcode = sol.retcode)
+    dual_soln = scalar_nlsolve_dual_soln(sol.u, partials, prob.p)
+    return SciMLBase.build_solution(prob, alg, dual_soln, sol.resid; sol.retcode)
 end
-function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number, StaticArraysCore.SVector},
-            iip,
-            <:AbstractArray{<:Dual{T, V, P}}},
-        alg::AbstractSimpleNonlinearSolveAlgorithm, args...;
-        kwargs...) where {iip, T, V, P}
+
+function SciMLBase.solve(prob::NonlinearProblem{<:Union{Number, SVector, <:AbstractArray},
+            false, <:AbstractArray{<:Dual{T, V, P}}},
+        alg::AbstractSimpleNonlinearSolveAlgorithm, args...; kwargs...) where {T, V, P}
     sol, partials = scalar_nlsolve_ad(prob, alg, args...; kwargs...)
-    return SciMLBase.build_solution(prob, alg, Dual{T, V, P}(sol.u, partials), sol.resid;
-        retcode = sol.retcode)
+    dual_soln = scalar_nlsolve_dual_soln(sol.u, partials, prob.p)
+    return SciMLBase.build_solution(prob, alg, dual_soln, sol.resid; sol.retcode)
 end
 
 # avoid ambiguities
 for Alg in [Bisection]
     @eval function SciMLBase.solve(prob::IntervalNonlinearProblem{uType, iip,
-                <:Dual{T, V, P}},
-            alg::$Alg, args...;
-            kwargs...) where {uType, iip, T, V, P}
+                <:Dual{T, V, P}}, alg::$Alg, args...; kwargs...) where {uType, iip, T, V, P}
         sol, partials = scalar_nlsolve_ad(prob, alg, args...; kwargs...)
-        return SciMLBase.build_solution(prob, alg, Dual{T, V, P}(sol.u, partials),
-            sol.resid; retcode = sol.retcode,
+        dual_soln = scalar_nlsolve_dual_soln(sol.u, partials, prob.p)
+        return SciMLBase.build_solution(prob, alg, dual_soln, sol.resid; sol.retcode,
             left = Dual{T, V, P}(sol.left, partials),
             right = Dual{T, V, P}(sol.right, partials))
-        #return BracketingSolution(Dual{T,V,P}(sol.left, partials), Dual{T,V,P}(sol.right, partials), sol.retcode, sol.resid)
     end
     @eval function SciMLBase.solve(prob::IntervalNonlinearProblem{uType, iip,
-                <:AbstractArray{
-                    <:Dual{T,
-                        V,
-                        P},
-                }},
-            alg::$Alg, args...;
+                <: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 SciMLBase.build_solution(prob, alg, Dual{T, V, P}(sol.u, partials),
-            sol.resid; retcode = sol.retcode,
+        dual_soln = scalar_nlsolve_dual_soln(sol.u, partials, prob.p)
+        return SciMLBase.build_solution(prob, alg, dual_soln, sol.resid; sol.retcode,
             left = Dual{T, V, P}(sol.left, partials),
             right = Dual{T, V, P}(sol.right, partials))
-        #return BracketingSolution(Dual{T,V,P}(sol.left, partials), Dual{T,V,P}(sol.right, partials), sol.retcode, sol.resid)
     end
 end