test unitary fid callback

src/callbacks.jl

@@ -34,7 +34,7 @@ function best_rollout_fidelity_callback(prob::QuantumControlProblem)
 end
 
 function best_unitary_rollout_fidelity_callback(prob::QuantumControlProblem)
-    return best_rollout_callback(prob, fidelity_unitary)
+    return best_rollout_callback(prob, unitary_fidelity)
 end
 
 function trajectory_history_callback(prob::QuantumControlProblem)
@@ -82,7 +82,7 @@ end
     @test length(trajectory_history) == 21
 end
 
-@testitem "Callback can get best trajectory" begin
+@testitem "Callback can get best state trajectory" begin
     using MathOptInterface
     using NamedTrajectories
     const MOI = MathOptInterface
@@ -110,6 +110,34 @@ end
     @test after ≤ best
 end
 
+@testitem "Callback can get best unitary trajectory" begin
+    using MathOptInterface
+    using NamedTrajectories
+    const MOI = MathOptInterface
+    include("../test/test_utils.jl")
+
+    prob, system = smooth_unitary_problem(return_system=true)
+
+    callback, best_trajs = best_unitary_rollout_fidelity_callback(prob)
+    @test length(best_trajs) == 0
+
+    # measure fidelity
+    before = unitary_fidelity(prob)
+    solve!(prob, max_iter=20, callback=callback)
+
+    # length must increase if iterations are made
+    @test length(best_trajs) > 0
+    @test best_trajs[end] isa NamedTrajectory
+
+    # fidelity ranking
+    after = unitary_fidelity(prob)
+    best = unitary_fidelity(best_trajs[end], system)
+
+    @test before < after
+    @test before < best
+    @test after ≤ best
+end
+
 @testitem "Callback with full parameter test" begin
     using MathOptInterface
     using NamedTrajectories

test/test_utils.jl

@@ -134,3 +134,20 @@ function smooth_quantum_state_problem(; return_system::Bool=false)
         return prob
     end
 end
+
+function smooth_unitary_problem(; return_system::Bool=false)
+    T = 50
+    Δt = 0.2
+    sys = QuantumSystem(0.1 * PAULIS[:Z], [PAULIS[:X], PAULIS[:Y]])
+    U_goal = GATES[:H]
+    prob = UnitarySmoothPulseProblem(
+        sys, U_goal, T, Δt;
+        ipopt_options=IpoptOptions(print_level=1), 
+        piccolo_options=PiccoloOptions(verbose=false)
+    )
+    if return_system
+        return prob, sys
+    else
+        return prob
+    end
+end