Skip to content

Instantly share code, notes, and snippets.

@jiahao
Last active September 24, 2020 04:43
Show Gist options
  • Save jiahao/8e1fc43be4a3ba9c511684a8fba1b125 to your computer and use it in GitHub Desktop.
Save jiahao/8e1fc43be4a3ba9c511684a8fba1b125 to your computer and use it in GitHub Desktop.
Implementation of the stabilized Type-I Anderson acceleration (AA-I-S-m) algorithm of Zhang, O'Donoghue and Boyd (2018). This implementation solves g(x) = 0 as opposed to f(x) = x, which you obtain from g(x) = f(x) - x.
using Dates: now
using DataFrames
import Base: *, push!
mutable struct AAUpdate{Tu,Tv} # Matrix-free representation of the H matrix
m::Int #:: Size of the AA subspace
u::Vector{Tu} #:: The quantities s-Hỹ (note typo in paper)
v::Vector{Tv} #:: The quantities (H'ŝ)/(ŝ'Hŷ)
end
function start!(H::AAUpdate{Tu,Tv}) where {Tu,Tv}
H.m = 0
H.u = Tu[]
H.v = Tv[]
end
AAUpdate{Tu,Tv}() where {Tu,Tv} = AAUpdate{Tu,Tv}(0, Tu[], Tv[])
AAUpdate() where {Tu,Tv} = AAUpdate{Vector{Float64},Vector{Float64}}(0, [], [])
function *(H::AAUpdate{Tu}, g0::Tu) where {Tu}
g = copy(g0)
for j = 1:H.m
g += H.u[j] * (H.v[j] ⋅ g)
end
return g
end
function push!(H::AAUpdate{Tu,Tv}, u::Tu, v::Tv) where {Tu,Tv}
push!(H.u, u)
push!(H.v, v)
H.m += 1
end
# Algorithm 3 - Stabilized Type-I Anderson Acceleration (AA-I-S-m)
# Modified for solving f(x) = 0
# ref: https://stanford.edu/~boyd/papers/pdf/scs_2.0_v_global.pdf
function anderson_zob( #matrix free version
g, # LHS of g(x) = 0
x, #initial point
;
θ̄=0.01, τ=0.001, α=0.1, #regularization constants
D=1e6, ϵ=1e-6, # safeguard constants
m = 5, #max memory
maxiter = 1000,
verbose = false,
reltol = 1e-5,
save_x = false, #save intermediate solutions into convergence history
)
gα(x; gx = g(x)) = (1-α)*x + α*gx
#Powell-type regularizer
ϕ(η, θ̄) = abs(η)≥θ̄ ? one(η) : (1-sign(η)*θ̄)/(1-η)
H = AAUpdate{typeof(x), typeof(x)}()
n = 0 # Count number of normal (non-safeguarded) steps taken (nAA in paper)
g0 = gx = g(x)
gnrm = U = norm(g0)
ŝs = typeof(x)[]
x̃ = xnew = gα(x; gx=gx)
abstol = reltol * U
ch = DataFrame(
iter=[0],
timestamp=[now()],
restart=[true],
safeguard=[false],
gnorm=[U],
Δgnorm=Union{Missing,typeof(U)}[missing],
x=Any[save_x ? x : missing]
#x=Union{Missing,typeof(x)}[save_x ? x : missing]
)
for k = 1:maxiter
if gnrm ≤ abstol
if verbose
@info "Termination criterion reached"
end
@goto finished
end
ŝ = s = x̃ - x
# Orthogonalize using classical Gram-Schmidt
# Running it twice guarantees numerical stability to within eps()
for _ in 1:2, t in ŝs
ŝ -= t⋅s * t
end
# Check if restart needed
do_restart = H.m == m || norm(ŝ) < τ * norm(s)
if do_restart
ŝ = s
ŝs = typeof(ŝ)[]
start!(H)
end
#Save basis
r = norm(ŝ)
ŝ = ŝ / r
push!(ŝs, ŝ)
gx̃ = g(x̃)
y = gx̃ - gx
# Perform Powell-type regularized update
γ = ŝ ⋅ (H * y)
θ = ϕ(γ, θ̄)
ỹ = θ * y - (1-θ) * gx
Hỹ = H * ỹ
push!(H, s - Hỹ, (H * ŝ)/(ŝ'Hỹ)) #Note typo in paper
x̃ = x - H * gx
x = xnew
# Check safeguarding strategy
threshold = D * U * (n + 1)^-(1+ϵ)
do_safeguard = !isfinite(gnrm) || gnrm ≥ threshold || norm(gx̃) ≥ threshold
if do_safeguard
if verbose
@info "Iteration $k: safeguard activated (‖gₖ‖ ≥ threshold = $threshold)"
end
xnew = gα(x)
gx = g(xnew)
else
xnew = x̃
gx = gx̃
n += 1
end
gnrm = norm(gx)
push!(ch, [k now() do_restart do_safeguard gnrm (gnrm-ch[:gnorm][end]) (save_x ? x : missing)])
if verbose
@info "Iteration $k: ‖gₖ‖ = $gnrm"
end
if !isfinite(gnrm)
@error "Norm of ‖gₖ‖ became infinite"
@goto finished
end
end
@label finished
return x, ch
end
#Simple tests
let
anderson_zob(x->0.5*(9.0 ./ x .- x), [1.0]; maxiter=100, verbose=false)
anderson_zob(x->0.5*(9.0 ./ x .- x), [3.0-eps()]; maxiter=100, verbose=false)
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment