free phase template

src/problem_templates/_problem_templates.jl

@@ -20,6 +20,7 @@ include("unitary_smooth_pulse_problem.jl")
 include("unitary_variational_problem.jl")
 include("unitary_minimum_time_problem.jl")
 include("unitary_sampling_problem.jl")
+include("unitary_free_phase_problem.jl")
 
 include("quantum_state_smooth_pulse_problem.jl")
 include("quantum_state_minimum_time_problem.jl")

src/problem_templates/unitary_free_phase_problem.jl

@@ -0,0 +1,100 @@
+export UnitaryFreePhaseProblem
+
+
+function UnitaryFreePhaseProblem(
+    system::AbstractQuantumSystem,
+    goal::Function,
+    T::Int,
+    Δt::Union{Float64, <:AbstractVector{Float64}},
+    initial_phases::AbstractVector{Float64};
+    unitary_integrator=UnitaryIntegrator,
+    state_name::Symbol = :Ũ⃗,
+    control_name::Symbol = :a,
+    timestep_name::Symbol = :Δt,
+    phase_name::Symbol = :θ,
+    init_trajectory::Union{NamedTrajectory, Nothing}=nothing,
+    a_guess::Union{Matrix{Float64}, Nothing}=nothing,
+    a_bound::Float64=1.0,
+    a_bounds=fill(a_bound, system.n_drives),
+    da_bound::Float64=Inf,
+    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),
+    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,
+    constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
+    piccolo_options::PiccoloOptions=PiccoloOptions(),
+)
+    if piccolo_options.verbose
+        println("    constructing UnitarySmoothPulseProblem...")
+        println("\tusing integrator: $(typeof(unitary_integrator))")
+        println("\tinitial free phases: $(phase_name) = $(initial_phases)")
+    end
+
+    # Trajectory
+    if !isnothing(init_trajectory)
+        traj = init_trajectory
+    else
+        traj = initialize_trajectory(
+            goal(initial_phases),
+            T,
+            Δt,
+            system.n_drives,
+            (a_bounds, da_bounds, dda_bounds);
+            state_name=state_name,
+            control_name=control_name,
+            timestep_name=timestep_name,
+            phase_name=phase_name,
+            phase_data=initial_phases,
+            Δt_bounds=(Δt_min, Δt_max),
+            zero_initial_and_final_derivative=piccolo_options.zero_initial_and_final_derivative,
+            geodesic=piccolo_options.geodesic,
+            bound_state=piccolo_options.bound_state,
+            a_guess=a_guess,
+            system=system,
+            rollout_integrator=piccolo_options.rollout_integrator,
+            verbose=piccolo_options.verbose
+        )
+    end
+
+    # Objective
+    J = UnitaryFreePhaseInfidelityObjective(
+        goal, 
+        state_name, 
+        phase_name,
+        traj
+    )
+
+    control_names = [
+        name for name ∈ traj.names
+            if endswith(string(name), string(control_name))
+    ]
+
+    J += QuadraticRegularizer(control_names[1], traj, R_a)
+    J += QuadraticRegularizer(control_names[2], traj, R_da)
+    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=goal isa EmbeddedOperator ? get_leakage_indices(goal) : nothing
+    )
+
+    integrators = [
+        unitary_integrator(system, traj, state_name, control_name),
+        DerivativeIntegrator(traj, control_name, control_names[2]),
+        DerivativeIntegrator(traj, control_names[2], control_names[3]),
+    ]
+
+    return DirectTrajOptProblem(
+        traj,
+        J,
+        integrators;
+        constraints=constraints
+    )
+end

src/quantum_objectives.jl

@@ -4,6 +4,7 @@ export KetInfidelityObjective
 export UnitaryInfidelityObjective
 export DensityMatrixPureStateInfidelityObjective
 export UnitarySensitivityObjective
+export UnitaryFreePhaseInfidelityObjective
 
 using LinearAlgebra
 using NamedTrajectories
@@ -39,7 +40,7 @@ end
 
 function unitary_fidelity_loss(
     Ũ⃗::AbstractVector{<:Real},
-    U_goal::AbstractMatrix{<:Complex{Float64}}
+    U_goal::AbstractMatrix{<:Complex{<:Real}}
 )
     U = iso_vec_to_operator(Ũ⃗)
     n = size(U, 1)
@@ -67,6 +68,21 @@ function UnitaryInfidelityObjective(
     return TerminalObjective(ℓ, Ũ⃗_name, traj; Q=Q)
 end
 
+function UnitaryFreePhaseInfidelityObjective(
+    U_goal::Function,
+    Ũ⃗_name::Symbol,
+    θ_name::Symbol,
+    traj::NamedTrajectory;
+    Q=100.0
+)
+    d = traj.global_dims[θ_name]
+    function ℓ(x)
+        Ũ⃗, θ = x[1:end-d], x[end-d+1:end]
+        return abs(1 - QuantumObjectives.unitary_fidelity_loss(Ũ⃗, U_goal(θ)))
+    end
+    return TerminalObjective(ℓ, Ũ⃗_name, traj; Q=Q, global_names=[θ_name])
+end
+
 # ---------------------------------------------------------
 #                        Density Matrices
 # ---------------------------------------------------------