Feature: Ipopt callbacks

docs/literate/man/ipopt_callbacks.jl

@@ -11,7 +11,7 @@ import ..QuantumStateSmoothPulseProblem
 import ..Callbacks
 
 # By default, IpOpt callbacks are called at each optimization step with the following signature:
-function full_callback_signature(
+function full_argument_list_callback(
     alg_mod::Cint,
     iter_count::Cint,
     obj_value::Float64,
@@ -33,18 +33,8 @@ end
 # ProblemSolvers.solve!(prob::QuantumControlProblem; callback=nothing)
 # ```
 
-# The callback function can be used to monitor the optimization progress, save intermediate results, or modify the optimization process.
-# For example, the following callback function saves the optimization trajectory at each iteration:
 
-trajectory_history = []
-function get_history_callback(
-    kwargs...
-)
-    push!(trajectory_history, QuantumCollocation.Problems.get_datavec(prob))
-    return true
-end
-
-# The callback function can also be used to stop the optimization early by returning `false`. The following callback when passed to `solve!` will stop the optimization after the first iteration:
+# The callback function can be used to stop the optimization early by returning `false`. The following callback when passed to `solve!` will stop the optimization after the first iteration:
 my_callback = (kwargs...) -> false
 
 T = 50
@@ -60,36 +50,37 @@ prob = QuantumStateSmoothPulseProblem(
     ipopt_options=IpoptOptions(print_level=1), 
     piccolo_options=PiccoloOptions(verbose=false)
 )
+
 
-trajectory_history = []
+# The callback function can be used to monitor the optimization progress, save intermediate results, or modify the optimization process.
+# For example, the following callback function saves the optimization trajectory at each iteration:
+callback, trajectory_history = Callbacks.callback_get_trajectory_history(prob)
 # using the get_history_callback function to save the optimization trajectory at each iteration into the trajectory_history 
-solve!(prob, max_iter=20, callback=get_history_callback)
+solve!(prob, max_iter=20, callback=callback)
 
+# save trajectory images into files
 for (iter, traj) in enumerate(trajectory_history)
-    # get the length of the trajectory history depending on length and left pad the index with leading zeros
     str_index = lpad(iter, length(string(length(trajectory_history))), "0")
-    # plot the trajectory but on fixed xaxis and yaxis
-    plot("./iteration-$str_index-trajectory.png", NamedTrajectory(traj, prob.trajectory),  [:ψ̃1, :a], xlims=(-Δt, (T+5)*Δt), plot_ylims=(ψ̃1 = (-2, 2), a = (-1.1, 1.1)))
+    plot("./iteration-$str_index-trajectory.png", traj, [:ψ̃, :a], xlims=(-Δt, (T+5)*Δt), ylims=(ψ̃1 = (-2, 2), a = (-1.1, 1.1)))
 end
 
 # Using a callback to get the best trajectory from all the optimization iterations
-T = 50
-Δt = 0.2
-sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]])
-ψ_init =  Vector{ComplexF64}([0.0, 1.0])
-ψ_target =  Vector{ComplexF64}([1.0, 0.0])
+sys2 = QuantumSystem(0.15 * GATES[:Z], [GATES[:X], GATES[:Y]])
+ψ_init2 =  Vector{ComplexF64}([0.0, 1.0])
+ψ_target2 =  Vector{ComplexF64}([1.0, 0.0])
 
 # Single initial and target states
 # --------------------------------
-prob = QuantumStateSmoothPulseProblem(
-    sys, ψ_init, ψ_target, T, Δt;
+prob2 = QuantumStateSmoothPulseProblem(
+    sys2, ψ_init2, ψ_target2, T, Δt;
     ipopt_options=IpoptOptions(print_level=1), 
     piccolo_options=PiccoloOptions(verbose=false)
 )
 
-(best_traj_callback, best_traj) = make_save_best_trajectory_callback(prob, prob.trajectory)
-solve!(prob, max_iter=20, callback=best_traj_callback)
+best_trajectory_callback, best_trajectory_list = callback_get_best_iterate(prob2)
+solve!(prob2, max_iter=20, callback=best_trajectory_callback)
 # fidelity of the last iterate
-@show fidelity(prob)
+@show Losses.fidelity(prob2)
+
 # fidelity of the best iterate
-@show fidelity(best_traj)
+@show QuantumCollocation.fidelity(best_trajectory_list[end], prob2.system)

src/callbacks.jl

@@ -126,4 +126,4 @@ end
     @test length(obj_vals) == 4   # problem init, iter 1, iter 2, iter 3 (terminate)
 end
 
-end
+end # module