Add wrapper for SpeedMapping

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "NonlinearSolve"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 authors = ["SciML"]
-version = "3.1.2"
+version = "3.2.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -32,18 +32,22 @@ UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 [weakdeps]
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 FastLevenbergMarquardt = "7a0df574-e128-4d35-8cbd-3d84502bf7ce"
+FixedPointAcceleration = "817d07cb-a79a-5c30-9a31-890123675176"
 LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891"
 MINPACK = "4854310b-de5a-5eb6-a2a5-c1dee2bd17f9"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
+SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 NonlinearSolveBandedMatricesExt = "BandedMatrices"
 NonlinearSolveFastLevenbergMarquardtExt = "FastLevenbergMarquardt"
+NonlinearSolveFixedPointAccelerationExt = "FixedPointAcceleration"
 NonlinearSolveLeastSquaresOptimExt = "LeastSquaresOptim"
 NonlinearSolveMINPACKExt = "MINPACK"
 NonlinearSolveNLsolveExt = "NLsolve"
+NonlinearSolveSpeedMappingExt = "SpeedMapping"
 NonlinearSolveSymbolicsExt = "Symbolics"
 NonlinearSolveZygoteExt = "Zygote"
 
@@ -60,6 +64,7 @@ Enzyme = "0.11.11"
 FastBroadcast = "0.2.8"
 FastLevenbergMarquardt = "0.1"
 FiniteDiff = "2.21"
+FixedPointAcceleration = "0.3"
 ForwardDiff = "0.10.36"
 LazyArrays = "1.8.2"
 LeastSquaresOptim = "0.8.5"
@@ -84,6 +89,7 @@ SciMLOperators = "0.3.7"
 SimpleNonlinearSolve = "1.0.2"
 SparseArrays = "<0.0.1, 1"
 SparseDiffTools = "2.14"
+SpeedMapping = "0.3"
 StableRNGs = "1"
 StaticArrays = "1.7"
 Symbolics = "5.13"
@@ -99,6 +105,7 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FastLevenbergMarquardt = "7a0df574-e128-4d35-8cbd-3d84502bf7ce"
+FixedPointAcceleration = "817d07cb-a79a-5c30-9a31-890123675176"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -112,11 +119,12 @@ Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
+SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices", "DiffEqBase", "StableRNGs", "MINPACK", "NLsolve", "OrdinaryDiffEq"]
+test = ["Aqua", "Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices", "DiffEqBase", "StableRNGs", "MINPACK", "NLsolve", "OrdinaryDiffEq", "SpeedMapping", "FixedPointAcceleration"]

docs/pages.jl

@@ -20,6 +20,7 @@ pages = ["index.md",
         "solvers/BracketingSolvers.md",
         "solvers/SteadyStateSolvers.md",
         "solvers/NonlinearLeastSquaresSolvers.md",
+        "solvers/FixedPointSolvers.md",
         "solvers/LineSearch.md"],
     "Detailed Solver APIs" => Any["api/nonlinearsolve.md",
         "api/simplenonlinearsolve.md",

docs/src/api/fixedpointacceleration.md

@@ -0,0 +1,17 @@
+# FixedPointAcceleration.jl
+
+This is a extension for importing solvers from FixedPointAcceleration.jl into the SciML
+interface. Note that these solvers do not come by default, and thus one needs to install
+the package before using these solvers:
+
+```julia
+using Pkg
+Pkg.add("FixedPointAcceleration")
+using FixedPointAcceleration, NonlinearSolve
+```
+
+## Solver API
+
+```@docs
+FixedPointAccelerationJL
+```

docs/src/api/speedmapping.md

@@ -0,0 +1,17 @@
+# SppedMapping.jl
+
+This is a extension for importing solvers from SppedMapping.jl into the SciML
+interface. Note that these solvers do not come by default, and thus one needs to install
+the package before using these solvers:
+
+```julia
+using Pkg
+Pkg.add("SppedMapping")
+using SppedMapping, NonlinearSolve
+```
+
+## Solver API
+
+```@docs
+SppedMappingJL
+```

docs/src/solvers/FixedPointSolvers.md

@@ -0,0 +1,42 @@
+# Fixed Point Solvers
+
+Currently we don't have an API to directly specify Fixed Point Solvers. However, a Fixed
+Point Problem can be triviall converted to a Root Finding Problem. Say we want to solve:
+
+```math
+f(u) = u
+```
+
+This can be written as:
+
+```math
+g(u) = f(u) - u = 0
+```
+
+Where ``g(u) = 0`` is a root finding problem. Note that we can use any root finding
+algorithm to solve this problem. However, this is often not the most efficient way to
+solve a fixed point problem. We provide a few algorithms available via extensions that
+are more efficient for fixed point problems.
+
+Note that even if you use one of the Fixed Point Solvers mentioned here, you must still
+use the `NonlinearProblem` API to specify the problem, i.e., ``g(u) = 0``.
+
+## Recommended Methods
+
+Using [native NonlinearSolve.jl methods](@ref nonlinearsystemsolvers) is the recommended
+approach. For systems where constructing Jacobian Matrices are expensive, we recommend
+using a Krylov Method with one of those solvers.
+
+## Full List of Methods
+
+We are only listing the methods that natively solve fixed point problems.
+
+### SpeedMapping.jl
+
+  - `SpeedMappingJL()`: accelerates the convergence of a mapping to a fixed point by the
+    Alternating cyclic extrapolation algorithm (ACX).
+
+### FixedPointAcceleration.jl
+
+  - `FixedPointAccelerationJL()`: accelerates the convergence of a mapping to a fixed point
+    by the Anderson acceleration algorithm and a few other methods.

docs/src/solvers/NonlinearSystemSolvers.md

@@ -1,6 +1,6 @@
 # [Nonlinear System Solvers](@id nonlinearsystemsolvers)
 
-`solve(prob::NonlinearProblem,alg;kwargs)`
+`solve(prob::NonlinearProblem, alg; kwargs)`
 
 Solves for ``f(u)=0`` in the problem defined by `prob` using the algorithm
 `alg`. If no algorithm is given, a default algorithm will be chosen.

ext/NonlinearSolveFixedPointAccelerationExt.jl

@@ -0,0 +1,3 @@
+module NonlinearSolveFixedPointAccelerationExt
+
+end

ext/NonlinearSolveSpeedMappingExt.jl

@@ -0,0 +1,49 @@
+module NonlinearSolveSpeedMappingExt
+
+using NonlinearSolve, SpeedMapping, DiffEqBase, SciMLBase
+import UnPack: @unpack
+
+function SciMLBase.__solve(prob::NonlinearProblem, alg::SpeedMappingJL, args...;
+        abstol = nothing, maxiters = 1000, alias_u0::Bool = false,
+        store_trace::Val{store_info} = Val(false), termination_condition = nothing,
+        kwargs...) where {store_info}
+    @assert (termination_condition ===
+             nothing)||(termination_condition isa AbsNormTerminationMode) "SpeedMappingJL does not support termination conditions!"
+
+    if typeof(prob.u0) <: Number
+        u0 = [prob.u0]
+    else
+        u0 = NonlinearSolve.__maybe_unaliased(prob.u0, alias_u0)
+    end
+
+    T = eltype(u0)
+    iip = isinplace(prob)
+    p = prob.p
+
+    if prob.u0 isa Number
+        resid = [NonlinearSolve.evaluate_f(prob, first(u0))]
+    else
+        resid = NonlinearSolve.evaluate_f(prob, u0)
+    end
+
+    if !iip && prob.u0 isa Number
+        m! = (du, u) -> (du .= prob.f(first(u), p) .+ first(u))
+    elseif !iip
+        m! = (du, u) -> (du .= prob.f(u, p) .+ u)
+    else
+        m! = (du, u) -> (prob.f(du, u, p); du .+= u)
+    end
+
+    tol = abstol === nothing ? real(oneunit(T)) * (eps(real(one(T))))^(4 // 5) : abstol
+
+    sol = speedmapping(u0; m!, tol, Lp = Inf, maps_limit = maxiters, alg.orders,
+        alg.check_obj, store_info, alg.σ_min, alg.stabilize)
+    res = prob.u0 isa Number ? first(sol.minimizer) : sol.minimizer
+    resid = NonlinearSolve.evaluate_f(prob, sol.minimizer)
+
+    return SciMLBase.build_solution(prob, alg, res, resid;
+        retcode = sol.converged ? ReturnCode.Success : ReturnCode.Failure,
+        stats = SciMLBase.NLStats(sol.maps, 0, 0, 0, sol.maps), original = sol)
+end
+
+end

src/NonlinearSolve.jl

@@ -236,7 +236,8 @@ export RadiusUpdateSchemes
 
 export NewtonRaphson, TrustRegion, LevenbergMarquardt, DFSane, GaussNewton, PseudoTransient,
     Broyden, Klement, LimitedMemoryBroyden
-export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL
+export LeastSquaresOptimJL,
+    FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, FixedPointAccelerationJL, SpeedMappingJL
 export NonlinearSolvePolyAlgorithm,
     RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg
 

src/extension_algs.jl

@@ -19,7 +19,7 @@ for solving `NonlinearLeastSquaresProblem`.
 
     This algorithm is only available if `LeastSquaresOptim.jl` is installed.
 """
-struct LeastSquaresOptimJL{alg, linsolve} <: AbstractNonlinearSolveAlgorithm
+struct LeastSquaresOptimJL{alg, linsolve} <: AbstractNonlinearAlgorithm
     autodiff::Symbol
 end
 
@@ -58,7 +58,7 @@ for solving `NonlinearLeastSquaresProblem`.
 
     This algorithm is only available if `FastLevenbergMarquardt.jl` is installed.
 """
-@concrete struct FastLevenbergMarquardtJL{linsolve} <: AbstractNonlinearSolveAlgorithm
+@concrete struct FastLevenbergMarquardtJL{linsolve} <: AbstractNonlinearAlgorithm
     autodiff
     factor
     factoraccept
@@ -206,3 +206,45 @@ function NLsolveJL(; method = :trust_region, autodiff = :central, store_trace =
     return NLsolveJL(method, autodiff, store_trace, extended_trace, linesearch, linsolve,
         factor, autoscale, m, beta, show_trace)
 end
+
+"""
+    SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false,
+        orders::Vector{Int} = [3, 3, 2], time_limit::Real = 1000)
+
+Wrapper over [SpeedMapping.jl](https://nicolasl-s.github.io/SpeedMapping.jl) for solving
+Fixed Point Problems. We allow using this algorithm to solve root finding problems as well.
+
+## Arguments:
+
+  - `σ_min`: Setting to `1` may avoid stalling (see paper).
+  - `stabilize`: performs a stabilization mapping before extrapolating. Setting to `true`
+    may improve the performance for applications like accelerating the EM or MM algorithms
+    (see paper).
+  - `check_obj`: In case of NaN or Inf values, the algorithm restarts at the best past
+    iterate.
+  - `orders`: determines ACX's alternating order. Must be between `1` and `3` (where `1`
+    means no extrapolation). The two recommended orders are `[3, 2]` and `[3, 3, 2]`, the
+    latter being potentially better for highly non-linear applications (see paper).
+  - `time_limit`: time limit for the algorithm.
+
+## References:
+
+  - N. Lepage-Saucier, Alternating cyclic extrapolation methods for optimization algorithms,
+    arXiv:2104.04974 (2021). https://arxiv.org/abs/2104.04974.
+"""
+@concrete struct SpeedMappingJL <: AbstractNonlinearAlgorithm
+    σ_min
+    stabilize::Bool
+    check_obj::Bool
+    orders::Vector{Int}
+    time_limit
+end
+
+function SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false,
+        orders::Vector{Int} = [3, 3, 2], time_limit::Real = 1000)
+    if Base.get_extension(@__MODULE__, :NonlinearSolveSpeedMappingExt) === nothing
+        error("SpeedMappingJL requires SpeedMapping.jl to be loaded")
+    end
+
+    return SpeedMappingJL(σ_min, stabilize, check_obj, orders, time_limit)
+end