Skip to content

Commit

Permalink
Add batched lbroyden
Browse files Browse the repository at this point in the history
  • Loading branch information
@avik-pal
avik-pal committed Mar 8, 2023
1 parent bc5244d commit 5b9548e
Show file tree
Hide file tree
Showing 4 changed files with 121 additions and 45 deletions.
2 changes: 1 addition & 1 deletion Project.toml
View file
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
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
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
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

0 comments on commit 5b9548e

Please sign in to comment.