update problem templates and bump PQO compat

Project.toml

@@ -1,13 +1,12 @@
 name = "QuantumCollocation"
 uuid = "0dc23a59-5ffb-49af-b6bd-932a8ae77adf"
+version = "0.9.0"
 authors = ["Aaron Trowbridge <aaron.j.trowbridge@gmail.com> and contributors"]
-version = "0.8.1"
 
 [deps]
 DirectTrajOpt = "c823fa1f-8872-4af5-b810-2b9b72bbbf56"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 ExponentialAction = "e24c0720-ea99-47e8-929e-571b494574d3"
-JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NamedTrajectories = "538bc3a1-5ab9-4fc3-b776-35ca1e893e08"
@@ -20,13 +19,12 @@ TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe"
 TrajectoryIndexingUtils = "6dad8b7f-dd9a-4c28-9b70-85b9a079bfc8"
 
 [compat]
-DirectTrajOpt = "0.4"
+DirectTrajOpt = "0.5"
 Distributions = "0.25"
 ExponentialAction = "0.2"
-JLD2 = "0.5"
 LinearAlgebra = "1.10, 1.11, 1.12"
-NamedTrajectories = "0.5"
-PiccoloQuantumObjects = "0.6"
+NamedTrajectories = "0.6"
+PiccoloQuantumObjects = "0.7"
 Random = "1.10, 1.11, 1.12"
 Reexport = "1.2"
 SparseArrays = "1.10, 1.11, 1.12"

docs/literate/examples/multilevel_transmon.jl

@@ -44,7 +44,7 @@ levels = 5
 δ = 0.2
 
 ## add a bound to the controls
-a_bound = 0.2
+u_bound = 0.2
 
 ## create the system
 sys = TransmonSystem(levels=levels, δ=δ)
@@ -75,7 +75,7 @@ get_subspace_identity(op) |> sparse
 # We can then pass this embedded operator to the `UnitarySmoothPulseProblem` template to create the problem
 
 ## create the problem
-prob = UnitarySmoothPulseProblem(sys, op, T, Δt; a_bound=a_bound)
+prob = UnitarySmoothPulseProblem(sys, op, T, Δt; u_bound=u_bound)
 
 ## solve the problem
 solve!(prob; max_iter=50)
@@ -98,8 +98,8 @@ plot_unitary_populations(prob.trajectory; fig_size=(900, 700))
 ## create the a leakage suppression problem, initializing with the previous solution
 
 prob_leakage = UnitarySmoothPulseProblem(sys, op, T, Δt;
-    a_bound=a_bound,
-    a_guess=prob.trajectory.a[:, :],
+    u_bound=u_bound,
+    u_guess=prob.trajectory.u[:, :],
     piccolo_options=PiccoloOptions(
         leakage_constraint=true,
         leakage_constraint_value=1e-2,

docs/literate/examples/two_qubit_gates.jl

@@ -13,15 +13,15 @@
 # Specifically, for a simple two-qubit system in a rotating frame, we have
 
 # ```math
-# H = J_{12} \sigma_1^x \sigma_2^x + \sum_{i \in {1,2}} a_i^R(t) {\sigma^x_i \over 2} + a_i^I(t) {\sigma^y_i \over 2}.
+# H = J_{12} \sigma_1^x \sigma_2^x + \sum_{i \in {1,2}} u_i^R(t) {\sigma^x_i \over 2} + u_i^I(t) {\sigma^y_i \over 2}.
 # ```
 
 # where
 
 # ```math
 # \begin{align*}
 # J_{12} &= 0.001 \text{ GHz}, \\
-# |a_i^R(t)| &\leq 0.1 \text{ GHz} \\
+# |u_i^R(t)| &\leq 0.1 \text{ GHz} \\
 # \end{align*}
 # ```
 
@@ -53,7 +53,7 @@ Id = GATES[:I]
 
 ## Define the parameters of the Hamiltonian
 J_12 = 0.001 # GHz
-a_bound = 0.100 # GHz
+u_bound = 0.100 # GHz
 
 ## Define the drift (coupling) Hamiltonian
 H_drift = J_12 * (σx ⊗ σx)
@@ -62,9 +62,9 @@ H_drift = J_12 * (σx ⊗ σx)
 H_drives = [σx_1 / 2, σy_1 / 2, σx_2 / 2, σy_2 / 2]
 
 ## Define control (and higher derivative) bounds
-a_bound = 0.1
-da_bound = 0.0005
-dda_bound = 0.0025
+u_bound = 0.1
+du_bound = 0.0005
+ddu_bound = 0.0025
 
 ## Scale the Hamiltonians by 2π
 H_drift *= 2π
@@ -95,11 +95,11 @@ prob = UnitarySmoothPulseProblem(
     U_goal,
     T,
     Δt;
-    a_bound=a_bound,
-    da_bound=da_bound,
-    dda_bound=dda_bound,
-    R_da=0.01,
-    R_dda=0.01,
+    u_bound=u_bound,
+    du_bound=du_bound,
+    ddu_bound=ddu_bound,
+    R_du=0.01,
+    R_ddu=0.01,
     Δt_max=Δt_max,
     piccolo_options=PiccoloOptions(bound_state=true),
 )
@@ -166,11 +166,11 @@ prob = UnitarySmoothPulseProblem(
     U_goal,
     T,
     Δt;
-    a_bound=a_bound,
-    da_bound=da_bound,
-    dda_bound=dda_bound,
-    R_da=0.01,
-    R_dda=0.01,
+    u_bound=u_bound,
+    du_bound=du_bound,
+    ddu_bound=ddu_bound,
+    R_du=0.01,
+    R_ddu=0.01,
     Δt_max=Δt_max,
     piccolo_options=PiccoloOptions(bound_state=true),
 )

docs/literate/man/ket_problem_templates.jl

@@ -39,7 +39,7 @@ solve!(state_prob, max_iter=100, verbose=true, print_level=1);
 println("After: ", rollout_fidelity(state_prob.trajectory, system))
 
 # _extract the control pulses_
-state_prob.trajectory.a |> size
+state_prob.trajectory.u |> size
 
 # -----
 

docs/literate/man/unitary_problem_templates.jl

@@ -41,7 +41,7 @@ The `NamedTrajectory` object stores the control pulse, state variables, and the
 =#
 
 # _extract the control pulses_
-prob.trajectory.a |> size
+prob.trajectory.u |> size
 
 # -----
 

docs/src/index.md

@@ -36,26 +36,26 @@ Unitary Problem Templates:
 
 *Problem Templates* are reusable design patterns for setting up and solving common quantum control problems. 
 
-For example, a *UnitarySmoothPulseProblem* is tasked with generating a *pulse* sequence $a_{1:T-1}$ in orderd to minimize infidelity, subject to constraints from the Schroedinger equation,
+For example, a *UnitarySmoothPulseProblem* is tasked with generating a *pulse* sequence $u_{1:T-1}$ in order to minimize infidelity, subject to constraints from the Schroedinger equation,
 ```math
     \begin{aligned}
         \arg \min_{\mathbf{Z}}\quad & |1 - \mathcal{F}(U_T, U_\text{goal})|  \\
         \nonumber \text{s.t.}
-        \qquad & U_{t+1} = \exp\{- i H(a_t) \Delta t_t \} U_t, \quad \forall\, t \\
+        \qquad & U_{t+1} = \exp\{- i H(u_t) \Delta t_t \} U_t, \quad \forall\, t \\
     \end{aligned}
 ```
 while a *UnitaryMinimumTimeProblem* minimizes time and constrains fidelity,
 ```math
     \begin{aligned}
         \arg \min_{\mathbf{Z}}\quad & \sum_{t=1}^T \Delta t_t \\
-        \qquad & U_{t+1} = \exp\{- i H(a_t) \Delta t_t \} U_t, \quad \forall\, t \\
+        \qquad & U_{t+1} = \exp\{- i H(u_t) \Delta t_t \} U_t, \quad \forall\, t \\
         \nonumber & \mathcal{F}(U_T, U_\text{goal}) \ge 0.9999
     \end{aligned}
 ```
 
 -----
 
-In each case, the dynamics between *knot points* $(U_t, a_t)$ and $(U_{t+1}, a_{t+1})$ are enforced as constraints on the states, which are free variables in the solver; this optimization framework is called *direct trajectory optimization*. 
+In each case, the dynamics between *knot points* $(U_t, u_t)$ and $(U_{t+1}, u_{t+1})$ are enforced as constraints on the states, which are free variables in the solver; this optimization framework is called *direct trajectory optimization*. 
 
 -----
 

src/problem_templates/_problem_templates.jl

@@ -12,7 +12,6 @@ using DirectTrajOpt
 using PiccoloQuantumObjects
 
 using ExponentialAction
-using JLD2
 using LinearAlgebra
 using SparseArrays
 using TestItems
@@ -53,7 +52,7 @@ function apply_piccolo_options!(
         end
 
         for (name, indices) ∈ zip(state_names, state_leakage_indices)
-            J += LeakageObjective(indices, name, traj, Qs=fill(piccolo_options.leakage_cost, traj.T))
+            J += LeakageObjective(indices, name, traj, Qs=fill(piccolo_options.leakage_cost, traj.N))
             push!(constraints, LeakageConstraint(val, indices, name, traj))
         end
     end
@@ -73,7 +72,7 @@ function apply_piccolo_options!(
             println("\tapplying complex control norm constraint: $(piccolo_options.complex_control_norm_constraint_name)")
         end
         norm_con = NonlinearKnotPointConstraint(
-            a -> [norm(a)^2 - piccolo_options.complex_control_norm_constraint_radius^2],
+            u -> [norm(u)^2 - piccolo_options.complex_control_norm_constraint_radius^2],
             piccolo_options.complex_control_norm_constraint_name,
             traj;
             equality=false,

src/problem_templates/quantum_state_minimum_time_problem.jl

@@ -89,14 +89,14 @@ end
     using NamedTrajectories
     using PiccoloQuantumObjects 
 
-    T = 51
+    N = 51
     Δt = 0.2
-    sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]])
+    sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]], 10.0, [1.0, 1.0])
     ψ_init = Vector{ComplexF64}[[1.0, 0.0]]
     ψ_target = Vector{ComplexF64}[[0.0, 1.0]]
 
     prob = QuantumStateSmoothPulseProblem(
-        sys, ψ_init, ψ_target, T, Δt;
+        sys, ψ_init, ψ_target, N, Δt;
         piccolo_options=PiccoloOptions(verbose=false)
     )
     initial = sum(get_timesteps(prob.trajectory))

src/problem_templates/quantum_state_sampling_problem.jl

@@ -1,7 +1,41 @@
 export QuantumStateSamplingProblem
 
 """
+    QuantumStateSamplingProblem(systems, ψ_inits, ψ_goals, T, Δt; kwargs...)
 
+Construct a quantum state sampling problem for multiple systems with shared controls.
+
+# Arguments
+- `systems::AbstractVector{<:AbstractQuantumSystem}`: A vector of quantum systems.
+- `ψ_inits::AbstractVector{<:AbstractVector{<:AbstractVector{<:ComplexF64}}}`: Initial states for each system.
+- `ψ_goals::AbstractVector{<:AbstractVector{<:AbstractVector{<:ComplexF64}}}`: Target states for each system.
+- `T::Int`: The number of time steps.
+- `Δt::Union{Float64, AbstractVector{Float64}}`: The time step value or vector of time steps.
+
+# Keyword Arguments
+- `ket_integrator=KetIntegrator`: The integrator to use for state dynamics.
+- `system_weights=fill(1.0, length(systems))`: The weights for each system.
+- `init_trajectory::Union{NamedTrajectory,Nothing}=nothing`: The initial trajectory.
+- `state_name::Symbol=:ψ̃`: The name of the state variable.
+- `control_name::Symbol=:u`: The name of the control variable.
+- `timestep_name::Symbol=:Δt`: The name of the timestep variable.
+- `u_bound::Float64=1.0`: The bound for the control amplitudes.
+- `u_bounds=fill(u_bound, systems[1].n_drives)`: The bounds for the control amplitudes.
+- `u_guess::Union{Matrix{Float64},Nothing}=nothing`: The initial guess for the control amplitudes.
+- `du_bound::Float64=Inf`: The bound for the control first derivatives.
+- `du_bounds=fill(du_bound, systems[1].n_drives)`: The bounds for the control first derivatives.
+- `ddu_bound::Float64=1.0`: The bound for the control second derivatives.
+- `ddu_bounds=fill(ddu_bound, systems[1].n_drives)`: The bounds for the control second derivatives.
+- `Δt_min::Float64=Δt isa Float64 ? 0.5 * Δt : 0.5 * minimum(Δt)`: The minimum time step size.
+- `Δt_max::Float64=Δt isa Float64 ? 2.0 * Δt : 2.0 * maximum(Δt)`: The maximum time step size.
+- `Q::Float64=100.0`: The fidelity weight.
+- `R::Float64=1e-2`: The regularization weight.
+- `R_u::Union{Float64,Vector{Float64}}=R`: The regularization weight for the control amplitudes.
+- `R_du::Union{Float64,Vector{Float64}}=R`: The regularization weight for the control first derivatives.
+- `R_ddu::Union{Float64,Vector{Float64}}=R`: The regularization weight for the control second derivatives.
+- `state_leakage_indices::Union{Nothing, AbstractVector{Int}}=nothing`: Indices of leakage states.
+- `constraints::Vector{<:AbstractConstraint}=AbstractConstraint[]`: The constraints.
+- `piccolo_options::PiccoloOptions=PiccoloOptions()`: The Piccolo options.
 """
 function QuantumStateSamplingProblem end
 
@@ -15,22 +49,22 @@ function QuantumStateSamplingProblem(
     system_weights=fill(1.0, length(systems)),
     init_trajectory::Union{NamedTrajectory,Nothing}=nothing,
     state_name::Symbol=:ψ̃,
-    control_name::Symbol=:a,
+    control_name::Symbol=:u,
     timestep_name::Symbol=:Δt,
-    a_bound::Float64=1.0,
-    a_bounds=fill(a_bound, systems[1].n_drives),
-    a_guess::Union{Matrix{Float64},Nothing}=nothing,
-    da_bound::Float64=Inf,
-    da_bounds=fill(da_bound, systems[1].n_drives),
-    dda_bound::Float64=1.0,
-    dda_bounds=fill(dda_bound, systems[1].n_drives),
-    Δt_min::Float64=0.5 * minimum(Δt),
-    Δt_max::Float64=2.0 * maximum(Δt),
+    u_bound::Float64=1.0,
+    u_bounds=fill(u_bound, systems[1].n_drives),
+    u_guess::Union{Matrix{Float64},Nothing}=nothing,
+    du_bound::Float64=Inf,
+    du_bounds=fill(du_bound, systems[1].n_drives),
+    ddu_bound::Float64=1.0,
+    ddu_bounds=fill(ddu_bound, systems[1].n_drives),
+    Δt_min::Float64=Δt isa Float64 ? 0.5 * Δt : 0.5 * minimum(Δt),
+    Δt_max::Float64=Δt isa Float64 ? 2.0 * Δt : 2.0 * maximum(Δ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,
+    R::Float64=1e-2,
+    R_u::Union{Float64,Vector{Float64}}=R,
+    R_du::Union{Float64,Vector{Float64}}=R,
+    R_ddu::Union{Float64,Vector{Float64}}=R,
     state_leakage_indices::Union{Nothing, AbstractVector{Int}}=nothing,
     constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
     piccolo_options::PiccoloOptions=PiccoloOptions(),
@@ -60,21 +94,21 @@ function QuantumStateSamplingProblem(
                 T,
                 Δt,
                 sys.n_drives,
-                (a_bounds, da_bounds, dda_bounds);
+                (u_bounds, du_bounds, ddu_bounds);
                 state_names=names,
                 control_name=control_name,
                 timestep_name=timestep_name,
                 Δt_bounds=(Δt_min, Δt_max),
                 zero_initial_and_final_derivative=piccolo_options.zero_initial_and_final_derivative,
                 bound_state=piccolo_options.bound_state,
-                a_guess=a_guess,
+                u_guess=u_guess,
                 system=sys,
                 rollout_integrator=piccolo_options.rollout_integrator,
                 verbose=false # loop
             )
         end
 
-        traj = merge(trajs, merge_names=(a=1, da=1, dda=1, Δt=1), timestep=timestep_name)
+        traj = merge(trajs, merge_names=(u=1, du=1, ddu=1, Δt=1), timestep=timestep_name)
     end
 
     control_names = [
@@ -83,9 +117,9 @@ function QuantumStateSamplingProblem(
     ]
 
     # Objective
-    J = QuadraticRegularizer(control_names[1], traj, R_a)
-    J += QuadraticRegularizer(control_names[2], traj, R_da)
-    J += QuadraticRegularizer(control_names[3], traj, R_dda)
+    J = QuadraticRegularizer(control_names[1], traj, R_u)
+    J += QuadraticRegularizer(control_names[2], traj, R_du)
+    J += QuadraticRegularizer(control_names[3], traj, R_ddu)
 
     for (weight, names) in zip(system_weights, state_names)
         for name in names
@@ -154,15 +188,15 @@ end
 @testitem "Sample systems with single initial, target" begin
     using PiccoloQuantumObjects
 
-    T = 50
+    N = 50
     Δt = 0.2
-    sys1 = QuantumSystem(0.3 * GATES[:Z], [GATES[:X], GATES[:Y]])
-    sys2 = QuantumSystem(-0.3 * GATES[:Z], [GATES[:X], GATES[:Y]])
+    sys1 = QuantumSystem(0.3 * GATES[:Z], [GATES[:X], GATES[:Y]], 10.0, [1.0, 1.0])
+    sys2 = QuantumSystem(-0.3 * GATES[:Z], [GATES[:X], GATES[:Y]], 10.0, [1.0, 1.0])
     ψ_init = Vector{ComplexF64}([1.0, 0.0])
     ψ_target = Vector{ComplexF64}([0.0, 1.0])
 
     prob = QuantumStateSamplingProblem(
-        [sys1, sys2], ψ_init, ψ_target, T, Δt;
+        [sys1, sys2], ψ_init, ψ_target, N, Δt;
         piccolo_options=PiccoloOptions(verbose=false)
     )
 
@@ -187,17 +221,17 @@ end
 @testitem "Sample systems with multiple initial, target" begin
     using PiccoloQuantumObjects
 
-    T = 50
+    N = 50
     Δt = 0.2
-    sys1 = QuantumSystem(0.3 * GATES[:Z], [GATES[:X], GATES[:Y]])
-    sys2 = QuantumSystem(-0.3 * GATES[:Z], [GATES[:X], GATES[:Y]])
+    sys1 = QuantumSystem(0.3 * GATES[:Z], [GATES[:X], GATES[:Y]], 10.0, [1.0, 1.0])
+    sys2 = QuantumSystem(-0.3 * GATES[:Z], [GATES[:X], GATES[:Y]], 10.0, [1.0, 1.0])
 
     # Multiple initial and target states
     ψ_inits = Vector{ComplexF64}.([[1.0, 0.0], [0.0, 1.0]])
     ψ_targets = Vector{ComplexF64}.([[0.0, 1.0], [1.0, 0.0]])
 
     prob = QuantumStateSamplingProblem(
-        [sys1, sys2], ψ_inits, ψ_targets, T, Δt;
+        [sys1, sys2], ψ_inits, ψ_targets, N, Δt;
         piccolo_options=PiccoloOptions(verbose=false)
     )
 
@@ -219,4 +253,5 @@ end
     end
 end
 
-# TODO: Test that a_guess can be used
+# TODO: Test that u_guess can be used
+