Add SIAMFANLEquations wrapper

ErikQQY · ErikQQY · commit fc437fd6372e · 2023-12-19T11:30:04.000+08:00
Signed-off-by: ErikQQY &lt;2283984853@qq.com&gt;
diff --git a/Project.toml b/Project.toml
@@ -35,6 +35,7 @@ FastLevenbergMarquardt = "7a0df574-e128-4d35-8cbd-3d84502bf7ce"
 LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891"
 MINPACK = "4854310b-de5a-5eb6-a2a5-c1dee2bd17f9"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
+SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
@@ -44,6 +45,7 @@ NonlinearSolveFastLevenbergMarquardtExt = "FastLevenbergMarquardt"
 NonlinearSolveLeastSquaresOptimExt = "LeastSquaresOptim"
 NonlinearSolveMINPACKExt = "MINPACK"
 NonlinearSolveNLsolveExt = "NLsolve"
+NonlinearSolveSIAMFANLEquationsExt = "SIAMFANLEquations"
 NonlinearSolveSymbolicsExt = "Symbolics"
 NonlinearSolveZygoteExt = "Zygote"
 
@@ -80,6 +82,7 @@ Reexport = "1.2"
 SafeTestsets = "0.1"
 SciMLBase = "2.11"
 SciMLOperators = "0.3.7"
+SIAMFANLEquations = "1.0.1"
 SimpleNonlinearSolve = "1.0.2"
 SparseArrays = "<0.0.1, 1"
 SparseDiffTools = "2.14"
@@ -109,6 +112,7 @@ NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4"
 SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -117,4 +121,4 @@ 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"]
+test = ["Aqua", "Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools", "NonlinearProblemLibrary", "LeastSquaresOptim", "FastLevenbergMarquardt", "NaNMath", "BandedMatrices", "DiffEqBase", "StableRNGs", "MINPACK", "NLsolve", "SIAMFANLEquations"]
diff --git a/docs/src/api/fastlevenbergmarquardt.md b/docs/src/api/fastlevenbergmarquardt.md
@@ -1,6 +1,6 @@
 # FastLevenbergMarquardt.jl
 
-This is a extension for importing solvers from FastLevenbergMarquardt.jl into the SciML
+This is an extension for importing solvers from FastLevenbergMarquardt.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:
 
diff --git a/docs/src/api/leastsquaresoptim.md b/docs/src/api/leastsquaresoptim.md
@@ -1,6 +1,6 @@
 # LeastSquaresOptim.jl
 
-This is a extension for importing solvers from LeastSquaresOptim.jl into the SciML
+This is an extension for importing solvers from LeastSquaresOptim.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:
 
diff --git a/docs/src/api/siamfanlequations.md b/docs/src/api/siamfanlequations.md
@@ -0,0 +1,17 @@
+# SIAMFANLEquations.jl
+
+This is an extension for importing solvers from [SIAMFANLEquations.jl](https://github.com/ctkelley/SIAMFANLEquations.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("SIAMFANLEquations")
+using SIAMFANLEquations, NonlinearSolve
+```
+
+## Solver API
+
+```@docs
+SIAMFANLEquationsJL
+```
diff --git a/ext/NonlinearSolveSIAMFANLEquationsExt.jl b/ext/NonlinearSolveSIAMFANLEquationsExt.jl
@@ -0,0 +1,176 @@
+module NonlinearSolveSIAMFANLEquationsExt
+
+using NonlinearSolve, SciMLBase
+using SIAMFANLEquations
+import ConcreteStructs: @concrete
+import UnPack: @unpack
+import FiniteDiff, ForwardDiff
+
+function SciMLBase.__solve(prob::NonlinearProblem, alg::SIAMFANLEquationsJL, args...; abstol = 1e-8,
+        reltol = 1e-8, alias_u0::Bool = false, maxiters = 1000, kwargs...)
+    @unpack method, autodiff, show_trace, delta, linsolve = alg
+
+    iip = SciMLBase.isinplace(prob)
+    if typeof(prob.u0) <: Number
+        f! = if iip
+            function (u)
+                du = similar(u)
+                prob.f(du, u, prob.p)
+                return du
+            end
+        else
+            u -> prob.f(u, prob.p)
+        end
+
+        if method == :newton
+            res = nsolsc(f!, prob.u0; maxit = maxiters, atol = abstol, rtol = reltol, printerr = show_trace)
+        elseif method == :pseudotransient
+            res = ptcsolsc(f!, prob.u0; delta0 = delta, maxit = maxiters, atol = abstol, rtol=reltol, printerr = show_trace)
+        elseif method == :secant
+            res = secant(f!, prob.u0; maxit = maxiters, atol = abstol, rtol = reltol, printerr = show_trace)
+        end
+
+        if res.errcode == 0
+            retcode = ReturnCode.Success
+        elseif res.errcode == 10
+            retcode = ReturnCode.MaxIters
+        elseif res.errcode == 1
+            retcode = ReturnCode.Failure
+            @error("Line search failed")
+        elseif res.errcode == -1
+            retcode = ReturnCode.Default
+            @info("Initial iterate satisfies the termination criteria")
+        end
+        stats = method == :pseudotransient ? nothing : (SciMLBase.NLStats(res.stats.ifun[1], res.stats.ijac[1], -1, -1, res.stats.iarm[1]))
+        return SciMLBase.build_solution(prob, alg, res.solution, res.history; retcode, stats)
+    else
+        u = NonlinearSolve.__maybe_unaliased(prob.u0, alias_u0)
+    end
+
+    fu = NonlinearSolve.evaluate_f(prob, u)
+
+    if iip
+        f! = function (du, u)
+            prob.f(du, u, prob.p)
+            return du
+        end
+    else
+        f! = function (du, u)
+            du .= prob.f(u, prob.p)
+            return du
+        end
+    end
+
+    # Allocate ahead for function and Jacobian
+    N = length(u)
+    FS = zeros(eltype(u), N)
+    FPS = zeros(eltype(u), N, N)
+    # Allocate ahead for Krylov basis
+
+    # Jacobian free Newton Krylov
+    if linsolve !== nothing
+        JVS = linsolve == :gmres ? zeros(eltype(u), N, 3) : zeros(eltype(u), N)
+        # `linsolve` as a Symbol to keep unified interface with other EXTs, SIAMFANLEquations directly use String to choose between linear solvers
+        linsolve_alg = strip(repr(linsolve), ':')
+
+        if method == :newton
+            res = nsoli(f!, u, FS, JVS; lsolver = linsolve_alg, maxit = maxiters, atol = abstol, rtol = reltol, printerr = show_trace)
+        elseif method == :pseudotransient
+            res = ptcsoli(f!, u, FS, JVS; lsolver = linsolve_alg, maxit = maxiters, atol = abstol, rtol = reltol, printerr = show_trace)
+        end
+        
+        if res.errcode == 0
+            retcode = ReturnCode.Success
+        elseif res.errcode == 10
+            retcode = ReturnCode.MaxIters
+        elseif res.errcode == 1
+            retcode = ReturnCode.Failure
+            @error("Line search failed")
+        elseif res.errcode == -1
+            retcode = ReturnCode.Default
+            @info("Initial iterate satisfies the termination criteria")
+        end
+        stats = method == :pseudotransient ? nothing : (SciMLBase.NLStats(res.stats.ifun[1], res.stats.ijac[1], -1, -1, res.stats.iarm[1]))
+        return SciMLBase.build_solution(prob, alg, res.solution, res.history; retcode, stats)
+    end
+
+    if prob.f.jac === nothing
+        use_forward_diff = if alg.autodiff === nothing
+            ForwardDiff.can_dual(eltype(u))
+        else
+            alg.autodiff isa AutoForwardDiff
+        end
+        uf = SciMLBase.JacobianWrapper{iip}(prob.f, prob.p)
+        if use_forward_diff
+            cache = iip ? ForwardDiff.JacobianConfig(uf, fu, u) :
+                    ForwardDiff.JacobianConfig(uf, u)
+        else
+            cache = FiniteDiff.JacobianCache(u, fu)
+        end
+        J! = if iip
+            if use_forward_diff
+                fu_cache = similar(fu)
+                function (J, x, p)
+                    uf.p = p
+                    ForwardDiff.jacobian!(J, uf, fu_cache, x, cache)
+                    return J
+                end
+            else
+                function (J, x, p)
+                    uf.p = p
+                    FiniteDiff.finite_difference_jacobian!(J, uf, x, cache)
+                    return J
+                end
+            end
+        else
+            if use_forward_diff
+                function (J, x, p)
+                    uf.p = p
+                    ForwardDiff.jacobian!(J, uf, x, cache)
+                    return J
+                end
+            else
+                function (J, x, p)
+                    uf.p = p
+                    J_ = FiniteDiff.finite_difference_jacobian(uf, x, cache)
+                    copyto!(J, J_)
+                    return J
+                end
+            end
+        end
+    else
+        J! = prob.f.jac
+    end
+
+    AJ!(J, u, x) = J!(J, x, prob.p)    
+
+    if method == :newton
+        res = nsol(f!, u, FS, FPS, AJ!;
+                    sham=1, rtol = reltol, atol = abstol, maxit = maxiters,
+                    printerr = show_trace)
+    elseif method == :pseudotransient
+        res = ptcsol(f!, u, FS, FPS, AJ!;
+                    rtol = reltol, atol = abstol, maxit = maxiters,
+                    delta0 = delta, printerr = show_trace)
+        
+    end
+
+    if res.errcode == 0
+        retcode = ReturnCode.Success
+    elseif res.errcode == 10
+        retcode = ReturnCode.MaxIters
+    elseif res.errcode == 1
+        retcode = ReturnCode.Failure
+        @error("Line search failed")
+    elseif res.errcode == -1
+        retcode = ReturnCode.Default
+        @info("Initial iterate satisfies the termination criteria")
+    end
+
+
+    # pseudo transient continuation has a fixed cost per iteration, iteration statistics are not interesting here.
+    stats = method == :pseudotransient ? nothing : (SciMLBase.NLStats(res.stats.ifun[1], res.stats.ijac[1], -1, -1, res.stats.iarm[1]))
+    return SciMLBase.build_solution(prob, alg, res.solution, res.history; retcode, stats)
+end
+
+end
diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl
@@ -236,7 +236,7 @@ export RadiusUpdateSchemes
 
 export NewtonRaphson, TrustRegion, LevenbergMarquardt, DFSane, GaussNewton, PseudoTransient,
     Broyden, Klement, LimitedMemoryBroyden
-export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL
+export LeastSquaresOptimJL, FastLevenbergMarquardtJL, CMINPACK, NLsolveJL, SIAMFANLEquationsJL
 export NonlinearSolvePolyAlgorithm,
     RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg
 
diff --git a/src/extension_algs.jl b/src/extension_algs.jl
@@ -206,3 +206,35 @@ 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
+
+"""
+    SIAMFANLEquationsJL(; method = :newton, autodiff = :central)
+
+### Keyword Arguments
+
+  - `method`: the choice of method for solving the nonlinear system.
+  - `autodiff`: the choice of method for generating the Jacobian. Defaults to `:central` or
+    central differencing via FiniteDiff.jl. The other choices are `:forward`.
+  - `show_trace`: whether to show the trace.
+  - `delta`: initial pseudo time step, default is 1e-3.
+  - `linsolve` : JFNK linear solvers, choices are `gmres` and `bicgstab`.
+
+### Submethod Choice
+
+  - `:newton`: Classical Newton method.
+  - `:pseudotransient`: 
+"""
+@concrete struct SIAMFANLEquationsJL <: AbstractNonlinearAlgorithm
+    method::Symbol
+    autodiff::Symbol
+    show_trace::Bool
+    delta
+    linsolve::Union{Symbol, Nothing}
+end
+
+function SIAMFANLEquationsJL(; method = :newton, autodiff = :central, show_trace = false, delta = 1e-3, linsolve = nothing)
+    if Base.get_extension(@__MODULE__, :NonlinearSolveSIAMFANLEquationsExt) === nothing
+        error("SIAMFANLEquationsJL requires SIAMFANLEquations.jl to be loaded")
+    end
+  return SIAMFANLEquationsJL(method, autodiff, show_trace, delta, linsolve)
+end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -23,6 +23,7 @@ end
     if GROUP == "All" || GROUP == "Wrappers"
         @time @safetestset "MINPACK" include("minpack.jl")
         @time @safetestset "NLsolve" include("nlsolve.jl")
+        @time @safetestset "SIAMFANLEquations" include("siamfanlequations.jl")
     end
 
     if GROUP == "All" || GROUP == "23TestProblems"
diff --git a/test/siamfanlequations.jl b/test/siamfanlequations.jl
