Reduce allocations in PANOC, PANOCplus, ZeroFPR, DRLS (#74)

Author: lostella

src/algorithms/drls.jl

@@ -89,7 +89,13 @@ Base.@kwdef mutable struct DRLSState{R,Tx,TH}
     temp_x2::Tx = similar(x)
 end
 
-DRE(f_u::Number, g_v::Number, x, u, res, gamma) = f_u + g_v - real(dot(x - u, res)) / gamma + 1 / (2 * gamma) * norm(res)^2
+function DRE(f_u::R, g_v, x, u, res, gamma) where R
+    dot_product = R(0)
+    for (x_i, u_i, res_i) in zip(x, u, res)
+        dot_product += (x_i - u_i) * res_i
+    end
+    return f_u + g_v - real(dot_product) / gamma + 1 / (2 * gamma) * norm(res)^2
+end
 
 DRE(state::DRLSState) = DRE(state.f_u, state.g_v, state.x, state.u, state.res, state.gamma)
 
@@ -107,18 +113,25 @@ function Base.iterate(iter::DRLSIteration)
     return state, state
 end
 
-set_next_direction!(::QuasiNewtonStyle, ::DRLSIteration, state::DRLSState) = mul!(state.d, state.H, -state.res)
+function set_next_direction!(::QuasiNewtonStyle, ::DRLSIteration, state::DRLSState)
+    mul!(state.d, state.H, state.res)
+    state.d .*= -1
+end
 set_next_direction!(::NesterovStyle, ::DRLSIteration, state::DRLSState) = state.d .= iterate(state.H)[1] .* (state.xbar .- state.xbar_prev) .+ (state.xbar .- state.x)
 set_next_direction!(::NoAccelerationStyle, ::DRLSIteration, state::DRLSState) = state.d .= state.xbar .- state.x
 set_next_direction!(iter::DRLSIteration, state::DRLSState) = set_next_direction!(acceleration_style(typeof(iter.directions)), iter, state)
 
-update_direction_state!(::QuasiNewtonStyle, ::DRLSIteration, state::DRLSState) = update!(state.H, state.d, state.res - state.res_prev)
+function update_direction_state!(::QuasiNewtonStyle, ::DRLSIteration, state::DRLSState)
+    state.res_prev .= state.res .- state.res_prev
+    update!(state.H, state.d, state.res_prev)
+end
 update_direction_state!(::NesterovStyle, ::DRLSIteration, state::DRLSState) = return
 update_direction_state!(::NoAccelerationStyle, ::DRLSIteration, state::DRLSState) = return
 update_direction_state!(iter::DRLSIteration, state::DRLSState) = update_direction_state!(acceleration_style(typeof(iter.directions)), iter, state)
 
 function Base.iterate(iter::DRLSIteration{R, Tx, Tf}, state::DRLSState) where {R, Tx, Tf}
     DRE_curr = DRE(state)
+    threshold = iter.dre_sign * DRE_curr - iter.c / iter.gamma * norm(state.res)^2
 
     set_next_direction!(iter, state)
 
@@ -139,7 +152,7 @@ function Base.iterate(iter::DRLSIteration{R, Tx, Tf}, state::DRLSState) where {R
     a, b, c = R(0), R(0), R(0)
 
     for k in 1:iter.max_backtracks
-        if iter.dre_sign * DRE(state) <= iter.dre_sign * DRE_curr - iter.c / iter.gamma * norm(state.res_prev)^2
+        if iter.dre_sign * DRE(state) <= threshold
             break
         end
 

src/algorithms/panoc.jl

@@ -98,11 +98,18 @@ function Base.iterate(iter::PANOCIteration{R}) where R
     return state, state
 end
 
-set_next_direction!(::QuasiNewtonStyle, ::PANOCIteration, state::PANOCState) = mul!(state.d, state.H, -state.res)
+function set_next_direction!(::QuasiNewtonStyle, ::PANOCIteration, state::PANOCState)
+    mul!(state.d, state.H, state.res)
+    state.d .*= -1
+end
 set_next_direction!(::NoAccelerationStyle, ::PANOCIteration, state::PANOCState) = state.d .= .-state.res
 set_next_direction!(iter::PANOCIteration, state::PANOCState) = set_next_direction!(acceleration_style(typeof(iter.directions)), iter, state)
 
-update_direction_state!(::QuasiNewtonStyle, ::PANOCIteration, state::PANOCState) = update!(state.H, state.x - state.x_prev, state.res - state.res_prev)
+function update_direction_state!(::QuasiNewtonStyle, ::PANOCIteration, state::PANOCState)
+    state.x_prev .= state.x .- state.x_prev
+    state.res_prev .= state.res .- state.res_prev
+    update!(state.H, state.x_prev, state.res_prev)
+end
 update_direction_state!(::NoAccelerationStyle, ::PANOCIteration, state::PANOCState) = return
 update_direction_state!(iter::PANOCIteration, state::PANOCState) = update_direction_state!(acceleration_style(typeof(iter.directions)), iter, state)
 

src/algorithms/panocplus.jl

@@ -103,11 +103,18 @@ function Base.iterate(iter::PANOCplusIteration{R}) where {R}
     return state, state
 end
 
-set_next_direction!(::QuasiNewtonStyle, ::PANOCplusIteration, state::PANOCplusState) = mul!(state.d, state.H, -state.res_prev)
+function set_next_direction!(::QuasiNewtonStyle, ::PANOCplusIteration, state::PANOCplusState)
+    mul!(state.d, state.H, state.res_prev)
+    state.d .*= -1
+end
 set_next_direction!(::NoAccelerationStyle, ::PANOCplusIteration, state::PANOCplusState) = state.d .= .-state.res_prev
 set_next_direction!(iter::PANOCplusIteration, state::PANOCplusState) = set_next_direction!(acceleration_style(typeof(iter.directions)), iter, state)
 
-update_direction_state!(::QuasiNewtonStyle, ::PANOCplusIteration, state::PANOCplusState) = update!(state.H, state.x - state.x_prev, state.res - state.res_prev)
+function update_direction_state!(::QuasiNewtonStyle, ::PANOCplusIteration, state::PANOCplusState)
+    state.x_prev .= state.x .- state.x_prev
+    state.res_prev .= state.res .- state.res_prev
+    update!(state.H, state.x_prev, state.res_prev)
+end
 update_direction_state!(::NoAccelerationStyle, ::PANOCplusIteration, state::PANOCplusState) = return
 update_direction_state!(iter::PANOCplusIteration, state::PANOCplusState) = update_direction_state!(acceleration_style(typeof(iter.directions)), iter, state)
 
@@ -116,8 +123,6 @@ reset_direction_state!(::NoAccelerationStyle, ::PANOCplusIteration, state::PANOC
 reset_direction_state!(iter::PANOCplusIteration, state::PANOCplusState) = reset_direction_state!(acceleration_style(typeof(iter.directions)), iter, state)
 
 function Base.iterate(iter::PANOCplusIteration{R}, state::PANOCplusState) where R
-    f_Az, a, b, c = R(Inf), R(Inf), R(Inf), R(Inf)
-
     # store iterate and residual for metric update later on
     state.x_prev .= state.x
     state.res_prev .= state.res

src/algorithms/zerofpr.jl

@@ -96,11 +96,18 @@ function Base.iterate(iter::ZeroFPRIteration{R}) where R
     return state, state
 end
 
-set_next_direction!(::QuasiNewtonStyle, ::ZeroFPRIteration, state::ZeroFPRState) = mul!(state.d, state.H, -state.res_xbar)
+function set_next_direction!(::QuasiNewtonStyle, ::ZeroFPRIteration, state::ZeroFPRState)
+    mul!(state.d, state.H, state.res_xbar)
+    state.d .*= -1
+end
 set_next_direction!(::NoAccelerationStyle, ::ZeroFPRIteration, state::ZeroFPRState) = state.d .= .-state.res
 set_next_direction!(iter::ZeroFPRIteration, state::ZeroFPRState) = set_next_direction!(acceleration_style(typeof(iter.directions)), iter, state)
 
-update_direction_state!(::QuasiNewtonStyle, ::ZeroFPRIteration, state::ZeroFPRState) = update!(state.H, state.xbar - state.xbar_prev, state.res_xbar - state.res_xbar_prev)
+function update_direction_state!(::QuasiNewtonStyle, ::ZeroFPRIteration, state::ZeroFPRState)
+    state.xbar_prev .= state.xbar .- state.xbar_prev
+    state.res_xbar_prev .= state.res_xbar .- state.res_xbar_prev
+    update!(state.H, state.xbar_prev, state.res_xbar_prev)
+end
 update_direction_state!(::NoAccelerationStyle, ::ZeroFPRIteration, state::ZeroFPRState) = return
 update_direction_state!(iter::ZeroFPRIteration, state::ZeroFPRState) = update_direction_state!(acceleration_style(typeof(iter.directions)), iter, state)