refactor: add quantum constraints module and update problem templates…

…, piccolo options for DTO

src/QuantumCollocation.jl

@@ -17,6 +17,9 @@ include("trajectory_interpolations.jl")
 include("quantum_objectives.jl")
 @reexport using .QuantumObjectives
 
+include("quantum_constraints.jl")
+@reexport using .QuantumConstraints
+
 include("quantum_integrators.jl")
 @reexport using .QuantumIntegrators
 

src/piccolo_options.jl

@@ -14,7 +14,6 @@ Options for the Piccolo quantum optimal control library.
 - `timesteps_all_equal::Bool = true`: Use equal timesteps
 - `rollout_integrator::Function = expv`: Integrator to use for rollout
 - `geodesic = true`: Use the geodesic to initialize the optimization.
-- `build_trajectory_constraints::Bool = true`: Build trajectory constraints.
 - `zero_initial_and_final_derivative::Bool=false`: Zero the initial and final control pulse derivatives.
 - `complex_control_norm_constraint_name::Union{Nothing, Symbol} = nothing`: Name of the complex control norm constraint.
 - `complex_control_norm_constraint_radius::Float64 = 1.0`: Radius of the complex control norm constraint.
@@ -27,12 +26,12 @@ Options for the Piccolo quantum optimal control library.
     timesteps_all_equal::Bool = true
     rollout_integrator::Function = expv
     geodesic::Bool = true
-    build_trajectory_constraints::Bool = true
     zero_initial_and_final_derivative::Bool = true
     complex_control_norm_constraint_name::Union{Nothing, Symbol} = nothing
     complex_control_norm_constraint_radius::Float64 = 1.0
     bound_state::Bool = false
     leakage_suppression::Bool = false
+    state_leakage_indices::AbstractVector{Int} = Int[]
     R_leakage::Float64 = 1.0
 end
 

src/problem_templates/_problem_templates.jl

@@ -4,10 +4,10 @@ using ..DirectSums
 using ..Rollouts
 using ..TrajectoryInitialization
 using ..QuantumObjectives
+using ..QuantumConstraints
 using ..QuantumIntegrators
 using ..Options
 
-using Distributions
 using TrajectoryIndexingUtils
 using NamedTrajectories
 using DirectTrajOpt
@@ -20,10 +20,7 @@ using TestItems
 
 include("unitary_smooth_pulse_problem.jl")
 include("unitary_minimum_time_problem.jl")
-include("unitary_robustness_problem.jl")
-include("unitary_direct_sum_problem.jl")
 include("unitary_sampling_problem.jl")
-include("unitary_bang_bang_problem.jl")
 
 include("quantum_state_smooth_pulse_problem.jl")
 include("quantum_state_minimum_time_problem.jl")
@@ -43,23 +40,23 @@ function apply_piccolo_options!(
     free_time::Bool=true,
 )
     if piccolo_options.leakage_suppression
-        if piccolo_options.verbose
-            println("\tapplying leakage suppression: $(state_names)")
-        end
+        throw(error("L1 is not implemented."))
+        # if piccolo_options.verbose
+        #     println("\tapplying leakage suppression: $(state_names)")
+        # end
 
-        if isnothing(state_leakage_indices)
-            error("You must provide leakage indices for leakage suppression.")
-        end
-        for (state_name, leakage_indices) ∈ zip(state_names, state_leakage_indices)
-            J += L1Regularizer!(
-                constraints,
-                state_name,
-                traj;
-                R_value=piccolo_options.R_leakage,
-                indices=leakage_indices,
-                eval_hessian=piccolo_options.eval_hessian
-            )
-        end
+        # if isnothing(state_leakage_indices)
+        #     error("You must provide leakage indices for leakage suppression.")
+        # end
+        # for (state_name, leakage_indices) ∈ zip(state_names, state_leakage_indices)
+        #     J += L1Regularizer!(
+        #         constraints,
+        #         state_name,
+        #         traj;
+        #         R_value=piccolo_options.R_leakage,
+        #         indices=leakage_indices,
+        #     )
+        # end
     end
 
     if free_time

src/problem_templates/quantum_state_minimum_time_problem.jl

@@ -7,122 +7,99 @@ export QuantumStateMinimumTimeProblem
 
 Construct a `DirectTrajOptProblem` for the minimum time problem of reaching a target state.
 
-# Arguments
-- `traj::NamedTrajectory`: The initial trajectory.
-- `sys::QuantumSystem`: The quantum system.
-- `obj::Objective`: The objective function.
-- `integrators::Vector{<:AbstractIntegrator}`: The integrators.
-- `constraints::Vector{<:AbstractConstraint}`: The constraints.
-or
-- `prob::DirectTrajOptProblem`: The quantum control problem.
-
 # Keyword Arguments
 - `state_name::Symbol=:ψ̃`: The symbol for the state variables.
 - `final_fidelity::Union{Real, Nothing}=nothing`: The final fidelity.
 - `D=1.0`: The cost weight on the time.
-- `ipopt_options::IpoptOptions=IpoptOptions()`: The Ipopt options.
 - `piccolo_options::PiccoloOptions=PiccoloOptions()`: The Piccolo options.
-- `kwargs...`: Additional keyword arguments, passed to `DirectTrajOptProblem`.
-
 """
 function QuantumStateMinimumTimeProblem end
 
 function QuantumStateMinimumTimeProblem(
-    traj::NamedTrajectory,
-    obj::Objective,
-    integrators::Vector{<:AbstractIntegrator},
+    trajectory::NamedTrajectory,
+    ψ_goals::AbstractVector{<:AbstractVector{<:ComplexF64}},
+    objective::Objective,
+    dynamics::TrajectoryDynamics,
     constraints::Vector{<:AbstractConstraint};
     state_name::Symbol=:ψ̃,
-    control_name::Symbol=:a,
-    final_fidelity::Union{Real, Nothing}=nothing,
+    final_fidelity::Float64=1.0,
     D=1.0,
-    ipopt_options::IpoptOptions=IpoptOptions(),
     piccolo_options::PiccoloOptions=PiccoloOptions(),
-    kwargs...
 )
-    state_names = [name for name in traj.names if startswith(name, state_name)]
+    state_names = [name for name in trajectory.names if startswith(string(name), string(state_name))]
     @assert length(state_names) ≥ 1 "No matching states found in trajectory"
+    @assert length(state_names) == length(ψ_goals) "Number of states and goals must match"
 
-    obj += MinimumTimeObjective(traj; D=D, eval_hessian=piccolo_options.eval_hessian)
-
-    # Default to average state fidelity
-    if isnothing(final_fidelity)
-        vals = [iso_fidelity(traj[n][:, end], traj.goal[n]) for n ∈ state_names]
-        final_fidelity = sum(vals) / length(vals)
+    if piccolo_options.verbose
+        println("    constructing QuantumStateMinimumTimeProblem...")
+        println("\tfinal fidelity: $(final_fidelity)")
     end
 
-    for state_name in state_names
-        fidelity_constraint = FinalQuantumStateFidelityConstraint(
+    objective += MinimumTimeObjective(
+        trajectory, D=D, timesteps_all_equal=piccolo_options.timesteps_all_equal
+    )
+
+    for (state_name, ψ_goal) in zip(state_names, ψ_goals)
+        fidelity_constraint = FinalKetFidelityConstraint(
+            ψ_goal,
             state_name,
             final_fidelity,
-            traj,
-            eval_hessian=piccolo_options.eval_hessian
+            trajectory,
         )
 
         push!(constraints, fidelity_constraint)
     end
 
     return DirectTrajOptProblem(
-        traj,
-        obj,
-        integrators;
-        constraints=constraints,
-        ipopt_options=ipopt_options,
-        piccolo_options=piccolo_options,
-        control_name=control_name,
-        kwargs...
+        trajectory,
+        objective,
+        dynamics,
+        constraints
     )
 end
 
+# TODO: goals from trajectory?
+
 function QuantumStateMinimumTimeProblem(
-    prob::DirectTrajOptProblem;
-    obj::Objective=get_objective(prob),
-    constraints::AbstractVector{<:AbstractConstraint}=get_constraints(prob),
-    ipopt_options::IpoptOptions=deepcopy(prob.ipopt_options),
-    piccolo_options::PiccoloOptions=deepcopy(prob.piccolo_options),
-    build_trajectory_constraints=false,
+    prob::DirectTrajOptProblem,
+    ψ_goals::AbstractVector{<:AbstractVector{<:ComplexF64}};
+    objective::Objective=prob.objective,
+    constraints::AbstractVector{<:AbstractConstraint}=prob.constraints,
     kwargs...
 )
-    piccolo_options.build_trajectory_constraints = build_trajectory_constraints
-
     return QuantumStateMinimumTimeProblem(
-        copy(prob.trajectory),
-        obj,
-        prob.integrators,
+        prob.trajectory,
+        ψ_goals,
+        objective,
+        prob.dynamics,
         constraints;
-        ipopt_options=ipopt_options,
-        piccolo_options=piccolo_options,
         kwargs...
     )
 end
 
 # *************************************************************************** #
 
 @testitem "Test quantum state minimum time" begin
-        using NamedTrajectories
-        using PiccoloQuantumObjects
-
-        # System
-        T = 50
-        Δt = 0.2
-        sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]])
-        ψ_init = Vector{ComplexF64}[[1.0, 0.0]]
-        ψ_target = Vector{ComplexF64}[[0.0, 1.0]]
-
-        prob = QuantumStateSmoothPulseProblem(
-            sys, ψ_init, ψ_target, T, Δt;
-            ipopt_options=IpoptOptions(print_level=1),
-            piccolo_options=PiccoloOptions(
-                verbose=false,
-                integrator=:pade
-            )
-        )
-        initial = sum(get_timesteps(prob.trajectory))
-        mintime_prob = QuantumStateMinimumTimeProblem(prob)
-        solve!(mintime_prob, max_iter=20)
-        final = sum(get_timesteps(mintime_prob.trajectory))
-        @test final < initial
-
-        # Test with final fidelity
-        QuantumStateMinimumTimeProblem(prob, final_fidelity=0.99)
+    using NamedTrajectories
+
+    T = 51
+    Δt = 0.2
+    sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]])
+    ψ_init = Vector{ComplexF64}[[1.0, 0.0]]
+    ψ_target = Vector{ComplexF64}[[0.0, 1.0]]
+
+    prob = QuantumStateSmoothPulseProblem(
+        sys, ψ_init, ψ_target, T, Δt;
+        piccolo_options=PiccoloOptions(verbose=false)
+    )
+    initial = sum(get_timesteps(prob.trajectory))
+
+    min_prob = QuantumStateMinimumTimeProblem(
+        prob, ψ_target,
+        piccolo_options=PiccoloOptions(verbose=false)
+    )
+    solve!(min_prob, max_iter=50, verbose=false, print_level=1)
+    final = sum(get_timesteps(min_prob.trajectory))
+
+    @test final < initial
 end