update fidelity constraints for globals

src/problem_templates/unitary_free_phase_problem.jl

@@ -33,7 +33,7 @@ function UnitaryFreePhaseProblem(
     piccolo_options::PiccoloOptions=PiccoloOptions(),
 )
     if piccolo_options.verbose
-        println("    constructing UnitarySmoothPulseProblem...")
+        println("    constructing UnitaryFreePhaseProblem...")
         println("\tusing integrator: $(typeof(unitary_integrator))")
         println("\tinitial free phases: $(phase_name) = $(initial_phases)")
     end
@@ -105,7 +105,7 @@ function UnitaryFreePhaseProblem(
         x, y = z[1:length(z) ÷ 2], z[length(z) ÷ 2 + 1:end]
         return x .^ 2 + y .^2 .- 1
     end
-    push!(constraints, GlobalNonlinearConstraint(phase_norm, phase_names, traj))
+    push!(constraints, NonlinearGlobalConstraint(phase_norm, phase_names, traj))
 
     integrators = [
         unitary_integrator(system, traj, state_name, control_name),

src/problem_templates/unitary_minimum_time_problem.jl

@@ -91,6 +91,68 @@ function UnitaryMinimumTimeProblem(
     )
 end
 
+# --------------------------------------------------------------------------- #
+# Free phases
+# --------------------------------------------------------------------------- #
+
+function UnitaryMinimumTimeProblem(
+    trajectory::NamedTrajectory,
+    goal::Function,
+    objective::Objective,
+    dynamics::TrajectoryDynamics,
+    constraints::AbstractVector{<:AbstractConstraint};
+    unitary_name::Symbol=:Ũ⃗,
+    phase_name::Symbol=:θ,
+    final_fidelity::Float64=1.0,
+    D::Float64=100.0,
+    piccolo_options::PiccoloOptions=PiccoloOptions(),
+)
+    # Collect phase names
+    phase_names = [
+        n for n in keys(trajectory.global_data) if endswith(string(n), string(phase_name))
+    ]
+
+    if piccolo_options.verbose
+        println("    constructing UnitaryMinimumTimeProblem...")
+        println("\tfinal fidelity: $(final_fidelity)")
+        println("\tphase names: $(phase_names)")
+    end
+
+    objective += MinimumTimeObjective(
+        trajectory; D=D, timesteps_all_equal=piccolo_options.timesteps_all_equal
+    )
+
+    fidelity_constraint = FinalUnitaryFidelityConstraint(
+        goal, unitary_name, phase_names, final_fidelity, trajectory
+    )
+
+    constraints = push!(constraints, fidelity_constraint)
+
+    return DirectTrajOptProblem(
+        trajectory,
+        objective,
+        dynamics,
+        constraints
+    )
+end
+
+function UnitaryMinimumTimeProblem(
+    prob::DirectTrajOptProblem,
+    goal::Function;
+    objective::Objective=prob.objective,
+    constraints::AbstractVector{<:AbstractConstraint}=prob.constraints,
+    kwargs...
+)
+    return UnitaryMinimumTimeProblem(
+        prob.trajectory,
+        goal,
+        objective,
+        prob.dynamics,
+        constraints;
+        kwargs...
+    )
+end
+
 # *************************************************************************** #
 
 @testitem "Minimum time Hadamard gate" begin

src/quantum_constraints.jl

@@ -21,7 +21,7 @@ function FinalKetFidelityConstraint(
     traj::NamedTrajectory
 )
     terminal_constraint = ψ̃ -> [
-        abs(QuantumObjectives.ket_fidelity_loss(ψ̃, ψ_goal) - final_fidelity)
+        final_fidelity - QuantumObjectives.ket_fidelity_loss(ψ̃, ψ_goal)
     ]
 
     return NonlinearKnotPointConstraint(
@@ -44,7 +44,7 @@ function FinalUnitaryFidelityConstraint(
     traj::NamedTrajectory
 )
     terminal_constraint = Ũ⃗ -> [
-        abs(QuantumObjectives.unitary_fidelity_loss(Ũ⃗, U_goal) - final_fidelity)
+        final_fidelity - QuantumObjectives.unitary_fidelity_loss(Ũ⃗, U_goal)
     ]
 
     return NonlinearKnotPointConstraint(
@@ -56,4 +56,29 @@ function FinalUnitaryFidelityConstraint(
     )
 end
 
+function FinalUnitaryFidelityConstraint(
+    U_goal::Function,
+    Ũ⃗_name::Symbol,
+    θ_names::AbstractVector{Symbol},
+    final_fidelity::Float64,
+    traj::NamedTrajectory
+)
+    d = sum(traj.global_dims[n] for n in θ_names)
+    function terminal_constraint(z)
+        Ũ⃗, θ = z[1:end-d], z[end-d+1:end]
+        return [
+            final_fidelity - QuantumObjectives.unitary_fidelity_loss(Ũ⃗, U_goal(θ))
+        ]
+    end
+
+    return NonlinearGlobalKnotPointConstraint(
+        terminal_constraint,
+        Ũ⃗_name,
+        θ_names,
+        traj,
+        equality=false,
+        times=[traj.T]
+    )
+end
+
 end

src/quantum_objectives.jl

@@ -80,7 +80,7 @@ function UnitaryFreePhaseInfidelityObjective(
         Ũ⃗, θ = z[1:end-d], z[end-d+1:end]
         return abs(1 - QuantumObjectives.unitary_fidelity_loss(Ũ⃗, U_goal(θ)))
     end
-    return TerminalObjective(ℓ, Ũ⃗_name, traj; Q=Q, global_names=θ_names)
+    return TerminalObjective(ℓ, Ũ⃗_name, θ_names, traj; Q=Q)
 end
 
 function UnitaryFreePhaseInfidelityObjective(