Prototype SimpleJNFK

Author: avik-pal

Project.toml

@@ -1,18 +1,20 @@
 name = "SimpleNonlinearSolve"
 uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 authors = ["SciML"]
-version = "0.1.20"
+version = "0.1.21"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 
 [weakdeps]

src/SimpleNonlinearSolve.jl

@@ -2,9 +2,10 @@ module SimpleNonlinearSolve
 
 using Reexport
 using FiniteDiff, ForwardDiff
-using ForwardDiff: Dual
+using ForwardDiff: Dual, Partials, Tag
 using StaticArraysCore
 using LinearAlgebra
+using LinearSolve
 import ArrayInterface
 using DiffEqBase
 
@@ -39,6 +40,7 @@ include("ad.jl")
 include("halley.jl")
 include("alefeld.jl")
 include("itp.jl")
+include("jnfk.jl")
 
 # Batched Solver Support
 include("batched/utils.jl")
@@ -77,7 +79,7 @@ end
 
 # DiffEq styled algorithms
 export Bisection, Brent, Broyden, LBroyden, SimpleDFSane, Falsi, Halley, Klement,
-    Ridder, SimpleNewtonRaphson, SimpleTrustRegion, Alefeld, ITP
+    Ridder, SimpleNewtonRaphson, SimpleTrustRegion, Alefeld, ITP, SimpleJFNK
 export BatchedBroyden, BatchedSimpleNewtonRaphson, BatchedSimpleDFSane
 
 end # module

src/jnfk.jl

@@ -0,0 +1,55 @@
+struct SimpleJNFKJacVecTag end
+
+function jvp_forwarddiff(f, x::AbstractArray{T}, v) where {T}
+    v_ = reshape(v, axes(x))
+    y = (Dual{Tag{SimpleJNFKJacVecTag, T}, T, 1}).(x, Partials.(tuple.(v_)))
+    return vec(ForwardDiff.partials.(vec(f(y)), 1))
+end
+
+struct JacVecOperator{F, X}
+    f::F
+    x::X
+end
+
+(jvp::JacVecOperator)(v, _, _) = jvp_forwarddiff(jvp.f, jvp.x, v)
+
+"""
+    SimpleJNFK()
+
+"""
+struct SimpleJFNK end
+
+function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleJFNK, args...;
+    abstol = nothing, reltol= nothing, maxiters = 1000, linsolve_kwargs = (;), kwargs...)
+    iip = SciMLBase.isinplace(prob)
+    @assert !iip "SimpleJFNK does not support inplace problems"
+
+    f = Base.Fix2(prob.f, prob.p)
+    x = float(prob.u0)
+    fx = f(x)
+    T = typeof(x)
+
+    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)
+
+    op = FunctionOperator(JacVecOperator(f, x), x)
+    linprob = LinearProblem(op, -fx)
+    lincache = init(linprob, SimpleGMRES(); abstol, reltol, maxiters, linsolve_kwargs...)
+
+    for i in 1:maxiters
+        iszero(fx) &&
+            return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.Success)
+
+        linsol = solve!(lincache)
+        x .-= linsol.u
+        lincache = linsol.cache
+
+        # FIXME: not nothing
+        if isapprox(x, nothing; atol, rtol)
+            return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.Success)
+        end
+    end
+
+    return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.MaxIters)
+end