Format

Author: avik-pal

ext/SimpleNonlinearSolveADLinearSolveExt.jl

@@ -1,8 +1,10 @@
 module SimpleNonlinearSolveADLinearSolveExt
 
-using AbstractDifferentiation, ArrayInterface, DiffEqBase, LinearAlgebra, LinearSolve,
+using AbstractDifferentiation,
+    ArrayInterface, DiffEqBase, LinearAlgebra, LinearSolve,
     SimpleNonlinearSolve, SciMLBase
-import SimpleNonlinearSolve: _construct_batched_problem_structure, _get_storage, _result_from_storage, _get_tolerance, @maybeinplace
+import SimpleNonlinearSolve: _construct_batched_problem_structure,
+    _get_storage, _result_from_storage, _get_tolerance, @maybeinplace
 
 const AD = AbstractDifferentiation
 
@@ -20,19 +22,18 @@ function SimpleNonlinearSolve.SimpleBatchedNewtonRaphson(; chunk_size = Val{0}()
     # TODO: Use `diff_type`. FiniteDiff.jl is currently not available in AD.jl
     chunksize = SciMLBase._unwrap_val(chunk_size) == 0 ? nothing : chunk_size
     ad = SciMLBase._unwrap_val(autodiff) ?
-        AD.ForwardDiffBackend(; chunksize) :
-        AD.FiniteDifferencesBackend()
-    return SimpleBatchedNewtonRaphson{typeof(ad), Nothing, typeof(termination_condition)}(
-        ad,
+         AD.ForwardDiffBackend(; chunksize) :
+         AD.FiniteDifferencesBackend()
+    return SimpleBatchedNewtonRaphson{typeof(ad), Nothing, typeof(termination_condition)}(ad,
         nothing,
         termination_condition)
 end
 
 function SciMLBase.__solve(prob::NonlinearProblem,
     alg::SimpleBatchedNewtonRaphson;
-    abstol=nothing,
-    reltol=nothing,
-    maxiters=1000,
+    abstol = nothing,
+    reltol = nothing,
+    maxiters = 1000,
     kwargs...)
     iip = isinplace(prob)
     @assert !iip "SimpleBatchedNewtonRaphson currently only supports out-of-place nonlinear problems."
@@ -57,9 +58,9 @@ function SciMLBase.__solve(prob::NonlinearProblem,
             alg,
             reconstruct(xₙ),
             reconstruct(fₙ);
-            retcode=ReturnCode.Success)
+            retcode = ReturnCode.Success)
 
-        solve(LinearProblem(𝓙, vec(fₙ); u0=vec(δx)), alg.linsolve; kwargs...)
+        solve(LinearProblem(𝓙, vec(fₙ); u0 = vec(δx)), alg.linsolve; kwargs...)
         xₙ .-= δx
 
         if termination_condition(fₙ, xₙ, xₙ₋₁, atol, rtol)
@@ -83,7 +84,7 @@ function SciMLBase.__solve(prob::NonlinearProblem,
         alg,
         reconstruct(xₙ),
         reconstruct(fₙ);
-        retcode=ReturnCode.MaxIters)
+        retcode = ReturnCode.MaxIters)
 end
 
 end

ext/SimpleNonlinearSolveNNlibExt.jl

@@ -1,7 +1,8 @@
 module SimpleNonlinearSolveNNlibExt
 
 using ArrayInterface, DiffEqBase, LinearAlgebra, NNlib, SimpleNonlinearSolve, SciMLBase
-import SimpleNonlinearSolve: _construct_batched_problem_structure, _get_storage, _init_𝓙, _result_from_storage, _get_tolerance, @maybeinplace
+import SimpleNonlinearSolve: _construct_batched_problem_structure,
+    _get_storage, _init_𝓙, _result_from_storage, _get_tolerance, @maybeinplace
 
 function __init__()
     SimpleNonlinearSolve.NNlibExtLoaded[] = true
@@ -10,9 +11,9 @@ end
 
 @views function SciMLBase.__solve(prob::NonlinearProblem,
     alg::BatchedBroyden;
-    abstol=nothing,
-    reltol=nothing,
-    maxiters=1000,
+    abstol = nothing,
+    reltol = nothing,
+    maxiters = 1000,
     kwargs...)
     iip = isinplace(prob)
 
@@ -74,7 +75,7 @@ end
         alg,
         reconstruct(xₙ),
         reconstruct(fₙ);
-        retcode=ReturnCode.MaxIters)
+        retcode = ReturnCode.MaxIters)
 end
 
 end

src/SimpleNonlinearSolve.jl

@@ -22,7 +22,8 @@ abstract type AbstractSimpleNonlinearSolveAlgorithm <: SciMLBase.AbstractNonline
 abstract type AbstractBracketingAlgorithm <: AbstractSimpleNonlinearSolveAlgorithm end
 abstract type AbstractNewtonAlgorithm{CS, AD, FDT} <: AbstractSimpleNonlinearSolveAlgorithm end
 abstract type AbstractImmutableNonlinearSolver <: AbstractSimpleNonlinearSolveAlgorithm end
-abstract type AbstractBatchedNonlinearSolveAlgorithm <: AbstractSimpleNonlinearSolveAlgorithm end
+abstract type AbstractBatchedNonlinearSolveAlgorithm <:
+              AbstractSimpleNonlinearSolveAlgorithm end
 
 include("utils.jl")
 include("bisection.jl")

src/batched/broyden.jl

@@ -1,5 +1,5 @@
 struct BatchedBroyden{TC <: NLSolveTerminationCondition} <:
-    AbstractBatchedNonlinearSolveAlgorithm
+       AbstractBatchedNonlinearSolveAlgorithm
     termination_condition::TC
 end
 

src/batched/dfsane.jl

@@ -1,4 +1,4 @@
-@kwdef struct SimpleBatchedDFSane{T, F, TC <: NLSolveTerminationCondition} <:
+Base.@kwdef struct SimpleBatchedDFSane{T, F, TC <: NLSolveTerminationCondition} <:
               AbstractBatchedNonlinearSolveAlgorithm
     σₘᵢₙ::T = 1.0f-10
     σₘₐₓ::T = 1.0f+10
@@ -10,17 +10,17 @@
     nₑₓₚ::Int = 2
     ηₛ::F = (f₍ₙₒᵣₘ₎₁, n, xₙ, fₙ) -> f₍ₙₒᵣₘ₎₁ ./ n .^ 2
     termination_condition::TC = NLSolveTerminationCondition(NLSolveTerminationMode.NLSolveDefault;
-        abstol=nothing,
-        reltol=nothing)
+        abstol = nothing,
+        reltol = nothing)
     max_inner_iterations::Int = 1000
 end
 
 function SciMLBase.__solve(prob::NonlinearProblem,
     alg::SimpleBatchedDFSane,
     args...;
-    abstol=nothing,
-    reltol=nothing,
-    maxiters=100,
+    abstol = nothing,
+    reltol = nothing,
+    maxiters = 100,
     kwargs...)
     iip = isinplace(prob)
 
@@ -60,7 +60,7 @@ function SciMLBase.__solve(prob::NonlinearProblem,
         return fₓ
     end
 
-    @maybeinplace iip fₙ₋₁ = ff!(f₍ₙₒᵣₘ₎ₙ₋₁, xₙ) xₙ
+    @maybeinplace iip fₙ₋₁=ff!(f₍ₙₒᵣₘ₎ₙ₋₁, xₙ) xₙ
     iip && (fₙ = similar(fₙ₋₁))
     ℋ = repeat(f₍ₙₒᵣₘ₎ₙ₋₁, M, 1)
     f̄ = similar(ℋ, 1, N)
@@ -79,7 +79,7 @@ function SciMLBase.__solve(prob::NonlinearProblem,
         fill!(α₋, α₁)
         @. xₙ = xₙ₋₁ + α₊ * 𝒹
 
-        @maybeinplace iip fₙ = ff!(f₍ₙₒᵣₘ₎ₙ, xₙ)
+        @maybeinplace iip fₙ=ff!(f₍ₙₒᵣₘ₎ₙ, xₙ)
 
         for _ in 1:(alg.max_inner_iterations)
             𝒸 = @. f̄ + η - γ * α₊^2 * f₍ₙₒᵣₘ₎ₙ₋₁
@@ -90,15 +90,15 @@ function SciMLBase.__solve(prob::NonlinearProblem,
                 τₘᵢₙ * α₊,
                 τₘₐₓ * α₊)
             @. xₙ = xₙ₋₁ - α₋ * 𝒹
-            @maybeinplace iip fₙ = ff!(f₍ₙₒᵣₘ₎ₙ, xₙ)
+            @maybeinplace iip fₙ=ff!(f₍ₙₒᵣₘ₎ₙ, xₙ)
 
             (sum(f₍ₙₒᵣₘ₎ₙ .≤ 𝒸) ≥ N ÷ 2) && break
 
             @. α₋ = clamp(α₋^2 * f₍ₙₒᵣₘ₎ₙ₋₁ / (f₍ₙₒᵣₘ₎ₙ + (T(2) * α₋ - T(1)) * f₍ₙₒᵣₘ₎ₙ₋₁),
                 τₘᵢₙ * α₋,
                 τₘₐₓ * α₋)
             @. xₙ = xₙ₋₁ + α₊ * 𝒹
-            @maybeinplace iip fₙ = ff!(f₍ₙₒᵣₘ₎ₙ, xₙ)
+            @maybeinplace iip fₙ=ff!(f₍ₙₒᵣₘ₎ₙ, xₙ)
         end
 
         if termination_condition(fₙ, xₙ, xₙ₋₁, atol, rtol)
@@ -129,12 +129,12 @@ function SciMLBase.__solve(prob::NonlinearProblem,
 
     if mode ∈ DiffEqBase.SAFE_BEST_TERMINATION_MODES
         xₙ = storage.u
-        @maybeinplace iip fₙ = f(xₙ)
+        @maybeinplace iip fₙ=f(xₙ)
     end
 
     return DiffEqBase.build_solution(prob,
         alg,
         reconstruct(xₙ),
         reconstruct(fₙ);
-        retcode=ReturnCode.MaxIters)
+        retcode = ReturnCode.MaxIters)
 end

src/batched/lbroyden.jl

@@ -0,0 +1 @@
+

src/batched/raphson.jl

@@ -1,5 +1,5 @@
 struct SimpleBatchedNewtonRaphson{AD, LS, TC <: NLSolveTerminationCondition} <:
-    AbstractBatchedNonlinearSolveAlgorithm
+       AbstractBatchedNonlinearSolveAlgorithm
     autodiff::AD
     linsolve::LS
     termination_condition::TC

src/batched/utils.jl

@@ -1,4 +1,4 @@
-macro maybeinplace(iip::Symbol, expr::Expr, u0::Union{Symbol, Nothing}=nothing)
+macro maybeinplace(iip::Symbol, expr::Expr, u0::Union{Symbol, Nothing} = nothing)
     @assert expr.head == :(=)
     x1, x2 = expr.args
     @assert x2.head == :call
@@ -64,7 +64,7 @@ function _result_from_storage(storage::NLSolveSafeTerminationResult, xₙ, fₙ,
         return ReturnCode.Success, xₙ, fₙ
     else
         if mode ∈ DiffEqBase.SAFE_BEST_TERMINATION_MODES
-            @maybeinplace iip fₙ = f(xₙ)
+            @maybeinplace iip fₙ=f(xₙ)
             return ReturnCode.Terminated, storage.u, fₙ
         else
             return ReturnCode.Terminated, xₙ, fₙ

src/raphson.jl

@@ -35,7 +35,7 @@ and static array problems.
 """
 struct SimpleNewtonRaphson{CS, AD, FDT} <: AbstractNewtonAlgorithm{CS, AD, FDT} end
 
-function SimpleNewtonRaphson(; batched=false,
+function SimpleNewtonRaphson(; batched = false,
     chunk_size = Val{0}(),
     autodiff = Val{true}(),
     diff_type = Val{:forward},
@@ -46,10 +46,10 @@ function SimpleNewtonRaphson(; batched=false,
     if batched
         @assert ADLinearSolveExtLoaded[] "Please install and load `LinearSolve.jl` and `AbstractDifferentiation.jl` to use batched Newton-Raphson."
         termination_condition = ismissing(termination_condition) ?
-            NLSolveTerminationCondition(NLSolveTerminationMode.NLSolveDefault;
-                abstol = nothing,
-                reltol = nothing) :
-            termination_condition
+                                NLSolveTerminationCondition(NLSolveTerminationMode.NLSolveDefault;
+            abstol = nothing,
+            reltol = nothing) :
+                                termination_condition
         return SimpleBatchedNewtonRaphson(; chunk_size,
             autodiff,
             diff_type,