Fix penalty parameter setting in augmented Lagrangian method

[JoshuaLampert]

In #630, the penalty parameter in the augmented Lagrangian method was hardcoded to 1/3 instead of inheriting it from alms.ρ. This looks unintended to me and lead to failing tests for me in JoshuaLampert/SummationByPartsOperatorsExtra.jl#100.

Note that I let an AI find what the issue is, but I reviewed the change it proposed manually. An MWE created by the AI is

# MWE: the sub-problem's cost in `augmented_Lagrangian_method` uses a hard-coded `ρ = 1/3`.
#
# In `step_solver!` (src/solvers/augmented_Lagrangian_method.jl) the penalty parameter of the
# sub-problem's *cost* is set to the literal `1 / 3`, while its *gradient* gets the actual `alms.ρ`:
#
#     set_parameter!(alms.sub_problem, Val(:Objective), Val(:Cost),     Val(:ρ), 1 / 3)
#     set_parameter!(alms.sub_problem, Val(:Objective), Val(:Gradient), Val(:ρ), alms.ρ)
#
# The sub-solver thus minimizes one function while following the gradient of another. Since ALM
# increases ρ geometrically, the mismatch grows with every outer iteration. Up to and including
# v0.6.3 both lines used `alms.ρ`.
#
# Requires only Manopt.jl and Manifolds.jl.

using Manopt
using Manifolds: Euclidean
using LinearAlgebra: norm

# minimize f(p) = p₁ subject to h(p) = ‖p‖² - 1 = 0.
# The minimizer is p* = (-1, 0); the constraint is active there, so ALM has to grow ρ.
M = Euclidean(2)
f(M, p) = p[1]
grad_f(M, p) = [1.0, 0.0]
h(M, p) = [p[1]^2 + p[2]^2 - 1.0]
grad_h(M, p) = [[2 * p[1], 2 * p[2]]]

p_star = [-1.0, 0.0]
p0 = [0.5, 0.5]

# Read the two penalty parameters that are supposed to agree (internal, hence the fallbacks)
function penalty_parameters(sub_problem)
    objective = Manopt.get_objective(sub_problem)
    cost, gradient = if hasproperty(objective, :functions)
        objective.functions.cost, objective.functions.gradient   # v0.6.4+
    else
        objective.cost, getproperty(objective, :gradient!!)      # up to v0.6.3
    end
    unwrap(fn) = hasproperty(fn, :ρ) ? fn : fn.f
    return unwrap(cost).ρ, unwrap(gradient).ρ
end

println("Manopt v", pkgversion(Manopt))
println()
println(" iterations |          minimizer          |  |p - p*|  |    |h(p)|   | ρ cost | ρ gradient")
println("------------+-----------------------------+------------+------------+--------+-----------")
for iterations in (5, 10, 20, 50)
    state = augmented_Lagrangian_method(M, f, grad_f, p0; h = h, grad_h = grad_h,
                                        equality_constraints = 1,
                                        stopping_criterion = StopAfterIteration(iterations),
                                        return_state = true)
    p = get_solver_result(state)
    rho_cost, rho_gradient = penalty_parameters(state.sub_problem)
    println(lpad(iterations, 11), " | ", lpad(string(round.(p, digits = 6)), 27), " | ",
            lpad(round(norm(p - p_star), sigdigits = 3), 10), " | ",
            lpad(round(abs(only(h(M, p))), sigdigits = 3), 10), " | ",
            lpad(round(rho_cost, sigdigits = 3), 6), " | ",
            lpad(round(rho_gradient, sigdigits = 3), 9))
end

Note that with the fix, we do not quite get the results from before v0.6.4, so it seems there is something else that changed with v0.6.4. But at least the results got better:

Manopt v0.6.3
       
 iterations |          minimizer          |  |p - p*|  |    |h(p)|   | ρ cost | ρ gradient
------------+-----------------------------+------------+------------+--------+-----------
          5 |        [-0.999567, -3.0e-6] |   0.000433 |   0.000866 |    1.0 |       1.0
         10 |             [-1.0, -1.0e-6] |    8.15e-7 |    3.61e-7 |   11.1 |      11.1
         20 |             [-1.0, -1.0e-6] |    7.42e-7 |    2.51e-8 | 1370.0 |    1370.0
         50 |        [-1.050304, -1.0e-6] |     0.0503 |      0.103 | 3.19e11 |   3.19e11

Manopt v0.6.5

 iterations |          minimizer          |  |p - p*|  |    |h(p)|   | ρ cost | ρ gradient
------------+-----------------------------+------------+------------+--------+-----------
          5 |       [-0.684193, 0.031784] |      0.317 |      0.531 |  0.333 |      37.0
         10 |       [-0.684208, 0.031785] |      0.317 |      0.531 |  0.333 |   15200.0
         20 |       [-0.993978, 0.046175] |     0.0466 |    0.00988 |  0.333 |    2.32e8
         50 |                  [NaN, NaN] |        NaN |        NaN |  0.333 |   3.54e12

Manopt v0.6.6

 iterations |          minimizer          |  |p - p*|  |    |h(p)|   | ρ cost | ρ gradient
------------+-----------------------------+------------+------------+--------+-----------
          5 |        [-0.981136, -3.5e-5] |     0.0189 |     0.0374 |   11.1 |      11.1
         10 |        [-0.998304, -3.5e-5] |     0.0017 |    0.00339 |  123.0 |     123.0
         20 |        [-1.000001, -3.5e-5] |    3.51e-5 |    2.71e-6 | 169000.0 |  169000.0
         50 |                  [NaN, NaN] |        NaN |        NaN | 3.93e13 |   3.93e13

[kellertuer]

Thanks for spotting! I am not sure how that happened – sorry.

Will take a closer look tomorrow at your PR, but at first glance it looks fine.

[kellertuer]

Ah – I see that we accidentally changed that (but I do not see why, sorry! We accidentally introduced a bug).
Yours looks fine here – will merge (after dinner, when the tests ran) and register a new version.

Doo you think adding a small test somewhere would help to avoid that this accidentally happens again? If so feel free to add it.

Again – very sorry, I am not sure how that happened. But it was also a super large PR over several weeks – and I tried for the first time a “find further typos” AI assisted check – hopefully that was not the reason it was introduced but really just me somewhere when refactoring.

[kellertuer]

So feel free to also set the date of 0.6.6 to today, again, I can merge this later.

[JoshuaLampert]

No worries. These kind of issues can happen from time to time. My AI was pretty quick in identifying the issue. So it didn't take terribly long for me to debug this. I will add a test for this once I am back home (currently on my phone).

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 100.00%. Comparing base (5e3a1f6) to head (107743e).

Additional details and impacted files

@@            Coverage Diff            @@
##            master      #637   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files           89        89           
  Lines        12536     12536           
=========================================
  Hits         12536     12536           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[kellertuer]

I carefully checked on the branch the PR came from, when that happened – and I can not even really blame an AI, I fear, since at the time I did that, I did not yet have proper access to it. So I have no clue why I ever changed that to 1/3.

Well – a good lesson to even be more careful.

Let me know when you feel this is ready to merge and register; independent of whether you want to add a test or not. But we could do something like a single step and make sure the parameter is set correctly or so.

[JoshuaLampert]

I added a test (basically the MWE above), again, with the help of AI, which also tests if the ρ is correct. While preparing this, the AI found another bug in AugmentedLagrangianCost (edit: nor sure anymore...). Additionally, the AI pointed to this comment

Manopt.jl/test/solvers/test_augmented_lagrangian.jl

Lines 17 to 18 in 5e3a1f6

	# This run (and the following as well) seem to stall after iteration 3 (previously 14)
	# not reaching the minimiser as before

(also introduced in #630) and says it has the same cause as the issue fixed by this PR. Should I look closer into this to see if this changed with this PR and, when appropriate, also update that part (e.g., the increased tolerance) accordingly?

[kellertuer]

Now I am confused. Very confused.

You changed the AbstractConstrainedFunction{CO} to T? But the parameter is the cost it contains not some real value type of. field. This does not make sense and would yield something super inconsistent. All “wrappers” of objectives (like the one abstract type here) should always have the cost they wrap as parameter.

[JoshuaLampert]

Still working on it. Give me a minute.

[kellertuer]

HM the Grad below is maybe then wrong? I am super confused how that came along.

Maybe the large rework PR was not a good idea then. and now Manopt is super buggy and should no longer be used.

I am super confused and sad, that there is so many bugs. Sorry.

[JoshuaLampert]

So what I can say is that with this additional change the MWE I posted above reproduces the exact same numbers as Manopt v0.6.3, which is why I was happy with what the AI found and hoped it was correct, but I have to say that I did not fully understand every part of it yet.

[JoshuaLampert]

The AIs explanation is (sorry for pasting AI messages here; I usually try to understand the answers myself first and rewrite it in my own words, but I need to eat something...):

[...] Second, AugmentedLagrangianCost declared its supertype as AbstractConstrainedFunction{CO} with the objective type CO, while AugmentedLagrangianGrad uses AbstractConstrainedFunction{T} with the type T of the dual variables — which is what the supertype's own docstring specifies. Hence set_parameter!(::AbstractConstrainedFunction{T}, ::Val{:μ}, ::T) and its :λ counterpart never applied to the cost. As set_parameter!(f, args...) falls back to a no-op, the corresponding calls in step_solver! failed silently and the cost's μ and λ kept their initial values for the entire run: the cost was effectively a fixed-multiplier penalty function while the gradient followed the actual augmented Lagrangian. The multiplier update is precisely what distinguishes ALM from a plain penalty method, so this disabled its central mechanism.

[kellertuer]

Yes, the first T is wrong and should be CO in set_parameter!(::AbstractConstrainedFunction{T}, ::Val{:μ}, ::T)

I am not sure how a T ever made it into the first parameter.

since CO is nowhere used in the signature, a correct fix is to just remove the {T} for that case there.

)yeah eating is important! I just had dinner. Guten Appetit!)

I you allow me to push here, I can take a careful look tomorrow, though only after my lecture, so around 1015 probably

[JoshuaLampert]

I applied your suggestion for an alternative fix, however, note that (every?) other subtype(s) of AbstractConstrainedFunction passes the dual/value type rather than the wrapped objective type, see, e.g.

Manopt.jl/src/commons/sub_functions.jl

Line 789 in 5e3a1f6

mutable struct LagrangianCost{CO, T} <: AbstractConstrainedFunction{T}

Which solution do you prefer? If the currently implemented one, then the other subtypes of AbstractConstrainedFunction should probably follow the same logic (maybe in another PR?). Feel free to push to this PR. Do you need permissions because this comes from my fork?

[kellertuer]

Thanks for the fix. I can also check the other cases / subtypes of the AbstractConstrainedFunction later today on a new PR.

So I would merge yours, check and unify on a new PR and register both as a new version.
If for example all subtypes have dual variables, the super type could also have two parameters :)

But yeah I can check that probably easiest on a follow up PR.

Will probably work on that after my lecture today then

[JoshuaLampert]

Sounds good.