Implementation of a trust-region solver

Author: Deltadahl

src/SimpleNonlinearSolve.jl

@@ -21,6 +21,7 @@ include("raphson.jl")
 include("ad.jl")
 include("broyden.jl")
 include("klement.jl")
+include("trustRegion.jl")
 
 import SnoopPrecompile
 
@@ -30,6 +31,10 @@ SnoopPrecompile.@precompile_all_calls begin for T in (Float32, Float64)
         solve(prob_no_brack, alg(), tol = T(1e-2))
     end
 
+    for alg in (TrustRegion(1.0),)
+        solve(prob_no_brack, alg, tol = T(1e-2))
+    end
+
     #=
     for alg in (SimpleNewtonRaphson,)
         for u0 in ([1., 1.], StaticArraysCore.SA[1.0, 1.0])
@@ -47,6 +52,7 @@ SnoopPrecompile.@precompile_all_calls begin for T in (Float32, Float64)
 end end
 
 # DiffEq styled algorithms
-export Bisection, Broyden, Falsi, Klement, SimpleNewtonRaphson
+export Bisection, Broyden, Falsi, Klement, SimpleNewtonRaphson,
+       TrustRegion
 
 end # module

src/trustRegion.jl

@@ -0,0 +1,125 @@
+"""
+```julia
+TrustRegion(max_trust_radius::Number; chunk_size = Val{0}(),
+                               autodiff = Val{true}(), diff_type = Val{:forward})
+```
+
+A low-overhead implementation of a
+[trust-region](https://optimization.mccormick.northwestern.edu/index.php/Trust-region_methods)
+solver
+
+
+### Keyword Arguments
+- `max_trust_radius`: the maximum radius of the trust region. The step size in the algorithm
+  will change dynamically. However, it will never be greater than the `max_trust_radius`.
+
+### Keyword Arguments
+
+- `chunk_size`: the chunk size used by the internal ForwardDiff.jl automatic differentiation
+  system. This allows for multiple derivative columns to be computed simultaneously,
+  improving performance. Defaults to `0`, which is equivalent to using ForwardDiff.jl's
+  default chunk size mechanism. For more details, see the documentation for
+  [ForwardDiff.jl](https://juliadiff.org/ForwardDiff.jl/stable/).
+- `autodiff`: whether to use forward-mode automatic differentiation for the Jacobian.
+  Note that this argument is ignored if an analytical Jacobian is passed; as that will be
+  used instead. Defaults to `Val{true}`, which means ForwardDiff.jl is used by default.
+  If `Val{false}`, then FiniteDiff.jl is used for finite differencing.
+- `diff_type`: the type of finite differencing used if `autodiff = false`. Defaults to
+  `Val{:forward}` for forward finite differences. For more details on the choices, see the
+  [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl) documentation.
+"""
+struct TrustRegion{CS, AD, FDT} <: AbstractNewtonAlgorithm{CS, AD, FDT}
+    max_trust_radius::Number
+    function TrustRegion(max_turst_radius::Number; chunk_size = Val{0}(),
+                                            autodiff = Val{true}(),
+                                            diff_type = Val{:forward})
+        new{SciMLBase._unwrap_val(chunk_size), SciMLBase._unwrap_val(autodiff),
+            SciMLBase._unwrap_val(diff_type)}(max_trust_radius)
+    end
+end
+
+function SciMLBase.solve(prob::NonlinearProblem,
+                         alg::TrustRegion, args...; abstol = nothing,
+                         reltol = nothing,
+                         maxiters = 1000, kwargs...)
+    f = Base.Fix2(prob.f, prob.p)
+    x = float(prob.u0)
+    T = typeof(x)
+    Δₘₐₓ = float(alg.max_trust_radius)  # The maximum trust region radius.
+    Δ = Δₘₐₓ / 5  # Initial trust region radius.
+    η₁ = 0.1   # Threshold for taking a step.
+    η₂ = 0.25  # Threshold for shrinking the trust region.
+    η₃ = 0.75  # Threshold for expanding the trust region.
+    t₁ = 0.25  # Factor to shrink the trust region with.
+    t₂ = 2.0   # Factor to expand the trust region with.
+
+    if SciMLBase.isinplace(prob)
+        error("TrustRegion currently only supports out-of-place nonlinear problems")
+    end
+
+    atol = abstol !== nothing ? abstol :
+           real(oneunit(eltype(T))) * (eps(real(one(eltype(T)))))^(4 // 5)
+    rtol = reltol !== nothing ? reltol : eps(real(one(eltype(T))))^(4 // 5)
+
+    if alg_autodiff(alg)
+        F, ∇f = value_derivative(f, x)
+    elseif x isa AbstractArray
+        F = f(x)
+        ∇f = FiniteDiff.finite_difference_jacobian(f, x, diff_type(alg), eltype(x), F)
+    else
+        F = f(x)
+        ∇f = FiniteDiff.finite_difference_derivative(f, x, diff_type(alg), eltype(x), F)
+    end
+
+    fₖ = 0.5 * norm(F)^2
+    H = ∇f * ∇f
+    g = ∇f * F
+
+    for k in 1:maxiters
+        # Solve the trust region subproblem.
+        δ = dogleg_method(H, g, Δ)
+        xₖ₊₁ = x + δ
+        Fₖ₊₁ = f(xₖ₊₁)
+        fₖ₊₁ = 0.5 * norm(Fₖ₊₁)^2
+
+        # Compute the ratio of the actual to predicted reduction.
+        model = -(δ' * g + 0.5 * δ' * H * δ)
+        r = (fₖ - fₖ₊₁) / model
+
+        # Update the trust region radius.
+        if r < η₂
+            Δ *= t₁
+        if r > η₁
+            if isapprox(x̂, x, atol = atol, rtol = rtol)
+                return SciMLBase.build_solution(prob, alg, x, F;
+                                                retcode = ReturnCode.Success)
+            end
+
+            x = xₖ₊₁
+            F = Fₖ₊₁
+            if alg_autodiff(alg)
+                F, ∇f = value_derivative(f, x)
+            elseif x isa AbstractArray
+                ∇f = FiniteDiff.finite_difference_jacobian(f, x, diff_type(alg), eltype(x),
+                                                            F)
+            else
+                ∇f = FiniteDiff.finite_difference_derivative(f, x, diff_type(alg),
+                                                              eltype(x),
+                                                              F)
+            end
+
+            iszero(F) &&
+                return SciMLBase.build_solution(prob, alg, x, F;
+                                                retcode = ReturnCode.Success)
+            # Update the trust region radius.
+            if r > η₃ && norm(δ) ≈ Δ
+                Δ = min(t₂ * Δ, Δₘₐₓ)
+            end
+            fₖ = f̂
+            H = ∇f * ∇f
+            g = ∇f * F
+        end
+    end
+
+    return SciMLBase.build_solution(prob, alg, x, F; retcode = ReturnCode.MaxIters)
+end

src/utils.jl

@@ -43,3 +43,28 @@ function init_J(x)
     end
     return J
 end
+
+function dogleg_method(H, g, Δ)
+    # Compute the Newton step.
+    δN = -H \ g
+    # Test if the full step is within the trust region.
+    if norm(δN) ≤ Δ
+        return δN
+    end
+
+    # Calcualte Cauchy point, optimum along the steepest descent direction.
+    δsd = -g
+    norm_δsd = norm(δsd)
+    if norm_δsd ≥ Δ
+        return δsd .* Δ / norm_δsd
+    end
+
+    # Find the intersection point on the boundary.
+    δN_δsd = δN - δsd
+    dot_δN_δsd = dot(δN_δsd, δN_δsd)
+    dot_δsd_δN_δsd = dot(δsd, δN_δsd)
+    dot_δsd = dot(δsd, δsd)
+    fact = dot_δsd_δN_δsd^2 - dot_δN_δsd * (dot_δsd - Δ^2)
+    tau = (-dot_δsd_δN_δsd + sqrt(fact)) / dot_δN_δsd
+    return δsd + tau * δN_δsd
+end

test/basictests.jl

@@ -46,13 +46,24 @@ sol = benchmark_scalar(sf, csu0)
 @test sol.u * sol.u - 2 < 1e-9
 @test (@ballocated benchmark_scalar(sf, csu0)) == 0
 
+# SimpleNewtonRaphsonTrustRegion
+function benchmark_scalar(f, u0)
+    probN = NonlinearProblem{false}(f, u0)
+    sol = (solve(probN, SimpleNewtonRaphsonTrustRegion(1.0)))
+end
+
+sol = benchmark_scalar(sf, csu0)
+@test sol.retcode === ReturnCode.Success
+@test sol.u * sol.u - 2 < 1e-9
+
 # AD Tests
 using ForwardDiff
 
 # Immutable
 f, u0 = (u, p) -> u .* u .- p, @SVector[1.0, 1.0]
 
-for alg in [SimpleNewtonRaphson(), Broyden(), Klement()]
+for alg in [SimpleNewtonRaphson(), Broyden(), Klement(),
+            SimpleNewtonRaphsonTrustRegion(10.0)]
     g = function (p)
         probN = NonlinearProblem{false}(f, csu0, p)
         sol = solve(probN, alg, tol = 1e-9)
@@ -67,7 +78,8 @@ end
 
 # Scalar
 f, u0 = (u, p) -> u * u - p, 1.0
-for alg in [SimpleNewtonRaphson(), Broyden(), Klement()]
+for alg in [SimpleNewtonRaphson(), Broyden(), Klement(),
+            SimpleNewtonRaphsonTrustRegion(10.0)]
     g = function (p)
         probN = NonlinearProblem{false}(f, oftype(p, u0), p)
         sol = solve(probN, alg)
@@ -108,7 +120,8 @@ for alg in [Bisection(), Falsi()]
     @test ForwardDiff.jacobian(g, p) ≈ ForwardDiff.jacobian(t, p)
 end
 
-for alg in [SimpleNewtonRaphson(), Broyden(), Klement()]
+for alg in [SimpleNewtonRaphson(), Broyden(), Klement(),
+            SimpleNewtonRaphsonTrustRegion(1.0)]
     global g, p
     g = function (p)
         probN = NonlinearProblem{false}(f, 0.5, p)
@@ -128,6 +141,11 @@ probN = NonlinearProblem(f, u0)
 @test solve(probN, SimpleNewtonRaphson(; autodiff = false)).u[end] ≈ sqrt(2.0)
 @test solve(probN, SimpleNewtonRaphson(; autodiff = false)).u[end] ≈ sqrt(2.0)
 
+@test solve(probN, SimpleNewtonRaphsonTrustRegion(1.0)).u[end] ≈ sqrt(2.0)
+@test solve(probN, SimpleNewtonRaphsonTrustRegion(1.0); immutable = false).u[end] ≈ sqrt(2.0)
+@test solve(probN, SimpleNewtonRaphsonTrustRegion(1.0; autodiff = false)).u[end] ≈ sqrt(2.0)
+@test solve(probN, SimpleNewtonRaphsonTrustRegion(1.0; autodiff = false)).u[end] ≈ sqrt(2.0)
+
 @test solve(probN, Broyden()).u[end] ≈ sqrt(2.0)
 @test solve(probN, Broyden(); immutable = false).u[end] ≈ sqrt(2.0)
 
@@ -144,6 +162,10 @@ for u0 in [1.0, [1, 1.0]]
     @test solve(probN, SimpleNewtonRaphson()).u ≈ sol
     @test solve(probN, SimpleNewtonRaphson(; autodiff = false)).u ≈ sol
 
+    @test solve(probN, SimpleNewtonRaphsonTrustRegion(1.0)).u ≈ sol
+    @test solve(probN, SimpleNewtonRaphsonTrustRegion(1.0)).u ≈ sol
+    @test solve(probN, SimpleNewtonRaphsonTrustRegion(1.0; autodiff = false)).u ≈ sol
+
     @test solve(probN, Broyden()).u ≈ sol
 
     @test solve(probN, Klement()).u ≈ sol