Merge/refactor commit '8c618b09804346b13c56883aafeda8423fd7302c' into…

… feature/free-phase

Project.toml

@@ -7,7 +7,6 @@ version = "0.8.0"
 DirectTrajOpt = "c823fa1f-8872-4af5-b810-2b9b72bbbf56"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 ExponentialAction = "e24c0720-ea99-47e8-929e-571b494574d3"
-Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -21,11 +20,13 @@ TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe"
 TrajectoryIndexingUtils = "6dad8b7f-dd9a-4c28-9b70-85b9a079bfc8"
 
 [compat]
+DirectTrajOpt = "0.3.0"
 Distributions = "0.25"
 ExponentialAction = "0.2"
-Interpolations = "0.15"
 JLD2 = "0.5"
 LinearAlgebra = "1.10, 1.11"
+NamedTrajectories = "0.4.0"
+PiccoloQuantumObjects = "0.6.0"
 Random = "1.10, 1.11"
 Reexport = "1.2"
 SparseArrays = "1.10, 1.11"

docs/src/lib.md

@@ -13,9 +13,4 @@ Modules = [QuantumCollocation.Options]
 ## Trajectory Initialization
 ```@autodocs
 Modules = [QuantumCollocation.TrajectoryInitialization]
-```
-
-## Trajectory Interpolations
-```@autodocs
-Modules = [QuantumCollocation.TrajectoryInterpolations]
 ```

src/QuantumCollocation.jl

@@ -10,9 +10,6 @@ include("piccolo_options.jl")
 include("trajectory_initialization.jl")
 @reexport using .TrajectoryInitialization
 
-include("trajectory_interpolations.jl")
-@reexport using .TrajectoryInterpolations
-
 include("quantum_objectives.jl")
 @reexport using .QuantumObjectives
 

src/problem_templates/quantum_state_minimum_time_problem.jl

@@ -35,9 +35,7 @@ function QuantumStateMinimumTimeProblem(
         println("\tfinal fidelity: $(final_fidelity)")
     end
 
-    objective += MinimumTimeObjective(
-        trajectory, D=D, timesteps_all_equal=piccolo_options.timesteps_all_equal
-    )
+    objective += MinimumTimeObjective(trajectory, D=D)
 
     for (state_name, ψ_goal) in zip(state_names, ψ_goals)
         fidelity_constraint = FinalKetFidelityConstraint(

src/problem_templates/quantum_state_sampling_problem.jl

@@ -73,7 +73,7 @@ function QuantumStateSamplingProblem(
             )
         end
 
-        traj = merge(trajs, merge_names=(; a=1, da=1, dda=1, Δt=1), free_time=true)
+        traj = merge(trajs, merge_names=(a=1, da=1, dda=1, Δt=1), timestep=timestep_name)
     end
 
     control_names = [

src/problem_templates/unitary_sampling_problem.jl

@@ -106,8 +106,10 @@ function UnitarySamplingProblem(
             )
         end
 
-        traj = merge(trajs, merge_names=(; a=1, da=1, dda=1, Δt=1), free_time=true)
-    end    
+        traj = merge(
+            trajs, merge_names=(a=1, da=1, dda=1, Δt=1), timestep=timestep_name
+        )
+    end
 
     control_names = [
         name for name ∈ traj.names

src/quantum_constraints.jl

@@ -1,14 +1,16 @@
 module QuantumConstraints
 
 using ..QuantumObjectives
+using ..QuantumObjectives: ket_fidelity_loss, unitary_fidelity_loss
 
+using DirectTrajOpt
+using LinearAlgebra
 using NamedTrajectories
 using PiccoloQuantumObjects
-using DirectTrajOpt
 
 export FinalKetFidelityConstraint
 export FinalUnitaryFidelityConstraint
-
+export LeakageConstraint
 
 # ---------------------------------------------------------
 #                        Kets
@@ -20,9 +22,7 @@ function FinalKetFidelityConstraint(
     final_fidelity::Float64,
     traj::NamedTrajectory
 )
-    terminal_constraint = ψ̃ -> [
-        final_fidelity - QuantumObjectives.ket_fidelity_loss(ψ̃, ψ_goal)
-    ]
+    terminal_constraint = ψ̃ -> [final_fidelity - ket_fidelity_loss(ψ̃, ψ_goal)]
 
     return NonlinearKnotPointConstraint(
         terminal_constraint,
@@ -43,9 +43,7 @@ function FinalUnitaryFidelityConstraint(
     final_fidelity::Float64,
     traj::NamedTrajectory
 )
-    terminal_constraint = Ũ⃗ -> [
-        final_fidelity - QuantumObjectives.unitary_fidelity_loss(Ũ⃗, U_goal)
-    ]
+    terminal_constraint = Ũ⃗ -> [final_fidelity - unitary_fidelity_loss(Ũ⃗, U_goal)]
 
     return NonlinearKnotPointConstraint(
         terminal_constraint,
@@ -63,12 +61,10 @@ function FinalUnitaryFidelityConstraint(
     final_fidelity::Float64,
     traj::NamedTrajectory
 )
-    d = sum(traj.global_dims[n] for n in θ_names)
+    θ_dim = 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(θ))
-        ]
+        Ũ⃗, θ = z[1:end-θ_dim], z[end-θ_dim+1:end]
+        return [final_fidelity - unitary_fidelity_loss(Ũ⃗, U_goal(θ))]
     end
 
     return NonlinearGlobalKnotPointConstraint(
@@ -81,4 +77,32 @@ function FinalUnitaryFidelityConstraint(
     )
 end
 
+# ---------------------------------------------------------
+# Leakage Constraint
+# ---------------------------------------------------------
+
+"""
+    LeakageConstraint(value, indices, name, traj::NamedTrajectory)
+
+Construct a `KnotPointConstraint` that bounds leakage of `name` at the knot points specified by `times` at any `indices` that are outside the computational subspace.
+
+"""
+function LeakageConstraint(
+    value::Float64,
+    indices::AbstractVector{Int},
+    name::Symbol,
+    traj::NamedTrajectory;
+    times=1:traj.T,
+)
+    leakage_constraint(x) = [sum(abs2.(x[indices])) - value]
+
+    return NonlinearKnotPointConstraint(
+        leakage_constraint,
+        name,
+        traj,
+        equality=false,
+        times=times,
+    )
+end
+
 end

src/quantum_objectives.jl

@@ -5,6 +5,7 @@ export UnitaryInfidelityObjective
 export DensityMatrixPureStateInfidelityObjective
 export UnitarySensitivityObjective
 export UnitaryFreePhaseInfidelityObjective
+export LeakageObjective
 
 using LinearAlgebra
 using NamedTrajectories
@@ -13,7 +14,7 @@ using DirectTrajOpt
 using TestItems
 
 # --------------------------------------------------------- 
-#                        Kets
+#                       Kets
 # ---------------------------------------------------------
 
 function ket_fidelity_loss(
@@ -36,7 +37,7 @@ end
 
 
 # ---------------------------------------------------------
-#                        Unitaries
+#                       Unitaries
 # ---------------------------------------------------------
 
 function unitary_fidelity_loss(
@@ -95,39 +96,7 @@ function UnitaryFreePhaseInfidelityObjective(
 end
 
 # ---------------------------------------------------------
-#                Leakage Suppression Objective
-# ---------------------------------------------------------
-
-"""
-        LeakageSuppressionObjective(leakage_inds::AbstractVector{<:Integer}, name::Symbol, traj::NamedTrajectory; kwargs...)
-
-Construct a `KnotPointObjective` that penalizes leakage outside the computational subspace.
-
-- `leakage_inds`: Vector of indices specifying the leakage subspace in the isomorphic vector representation (e.g., `get_iso_vec_leakage_indices`).
-- `name`: The variable name in the trajectory (e.g., `:u`).
-- `traj`: The `NamedTrajectory` containing the variable.
-- `kwargs...`: Passed to `KnotPointObjective`.
-
-The objective is: 
-    sum_{i ∈ I_leakage} abs(U[i])
-where `I_leakage` is given by `leakage_inds` and `U` is the vectorized variable.
-"""
-function LeakageSuppressionObjective(
-    leakage_inds::AbstractVector{<:Integer},
-    name::Symbol,
-    traj::NamedTrajectory;
-    kwargs...
-)
-    return KnotPointObjective(
-        (U, args...) -> sum(abs(U[i]) for i in leakage_inds),
-        name,
-        traj;
-        kwargs...
-    )
-end
-
-# ---------------------------------------------------------
-#                        Density Matrices
+#                       Density Matrices
 # ---------------------------------------------------------
 
 function density_matrix_pure_state_infidelity_loss(
@@ -150,7 +119,7 @@ function DensityMatrixPureStateInfidelityObjective(
 end
 
 # ---------------------------------------------------------
-#                        Sensitivity
+#                       Sensitivity
 # ---------------------------------------------------------
 
 function unitary_fidelity_loss(
@@ -179,22 +148,33 @@ function UnitarySensitivityObjective(
     )
 end
 
-@testitem "LeakageSuppressionObjective basic functionality" begin
-    using PiccoloQuantumObjects
-    using NamedTrajectories
-
-    # U with zeros everywhere (no leakage)
-    U = zeros(3,3)
-    traj = NamedTrajectory((; u = reshape(U, 9, 1)); controls=:u, timestep=1.0)
-    leakage_inds = [4, 8]
-    obj = QuantumObjectives.LeakageSuppressionObjective(leakage_inds, :u, traj)
-    @test obj.L(traj.datavec) == 0.0
-
-    # U with leakage at (1,2) and (2,3)
-    U[1,2] = 2.0
-    U[2,3] = 3.0
-    traj = NamedTrajectory((; u = reshape(U, 9, 1)); controls=:u, timestep=1.0)
-    @test obj.L(traj.datavec) ≈ 2.0 + 3.0
+# ---------------------------------------------------------
+#                       Leakage
+# ---------------------------------------------------------
+
+"""
+    LeakageObjective(indices, name, traj::NamedTrajectory)
+
+Construct a `KnotPointObjective` that penalizes leakage of `name` at the knot points specified by `times` at any `indices` that are outside the computational subspace.
+
+"""
+function LeakageObjective(
+    indices::AbstractVector{Int},
+    name::Symbol,
+    traj::NamedTrajectory;
+    times=1:traj.T,
+    Qs::AbstractVector{<:Float64}=fill(1.0, length(times)),
+)
+    leakage_objective(x) = sum(abs2.(x[indices]))
+
+    return KnotPointObjective(
+        leakage_objective,
+        name,
+        traj;
+        Qs=Qs,
+        times=times,
+    )
 end
 
+
 end

src/quantum_system_templates/cats.jl

@@ -62,3 +62,6 @@ function get_cat_controls(sys::AbstractQuantumSystem, α::Real, T::Int)
         fill(cat_kerr_correction, T)
     ], dims=1)
 end
+
+# *************************************************************************** #
+

src/quantum_system_templates/rydberg.jl

@@ -119,6 +119,8 @@ function RydbergChainSystem(;
     )
 end
 
+# *************************************************************************** #
+
 @testitem "Rydberg system test" begin
     using PiccoloQuantumObjects