ChambollePock.jl

@doc raw"""
    ChambollePockState <: AbstractPrimalDualSolverState

stores all options and variables within a linearized or exact Chambolle Pock.
The following list provides the order for the constructor, where the previous iterates are
initialized automatically and values with a default may be left out.

* `m`:                              base point on ``\mathcal M``
* `n`:                              base point on ``\mathcal N``
* `p`:                              an initial point on ``x^{(0)} ∈\mathcal M`` (and its previous iterate)
* `X`:                              an initial tangent vector ``X^{(0)}∈T^*\mathcal N`` (and its previous iterate)
* `pbar`:                           the relaxed iterate used in the next dual update step (when using `:primal` relaxation)
* `Xbar`:                           the relaxed iterate used in the next primal update step (when using `:dual` relaxation)
* `primal_stepsize`:                (`1/sqrt(8)`) proximal parameter of the primal prox
* `dual_stepsize`:                  (`1/sqrt(8)`) proximal parameter of the dual prox
* `acceleration`:                   (`0.`) acceleration factor due to Chambolle & Pock
* `relaxation`:                     (`1.`) relaxation in the primal relaxation step (to compute `pbar`)
* `relax`:                          (`:primal`) which variable to relax (`:primal` or `:dual`)
* `stop`:                           a [`StoppingCriterion`](@ref)
* `variant`:                        (`exact`) whether to perform an `:exact` or `:linearized` Chambolle-Pock
* `update_primal_base`:             (`(p,o,i) -> o.m`) function to update the primal base
* `update_dual_base`:               (`(p,o,i) -> o.n`) function to update the dual base
* `retraction_method`:              (`default_retraction_method(M, typeof(p))`) the retraction to use
* `inverse_retraction_method`:      (`default_inverse_retraction_method(M, typeof(p))`) an inverse
  retraction to use on the manifold ``\mathcal M``.
* `inverse_retraction_method_dual`: (`default_inverse_retraction_method(N, typeof(n))`)
  an inverse retraction to use on manifold ``\mathcal N``.
* `vector_transport_method`:        (`default_vector_transport_method(M, typeof(p))`) a vector transport to
  use on the manifold ``\mathcal M``.
* `vector_transport_method_dual`:   (`default_vector_transport_method(N, typeof(n))`) a
  vector transport to use on manifold ``\mathcal N``.

where for the last two the functions a [`AbstractManoptProblem`](@ref)` p`,
[`AbstractManoptSolverState`](@ref)` o` and the current iterate `i` are the arguments.
If you activate these to be different from the default identity, you have to provide
`p.Λ` for the algorithm to work (which might be `missing` in the linearized case).

# Constructor

    ChambollePockState(M::AbstractManifold, N::AbstractManifold,
        m::P, n::Q, p::P, X::T, primal_stepsize::Float64, dual_stepsize::Float64;
        kwargs...
    )

where all other fields are keyword arguments with their default values given in brackets.

if `Manifolds.jl` is loaded, `N` is also a keyword argument and set to `TangentBundle(M)` by default.
"""
mutable struct ChambollePockState{
    P,
    Q,
    T,
    RM<:AbstractRetractionMethod,
    IRM<:AbstractInverseRetractionMethod,
    IRM_Dual<:AbstractInverseRetractionMethod,
    VTM<:AbstractVectorTransportMethod,
    VTM_Dual<:AbstractVectorTransportMethod,
} <: AbstractPrimalDualSolverState
    m::P
    n::Q
    p::P
    pbar::P
    X::T
    Xbar::T
    primal_stepsize::Float64
    dual_stepsize::Float64
    acceleration::Float64
    relaxation::Float64
    relax::Symbol
    stop::StoppingCriterion
    variant::Symbol
    update_primal_base::Union{Function,Missing}
    update_dual_base::Union{Function,Missing}
    retraction_method::RM
    inverse_retraction_method::IRM
    inverse_retraction_method_dual::IRM_Dual
    vector_transport_method::VTM
    vector_transport_method_dual::VTM_Dual
end
function Manopt.ChambollePockState(
    M::AbstractManifold,
    N::AbstractManifold,
    m::P,
    n::Q,
    p::P,
    X::T;
    primal_stepsize::Float64=1 / sqrt(8),
    dual_stepsize::Float64=1 / sqrt(8),
    acceleration::Float64=0.0,
    relaxation::Float64=1.0,
    relax::Symbol=:primal,
    stopping_criterion::StoppingCriterion=StopAfterIteration(300),
    variant::Symbol=:exact,
    update_primal_base::Union{Function,Missing}=missing,
    update_dual_base::Union{Function,Missing}=missing,
    retraction_method::RM=default_retraction_method(M, typeof(p)),
    inverse_retraction_method::IRM=default_inverse_retraction_method(M, typeof(p)),
    inverse_retraction_method_dual::IRM_Dual=default_inverse_retraction_method(
        N, typeof(p)
    ),
    vector_transport_method::VTM=default_vector_transport_method(M, typeof(n)),
    vector_transport_method_dual::VTM_Dual=default_vector_transport_method(N, typeof(n)),
) where {
    P,
    Q,
    T,
    RM<:AbstractRetractionMethod,
    IRM<:AbstractInverseRetractionMethod,
    IRM_Dual<:AbstractInverseRetractionMethod,
    VTM<:AbstractVectorTransportMethod,
    VTM_Dual<:AbstractVectorTransportMethod,
}
    return ChambollePockState{P,Q,T,RM,IRM,IRM_Dual,VTM,VTM_Dual}(
        m,
        n,
        p,
        copy(M, p),
        X,
        copy(N, X),
        primal_stepsize,
        dual_stepsize,
        acceleration,
        relaxation,
        relax,
        stopping_criterion,
        variant,
        update_primal_base,
        update_dual_base,
        retraction_method,
        inverse_retraction_method,
        inverse_retraction_method_dual,
        vector_transport_method,
        vector_transport_method_dual,
    )
end
function show(io::IO, cps::ChambollePockState)
    i = get_count(cps, :Iterations)
    Iter = (i > 0) ? "After $i iterations\n" : ""
    Conv = indicates_convergence(cps.stop) ? "Yes" : "No"
    s = """
    # Solver state for `Manopt.jl`s Chambolle-Pock Algorithm
    $Iter
    ## Parameters
    * primal_stepsize:  $(cps.primal_stepsize)
    * dual_stepsize:    $(cps.dual_stepsize)
    * acceleration:     $(cps.acceleration)
    * relaxation:       $(cps.relaxation)
    * relax:            $(cps.relax)
    * variant:          :$(cps.variant)
    * retraction_method:              $(cps.retraction_method)
    * inverse_retraction_method:      $(cps.inverse_retraction_method)
    * vector_transport_method:        $(cps.vector_transport_method)
    * inverse_retraction_method_dual: $(cps.inverse_retraction_method_dual)
    * vector_transport_method_dual:   $(cps.vector_transport_method_dual)

    ## Stopping criterion

    $(status_summary(cps.stop))
    This indicates convergence: $Conv"""
    return print(io, s)
end
get_solver_result(apds::AbstractPrimalDualSolverState) = get_iterate(apds)
get_iterate(apds::AbstractPrimalDualSolverState) = apds.p
function set_iterate!(apds::AbstractPrimalDualSolverState, p)
    apds.p = p
    return apds
end

@doc raw"""
    ChambollePock(
        M, N, cost, x0, ξ0, m, n, prox_F, prox_G_dual, adjoint_linear_operator;
        forward_operator=missing,
        linearized_forward_operator=missing,
        evaluation=AllocatingEvaluation()
    )

Perform the Riemannian Chambolle—Pock algorithm.

Given a `cost` function ``\mathcal E:\mathcal M → ℝ`` of the form
```math
\mathcal E(p) = F(p) + G( Λ(p) ),
```
where ``F:\mathcal M → ℝ``, ``G:\mathcal N → ℝ``,
and ``Λ:\mathcal M → \mathcal N``. The remaining input parameters are

* `p, X`:                         primal and dual start points ``x∈\mathcal M`` and ``ξ∈T_n\mathcal N``
* `m,n`:                          base points on ``\mathcal M`` and ``\mathcal N``, respectively.
* `adjoint_linearized_operator`:  the adjoint ``DΛ^*`` of the linearized operator ``DΛ(m): T_{m}\mathcal M → T_{Λ(m)}\mathcal N``
* `prox_F, prox_G_Dual`:          the proximal maps of ``F`` and ``G^\ast_n``

note that depending on the [`AbstractEvaluationType`](@ref) `evaluation` the last three parameters
as well as the forward operator `Λ` and the `linearized_forward_operator` can be given as
allocating functions `(Manifolds, parameters) -> result`  or as mutating functions
`(Manifold, result, parameters)` -> result` to spare allocations.

By default, this performs the exact Riemannian Chambolle Pock algorithm, see the optional parameter
`DΛ` for their linearized variant.

For more details on the algorithm, see [BergmannHerzogSilvaLouzeiroTenbrinckVidalNunez:2021](@cite).

# Optional parameters

* `acceleration`:                (`0.05`)
* `dual_stepsize`:               (`1/sqrt(8)`) proximal parameter of the primal prox
* `evaluation`:                  ([`AllocatingEvaluation`](@ref)`()) specify whether the proximal maps
  and operators are allocating functions `(Manifolds, parameters) -> result`  or
   given as mutating functions `(Manifold, result, parameters)` -> result`
* `Λ`:                           (`missing`) the (forward) operator ``Λ(⋅)`` (required for the `:exact` variant)
* `linearized_forward_operator`: (`missing`) its linearization ``DΛ(⋅)[⋅]`` (required for the `:linearized` variant)
* `primal_stepsize`:             (`1/sqrt(8)`) proximal parameter of the dual prox
* `relaxation`:                  (`1.`) the relaxation parameter ``γ``
* `relax`:                       (`:primal`) whether to relax the primal or dual
* `variant`:                     (`:exact` if `Λ` is missing, otherwise `:linearized`) variant to use.
  Note that this changes the arguments the `forward_operator` is called with.
* `stopping_criterion`:          (`[StopAfterIteration`](@ref)`(100)`) a [`StoppingCriterion`](@ref)
* `update_primal_base`:          (`missing`) function to update `m` (identity by default/missing)
* `update_dual_base`:            (`missing`) function to update `n` (identity by default/missing)
* `retraction_method`:           (`default_retraction_method(M, typeof(p))`) the retraction to use
* `inverse_retraction_method`    (`default_inverse_retraction_method(M, typeof(p))`) an inverse retraction to use.
* `vector_transport_method`      (`default_vector_transport_method(M, typeof(p))`) a vector transport to use

# Output

the obtained (approximate) minimizer ``p^*``, see [`get_solver_return`](@ref) for details.
"""
function ChambollePock(
    M::AbstractManifold,
    N::AbstractManifold,
    cost::TF,
    p::P,
    X::T,
    m::P,
    n::Q,
    prox_F::Function,
    prox_G_dual::Function,
    adjoint_linear_operator::Function;
    Λ::Union{Function,Missing}=missing,
    linearized_forward_operator::Union{Function,Missing}=missing,
    kwargs...,
) where {TF,P,T,Q}
    q = copy(M, p)
    Y = copy(N, n, X)
    m2 = copy(M, m)
    n2 = copy(N, n)
    return ChambollePock!(
        M,
        N,
        cost,
        q,
        Y,
        m2,
        n2,
        prox_F,
        prox_G_dual,
        adjoint_linear_operator;
        Λ=Λ,
        linearized_forward_operator=linearized_forward_operator,
        kwargs...,
    )
end
@doc raw"""
    ChambollePock(M, N, cost, x0, ξ0, m, n, prox_F, prox_G_dual, adjoint_linear_operator)

Perform the Riemannian Chambolle—Pock algorithm in place of `x`, `ξ`, and potentially `m`,
`n` if they are not fixed. See [`ChambollePock`](@ref) for details and optional parameters.
"""
function ChambollePock!(
    M::AbstractManifold,
    N::AbstractManifold,
    cost::TF,
    p::P,
    X::T,
    m::P,
    n::Q,
    prox_F::Function,
    prox_G_dual::Function,
    adjoint_linear_operator::Function;
    Λ::Union{Function,Missing}=missing,
    linearized_forward_operator::Union{Function,Missing}=missing,
    acceleration=0.05,
    dual_stepsize=1 / sqrt(8),
    primal_stepsize=1 / sqrt(8),
    relaxation=1.0,
    relax::Symbol=:primal,
    stopping_criterion::StoppingCriterion=StopAfterIteration(200),
    update_primal_base::Union{Function,Missing}=missing,
    update_dual_base::Union{Function,Missing}=missing,
    retraction_method::RM=default_retraction_method(M, typeof(p)),
    inverse_retraction_method::IRM=default_inverse_retraction_method(M, typeof(p)),
    vector_transport_method::VTM=default_vector_transport_method(M, typeof(p)),
    variant=ismissing(Λ) ? :exact : :linearized,
    kwargs...,
) where {
    TF,
    P,
    Q,
    T,
    RM<:AbstractRetractionMethod,
    IRM<:AbstractInverseRetractionMethod,
    VTM<:AbstractVectorTransportMethod,
}
    pdmo = PrimalDualManifoldObjective(
        cost,
        prox_F,
        prox_G_dual,
        adjoint_linear_operator;
        linearized_forward_operator=linearized_forward_operator,
        Λ=Λ,
    )
    dpdmo = decorate_objective!(M, pdmo; kwargs...)
    tmp = TwoManifoldProblem(M, N, dpdmo)
    cps = ChambollePockState(
        M,
        m,
        n,
        p,
        X;
        N=N,
        primal_stepsize=primal_stepsize,
        dual_stepsize=dual_stepsize,
        acceleration=acceleration,
        relaxation=relaxation,
        stopping_criterion=stopping_criterion,
        relax=relax,
        update_primal_base=update_primal_base,
        update_dual_base=update_dual_base,
        variant=variant,
        retraction_method=retraction_method,
        inverse_retraction_method=inverse_retraction_method,
        vector_transport_method=vector_transport_method,
    )
    dcps = decorate_state!(cps; kwargs...)
    solve!(tmp, dcps)
    return get_solver_return(get_objective(tmp), dcps)
end

function initialize_solver!(::TwoManifoldProblem, ::ChambollePockState) end

function step_solver!(tmp::TwoManifoldProblem, cps::ChambollePockState, iter)
    N = get_manifold(tmp, 2)
    primal_dual_step!(tmp, cps, Val(cps.relax))
    cps.m =
        ismissing(cps.update_primal_base) ? cps.m : cps.update_primal_base(tmp, cps, iter)
    if !ismissing(cps.update_dual_base)
        n_old = deepcopy(cps.n)
        cps.n = cps.update_dual_base(tmp, cps, iter)
        vector_transport_to!(
            N, cps.X, n_old, cps.X, cps.n, cps.vector_transport_method_dual
        )
        vector_transport_to!(
            N, cps.Xbar, n_old, cps.Xbar, cps.n, cps.vector_transport_method_dual
        )
    end
    return cps
end
#
# Variant 1: primal relax
#
function primal_dual_step!(tmp::TwoManifoldProblem, cps::ChambollePockState, ::Val{:primal})
    dual_update!(tmp, cps, cps.pbar, Val(cps.variant))
    obj = get_objective(tmp)
    M = get_manifold(tmp, 1)
    N = get_manifold(tmp, 2)
    if !hasproperty(obj, :Λ!!) || ismissing(obj.Λ!!)
        ptXn = cps.X
    else
        ptXn = vector_transport_to(
            N, cps.n, cps.X, forward_operator(tmp, cps.m), cps.vector_transport_method_dual
        )
    end
    p_old = cps.p
    cps.p = get_primal_prox!(
        tmp,
        cps.p,
        cps.primal_stepsize,
        retract(
            M,
            cps.p,
            vector_transport_to(
                M,
                cps.m,
                -cps.primal_stepsize *
                (adjoint_linearized_operator(tmp, cps.m, cps.n, ptXn)),
                cps.p,
                cps.vector_transport_method,
            ),
            cps.retraction_method,
        ),
    )
    update_prox_parameters!(cps)
    retract!(
        M,
        cps.pbar,
        cps.p,
        -cps.relaxation * inverse_retract(M, cps.p, p_old, cps.inverse_retraction_method),
        cps.retraction_method,
    )
    return cps
end
#
# Variant 2: dual relax
#
function primal_dual_step!(tmp::TwoManifoldProblem, cps::ChambollePockState, ::Val{:dual})
    obj = get_objective(tmp)
    M = get_manifold(tmp, 1)
    N = get_manifold(tmp, 2)
    if !hasproperty(obj, :Λ!!) || ismissing(obj.Λ!!)
        ptXbar = cps.Xbar
    else
        ptXbar = vector_transport_to(
            N,
            cps.n,
            cps.Xbar,
            forward_operator(tmp, cps.m),
            cps.vector_transport_method_dual,
        )
    end
    get_primal_prox!(
        tmp,
        cps.p,
        cps.primal_stepsize,
        retract(
            M,
            cps.p,
            vector_transport_to(
                M,
                cps.m,
                -cps.primal_stepsize *
                (adjoint_linearized_operator(tmp, cps.m, cps.n, ptXbar)),
                cps.p,
                cps.vector_transport_method,
            ),
            cps.retraction_method,
        ),
    )
    X_old = deepcopy(cps.X)
    dual_update!(tmp, cps, cps.p, Val(cps.variant))
    update_prox_parameters!(cps)
    cps.Xbar = cps.X + cps.relaxation * (cps.X - X_old)
    return cps
end
#
# Dual step: linearized
# depending on whether its primal relaxed or dual relaxed, start from start=o.x or start=o.xbar
#
function dual_update!(
    tmp::TwoManifoldProblem, cps::ChambollePockState, start::P, ::Val{:linearized}
) where {P}
    M = get_manifold(tmp, 1)
    N = get_manifold(tmp, 2)
    obj = get_objective(tmp)
    # (1) compute update direction
    X_update = linearized_forward_operator(
        tmp, cps.m, inverse_retract(M, cps.m, start, cps.inverse_retraction_method), cps.n
    )
    # (2) if p.Λ is missing, if n = Λ(m) and do not PT, otherwise do
    (hasproperty(obj, :Λ!!) && !ismissing(obj.Λ!!)) && vector_transport_to!(
        N,
        X_update,
        forward_operator(tmp, cps.m),
        X_update,
        cps.n,
        cps.vector_transport_method_dual,
    )
    # (3) to the dual update
    get_dual_prox!(
        tmp, cps.X, cps.n, cps.dual_stepsize, cps.X + cps.dual_stepsize * X_update
    )
    return cps
end
#
# Dual step: exact
# depending on whether its primal relaxed or dual relaxed start from start=o.x or start=o.xbar here
#
function dual_update!(
    tmp::TwoManifoldProblem, cps::ChambollePockState, start::P, ::Val{:exact}
) where {P}
    N = get_manifold(tmp, 2)
    ξ_update = inverse_retract(
        N, cps.n, forward_operator(tmp, start), cps.inverse_retraction_method_dual
    )
    get_dual_prox!(
        tmp, cps.X, cps.n, cps.dual_stepsize, cps.X + cps.dual_stepsize * ξ_update
    )
    return cps
end

@doc raw"""
    update_prox_parameters!(o)
update the prox parameters as described in Algorithm 2 of [ChambollePock:2011](@cite),

1. ``θ_{n} = \frac{1}{\sqrt{1+2γτ_n}}``
2. ``τ_{n+1} = θ_nτ_n``
3. ``σ_{n+1} = \frac{σ_n}{θ_n}``
"""
function update_prox_parameters!(pds::S) where {S<:AbstractPrimalDualSolverState}
    if pds.acceleration > 0
        pds.relaxation = 1 / sqrt(1 + 2 * pds.acceleration * pds.primal_stepsize)
        pds.primal_stepsize = pds.primal_stepsize * pds.relaxation
        pds.dual_stepsize = pds.dual_stepsize / pds.relaxation
    end
    return pds
end