add leakage constraint to piccolo options

src/piccolo_options.jl

@@ -19,7 +19,7 @@ Options for the Piccolo quantum optimal control library.
 - `complex_control_norm_constraint_radius::Float64 = 1.0`: Radius of the complex control norm constraint.
 - `bound_state::Bool = false`: Bound the state.
 - `leakage_suppression::Bool = false`: Suppress leakage.
-- `R_leakage::Float64 = 1.0`: Leakage suppression parameter.
+- `leakage_cost::Float64 = 1.0`: Leakage suppression parameter.
 """
 @kwdef mutable struct PiccoloOptions
     verbose::Bool = true
@@ -30,9 +30,9 @@ Options for the Piccolo quantum optimal control library.
     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
+    leakage_constraint::Bool = false
+    leakage_constraint_value::Float64 = 1e-2
+    leakage_cost::Float64 = 1e-2
 end
 
 end

src/problem_templates/_problem_templates.jl

@@ -17,9 +17,6 @@ using LinearAlgebra
 using SparseArrays
 using TestItems
 
-# using Statistics
-mean(x) = sum(x) / length(x)
-
 include("unitary_smooth_pulse_problem.jl")
 include("unitary_variational_problem.jl")
 include("unitary_minimum_time_problem.jl")
@@ -30,47 +27,45 @@ include("quantum_state_smooth_pulse_problem.jl")
 include("quantum_state_minimum_time_problem.jl")
 include("quantum_state_sampling_problem.jl")
 
-# TODO: refactor
+
 function apply_piccolo_options!(
-    J::Objective,
-    constraints::AbstractVector{<:AbstractConstraint},
     piccolo_options::PiccoloOptions,
-    traj::NamedTrajectory,
-    state_names::AbstractVector{Symbol},
-    timestep_name::Symbol;
-    state_leakage_indices::Union{Nothing, AbstractVector{<:AbstractVector{Int}}}=nothing,
-    free_time::Bool=true,
+    constraints::AbstractVector{<:AbstractConstraint},
+    traj::NamedTrajectory;
+    state_names::Union{Nothing, Symbol, AbstractVector{Symbol}}=nothing,
+    state_leakage_indices::Union{Nothing, AbstractVector{Int}, AbstractVector{<:AbstractVector{Int}}}=nothing,
 )
-    if piccolo_options.leakage_suppression
-        throw(error("L1 is not implemented."))
-        # if piccolo_options.verbose
-        #     println("\tapplying leakage suppression: $(state_names)")
-        # end
+    J = NullObjective(traj)
+
+    if piccolo_options.leakage_constraint
+        val = piccolo_options.leakage_constraint_value
+        if piccolo_options.verbose
+            println("\tapplying leakage suppression: $(state_names) < $(val)")
+        end
+
+        if isnothing(state_leakage_indices)
+            throw(ValueError("Leakage indices are required for leakage suppression."))
+        end
+
+        if state_names isa Symbol
+            state_names = [state_names]
+            state_leakage_indices = [state_leakage_indices]
+        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
+        for (name, indices) ∈ zip(state_names, state_leakage_indices)
+            J += LeakageObjective(indices, name, traj, Qs=fill(piccolo_options.leakage_cost, traj.T))
+            push!(constraints, LeakageConstraint(val, indices, name, traj))
+        end
     end
 
-    if free_time
-        if piccolo_options.timesteps_all_equal
-            if piccolo_options.verbose
-                println("\tapplying timesteps_all_equal constraint: $(traj.timestep)")
-            end
-            push!(
-                constraints,
-                TimeStepsAllEqualConstraint(traj)
-            )
+    if piccolo_options.timesteps_all_equal
+        if piccolo_options.verbose
+            println("\tapplying timesteps_all_equal constraint: $(traj.timestep)")
         end
+        push!(
+            constraints,
+            TimeStepsAllEqualConstraint(traj)
+        )
     end
 
     if !isnothing(piccolo_options.complex_control_norm_constraint_name)
@@ -86,35 +81,7 @@ function apply_piccolo_options!(
         push!(constraints, norm_con)
     end
 
-    return
+    return J
 end
 
-function apply_piccolo_options!(
-    J::Objective,
-    constraints::AbstractVector{<:AbstractConstraint},
-    piccolo_options::PiccoloOptions,
-    traj::NamedTrajectory,
-    state_name::Symbol,
-    timestep_name::Symbol;
-    state_leakage_indices::Union{Nothing, AbstractVector{Int}}=nothing,
-    kwargs...
-)
-    state_names = [
-        name for name ∈ traj.names
-            if startswith(string(name), string(state_name))
-    ]
-
-    return apply_piccolo_options!(
-        J,
-        constraints,
-        piccolo_options,
-        traj,
-        state_names,
-        timestep_name;
-        state_leakage_indices=isnothing(state_leakage_indices) ? nothing : fill(state_leakage_indices, length(state_names)),
-        kwargs...
-    )
-end
-
-
 end

src/problem_templates/quantum_state_sampling_problem.jl

@@ -10,7 +10,7 @@ function QuantumStateSamplingProblem(
     ψ_inits::AbstractVector{<:AbstractVector{<:AbstractVector{<:ComplexF64}}},
     ψ_goals::AbstractVector{<:AbstractVector{<:AbstractVector{<:ComplexF64}}},
     T::Int,
-    Δt::Union{Float64,Vector{Float64}};
+    Δt::Union{Float64, AbstractVector{Float64}};
     ket_integrator=KetIntegrator,
     system_weights=fill(1.0, length(systems)),
     init_trajectory::Union{NamedTrajectory,Nothing}=nothing,
@@ -24,13 +24,14 @@ function QuantumStateSamplingProblem(
     da_bounds::Vector{Float64}=fill(da_bound, systems[1].n_drives),
     dda_bound::Float64=1.0,
     dda_bounds::Vector{Float64}=fill(dda_bound, systems[1].n_drives),
-    Δt_min::Float64=0.5 * Δt,
-    Δt_max::Float64=1.5 * Δt,
+    Δt_min::Float64=0.5 * minimum(Δt),
+    Δt_max::Float64=2.0 * maximum(Δt),
     Q::Float64=100.0,
     R=1e-2,
     R_a::Union{Float64,Vector{Float64}}=R,
     R_da::Union{Float64,Vector{Float64}}=R,
     R_dda::Union{Float64,Vector{Float64}}=R,
+    state_leakage_indices::Union{Nothing, AbstractVector{Int}}=nothing,
     constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
     piccolo_options::PiccoloOptions=PiccoloOptions(),
 )
@@ -93,9 +94,10 @@ function QuantumStateSamplingProblem(
     end
 
     # Optional Piccolo constraints and objectives
-    apply_piccolo_options!(
-        J, constraints, piccolo_options, traj, state_name, timestep_name;
-        state_leakage_indices=piccolo_options.state_leakage_indices
+    J += apply_piccolo_options!(
+        piccolo_options, constraints, traj;
+        state_names=vcat(state_names...),
+        state_leakage_indices=isnothing(state_leakage_indices) ? nothing : fill(state_leakage_indices, length(vcat(state_names...))),
     )
 
     # Integrators

src/problem_templates/quantum_state_smooth_pulse_problem.jl

@@ -69,13 +69,14 @@ function QuantumStateSmoothPulseProblem(
     da_bounds::Vector{Float64}=fill(da_bound, sys.n_drives),
     dda_bound::Float64=1.0,
     dda_bounds::Vector{Float64}=fill(dda_bound, sys.n_drives),
-    Δt_min::Float64=0.001 * Δt,
-    Δt_max::Float64=2.0 * Δt,
+    Δt_min::Float64=0.5 * minimum(Δt),
+    Δt_max::Float64=2.0 * maximum(Δt),
     Q::Float64=100.0,
     R=1e-2,
     R_a::Union{Float64, Vector{Float64}}=R,
     R_da::Union{Float64, Vector{Float64}}=R,
     R_dda::Union{Float64, Vector{Float64}}=R,
+    state_leakage_indices::Union{Nothing, AbstractVector{Int}}=nothing,
     constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
     piccolo_options::PiccoloOptions=PiccoloOptions(),
 )
@@ -131,9 +132,10 @@ function QuantumStateSmoothPulseProblem(
     end
 
     # Optional Piccolo constraints and objectives
-    apply_piccolo_options!(
-        J, constraints, piccolo_options, traj, state_name, timestep_name;
-        state_leakage_indices=piccolo_options.state_leakage_indices
+    J += apply_piccolo_options!(
+        piccolo_options, constraints, traj;
+        state_names=state_names,
+        state_leakage_indices= isnothing(state_leakage_indices) ? nothing : fill(state_leakage_indices, length(state_names))
     )
 
     state_names = [

src/problem_templates/unitary_free_phase_problem.jl

@@ -13,7 +13,7 @@ function UnitaryFreePhaseProblem(
     system::AbstractQuantumSystem,
     goal::Function,
     T::Int,
-    Δt::Union{Float64, <:AbstractVector{Float64}};
+    Δt::Union{Float64, AbstractVector{Float64}};
     unitary_integrator=UnitaryIntegrator,
     state_name::Symbol = :Ũ⃗,
     control_name::Symbol = :a,
@@ -28,8 +28,8 @@ function UnitaryFreePhaseProblem(
     da_bounds::Vector{Float64}=fill(da_bound, system.n_drives),
     dda_bound::Float64=1.0,
     dda_bounds::Vector{Float64}=fill(dda_bound, system.n_drives),
-    Δt_min::Float64=Δt isa Float64 ? 0.5 * Δt : 0.5 * mean(Δt),
-    Δt_max::Float64=Δt isa Float64 ? 1.5 * Δt : 1.5 * mean(Δt),
+    Δt_min::Float64=0.5 * minimum(Δt),
+    Δt_max::Float64=2.0 * maximum(Δt),
     Q::Float64=100.0,
     R=1e-2,
     R_a::Union{Float64, Vector{Float64}}=R,
@@ -103,9 +103,12 @@ function UnitaryFreePhaseProblem(
     J += QuadraticRegularizer(control_names[3], traj, R_dda)
 
     # Optional Piccolo constraints and objectives
-    ProblemTemplates.apply_piccolo_options!(
-        J, constraints, piccolo_options, traj, state_name, timestep_name;
-        state_leakage_indices=eval_goal isa EmbeddedOperator ? get_leakage_indices(eval_goal) : nothing
+    J += apply_piccolo_options!(
+        piccolo_options, constraints, traj;
+        state_names=state_name,
+        state_leakage_indices=eval_goal isa EmbeddedOperator ? 
+            get_iso_vec_leakage_indices(eval_goal) :
+            nothing
     )
 
     # Phase constraint

src/problem_templates/unitary_sampling_problem.jl

@@ -61,8 +61,8 @@ function UnitarySamplingProblem(
     da_bounds::Vector{Float64}=fill(da_bound, systems[1].n_drives),
     dda_bound::Float64=1.0,
     dda_bounds::Vector{Float64}=fill(dda_bound, systems[1].n_drives),
-    Δt_min::Float64=0.5 * Δt,
-    Δt_max::Float64=1.5 * Δt,
+    Δt_min::Float64=0.5 * minimum(Δt),
+    Δt_max::Float64=2.0 * maximum(Δt),
     Q::Float64=100.0,
     R=1e-2,
     R_a::Union{Float64,Vector{Float64}}=R,
@@ -126,9 +126,12 @@ function UnitarySamplingProblem(
     end
 
     # Optional Piccolo constraints and objectives
-    apply_piccolo_options!(
-        J, constraints, piccolo_options, traj, state_names, timestep_name;
-        state_leakage_indices=all(op -> op isa EmbeddedOperator, operators) ? get_leakage_indices.(operators) : nothing
+    J += apply_piccolo_options!(
+        piccolo_options, constraints, traj;
+        state_names=state_names,
+        state_leakage_indices=all(op -> op isa EmbeddedOperator, operators) ?       
+            get_iso_vec_leakage_indices.(operators) : 
+            nothing
     )
 
     # Integrators

src/problem_templates/unitary_smooth_pulse_problem.jl

@@ -82,8 +82,8 @@ function UnitarySmoothPulseProblem(
     da_bounds::Vector{Float64}=fill(da_bound, system.n_drives),
     dda_bound::Float64=1.0,
     dda_bounds::Vector{Float64}=fill(dda_bound, system.n_drives),
-    Δt_min::Float64=Δt isa Float64 ? 0.5 * Δt : 0.5 * mean(Δt),
-    Δt_max::Float64=Δt isa Float64 ? 1.5 * Δt : 1.5 * mean(Δt),
+    Δt_min::Float64=0.5 * minimum(Δt),
+    Δt_max::Float64=2.0 * maximum(Δt),
     Q::Float64=100.0,
     R=1e-2,
     R_a::Union{Float64, Vector{Float64}}=R,
@@ -134,9 +134,12 @@ function UnitarySmoothPulseProblem(
     J += QuadraticRegularizer(control_names[3], traj, R_dda)
 
     # Optional Piccolo constraints and objectives
-    apply_piccolo_options!(
-        J, constraints, piccolo_options, traj, state_name, timestep_name;
-        state_leakage_indices=goal isa EmbeddedOperator ? get_leakage_indices(goal) : nothing
+    J += apply_piccolo_options!(
+        piccolo_options, constraints, traj; 
+        state_names=state_name,
+        state_leakage_indices=goal isa EmbeddedOperator ? 
+            get_iso_vec_leakage_indices(goal) : 
+            nothing
     )
 
     integrators = [
@@ -165,7 +168,7 @@ end
 
 # *************************************************************************** #
 
-@testitem "Hadamard gate" begin
+@testitem "Hadamard gate improvement" begin
     using PiccoloQuantumObjects 
 
     sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]])
@@ -181,11 +184,10 @@ end
 
     initial = unitary_rollout_fidelity(prob.trajectory, sys)
     solve!(prob, max_iter=100, verbose=false, print_level=1)
-    final = unitary_rollout_fidelity(prob.trajectory, sys)
-    @test final > initial
+    @test unitary_rollout_fidelity(prob.trajectory, sys) > initial
 end
 
-@testitem "Hadamard gate with bound states and control norm constraint" begin
+@testitem "Bound states and control norm constraint" begin
     using PiccoloQuantumObjects 
 
     sys = QuantumSystem(GATES[:Z], [GATES[:X], GATES[:Y]])
@@ -202,14 +204,12 @@ end
         )
     )
 
-    initial = unitary_rollout_fidelity(prob.trajectory, sys)
-    solve!(prob, max_iter=100, verbose=false, print_level=1)
-    final = unitary_rollout_fidelity(prob.trajectory, sys)
-    @test_broken false
-    # @test final > initial
+    initial = copy(prob.trajectory.data)
+    solve!(prob, max_iter=5, verbose=false, print_level=1)
+    @test prob.trajectory.data != initial
 end
 
-@testitem "EmbeddedOperator Hadamard gate" begin
+@testitem "EmbeddedOperator tests" begin
     using PiccoloQuantumObjects 
 
     a = annihilate(3)
@@ -218,15 +218,32 @@ end
     T = 51
     Δt = 0.2
 
-    prob = UnitarySmoothPulseProblem(
-        sys, U_goal, T, Δt,
-        piccolo_options=PiccoloOptions(verbose=false)
-    )
+    @testset "EmbeddedOperator: solve gate" begin
+        prob = UnitarySmoothPulseProblem(
+            sys, U_goal, T, Δt,
+            piccolo_options=PiccoloOptions(verbose=false)
+        )
 
-    initial = unitary_rollout_fidelity(prob.trajectory, sys, subspace=U_goal.subspace)
-    solve!(prob, max_iter=100, verbose=false, print_level=1)
-    final = unitary_rollout_fidelity(prob.trajectory, sys, subspace=U_goal.subspace)
-    @test final > initial
+        initial = copy(prob.trajectory.data)
+        solve!(prob, max_iter=5, verbose=false, print_level=1)
+        @test prob.trajectory.data != initial
+    end
+
+    @testset "EmbeddedOperator: leakage constraint" begin
+        prob = UnitarySmoothPulseProblem(
+            sys, U_goal, T, Δt;
+            da_bound=1.0,
+            piccolo_options=PiccoloOptions(
+                leakage_constraint=true, 
+                leakage_constraint_value=5e-2, 
+                leakage_cost=1e-1,
+                verbose=false
+            )
+        )
+        initial = copy(prob.trajectory.data)
+        solve!(prob, max_iter=5, verbose=false, print_level=1)
+        @test prob.trajectory.data != initial
+    end
 end
 
 # TODO: Test changing names of control, state, and timestep