Merge pull request #177 from kestrelquantum/dev/aaron

leakage suppression refactor

Project.toml

@@ -1,7 +1,7 @@
 name = "QuantumCollocation"
 uuid = "0dc23a59-5ffb-49af-b6bd-932a8ae77adf"
 authors = ["Aaron Trowbridge <[email protected]> and contributors"]
-version = "0.5.0"
+version = "0.5.1"
 
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
@@ -30,7 +30,7 @@ TrajectoryIndexingUtils = "6dad8b7f-dd9a-4c28-9b70-85b9a079bfc8"
 
 [compat]
 BenchmarkTools = "1.6"
-CairoMakie = "0.12"
+CairoMakie = "0.13"
 Distributions = "0.25"
 Einsum = "0.4"
 ExponentialAction = "0.2"
@@ -39,7 +39,7 @@ Interpolations = "0.15"
 JLD2 = "0.5"
 LinearAlgebra = "1.10, 1.11"
 MathOptInterface = "1.35"
-NamedTrajectories = "0.2"
+NamedTrajectories = "0.2.6"
 PiccoloQuantumObjects = "0.1"
 PkgTemplates = "0.7"
 ProgressMeter = "1.10"

src/problem_templates/_problem_templates.jl

@@ -34,27 +34,23 @@ function apply_piccolo_options!(
     constraints::AbstractVector{<:AbstractConstraint},
     piccolo_options::PiccoloOptions,
     traj::NamedTrajectory,
-    operators::Union{Nothing, AbstractVector{<:AbstractPiccoloOperator}},
     state_names::AbstractVector{Symbol},
-    timestep_name::Symbol
+    timestep_name::Symbol;
+    state_leakage_indices::Union{Nothing, AbstractVector{<:AbstractVector{Int}}}=nothing,
 )
-    # TODO: This should be changed to leakage indices (more general, for states)
-    # (also, it should be looped over relevant state names, not broadcast)
-    if piccolo_options.leakage_suppression && !isnothing(operators)
-        for (operator, state_name) in zip(operators, state_names)
-            if operator isa EmbeddedOperator
-                leakage_indices = get_iso_vec_leakage_indices(operator)
-                J += L1Regularizer!(
-                    constraints,
-                    state_name,
-                    traj;
-                    R_value=piccolo_options.R_leakage,
-                    indices=leakage_indices,
-                    eval_hessian=piccolo_options.eval_hessian
-                )
-            else
-                @warn "leakage_suppression is only supported for embedded operators, ignoring."
-            end
+    if piccolo_options.leakage_suppression
+        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
     end
 
@@ -84,29 +80,23 @@ function apply_piccolo_options!(
     constraints::AbstractVector{<:AbstractConstraint},
     piccolo_options::PiccoloOptions,
     traj::NamedTrajectory,
-    operator::Union{Nothing, AbstractPiccoloOperator},
     state_name::Symbol,
-    timestep_name::Symbol
+    timestep_name::Symbol;
+    state_leakage_indices::Union{Nothing, AbstractVector{Int}}=nothing,
 )
     state_names = [
         name for name ∈ traj.names
             if startswith(string(name), string(state_name))
     ]
 
-    if !isnothing(operator)
-        operators = fill(operator, length(state_names))
-    else
-        operators = nothing
-    end
-
     return apply_piccolo_options!(
         J,
         constraints,
         piccolo_options,
         traj,
-        operators,
         state_names,
-        timestep_name
+        timestep_name;
+        state_leakage_indices=isnothing(state_leakage_indices) ? nothing : fill(state_leakage_indices, length(state_names)),
     )
 end
 

src/problem_templates/quantum_state_sampling_problem.jl

@@ -29,7 +29,7 @@ function QuantumStateSamplingProblem(
     R_a::Union{Float64,Vector{Float64}}=R,
     R_da::Union{Float64,Vector{Float64}}=R,
     R_dda::Union{Float64,Vector{Float64}}=R,
-    leakage_operator::Union{Nothing, EmbeddedOperator}=nothing,
+    leakage_indices::Union{Nothing, <:AbstractVector{Int}}=nothing,
     constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
     kwargs...
 )
@@ -117,7 +117,8 @@ function QuantumStateSamplingProblem(
 
     # Optional Piccolo constraints and objectives
     apply_piccolo_options!(
-        J, constraints, piccolo_options, traj, leakage_operator, state_name, timestep_name
+        J, constraints, piccolo_options, traj, state_name, timestep_name;
+        state_leakage_indices=leakage_indices
     )
 
     return QuantumControlProblem(

src/problem_templates/quantum_state_smooth_pulse_problem.jl

@@ -59,7 +59,7 @@ function QuantumStateSmoothPulseProblem(
     ψ_inits::Vector{<:AbstractVector{<:ComplexF64}},
     ψ_goals::Vector{<:AbstractVector{<:ComplexF64}},
     T::Int,
-    Δt::Float64;
+    Δt::Union{Float64, <:AbstractVector{Float64}};
     ipopt_options::IpoptOptions=IpoptOptions(),
     piccolo_options::PiccoloOptions=PiccoloOptions(),
     state_name::Symbol=:ψ̃,
@@ -82,7 +82,7 @@ function QuantumStateSmoothPulseProblem(
     R_a::Union{Float64, Vector{Float64}}=R,
     R_da::Union{Float64, Vector{Float64}}=R,
     R_dda::Union{Float64, Vector{Float64}}=R,
-    leakage_operator::Union{Nothing, EmbeddedOperator}=nothing,
+    leakage_indices::Union{Nothing, <:AbstractVector{Int}}=nothing,
     constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
     kwargs...
 )
@@ -191,7 +191,8 @@ function QuantumStateSmoothPulseProblem(
 
     # Optional Piccolo constraints and objectives
     apply_piccolo_options!(
-        J, constraints, piccolo_options, traj, leakage_operator, state_name, timestep_name
+        J, constraints, piccolo_options, traj, state_name, timestep_name;
+        state_leakage_indices=leakage_indices
     )
 
     return QuantumControlProblem(

src/problem_templates/unitary_bang_bang_problem.jl

@@ -170,7 +170,9 @@ function UnitaryBangBangProblem(
     ]
 
     # Optional Piccolo constraints and objectives
-    apply_piccolo_options!(J, constraints, piccolo_options, traj, operator, state_name, timestep_name)
+    apply_piccolo_options!(J, constraints, piccolo_options, traj, state_name, timestep_name;
+        state_leakage_indices=operator isa EmbeddedOperator ? get_leakage_indices(operator) : nothing
+    )
 
     return QuantumControlProblem(
         traj,

src/problem_templates/unitary_sampling_problem.jl

@@ -150,7 +150,8 @@ function UnitarySamplingProblem(
 
     # Optional Piccolo constraints and objectives
     apply_piccolo_options!(
-        J, constraints, piccolo_options, traj, operators, state_names, timestep_name
+        J, constraints, piccolo_options, traj, state_names, timestep_name;
+        state_leakage_indices=all(op -> op isa EmbeddedOperator, operators) ? get_leakage_indices.(operators) : nothing
     )
 
     return QuantumControlProblem(

src/problem_templates/unitary_smooth_pulse_problem.jl

@@ -71,7 +71,7 @@ function UnitarySmoothPulseProblem(
     system::AbstractQuantumSystem,
     operator::AbstractPiccoloOperator,
     T::Int,
-    Δt::Union{Float64, Vector{Float64}};
+    Δt::Union{Float64, <:AbstractVector{Float64}};
     ipopt_options::IpoptOptions=IpoptOptions(),
     piccolo_options::PiccoloOptions=PiccoloOptions(),
     state_name::Symbol = :Ũ⃗,
@@ -174,7 +174,8 @@ function UnitarySmoothPulseProblem(
 
     # Optional Piccolo constraints and objectives
     apply_piccolo_options!(
-        J, constraints, piccolo_options, traj, operator, state_name, timestep_name
+        J, constraints, piccolo_options, traj, state_name, timestep_name;
+        state_leakage_indices=operator isa EmbeddedOperator ? get_leakage_indices(operator) : nothing
     )
 
     return QuantumControlProblem(