LBroyden

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,18 +1,20 @@
 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
     return
 end
 
+# Broyden's method
 @views function SciMLBase.__solve(prob::NonlinearProblem,
     alg::BatchedBroyden;
-    abstol=nothing,
-    reltol=nothing,
-    maxiters=1000,
+    abstol = nothing,
+    reltol = nothing,
+    maxiters = 1000,
     kwargs...)
     iip = isinplace(prob)
 
@@ -24,7 +26,7 @@ end
 
     storage = _get_storage(mode, u)
 
-    xₙ, xₙ₋₁, δx, δf = ntuple(_ -> copy(u), 4)
+    xₙ, xₙ₋₁, δxₙ, δf = ntuple(_ -> copy(u), 4)
     T = eltype(u)
 
     atol = _get_tolerance(abstol, tc.abstol, T)
@@ -41,16 +43,16 @@ end
         xₙ .= xₙ₋₁ .- 𝓙⁻¹f
 
         @maybeinplace iip fₙ=f(xₙ)
-        δx .= xₙ .- xₙ₋₁
+        δxₙ .= xₙ .- xₙ₋₁
         δf .= fₙ .- fₙ₋₁
 
         batched_mul!(reshape(𝓙⁻¹f, L, 1, N), 𝓙⁻¹, reshape(δf, L, 1, N))
-        δxᵀ = reshape(δx, 1, L, N)
+        δxₙᵀ = reshape(δxₙ, 1, L, N)
 
-        batched_mul!(reshape(xᵀ𝓙⁻¹δf, 1, 1, N), δxᵀ, reshape(𝓙⁻¹f, L, 1, N))
-        batched_mul!(xᵀ𝓙⁻¹, δxᵀ, 𝓙⁻¹)
-        δx .= (δx .- 𝓙⁻¹f) ./ (xᵀ𝓙⁻¹δf .+ T(1e-5))
-        batched_mul!(𝓙⁻¹, reshape(δx, L, 1, N), xᵀ𝓙⁻¹, one(T), one(T))
+        batched_mul!(reshape(xᵀ𝓙⁻¹δf, 1, 1, N), δxₙᵀ, reshape(𝓙⁻¹f, L, 1, N))
+        batched_mul!(xᵀ𝓙⁻¹, δxₙᵀ, 𝓙⁻¹)
+        δxₙ .= (δxₙ .- 𝓙⁻¹f) ./ (xᵀ𝓙⁻¹δf .+ T(1e-5))
+        batched_mul!(𝓙⁻¹, reshape(δxₙ, L, 1, N), xᵀ𝓙⁻¹, one(T), one(T))
 
         if termination_condition(fₙ, xₙ, xₙ₋₁, atol, rtol)
             retcode, xₙ, fₙ = _result_from_storage(storage, xₙ, fₙ, f, mode, iip)
@@ -74,7 +76,103 @@ end
         alg,
         reconstruct(xₙ),
         reconstruct(fₙ);
-        retcode=ReturnCode.MaxIters)
+        retcode = ReturnCode.MaxIters)
+end
+
+# Limited Memory Broyden's method
+@views function SciMLBase.__solve(prob::NonlinearProblem,
+    alg::BatchedLBroyden;
+    abstol = nothing,
+    reltol = nothing,
+    maxiters = 1000,
+    kwargs...)
+    iip = isinplace(prob)
+
+    u, f, reconstruct = _construct_batched_problem_structure(prob)
+    L, N = size(u)
+    T = eltype(u)
+
+    tc = alg.termination_condition
+    mode = DiffEqBase.get_termination_mode(tc)
+
+    storage = _get_storage(mode, u)
+
+    η = min(maxiters, alg.threshold)
+    U = fill!(similar(u, (η, L, N)), zero(T))
+    Vᵀ = fill!(similar(u, (L, η, N)), zero(T))
+
+    xₙ, xₙ₋₁, δfₙ = ntuple(_ -> copy(u), 3)
+
+    atol = _get_tolerance(abstol, tc.abstol, T)
+    rtol = _get_tolerance(reltol, tc.reltol, T)
+    termination_condition = tc(storage)
+
+    @maybeinplace iip fₙ₋₁=f(xₙ) u
+    iip && (fₙ = copy(fₙ₋₁))
+    δxₙ = -copy(fₙ₋₁)
+    ηNx = similar(xₙ, η, N)
+
+    for i in 1:maxiters
+        @. xₙ = xₙ₋₁ - δxₙ
+        @maybeinplace iip fₙ=f(xₙ)
+        @. δxₙ = xₙ - xₙ₋₁
+        @. δfₙ = fₙ - fₙ₋₁
+
+        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
+
+        _L = min(i, η)
+        _U = U[1:_L, :, :]
+        _Vᵀ = Vᵀ[:, 1:_L, :]
+
+        idx = mod1(i, η)
+
+        if i > 1
+            partial_ηNx = ηNx[1:_L, :]
+
+            _ηNx = reshape(partial_ηNx, 1, :, N)
+            batched_mul!(_ηNx, reshape(δxₙ, 1, L, N), _Vᵀ)
+            batched_mul!(Vᵀ[:, idx:idx, :], _ηNx, _U)
+            Vᵀ[:, idx, :] .-= δxₙ
+
+            _ηNx = reshape(partial_ηNx, :, 1, N)
+            batched_mul!(_ηNx, _U, reshape(δfₙ, L, 1, N))
+            batched_mul!(U[idx:idx, :, :], _Vᵀ, _ηNx)
+            U[idx, :, :] .-= δfₙ
+        else
+            Vᵀ[:, idx, :] .= -δxₙ
+            U[idx, :, :] .= -δfₙ
+        end
+
+        U[idx, :, :] .= (δxₙ .- U[idx, :, :]) ./
+                        (sum(Vᵀ[:, idx, :] .* δfₙ; dims = 1) .+
+                         convert(T, 1e-5))
+
+        _L = min(i + 1, η)
+        _ηNx = reshape(ηNx[1:_L, :], :, 1, N)
+        batched_mul!(_ηNx, U[1:_L, :, :], reshape(δfₙ, L, 1, N))
+        batched_mul!(reshape(δxₙ, L, 1, N), Vᵀ[:, 1:_L, :], _ηNx)
+
+        xₙ₋₁ .= xₙ
+        fₙ₋₁ .= fₙ
+    end
+
+    if mode ∈ DiffEqBase.SAFE_BEST_TERMINATION_MODES
+        xₙ = storage.u
+        @maybeinplace iip fₙ=f(xₙ)
+    end
+
+    return DiffEqBase.build_solution(prob,
+        alg,
+        reconstruct(xₙ),
+        reconstruct(fₙ);
+        retcode = ReturnCode.MaxIters)
 end
 
 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

src/batched/lbroyden.jl

@@ -0,0 +1,7 @@
+struct BatchedLBroyden{TC <: NLSolveTerminationCondition} <:
+    AbstractBatchedNonlinearSolveAlgorithm
+    termination_condition::TC
+    threshold::Int
+end
+
+# Implementation of solve using Package Extensions

src/broyden.jl

@@ -30,6 +30,9 @@ end
 
 function SciMLBase.__solve(prob::NonlinearProblem, alg::Broyden, args...;
     abstol = nothing, reltol = nothing, maxiters = 1000, kwargs...)
+    if SciMLBase.isinplace(prob)
+        error("Broyden currently only supports out-of-place nonlinear problems")
+    end
     tc = alg.termination_condition
     mode = DiffEqBase.get_termination_mode(tc)
     f = Base.Fix2(prob.f, prob.p)
@@ -39,19 +42,14 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::Broyden, args...;
     T = eltype(x)
     J⁻¹ = init_J(x)
 
-    if SciMLBase.isinplace(prob)
-        error("Broyden currently only supports out-of-place nonlinear problems")
-    end
-
     atol = _get_tolerance(abstol, tc.abstol, T)
     rtol = _get_tolerance(reltol, tc.reltol, T)
 
     if mode ∈ DiffEqBase.SAFE_BEST_TERMINATION_MODES
         error("Broyden currently doesn't support SAFE_BEST termination modes")
     end
 
-    storage = mode ∈ DiffEqBase.SAFE_TERMINATION_MODES ? NLSolveSafeTerminationResult() :
-              nothing
+    storage = _get_storage(mode, x)
     termination_condition = tc(storage)
 
     xₙ = x