Skip to content
This repository was archived by the owner on May 15, 2025. It is now read-only.

Add batched lbroyden #50

Merged
merged 1 commit into from
Mar 8, 2023
Merged

Add batched lbroyden #50

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "SimpleNonlinearSolve"
uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7"
authors = ["SciML"]
version = "0.1.13"
version = "0.1.14"

[deps]
ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
Expand Down
15 changes: 10 additions & 5 deletions src/halley.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,14 +80,19 @@ function SciMLBase.__solve(prob::NonlinearProblem,
else
if isa(x, Number)
fx = f(x)
dfx = FiniteDiff.finite_difference_derivative(f, x, diff_type(alg), eltype(x))
d2fx = FiniteDiff.finite_difference_derivative(x -> FiniteDiff.finite_difference_derivative(f, x), x,
diff_type(alg), eltype(x))
dfx = FiniteDiff.finite_difference_derivative(f, x, diff_type(alg),
eltype(x))
d2fx = FiniteDiff.finite_difference_derivative(x -> FiniteDiff.finite_difference_derivative(f,
x),
x,
diff_type(alg), eltype(x))
else
fx = f(x)
dfx = FiniteDiff.finite_difference_jacobian(f, x, diff_type(alg), eltype(x))
d2fx = FiniteDiff.finite_difference_jacobian(x -> FiniteDiff.finite_difference_jacobian(f, x), x,
diff_type(alg), eltype(x))
d2fx = FiniteDiff.finite_difference_jacobian(x -> FiniteDiff.finite_difference_jacobian(f,
x),
x,
diff_type(alg), eltype(x))
ai = -(dfx \ fx)
A = reshape(d2fx * ai, (n, n))
bi = (dfx) \ (A * ai)
Expand Down
101 changes: 78 additions & 23 deletions src/lbroyden.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,19 +1,40 @@
"""
LBroyden(threshold::Int = 27)
LBroyden(; batched = false,
termination_condition = NLSolveTerminationCondition(NLSolveTerminationMode.NLSolveDefault;
abstol = nothing, reltol = nothing),
threshold::Int = 27)

A limited memory implementation of Broyden. This method applies the L-BFGS scheme to
Broyden's method.

!!! warn

This method is not very stable and can diverge even for very simple problems. This has mostly been
tested for neural networks in DeepEquilibriumNetworks.jl.
"""
Base.@kwdef struct LBroyden <: AbstractSimpleNonlinearSolveAlgorithm
threshold::Int = 27
struct LBroyden{batched, TC <: NLSolveTerminationCondition} <:
AbstractSimpleNonlinearSolveAlgorithm
termination_condition::TC
threshold::Int

function LBroyden(; batched = false, threshold::Int = 27,
termination_condition = NLSolveTerminationCondition(NLSolveTerminationMode.NLSolveDefault;
abstol = nothing,
reltol = nothing))
return new{batched, typeof(termination_condition)}(termination_condition, threshold)
end
end

@views function SciMLBase.__solve(prob::NonlinearProblem, alg::LBroyden, args...;
@views function SciMLBase.__solve(prob::NonlinearProblem, alg::LBroyden{batched}, args...;
abstol = nothing, reltol = nothing, maxiters = 1000,
batch = false, kwargs...)
kwargs...) where {batched}
tc = alg.termination_condition
mode = DiffEqBase.get_termination_mode(tc)
threshold = min(maxiters, alg.threshold)
x = float(prob.u0)

batched && @assert ndims(x)==2 "Batched LBroyden only supports 2D arrays"

if x isa Number
restore_scalar = true
x = [x]
Expand All @@ -30,12 +51,20 @@ end
error("LBroyden currently only supports out-of-place nonlinear problems")
end

U = fill!(similar(x, (threshold, length(x))), zero(T))
Vᵀ = fill!(similar(x, (length(x), threshold)), zero(T))
U, Vᵀ = _init_lbroyden_state(batched, x, threshold)

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)
(tc.abstol !== nothing ? tc.abstol :
real(oneunit(eltype(T))) * (eps(real(one(eltype(T)))))^(4 // 5))
rtol = reltol !== nothing ? reltol :
(tc.reltol !== nothing ? tc.reltol : eps(real(one(eltype(T))))^(4 // 5))

if mode ∈ DiffEqBase.SAFE_BEST_TERMINATION_MODES
error("LBroyden currently doesn't support SAFE_BEST termination modes")
end

storage = mode ∈ DiffEqBase.SAFE_TERMINATION_MODES ? Dict() : nothing
termination_condition = tc(storage)

xₙ = x
xₙ₋₁ = x
Expand All @@ -47,27 +76,23 @@ end
Δxₙ = xₙ .- xₙ₋₁
Δfₙ = fₙ .- fₙ₋₁

if iszero(fₙ)
xₙ = restore_scalar ? xₙ[] : xₙ
return SciMLBase.build_solution(prob, alg, xₙ, fₙ; retcode = ReturnCode.Success)
end

if isapprox(xₙ, xₙ₋₁; atol, rtol)
if termination_condition(restore_scalar ? [fₙ] : fₙ, xₙ, xₙ₋₁, atol, rtol)
xₙ = restore_scalar ? xₙ[] : xₙ
return SciMLBase.build_solution(prob, alg, xₙ, fₙ; retcode = ReturnCode.Success)
end

_U = U[1:min(threshold, i), :]
_Vᵀ = Vᵀ[:, 1:min(threshold, i)]
_U = selectdim(U, 1, 1:min(threshold, i))
_Vᵀ = selectdim(Vᵀ, 2, 1:min(threshold, i))

vᵀ = _rmatvec(_U, _Vᵀ, Δxₙ)
mvec = _matvec(_U, _Vᵀ, Δfₙ)
Δxₙ = (Δxₙ .- mvec) ./ (sum(vᵀ .* Δfₙ) .+ convert(T, 1e-5))
u = (Δxₙ .- mvec) ./ (sum(vᵀ .* Δfₙ) .+ convert(T, 1e-5))

Vᵀ[:, mod1(i, threshold)] .= vᵀ
U[mod1(i, threshold), :] .= Δxₙ
selectdim(Vᵀ, 2, mod1(i, threshold)) .= vᵀ
selectdim(U, 1, mod1(i, threshold)) .= u

update = -_matvec(U[1:min(threshold, i + 1), :], Vᵀ[:, 1:min(threshold, i + 1)], fₙ)
update = -_matvec(selectdim(U, 1, 1:min(threshold, i + 1)),
selectdim(Vᵀ, 2, 1:min(threshold, i + 1)), fₙ)

xₙ₋₁ = xₙ
fₙ₋₁ = fₙ
Expand All @@ -77,12 +102,42 @@ end
return SciMLBase.build_solution(prob, alg, xₙ, fₙ; retcode = ReturnCode.MaxIters)
end

function _init_lbroyden_state(batched::Bool, x, threshold)
T = eltype(x)
if batched
U = fill!(similar(x, (threshold, size(x, 1), size(x, 2))), zero(T))
Vᵀ = fill!(similar(x, (size(x, 1), threshold, size(x, 2))), zero(T))
else
U = fill!(similar(x, (threshold, length(x))), zero(T))
Vᵀ = fill!(similar(x, (length(x), threshold)), zero(T))
end
return U, Vᵀ
end

function _rmatvec(U::AbstractMatrix, Vᵀ::AbstractMatrix,
x::Union{<:AbstractVector, <:Number})
return -x .+ dropdims(sum(U .* sum(Vᵀ .* x; dims = 1)'; dims = 1); dims = 1)
length(U) == 0 && return x
return -x .+ vec((x' * Vᵀ) * U)
end

function _rmatvec(U::AbstractArray{T1, 3}, Vᵀ::AbstractArray{T2, 3},
x::AbstractMatrix) where {T1, T2}
length(U) == 0 && return x
Vᵀx = sum(Vᵀ .* reshape(x, size(x, 1), 1, size(x, 2)); dims = 1)
return -x .+ _drdims_sum(U .* permutedims(Vᵀx, (2, 1, 3)); dims = 1)
end

function _matvec(U::AbstractMatrix, Vᵀ::AbstractMatrix,
x::Union{<:AbstractVector, <:Number})
return -x .+ dropdims(sum(sum(x .* U'; dims = 1) .* Vᵀ; dims = 2); dims = 2)
length(U) == 0 && return x
return -x .+ vec(Vᵀ * (U * x))
end

function _matvec(U::AbstractArray{T1, 3}, Vᵀ::AbstractArray{T2, 3},
x::AbstractMatrix) where {T1, T2}
length(U) == 0 && return x
xUᵀ = sum(reshape(x, size(x, 1), 1, size(x, 2)) .* permutedims(U, (2, 1, 3)); dims = 1)
return -x .+ _drdims_sum(xUᵀ .* Vᵀ; dims = 2)
end

_drdims_sum(args...; dims = :) = dropdims(sum(args...; dims); dims)
48 changes: 32 additions & 16 deletions test/basictests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@ using Test

const BATCHED_BROYDEN_SOLVERS = Broyden[]
const BROYDEN_SOLVERS = Broyden[]
const BATCHED_LBROYDEN_SOLVERS = LBroyden[]
const LBROYDEN_SOLVERS = LBroyden[]

for mode in instances(NLSolveTerminationMode.T)
if mode ∈
Expand All @@ -18,6 +20,8 @@ for mode in instances(NLSolveTerminationMode.T)
reltol = nothing)
push!(BROYDEN_SOLVERS, Broyden(; batched = false, termination_condition))
push!(BATCHED_BROYDEN_SOLVERS, Broyden(; batched = true, termination_condition))
push!(LBROYDEN_SOLVERS, LBroyden(; batched = false, termination_condition))
push!(BATCHED_LBROYDEN_SOLVERS, LBroyden(; batched = true, termination_condition))
end

# SimpleNewtonRaphson
Expand Down Expand Up @@ -134,24 +138,38 @@ for alg in (SimpleNewtonRaphson(), LBroyden(), Klement(), SimpleTrustRegion(),
end

for p in 1.1:0.1:100.0
@test abs.(g(p)) ≈ sqrt(p)
@test abs.(ForwardDiff.derivative(g, p)) ≈ 1 / (2 * sqrt(p))
res = abs.(g(p))
# Not surprising if LBrouden fails to converge
if any(x -> isnan(x) || x <= 1e-5 || x >= 1e5, res) && alg isa LBroyden
@test_broken res ≈ sqrt(p)
@test_broken abs.(ForwardDiff.derivative(g, p)) ≈ 1 / (2 * sqrt(p))
else
@test res ≈ sqrt(p)
@test abs.(ForwardDiff.derivative(g, p)) ≈ 1 / (2 * sqrt(p))
end
end
end

# Scalar
f, u0 = (u, p) -> u * u - p, 1.0
for alg in (SimpleNewtonRaphson(), LBroyden(), Klement(), SimpleTrustRegion(),
SimpleDFSane(), Halley(), BROYDEN_SOLVERS...)
for alg in (SimpleNewtonRaphson(), Klement(), SimpleTrustRegion(),
SimpleDFSane(), Halley(), BROYDEN_SOLVERS..., LBROYDEN_SOLVERS...)
g = function (p)
probN = NonlinearProblem{false}(f, oftype(p, u0), p)
sol = solve(probN, alg)
return sol.u
end

for p in 1.1:0.1:100.0
@test abs(g(p)) ≈ sqrt(p)
@test abs(ForwardDiff.derivative(g, p)) ≈ 1 / (2 * sqrt(p))
res = abs.(g(p))
# Not surprising if LBrouden fails to converge
if any(x -> isnan(x) || x <= 1e-5 || x >= 1e5, res) && alg isa LBroyden
@test_broken res ≈ sqrt(p)
@test_broken abs.(ForwardDiff.derivative(g, p)) ≈ 1 / (2 * sqrt(p))
else
@test res ≈ sqrt(p)
@test abs.(ForwardDiff.derivative(g, p)) ≈ 1 / (2 * sqrt(p))
end
end
end

Expand Down Expand Up @@ -207,8 +225,8 @@ for alg in [Bisection(), Falsi(), Ridder(), Brent()]
@test ForwardDiff.jacobian(g, p) ≈ ForwardDiff.jacobian(t, p)
end

for alg in (SimpleNewtonRaphson(), LBroyden(), Klement(), SimpleTrustRegion(),
SimpleDFSane(), Halley(), BROYDEN_SOLVERS...)
for alg in (SimpleNewtonRaphson(), Klement(), SimpleTrustRegion(),
SimpleDFSane(), Halley(), BROYDEN_SOLVERS..., LBROYDEN_SOLVERS...)
global g, p
g = function (p)
probN = NonlinearProblem{false}(f, 0.5, p)
Expand All @@ -225,26 +243,24 @@ probN = NonlinearProblem(f, u0)

for alg in (SimpleNewtonRaphson(), SimpleNewtonRaphson(; autodiff = false),
SimpleTrustRegion(),
SimpleTrustRegion(; autodiff = false), Halley(), Halley(; autodiff = false), LBroyden(), Klement(), SimpleDFSane(),
BROYDEN_SOLVERS...)
SimpleTrustRegion(; autodiff = false), Halley(), Halley(; autodiff = false),
Klement(), SimpleDFSane(),
BROYDEN_SOLVERS..., LBROYDEN_SOLVERS...)
sol = solve(probN, alg)

@test sol.retcode == ReturnCode.Success
@test sol.u[end] ≈ sqrt(2.0)
end


for u0 in [1.0, [1, 1.0]]
local f, probN, sol
f = (u, p) -> u .* u .- 2.0
probN = NonlinearProblem(f, u0)
sol = sqrt(2) * u0

for alg in (SimpleNewtonRaphson(), SimpleNewtonRaphson(; autodiff = false),
SimpleTrustRegion(),
SimpleTrustRegion(; autodiff = false), LBroyden(), Klement(),
SimpleDFSane(),
BROYDEN_SOLVERS...)
SimpleTrustRegion(), SimpleTrustRegion(; autodiff = false), Klement(),
SimpleDFSane(), BROYDEN_SOLVERS..., LBROYDEN_SOLVERS...)
sol2 = solve(probN, alg)

@test sol2.retcode == ReturnCode.Success
Expand Down Expand Up @@ -430,7 +446,7 @@ sol = solve(probN, Broyden(batched = true))

@test abs.(sol.u) ≈ sqrt.(p)

for alg in BATCHED_BROYDEN_SOLVERS
for alg in (BATCHED_BROYDEN_SOLVERS..., BATCHED_LBROYDEN_SOLVERS...)
sol = solve(probN, alg)

@test sol.retcode == ReturnCode.Success
Expand Down