SimpleJNFK working version

Author: avik-pal

Diff for: ext/SimpleNonlinearSolveNNlibExt.jl

@@ -4,10 +4,7 @@ using ArrayInterface, DiffEqBase, LinearAlgebra, NNlib, SimpleNonlinearSolve, Sc
 import SimpleNonlinearSolve: _construct_batched_problem_structure,
     _get_storage, _init_𝓙, _result_from_storage, _get_tolerance, @maybeinplace
 
-function __init__()
-    SimpleNonlinearSolve.NNlibExtLoaded[] = true
-    return
-end
+SimpleNonlinearSolve.extension_loaded(::Val{NNlib}) = true
 
 @views function SciMLBase.__solve(prob::NonlinearProblem,
     alg::BatchedBroyden;

Diff for: src/SimpleNonlinearSolve.jl

@@ -16,7 +16,7 @@ function __init__()
     @require_extensions
 end
 
-const NNlibExtLoaded = Ref{Bool}(false)
+extension_loaded(::Val) = false
 
 abstract type AbstractSimpleNonlinearSolveAlgorithm <: SciMLBase.AbstractNonlinearAlgorithm end
 abstract type AbstractBracketingAlgorithm <: AbstractSimpleNonlinearSolveAlgorithm end

Diff for: src/broyden.jl

@@ -22,7 +22,7 @@ function Broyden(; batched = false,
         abstol = nothing,
         reltol = nothing))
     if batched
-        @assert NNlibExtLoaded[] "Please install and load `NNlib.jl` to use batched Broyden."
+        @assert extension_loaded(Val(:NNlib)) "Please install and load `NNlib.jl` to use batched Broyden."
         return BatchedBroyden(termination_condition)
     end
     return Broyden(termination_condition)

Diff for: src/jnfk.jl

@@ -5,22 +5,32 @@ function jvp_forwarddiff(f, x::AbstractArray{T}, v) where {T}
     y = (Dual{Tag{SimpleJNFKJacVecTag, T}, T, 1}).(x, Partials.(tuple.(v_)))
     return vec(ForwardDiff.partials.(vec(f(y)), 1))
 end
+jvp_forwarddiff!(r, f, x, v) = copyto!(r, jvp_forwarddiff(f, x, v))
 
 struct JacVecOperator{F, X}
     f::F
     x::X
 end
 
 (jvp::JacVecOperator)(v, _, _) = jvp_forwarddiff(jvp.f, jvp.x, v)
+(jvp::JacVecOperator)(r, v, _, _) = jvp_forwarddiff!(r, jvp.f, jvp.x, v)
 
 """
-    SimpleJNFK()
+    SimpleJNFK(; batched::Bool = false)
 
+A low overhead Jacobian-free Newton-Krylov method. This method internally uses `GMRES` to
+avoid computing the Jacobian Matrix.
+
+!!! warning
+
+    JNFK doesn't work well without preconditioning, which is currently not supported. We
+    recommend using `NewtonRaphson(linsolve = KrylovJL_GMRES())` for preconditioning
+    support.
 """
 struct SimpleJFNK end
 
 function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleJFNK, args...;
-    abstol = nothing, reltol= nothing, maxiters = 1000, linsolve_kwargs = (;), kwargs...)
+    abstol = nothing, reltol = nothing, maxiters = 1000, linsolve_kwargs = (;), kwargs...)
     iip = SciMLBase.isinplace(prob)
     @assert !iip "SimpleJFNK does not support inplace problems"
 
@@ -29,26 +39,30 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleJFNK, args...;
     fx = f(x)
     T = typeof(x)
 
+    iszero(fx) &&
+        return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.Success)
+
     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...)
+    linprob = LinearProblem(op, vec(fx))
+    lincache = init(linprob, KrylovJL_GMRES(); abstol = atol, reltol = rtol, 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
+        axpy!(-1, linsol.u, x)
         lincache = linsol.cache
 
-        # FIXME: not nothing
-        if isapprox(x, nothing; atol, rtol)
+        fx = f(x)
+
+        norm(fx, Inf) ≤ atol &&
             return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.Success)
-        end
+
+        lincache.b = vec(fx)
+        lincache.A = FunctionOperator(JacVecOperator(f, x), x)
     end
 
     return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.MaxIters)

Diff for: src/raphson.jl

@@ -42,7 +42,6 @@ function SimpleNewtonRaphson(; batched = false,
         throw(ArgumentError("`termination_condition` is currently only supported for batched problems"))
     end
     if batched
-        # @assert ADLinearSolveFDExtLoaded[] "Please install and load `LinearSolve.jl`, `FiniteDifferences.jl` and `AbstractDifferentiation.jl` to use batched Newton-Raphson."
         termination_condition = ismissing(termination_condition) ?
                                 NLSolveTerminationCondition(NLSolveTerminationMode.NLSolveDefault;
             abstol = nothing,