Update batched Raphson to use same parameters as unbatched

Author: avik-pal

Project.toml

@@ -16,21 +16,16 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 
 [weakdeps]
-AbstractDifferentiation = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"
-LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 
 [extensions]
-SimpleNonlinearSolveADLinearSolveExt = ["AbstractDifferentiation", "LinearSolve"]
 SimpleNonlinearSolveNNlibExt = "NNlib"
 
 [compat]
-AbstractDifferentiation = "0.5"
 ArrayInterface = "6, 7"
 DiffEqBase = "6.126"
 FiniteDiff = "2"
 ForwardDiff = "0.10.3"
-LinearSolve = "2"
 NNlib = "0.8, 0.9"
 PackageExtensionCompat = "1"
 PrecompileTools = "1"
@@ -40,6 +35,4 @@ StaticArraysCore = "1.4"
 julia = "1.6"
 
 [extras]
-AbstractDifferentiation = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"
-LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"

ext/SimpleNonlinearSolveADLinearSolveExt.jl

@@ -15,7 +15,6 @@ function __init__()
     @require_extensions
 end
 
-const ADLinearSolveExtLoaded = Ref{Bool}(false)
 const NNlibExtLoaded = Ref{Bool}(false)
 
 abstract type AbstractSimpleNonlinearSolveAlgorithm <: SciMLBase.AbstractNonlinearAlgorithm end
@@ -78,6 +77,6 @@ end
 # DiffEq styled algorithms
 export Bisection, Brent, Broyden, LBroyden, SimpleDFSane, Falsi, Halley, Klement,
     Ridder, SimpleNewtonRaphson, SimpleTrustRegion, Alefeld
-export BatchedBroyden, SimpleBatchedNewtonRaphson, SimpleBatchedDFSane
+export BatchedBroyden, BatchedSimpleNewtonRaphson, BatchedSimpleDFSane
 
 end # module

src/SimpleNonlinearSolve.jl

@@ -1,4 +1,4 @@
-Base.@kwdef struct SimpleBatchedDFSane{T, F, TC <: NLSolveTerminationCondition} <:
+Base.@kwdef struct BatchedSimpleDFSane{T, F, TC <: NLSolveTerminationCondition} <:
                    AbstractBatchedNonlinearSolveAlgorithm
     σₘᵢₙ::T = 1.0f-10
     σₘₐₓ::T = 1.0f+10
@@ -16,7 +16,7 @@ Base.@kwdef struct SimpleBatchedDFSane{T, F, TC <: NLSolveTerminationCondition}
 end
 
 function SciMLBase.__solve(prob::NonlinearProblem,
-    alg::SimpleBatchedDFSane,
+    alg::BatchedSimpleDFSane,
     args...;
     abstol = nothing,
     reltol = nothing,

src/batched/dfsane.jl

@@ -1,8 +1,77 @@
-struct SimpleBatchedNewtonRaphson{AD, LS, TC <: NLSolveTerminationCondition} <:
+struct BatchedSimpleNewtonRaphson{CS, AD, FDT, TC <: NLSolveTerminationCondition} <:
        AbstractBatchedNonlinearSolveAlgorithm
-    autodiff::AD
-    linsolve::LS
     termination_condition::TC
 end
 
-# Implementation of solve using Package Extensions
+alg_autodiff(alg::BatchedSimpleNewtonRaphson{CS, AD, FDT}) where {CS, AD, FDT} = AD
+diff_type(alg::BatchedSimpleNewtonRaphson{CS, AD, FDT}) where {CS, AD, FDT} = FDT
+
+function BatchedSimpleNewtonRaphson(; chunk_size = Val{0}(),
+    autodiff = Val{true}(),
+    diff_type = Val{:forward},
+    termination_condition = NLSolveTerminationCondition(NLSolveTerminationMode.NLSolveDefault;
+        abstol = nothing,
+        reltol = nothing))
+    return BatchedSimpleNewtonRaphson{SciMLBase._unwrap_val(chunk_size),
+        SciMLBase._unwrap_val(autodiff),
+        SciMLBase._unwrap_val(diff_type), typeof(termination_condition)}(termination_condition)
+end
+
+function SciMLBase.__solve(prob::NonlinearProblem, alg::BatchedSimpleNewtonRaphson;
+    abstol = nothing, reltol = nothing, maxiters = 1000, kwargs...)
+    iip = SciMLBase.isinplace(prob)
+    @assert !iip "BatchedSimpleNewtonRaphson currently only supports out-of-place nonlinear problems."
+    u, f, reconstruct = _construct_batched_problem_structure(prob)
+
+    tc = alg.termination_condition
+    mode = DiffEqBase.get_termination_mode(tc)
+
+    storage = _get_storage(mode, u)
+
+    xₙ, xₙ₋₁ = copy(u), copy(u)
+    T = eltype(u)
+
+    atol = _get_tolerance(abstol, tc.abstol, T)
+    rtol = _get_tolerance(reltol, tc.reltol, T)
+    termination_condition = tc(storage)
+
+    for i in 1:maxiters
+        if alg_autodiff(alg)
+            fₙ, 𝓙 = value_derivative(f, xₙ)
+        else
+            fₙ = f(xₙ)
+            𝓙 = FiniteDiff.finite_difference_jacobian(f, xₙ, diff_type(alg), eltype(xₙ), fₙ)
+        end
+
+        iszero(fₙ) && return DiffEqBase.build_solution(prob,
+            alg,
+            reconstruct(xₙ),
+            reconstruct(fₙ);
+            retcode = ReturnCode.Success)
+
+        δx = reshape(𝓙 \ vec(fₙ), size(xₙ))
+        xₙ .-= δx
+
+        if termination_condition(fₙ, xₙ, xₙ₋₁, atol, rtol)
+            retcode, xₙ, fₙ = _result_from_storage(storage, xₙ, fₙ, f, mode, iip)
+            return DiffEqBase.build_solution(prob,
+                alg,
+                reconstruct(xₙ),
+                reconstruct(fₙ);
+                retcode)
+        end
+
+        xₙ₋₁ .= xₙ
+    end
+
+    if mode ∈ DiffEqBase.SAFE_BEST_TERMINATION_MODES
+        xₙ = storage.u
+        fₙ = f(xₙ)
+    end
+
+    return DiffEqBase.build_solution(prob,
+        alg,
+        reconstruct(xₙ),
+        reconstruct(fₙ);
+        retcode = ReturnCode.MaxIters)
+end

src/batched/raphson.jl

@@ -73,7 +73,7 @@ function SimpleDFSane(; σ_min::Real = 1e-10, σ_max::Real = 1e10, σ_1::Real =
     batched::Bool = false,
     max_inner_iterations = 1000)
     if batched
-        return SimpleBatchedDFSane(; σₘᵢₙ = σ_min,
+        return BatchedSimpleDFSane(; σₘᵢₙ = σ_min,
             σₘₐₓ = σ_max,
             σ₁ = σ_1,
             M,

src/dfsane.jl

@@ -2,7 +2,8 @@
     SimpleNewtonRaphson(; batched = false,
         chunk_size = Val{0}(),
         autodiff = Val{true}(),
-        diff_type = Val{:forward})
+        diff_type = Val{:forward},
+        termination_condition = missing)
 
 A low-overhead implementation of Newton-Raphson. This method is non-allocating on scalar
 and static array problems.
@@ -27,11 +28,8 @@ and static array problems.
 - `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.
-
-!!! note
-
-    To use the `batched` version, remember to load `AbstractDifferentiation` and
-    `LinearSolve`.
+- `termination_condition`: control the termination of the algorithm. (Only works for batched
+  problems)
 """
 struct SimpleNewtonRaphson{CS, AD, FDT} <: AbstractNewtonAlgorithm{CS, AD, FDT} end
 
@@ -44,16 +42,19 @@ function SimpleNewtonRaphson(; batched = false,
         throw(ArgumentError("`termination_condition` is currently only supported for batched problems"))
     end
     if batched
-        @assert ADLinearSolveExtLoaded[] "Please install and load `LinearSolve.jl` and `AbstractDifferentiation.jl` to use batched Newton-Raphson."
+        # @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,
             reltol = nothing) :
                                 termination_condition
-        return SimpleBatchedNewtonRaphson(; chunk_size,
+        return BatchedSimpleNewtonRaphson(; chunk_size,
             autodiff,
             diff_type,
             termination_condition)
+        return SimpleNewtonRaphson{SciMLBase._unwrap_val(chunk_size),
+            SciMLBase._unwrap_val(autodiff),
+            SciMLBase._unwrap_val(diff_type)}()
     end
     return SimpleNewtonRaphson{SciMLBase._unwrap_val(chunk_size),
         SciMLBase._unwrap_val(autodiff),

src/raphson.jl

@@ -1,10 +1,8 @@
 [deps]
-AbstractDifferentiation = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"

test/Project.toml

@@ -1,12 +1,13 @@
 using SimpleNonlinearSolve, StaticArrays, BenchmarkTools, DiffEqBase, LinearAlgebra, Test,
-    NNlib, AbstractDifferentiation, LinearSolve
+    NNlib
 
 const BATCHED_BROYDEN_SOLVERS = []
 const BROYDEN_SOLVERS = []
 const BATCHED_LBROYDEN_SOLVERS = []
 const LBROYDEN_SOLVERS = []
 const BATCHED_DFSANE_SOLVERS = []
 const DFSANE_SOLVERS = []
+const BATCHED_RAPHSON_SOLVERS = []
 
 for mode in instances(NLSolveTerminationMode.T)
     if mode ∈
@@ -23,6 +24,12 @@ for mode in instances(NLSolveTerminationMode.T)
     push!(BATCHED_LBROYDEN_SOLVERS, LBroyden(; batched = true, termination_condition))
     push!(DFSANE_SOLVERS, SimpleDFSane(; batched = false, termination_condition))
     push!(BATCHED_DFSANE_SOLVERS, SimpleDFSane(; batched = true, termination_condition))
+    push!(BATCHED_RAPHSON_SOLVERS,
+        SimpleNewtonRaphson(; batched = true,
+            termination_condition))
+    push!(BATCHED_RAPHSON_SOLVERS,
+        SimpleNewtonRaphson(; batched = true, autodiff = false,
+            termination_condition))
 end
 
 # SimpleNewtonRaphson
@@ -483,7 +490,8 @@ sol = solve(probN, Broyden(batched = true))
 
 @testset "Batched Solver: $(nameof(typeof(alg)))" for alg in (BATCHED_BROYDEN_SOLVERS...,
     BATCHED_LBROYDEN_SOLVERS...,
-    BATCHED_DFSANE_SOLVERS...)
+    BATCHED_DFSANE_SOLVERS...,
+    BATCHED_RAPHSON_SOLVERS...)
     sol = solve(probN, alg; abstol = 1e-3, reltol = 1e-3)
 
     @test sol.retcode == ReturnCode.Success

test/basictests.jl

@@ -1,7 +1,7 @@
 using SimpleNonlinearSolve, StaticArrays, BenchmarkTools, DiffEqBase, LinearAlgebra, Test,
-    NNlib, AbstractDifferentiation, LinearSolve
+    NNlib
 
-# Supported Solvers: BatchedBroyden, SimpleBatchedDFSane
+# Supported Solvers: BatchedBroyden, BatchedSimpleDFSane
 function f!(du::AbstractArray{<:Number, N},
     u::AbstractArray{<:Number, N},
     p::AbstractVector) where {N}