How to use callbacks

Ronny Bergmann

This tutorial illustrates how to use custom callbacks. Manopt.jl 0.6.2 introduces a generic way to work with callbacks within solvers. While simple debug and recording are available, callbacks offer a more flexible way to interact with a solver run.

The framework of callbacks allows to call user-defined functions to be called at specific places, referred to as “hooks” in the following, during the run of a solver.

using Manopt, Manifolds, Random, LinearAlgebra, ManifoldDiff
using ManifoldDiff: grad_distance
Random.seed!(42);

We consider the same example as in the Introductory tutorial to compute the Riemannian center of mass

n = 100
σ = π / 8
M = Sphere(2)
p = 1 / sqrt(2) * [1.0, 0.0, 1.0]
data = [exp(M, p,  σ * rand(M; vector_at=p)) for i in 1:n];
f(M, p) = sum(1 / (2 * n) * distance.(Ref(M), Ref(p), data) .^ 2)
grad_f(M, p) = sum(1 / n * grad_distance.(Ref(M), data, Ref(p)));

And we also use gradient_descent as the example solver for this tutorial.

The general callback function is called at every hook a solver has. By default there are at least

• :BeforeInit and :Init before and after initialize_solver! is executed,
• :BeforeStep and :Step before and after a step_solver! is performed.
• :BeforeStop before the stopping criterion is evaluated, and
• :Stop when the stopping criterion indicated to stop.

There are two different variants to introduce callbacks: A single callback that is called at every hook and specific callbacks for single hooks.

A general callback for all hooks

The general callback function has to have the signature of the form

cb(hook::Symbol, problem::AbstractManoptProblem, state::AbstractManoptSolverState, k::Int)

where the hook is the specific hook name currently calling cb, and problem and state can of course be more precisely specified to fit the solver you run.

As an example here, we generate a list of the Hooks visited when only performing the first step of the solver run. We collect them in a string, that starts before the initialisation, and generates a list for every step. All this is collected in a globak string

vhs = ""
function visited_hooks(hook, problem, state, k)
    if hook === :BeforeInit # Reset string on start
        global vhs = """
            # Visited Hooks during this solver run

            ## Initialisation
            """
    end
    if hook === :BeforeStep # Start new list every step
        vhs *= """

            ## Step #$(k)
            """
    end # For all: Add an item of the current hook name
    vhs *= "* :$(hook)\n"
    return nothing
end

To pass this to the solver, the general callbacks = can be used. For the callback for all hooks, we can just pass the function we just defined. This fills the string vhs as follows

gradient_descent(
    M, f, grad_f, data[1];
    callbacks = visited_hooks,
    stopping_criterion = StopAfterIteration(1),
)
vhs |> print

# Visited Hooks during this solver run

## Initialisation
* :BeforeInit
* :Init
* :BeforeStop

## Step #1
* :BeforeStep
* :Stepsize
* :Step
* :BeforeStop
* :Stop

And we see that an additional hook is available: :Stepsize. This is placed after the step size is computed but before the actual step is taken.

Solvers might offer several further hooks. Given a state type, these can be seen calling

provided_callbacks(GradientDescentState)

8-element Vector{Symbol}:
 :Any
 :BeforeInit
 :BeforeStep
 :BeforeStop
 :Init
 :Step
 :Stop
 :Stepsize

We further see the :Any hook. This is used internally to store the general fallback we constructed in this section.

Specific callbacks for single hooks.

If there is some computation that shall only be performed on a single hook, one can of course use the method from the last section and use the if variant as already used there as well. Alternatively, one can also specify a callback with a signature not requiring the symbol as before, that is

cb(problem::AbstractManoptProblem, state::AbstractManoptSolverState, k::Int)

and add it to just one of the hooks. As an example, while recording allows to record fields, we maybe want to modify the state as well. Here – as a toy example – we want to modify the descent direction after the step size is computed but before the step is taken. We could for example shorten it a bit, but also only for the first 10 steps.

function stepsize_callback(problem, state, k)
    if k <= 10
        state.X ./= 2
        (k==10) && (@info "Stop shortening direction at k=10")
    end
end

Adding this to a certain hook is done by passing a vector of Pairs to the callbacks = keyword.

gradient_descent(
    M, f, grad_f, data[1]; callbacks = [:Stepsize => stepsize_callback]
);

[ Info: Stop shortening direction at k=10

Both approaches can also be mixed, by just adding visited_hooks into the callbacks = array. This is equivalent to passing :Any => visited_hooks as a pair therein as well.

Status `~/work/Manopt.jl/Manopt.jl/tutorials/Project.toml`
  [47edcb42] ADTypes v1.22.2
  [6e4b80f9] BenchmarkTools v1.8.0
  [5ae59095] Colors v0.13.1
  [a0c0ee7d] DifferentiationInterface v0.7.20
  [31c24e10] Distributions v0.25.129
  [26cc04aa] FiniteDifferences v0.12.34
  [f6369f11] ForwardDiff v1.4.1
  [8ac3fa9e] LRUCache v1.6.2
  [af67fdf4] ManifoldDiff v0.4.5
  [1cead3c2] Manifolds v0.11.28
  [3362f125] ManifoldsBase v2.5.0
  [0fc0a36d] Manopt v0.6.2 `.`
  [91a5bcdd] Plots v1.41.6
  [731186ca] RecursiveArrayTools v4.3.4
  [37e2e46d] LinearAlgebra v1.12.0
  [9a3f8284] Random v1.11.0