feat: update version to 0.3.0 and add callbacks module; implement suf…

…fix utilities for integrators

Project.toml

@@ -1,7 +1,7 @@
 name = "QuantumCollocationCore"
 uuid = "2b384925-53cb-4042-a8d2-6faa627467e1"
 authors = ["Aaron Trowbridge <[email protected]> and contributors"]
-version = "0.2.1"
+version = "0.3.0"
 
 [deps]
 Einsum = "b7d42ee7-0b51-5a75-98ca-779d3107e4c0"
@@ -30,7 +30,7 @@ Ipopt = "1.7"
 JLD2 = "0.5"
 MathOptInterface = "1.35"
 NamedTrajectories = "0.2"
-PiccoloQuantumObjects = "0.2"
+PiccoloQuantumObjects = "0.3"
 Reexport = "1.2"
 TestItemRunner = "1.1"
 TestItems = "1.0"

src/QuantumCollocationCore.jl

@@ -32,4 +32,7 @@ include("save_load_utils.jl")
 include("problem_solvers.jl")
 @reexport using .ProblemSolvers
 
+include("callbacks.jl")
+@reexport using .Callbacks
+
 end

src/callbacks.jl

@@ -0,0 +1,173 @@
+module Callbacks
+
+export best_rollout_fidelity_callback
+export best_unitary_rollout_fidelity_callback
+export trajectory_history_callback
+
+using NamedTrajectories
+using TestItems
+
+using PiccoloQuantumObjects
+
+using ..Problems
+
+function best_rollout_callback(
+    prob::QuantumControlProblem,
+    system::Union{AbstractQuantumSystem, AbstractVector{<:AbstractQuantumSystem}}, rollout_fidelity::Function
+)
+    best_value = 0.0
+    best_trajectories = []
+
+    function callback(args...)
+        traj = NamedTrajectory(Problems.get_datavec(prob), prob.trajectory)
+        value = rollout_fidelity(traj, system)
+        if value > best_value
+            best_value = value
+            push!(best_trajectories, traj)
+        end
+        return true
+    end
+
+    return callback, best_trajectories
+end
+
+function trajectory_history_callback(prob::QuantumControlProblem)
+    trajectory_history = []
+    function callback(args...)
+        push!(trajectory_history, NamedTrajectory(Problems.get_datavec(prob), prob.trajectory))
+        return true
+    end
+
+    return callback, trajectory_history
+end
+
+# *************************************************************************** #
+
+# TODO: figure out a way to test callback from QCC
+#     - maybe figure out a way to set up empty ipopt Problems
+#     - or an abstract callback struct anticipating other interfaces
+#     - or something with this: https://pkgdocs.julialang.org/v1/creating-packages/index.html#Test-specific-dependencies
+
+# @testitem "Callback returns false early stops" begin
+#     using MathOptInterface
+#     const MOI = MathOptInterface
+#     using LinearAlgebra
+
+#     include("../test/test_utils.jl")
+
+#     prob = smooth_quantum_state_problem()
+
+#     my_callback = (kwargs...) -> false
+
+#     solve!(prob, max_iter=20, callback=my_callback)
+
+#     # callback forces problem to exit early as per Ipopt documentation
+#     @test MOI.get(prob.optimizer, MOI.TerminationStatus()) == MOI.INTERRUPTED
+# end
+
+
+# @testitem "Callback can get internal history" begin
+#     using MathOptInterface
+#     using NamedTrajectories
+#     const MOI = MathOptInterface
+#     include("../test/test_utils.jl")
+
+#     prob = smooth_quantum_state_problem()
+
+#     callback, trajectory_history = trajectory_history_callback(prob)
+
+#     solve!(prob, max_iter=20, callback=callback)
+#     @test length(trajectory_history) == 21
+# end
+
+# @testitem "Callback can get best state trajectory" begin
+#     using MathOptInterface
+#     using NamedTrajectories
+#     const MOI = MathOptInterface
+#     include("../test/test_utils.jl")
+
+#     prob, system = smooth_quantum_state_problem(return_system=true)
+
+#     callback, best_trajs = best_rollout_fidelity_callback(prob, system)
+#     @test length(best_trajs) == 0
+
+#     # measure fidelity
+#     before = rollout_fidelity(prob, system)
+#     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 = rollout_fidelity(prob, system)
+#     best = rollout_fidelity(best_trajs[end], system)
+
+#     @test before < after
+#     @test before < best
+#     @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, system)
+#     @test length(best_trajs) == 0
+
+#     # measure fidelity
+#     before = unitary_rollout_fidelity(prob.trajectory, system)
+#     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_rollout_fidelity(prob.trajectory, system)
+#     best = unitary_rollout_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
+#     const MOI = MathOptInterface
+#     include("../test/test_utils.jl")
+
+#     prob = smooth_quantum_state_problem()
+
+#     obj_vals = []
+#     function get_history_callback(
+#         alg_mod::Cint,
+#         iter_count::Cint,
+#         obj_value::Float64,
+#         inf_pr::Float64,
+#         inf_du::Float64,
+#         mu::Float64,
+#         d_norm::Float64,
+#         regularization_size::Float64,
+#         alpha_du::Float64,
+#         alpha_pr::Float64,
+#         ls_trials::Cint,
+#     )
+#         push!(obj_vals, obj_value)
+#         return iter_count < 3
+#     end
+
+#     solve!(prob, max_iter=20, callback=get_history_callback)
+
+#     @test MOI.get(prob.optimizer, MOI.TerminationStatus()) == MOI.INTERRUPTED
+#     @test length(obj_vals) == 4   # problem init, iter 1, iter 2, iter 3 (terminate)
+# end
+
+
+
+end

src/integrators/_integrators.jl

@@ -25,6 +25,11 @@ export sixth_order_pade
 export eighth_order_pade
 export tenth_order_pade
 
+export get_suffix
+export add_suffix
+export remove_suffix
+export modify_integrator_suffix
+
 using NamedTrajectories
 using TrajectoryIndexingUtils
 using PiccoloQuantumObjects
@@ -33,6 +38,8 @@ using SparseArrays
 using ForwardDiff
 using TestItems
 
+import PiccoloQuantumObjects
+
 const ⊗ = kron
 
 abstract type AbstractIntegrator end
@@ -67,4 +74,115 @@ include("pade_integrators.jl")
 include("exponential_integrators.jl")
 
 
+# ----------------------------------------------------------------------------
+# Integrator direct sum methods
+# ----------------------------------------------------------------------------
+
+function modify_integrator_suffix(
+    integrator::AbstractIntegrator,
+    modifier::Function,
+    suffix::String,
+    traj::NamedTrajectory,
+    mod_traj::NamedTrajectory
+)
+    mod_integrator = deepcopy(integrator)
+
+    if integrator isa QuantumIntegrator
+        state_name = get_component_names(traj, integrator.state_components)
+        drive_name = get_component_names(traj, integrator.drive_components)
+        mod_integrator.state_components = mod_traj.components[modifier(state_name, suffix)]
+        mod_integrator.drive_components = mod_traj.components[modifier(drive_name, suffix)]
+        return mod_integrator
+    elseif integrator isa DerivativeIntegrator
+        var_name = get_component_names(traj, integrator.variable_components)
+        der_name = get_component_names(traj, integrator.derivative_components)
+        mod_integrator.variable_components = mod_traj.components[modifier(var_name, suffix)]
+        mod_integrator.derivative_components = mod_traj.components[modifier(der_name, suffix)]
+        return mod_integrator
+    else
+        error("Integrator type not recognized")
+    end
+end
+
+function PiccoloQuantumObjects.add_suffix(
+    integrator::AbstractIntegrator,
+    suffix::String,
+    traj::NamedTrajectory,
+    mod_traj::NamedTrajectory
+)
+    return modify_integrator_suffix(integrator, add_suffix, suffix, traj, mod_traj)
+end
+
+function PiccoloQuantumObjects.add_suffix(
+    integrators::AbstractVector{<:AbstractIntegrator},
+    suffix::String,
+    traj::NamedTrajectory,
+    mod_traj::NamedTrajectory
+)
+    return [
+        add_suffix(integrator, suffix, traj, mod_traj)
+            for integrator ∈ integrators
+    ]
+end
+
+function PiccoloQuantumObjects.remove_suffix(
+    integrator::AbstractIntegrator,
+    suffix::String,
+    traj::NamedTrajectory,
+    mod_traj::NamedTrajectory
+)
+    return modify_integrator_suffix(integrator, remove_suffix, suffix, traj, mod_traj)
+end
+
+function PiccoloQuantumObjects.remove_suffix(
+    integrators::AbstractVector{<:AbstractIntegrator},
+    suffix::String,
+    traj::NamedTrajectory,
+    mod_traj::NamedTrajectory
+)
+    return [remove_suffix(intg, suffix, traj, mod_traj) for intg in integrators]
+end
+
+
+# Get suffix utilities
+# --------------------
+
+Base.endswith(symb::Symbol, suffix::AbstractString) = endswith(String(symb), suffix)
+Base.endswith(integrator::UnitaryPadeIntegrator, suffix::String) = endswith(integrator.unitary_symb, suffix)
+Base.endswith(integrator::DerivativeIntegrator, suffix::String) = endswith(integrator.variable, suffix)
+
+function Base.endswith(integrator::AbstractIntegrator, traj::NamedTrajectory, suffix::String)
+    if integrator isa UnitaryExponentialIntegrator
+        name = get_component_names(traj, integrator.state_components)
+    elseif integrator isa QuantumStateExponentialIntegrator
+        name = get_component_names(traj, integrator.state_components)
+    elseif integrator isa UnitaryPadeIntegrator
+        name = get_component_names(traj, integrator.state_components)
+    elseif integrator isa QuantumStatePadeIntegrator
+        name = get_component_names(traj, integrator.state_components)
+    elseif integrator isa DerivativeIntegrator
+        name = get_component_names(traj, integrator.variable_components)
+    else
+        error("Integrator type not recognized")
+    end
+    return endswith(name, suffix)
+end
+
+function PiccoloQuantumObjects.get_suffix(
+    integrators::AbstractVector{<:AbstractIntegrator},
+    sys::AbstractQuantumSystem,
+    traj::NamedTrajectory,
+    mod_traj::NamedTrajectory,
+    suffix::String
+)
+    found = AbstractIntegrator[]
+    for integrator ∈ integrators
+        if endswith(integrator, traj, suffix)
+            push!(found, remove_suffix(integrator, sys, traj, mod_traj, suffix))
+        end
+    end
+    return found
+end
+
+
 end

src/losses/unitary_infidelity_loss.jl

@@ -232,16 +232,6 @@ end
 ### UnitaryFreePhaseInfidelityLoss
 ###
 
-function free_phase(
-    ϕs::AbstractVector,
-    Hs::AbstractVector{<:AbstractMatrix}
-)
-    # NOTE: switch to expv for ForwardDiff
-    # return reduce(kron, [exp(im * ϕ * H) for (ϕ, H) ∈ zip(ϕs, Hs)])
-    Id = Matrix{eltype(Hs[1])}(I, size(Hs[1]))
-    return reduce(kron, [expv(im * ϕ, H, Id) for (ϕ, H) ∈ zip(ϕs, Hs)])
-end
-
 function free_phase_gradient(
     ϕs::AbstractVector,
     Hs::AbstractVector{<:AbstractMatrix}
@@ -531,7 +521,7 @@ end
     phase_operators = [PAULIS[:Z], PAULIS[:Z]]
     subspace = get_subspace_indices([1:2, 1:2], [n_levels, n_levels])
 
-    R = Losses.free_phase(phase_data, phase_operators)
+    R = free_phase(phase_data, phase_operators)
     @test R'R ≈ [1 0 0 0; 0 1 0 0; 0 0 1 0; 0 0 0 1]
     @test size(R) == (2^2, 2^2)