Migrate Manopt solver callbacks

Project.toml

@@ -18,7 +18,7 @@ TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 JuliaFormatter = "=2.10.1"
 Manifolds = "0.11.28"
 ManifoldsBase = "2.3.5"
-Manopt = "0.6"
+Manopt = "0.6.2"
 PrecompileTools = "1"
 ProgressMeter = "1.11.0"
 RecursiveArrayTools = "4.3"

src/solvers/lbfgs.jl

@@ -143,6 +143,7 @@ function solve_lbfgs(
         stepsize = _lbfgs_linesearch(linesearch),
         stopping_criterion = stopping,
         debug = callbacks.debug_actions,
+        callbacks = callbacks.solver_callbacks,
         count = [:Cost, :Gradient],
         return_state = true,
     )

src/solvers/lm.jl

@@ -218,6 +218,7 @@ function solve_lm(
         sub_objective = sub_objective,
         sub_state = sub_state,
         debug = callbacks.debug_actions,
+        callbacks = callbacks.solver_callbacks,
         return_state = true,
     )
 

src/solvers/manopt_helpers.jl

@@ -421,7 +421,7 @@ function _manopt_stopping(maxiter::Int, grad_stop_tol, dual_stop; extra = ())
     )
 end
 
-# Create progress and debug callbacks shared by Manopt-backed solvers.
+# Create progress, lifecycle callbacks, and debug actions shared by Manopt-backed solvers.
 function _manopt_callbacks(
     make_progress::Function,
     maxiter::Int,
@@ -444,18 +444,21 @@ function _manopt_callbacks(
         M;
         diagnostics_recorder,
     )
-    debug_actions = _solver_debug_actions(
-        verbose,
+    step_callback = _solver_callback_group(
         post_step_callback,
         diagnostics_callback,
         progress_callback,
         iteration_callbacks...,
     )
+    solver_callbacks =
+        isnothing(step_callback) ? Any[] : Any[:Init=>step_callback, :Step=>step_callback]
+    debug_actions = _solver_debug_actions(verbose)
     return (
         progress = progress,
         diagnostics_callback = diagnostics_callback,
         progress_callback = progress_callback,
         debug_actions = debug_actions,
+        solver_callbacks = solver_callbacks,
     )
 end
 
@@ -595,12 +598,20 @@ function _solver_stats(
     )
 end
 
-# Collect only active Manopt debug callbacks, dropping omitted hooks.
-_solver_debug_callbacks(callbacks...) = Any[cb for cb in callbacks if !isnothing(cb)]
+# Combine active solver callbacks in their declared execution order.
+function _solver_callback_group(callbacks...)
+    active_callbacks = Any[callback for callback in callbacks if !isnothing(callback)]
+    isempty(active_callbacks) && return nothing
+    return function (problem, state, k)
+        for callback in active_callbacks
+            callback(problem, state, k)
+        end
+        return nothing
+    end
+end
 
-# Build Manopt debug actions and attach TensorKitchen callback hooks.
-function _solver_debug_actions(verbose::Union{Nothing,Bool}, callbacks...)
-    callback_actions = _solver_debug_callbacks(callbacks...)
+# Build only Manopt display/debug actions; solver callbacks use `callbacks=`.
+function _solver_debug_actions(verbose::Union{Nothing,Bool})
     if verbose === true
         io = _SOLVER_DEBUG_SINK
         init_group = Manopt.DebugGroup([
@@ -619,11 +630,9 @@ function _solver_debug_actions(verbose::Union{Nothing,Bool}, callbacks...)
             ]),
             100,
         )
-        iteration_actions = Any[iter_group]
-        append!(iteration_actions, callback_actions)
-        return Any[:Start=>Any[init_group], :Iteration=>iteration_actions]
+        return Any[:Start=>Any[init_group], :Iteration=>Any[iter_group]]
     end
-    return callback_actions
+    return Any[]
 end
 
 # Create a Manopt iteration callback that updates TensorKitchen progress output.

src/solvers/rcg.jl

@@ -131,6 +131,7 @@ function solve_rcg(
         restart_condition = restart_rule,
         stopping_criterion = stopping,
         debug = callbacks.debug_actions,
+        callbacks = callbacks.solver_callbacks,
         count = [:Cost, :Gradient],
         return_state = true,
     )

src/solvers/rgd.jl

@@ -105,6 +105,7 @@ function solve_rgd(
         stepsize = armijo,
         stopping_criterion = stopping,
         debug = callbacks.debug_actions,
+        callbacks = callbacks.solver_callbacks,
         count = [:Cost, :Gradient],
         return_state = true,
     )
@@ -189,6 +190,7 @@ function solve_rgd_fixed(
         stepsize = Manopt.ConstantStepsize(M, T(stepsize) * setup.objective_scale),
         stopping_criterion = stopping,
         debug = callbacks.debug_actions,
+        callbacks = callbacks.solver_callbacks,
         count = [:Cost, :Gradient],
         return_state = true,
     )

test/basic_tests.jl

@@ -1169,7 +1169,7 @@ end
     @test res_init_sym.solver_info.function_evaluations >= 0
     @test res_init_sym.solver_info.gradient_evaluations >= 1
 
-    res_trace = cpd(
+    res_trace = @test_logs min_level = Base.CoreLogging.Warn cpd(
         A,
         r;
         solver = :rgd,
@@ -1202,6 +1202,9 @@ end
     @test hasproperty(trace_info, :component_trace_rgrad_failed_count)
     @test isfinite(trace_info.component_trace_start_rel_error)
     @test trace_info.component_trace_rgrad_failed_count == 0
+    @test length(trace_info.component_trace_iterations) == res_trace.iterations
+    @test length(trace_info.component_trace_cost_history) == res_trace.iterations
+    @test length(trace_info.component_trace_max_delta_history) == res_trace.iterations
     @test length(trace_info.component_trace_iterations) ==
           length(trace_info.component_trace_max_delta_history)
     @test length(trace_info.component_trace_delta_history) ==