free phase initialization

src/constraints/fidelity_constraint.jl

@@ -286,8 +286,8 @@ function FinalUnitaryFreePhaseFidelityConstraint(;
 end
 
 function FinalUnitaryFreePhaseFidelityConstraint(
-    state_symbol::Symbol,
-    global_symbol::Symbol,
+    state_name::Symbol,
+    phase_name::Symbol,
     phase_operators::AbstractVector{<:AbstractMatrix{<:Complex}},
     val::Float64,
     traj::NamedTrajectory;
@@ -296,9 +296,9 @@ function FinalUnitaryFreePhaseFidelityConstraint(
 )
     return FinalUnitaryFreePhaseFidelityConstraint(;
         value=val,
-        state_slice=slice(traj.T, traj.components[state_symbol], traj.dim),
-        phase_slice=trajectory.global_components[global_symbol],
-        goal=traj.goal[state_symbol],
+        state_slice=slice(traj.T, traj.components[state_name], traj.dim),
+        phase_slice=traj.global_components[phase_name],
+        goal=traj.goal[state_name],
         phase_operators=phase_operators,
         zdim=length(traj),
         subspace=subspace,

src/objectives/unitary_infidelity_objective.jl

@@ -119,8 +119,8 @@ Fields:
 """
 function UnitaryFreePhaseInfidelityObjective(;
     name::Union{Nothing,Symbol}=nothing,
-    phase_name::Union{Nothing,Symbol}=nothing,
     goal::Union{Nothing,AbstractVector{<:R}}=nothing,
+    phase_name::Union{Nothing,Symbol}=nothing,
     phase_operators::Union{Nothing,AbstractVector{<:AbstractMatrix{<:Complex{R}}}}=nothing,
 	Q::R=1.0,
 	eval_hessian::Bool=false,
@@ -168,3 +168,23 @@ function UnitaryFreePhaseInfidelityObjective(;
 
 	return Objective(L, ∇L, ∂²L, ∂²L_structure, Dict[params])
 end
+
+function UnitaryFreePhaseInfidelityObjective(
+    name::Symbol,
+    phase_name::Symbol,
+    phase_operators::AbstractVector{<:AbstractMatrix{<:Complex}},
+    traj::NamedTrajectory,
+    Q::Float64;
+    subspace=nothing,
+	eval_hessian::Bool=true
+)
+    return UnitaryFreePhaseInfidelityObjective(
+        name=name,
+        goal=traj.goal[name],
+        phase_name=phase_name,
+        phase_operators=phase_operators,
+        Q=Q,
+        subspace=subspace,
+        eval_hessian=eval_hessian,
+    )
+end

src/problem_templates/unitary_minimum_time_problem.jl

@@ -52,11 +52,12 @@ function UnitaryMinimumTimeProblem(
     integrators::Vector{<:AbstractIntegrator},
     constraints::Vector{<:AbstractConstraint};
     unitary_name::Symbol=:Ũ⃗,
-    global_name::Union{Nothing, Symbol}=nothing,
     final_fidelity::Union{Real, Nothing}=nothing,
     D=1.0,
     ipopt_options::IpoptOptions=IpoptOptions(),
     piccolo_options::PiccoloOptions=PiccoloOptions(),
+    phase_name::Symbol=:ϕ,
+    phase_operators::Union{AbstractVector{<:AbstractMatrix}, Nothing}=nothing,
     subspace=nothing,
     kwargs...
 )
@@ -70,7 +71,7 @@ function UnitaryMinimumTimeProblem(
 
     objective += MinimumTimeObjective(trajectory; D=D, eval_hessian=piccolo_options.eval_hessian)
 
-    if isnothing(global_name)
+    if isnothing(phase_operators)
         fidelity_constraint = FinalUnitaryFidelityConstraint(
             unitary_name,
             final_fidelity,
@@ -79,12 +80,9 @@ function UnitaryMinimumTimeProblem(
             eval_hessian=piccolo_options.eval_hessian
         )
     else
-        phase_operators= [
-            GATES[:Z] for _ in eachindex(trajectory.global_components[global_name])
-        ]
         fidelity_constraint = FinalUnitaryFreePhaseFidelityConstraint(
             unitary_name,
-            global_name,
+            phase_name,
             phase_operators,
             final_fidelity,
             trajectory;

src/problem_templates/unitary_robustness_problem.jl

@@ -17,7 +17,7 @@ export UnitaryRobustnessProblem
 Create a quantum control problem for robustness optimization of a unitary trajectory.
 
 # Keyword Arguments
-- `unitary_symbol::Symbol=:Ũ⃗`: The symbol for the unitary trajectory in `trajectory`.
+- `unitary_name::Symbol=:Ũ⃗`: The symbol for the unitary trajectory in `trajectory`.
 - `final_fidelity::Union{Real, Nothing}=nothing`: The target fidelity for the final unitary.
 - `ipopt_options::IpoptOptions=IpoptOptions()`: Options for the Ipopt solver.
 - `piccolo_options::PiccoloOptions=PiccoloOptions()`: Options for the Piccolo solver.
@@ -33,17 +33,19 @@ function UnitaryRobustnessProblem(
     objective::Objective,
     integrators::Vector{<:AbstractIntegrator},
     constraints::Vector{<:AbstractConstraint};
-    unitary_symbol::Symbol=:Ũ⃗,
+    unitary_name::Symbol=:Ũ⃗,
     final_fidelity::Union{Real, Nothing}=nothing,
+    phase_name::Symbol=:ϕ,
+    phase_operators::Union{AbstractVector{<:AbstractMatrix}, Nothing}=nothing,
     ipopt_options::IpoptOptions=IpoptOptions(),
     piccolo_options::PiccoloOptions=PiccoloOptions(),
     kwargs...
 )
-    @assert unitary_symbol ∈ trajectory.names
+    @assert unitary_name ∈ trajectory.names
 
     if isnothing(final_fidelity)
         final_fidelity = unitary_fidelity(
-            trajectory[unitary_symbol][:, end], trajectory.goal[unitary_symbol]
+            trajectory[unitary_name][:, end], trajectory.goal[unitary_name]
         )
     end
 
@@ -52,12 +54,28 @@ function UnitaryRobustnessProblem(
         eval_hessian=piccolo_options.eval_hessian,
     )
 
-    fidelity_constraint = FinalUnitaryFidelityConstraint(
-        unitary_symbol,
-        final_fidelity,
-        trajectory;
-        subspace=H_error isa EmbeddedOperator ? H_error.subspace_indices : nothing
-    )
+    subspace = H_error isa EmbeddedOperator ? H_error.subspace_indices : nothing
+
+    if isnothing(phase_operators)
+        fidelity_constraint = FinalUnitaryFidelityConstraint(
+            unitary_name,
+            final_fidelity,
+            trajectory;
+            subspace=subspace,
+            eval_hessian=piccolo_options.eval_hessian
+        )
+    else
+        fidelity_constraint = FinalUnitaryFreePhaseFidelityConstraint(
+            unitary_name,
+            phase_name,
+            phase_operators,
+            final_fidelity,
+            trajectory;
+            subspace=subspace,
+            eval_hessian=piccolo_options.eval_hessian
+        )
+    end
+
     push!(constraints, fidelity_constraint)
 
     return QuantumControlProblem(

src/problem_templates/unitary_smooth_pulse_problem.jl

@@ -62,7 +62,8 @@ with
 - `R_a::Union{Float64, Vector{Float64}}=R`: the weight on the regularization term for the control pulses
 - `R_da::Union{Float64, Vector{Float64}}=R`: the weight on the regularization term for the control pulse derivatives
 - `R_dda::Union{Float64, Vector{Float64}}=R`: the weight on the regularization term for the control pulse second derivatives
-- `global_data::Union{NamedTuple, Nothing}=nothing`: global data for the problem
+- `phase_name::Symbol=:ϕ`: the name of the phase
+- `phase_operators::Union{AbstractVector{<:AbstractMatrix}, Nothing}=nothing`: the phase operators for free phase corrections
 - `constraints::Vector{<:AbstractConstraint}=AbstractConstraint[]`: the constraints to enforce
 
 """
@@ -92,7 +93,8 @@ function UnitarySmoothPulseProblem(
     R_a::Union{Float64, Vector{Float64}}=R,
     R_da::Union{Float64, Vector{Float64}}=R,
     R_dda::Union{Float64, Vector{Float64}}=R,
-    global_data::Union{NamedTuple, Nothing}=nothing,
+    phase_name::Symbol=:ϕ,
+    phase_operators::Union{AbstractVector{<:AbstractMatrix}, Nothing}=nothing,
     constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
     kwargs...
 )
@@ -120,31 +122,26 @@ function UnitarySmoothPulseProblem(
             a_guess=a_guess,
             system=system,
             rollout_integrator=piccolo_options.rollout_integrator,
-            global_data=global_data
+            phase_name=phase_name,
+            phase_operators=phase_operators
         )
     end
 
     # Subspace
     subspace = operator isa EmbeddedOperator ? operator.subspace_indices : nothing
-    goal = operator_to_iso_vec(operator isa EmbeddedOperator ? operator.operator : operator)
 
     # Objective
-    if isnothing(global_data)
+    if isnothing(phase_operators)
         J = UnitaryInfidelityObjective(
             state_name, traj, Q; 
             subspace=subspace,
             eval_hessian=piccolo_options.eval_hessian,
         )
     else
-        # TODO: remove hardcoded args
         J = UnitaryFreePhaseInfidelityObjective(
-            name=state_name,
-            phase_name=:ϕ,
-            goal=goal,
-            phase_operators=fill(GATES[:Z], length(traj.global_components[:ϕ])),
-            Q=Q,
+            state_name, phase_name, phase_operators, traj, Q;
+            subspace=subspace,
             eval_hessian=piccolo_options.eval_hessian,
-            subspace=subspace
         )
     end
 

src/trajectory_initialization.jl

@@ -276,7 +276,8 @@ function initialize_trajectory(
     Δt_bounds::ScalarBound=(0.5 * Δt, 1.5 * Δt),
     drive_derivative_σ::Float64=0.1,
     a_guess::Union{AbstractMatrix{<:Float64}, Nothing}=nothing,
-    global_data::Union{NamedTuple, Nothing}=nothing,
+    phase_name::Symbol=:ϕ,
+    phase_data::Union{AbstractVector{<:Real}, Nothing}=nothing,
     verbose=false,
 )
     @assert length(state_data) == length(state_names) == length(state_inits) == length(state_goals) "state_data, state_names, state_inits, and state_goals must have the same length"
@@ -361,6 +362,9 @@ function initialize_trajectory(
         controls = (control_names[end],)
     end
 
+    # Construct global data for free phases
+    global_data = isnothing(phase_data) ? (;) : (; phase_name => phase_data)
+
     return NamedTrajectory(
         (; (names .=> values)...);
         controls=controls,
@@ -369,7 +373,7 @@ function initialize_trajectory(
         initial=initial,
         final=final,
         goal=goal,
-        global_data= isnothing(global_data) ? (;) : global_data
+        global_data=global_data
     )
 end
 
@@ -422,6 +426,7 @@ function initialize_trajectory(
     system::Union{AbstractQuantumSystem, AbstractVector{<:AbstractQuantumSystem}, Nothing}=nothing,
     rollout_integrator::Function=expv,
     geodesic=true,
+    phase_operators::Union{AbstractVector{<:AbstractMatrix}, Nothing}=nothing,
     verbose=false,
     kwargs...
 )
@@ -475,6 +480,9 @@ function initialize_trajectory(
         state_goals = [Ũ⃗_goal]
     end
 
+    # Construct phase data
+    phase_data = isnothing(phase_operators) ? nothing : π * randn(length(phase_operators))
+
     return initialize_trajectory(
         state_data,
         state_inits,
@@ -483,6 +491,7 @@ function initialize_trajectory(
         T,
         Δt,
         args...;
+        phase_data=phase_data,
         a_guess=a_guess,
         verbose=verbose,
         kwargs...