interface changes

Project.toml

@@ -4,6 +4,7 @@ authors = ["Aaron Trowbridge <[email protected]> and contributors"]
 version = "0.6.0"
 
 [deps]
+DirectTrajOpt = "82ab4c96-bdfe-4573-b0a2-347b6440cc7b"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 ExponentialAction = "e24c0720-ea99-47e8-929e-571b494574d3"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
@@ -12,7 +13,6 @@ Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NamedTrajectories = "538bc3a1-5ab9-4fc3-b776-35ca1e893e08"
 PiccoloQuantumObjects = "5a402ddf-f93c-42eb-975e-5582dcda653d"
-QuantumCollocationCore = "2b384925-53cb-4042-a8d2-6faa627467e1"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
@@ -26,9 +26,8 @@ ExponentialAction = "0.2"
 Interpolations = "0.15"
 JLD2 = "0.5"
 LinearAlgebra = "1.10, 1.11"
-NamedTrajectories = "0.2"
+NamedTrajectories = "0.3"
 PiccoloQuantumObjects = "0.3"
-QuantumCollocationCore = "0.3"
 Random = "1.10, 1.11"
 Reexport = "1.2"
 SparseArrays = "1.10, 1.11"

docs/Project.toml

@@ -4,7 +4,7 @@ Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 LiveServer = "16fef848-5104-11e9-1b77-fb7a48bbb589"
 NamedTrajectories = "538bc3a1-5ab9-4fc3-b776-35ca1e893e08"
 QuantumCollocation = "0dc23a59-5ffb-49af-b6bd-932a8ae77adf"
-QuantumCollocationCore = "2b384925-53cb-4042-a8d2-6faa627467e1"
+DirectTrajOpt = "2b384925-53cb-4042-a8d2-6faa627467e1"
 Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 
 [compat]

docs/literate/man/problem_templates.jl

@@ -4,7 +4,7 @@
 
 # # Problem Templates
 
-# We provide a number of problem templates for making it simple and easy to set up and solve certain types of quantum optimal control problems. These templates all construct a `QuantumControlProblem` object.  The problem templates are:
+# We provide a number of problem templates for making it simple and easy to set up and solve certain types of quantum optimal control problems. These templates all construct a `DirectTrajOptProblem` object.  The problem templates are:
 
 # ## Unitary Smooth Pulse Problem
 

src/QuantumCollocation.jl

@@ -2,15 +2,24 @@ module QuantumCollocation
 
 using Reexport
 
-@reexport using QuantumCollocationCore
+@reexport using DirectTrajOpt
 @reexport using PiccoloQuantumObjects
 
+include("piccolo_options.jl")
+@reexport using .Options
+
 include("trajectory_initialization.jl")
 @reexport using .TrajectoryInitialization
 
 include("trajectory_interpolations.jl")
 @reexport using .TrajectoryInterpolations
 
+include("quantum_objectives.jl")
+@reexport using .QuantumObjectives
+
+include("quantum_integrators.jl")
+@reexport using .QuantumIntegrators
+
 include("problem_templates/_problem_templates.jl")
 @reexport using .ProblemTemplates
 

src/piccolo_options.jl

@@ -0,0 +1,55 @@
+module Options 
+
+export PiccoloOptions
+
+using ExponentialAction
+
+"""
+    PiccoloOptions
+
+Options for the Piccolo quantum optimal control library.
+
+# Fields
+- `verbose::Bool = true`: Print verbose output
+- `verbose_evaluator::Bool = false`: Print verbose output from the evaluator
+- `free_time::Bool = true`: Allow free time optimization
+- `timesteps_all_equal::Bool = true`: Use equal timesteps
+- `integrator::Symbol = :pade`: Integrator to use
+- `pade_order::Int = 4`: Order of the Pade approximation
+- `rollout_integrator::Function = expv`: Integrator to use for rollout
+- `eval_hessian::Bool = false`: Evaluate the Hessian
+- `geodesic = true`: Use the geodesic to initialize the optimization.
+- `blas_multithreading::Bool = true`: Use BLAS multithreading.
+- `build_trajectory_constraints::Bool = true`: Build trajectory constraints.
+- `complex_control_norm_constraint_name::Union{Nothing, Symbol} = nothing`: Name of the complex control norm constraint.
+- `complex_control_norm_constraint_radius::Float64 = 1.0`: Radius of the complex control norm constraint.
+- `bound_state::Bool = false`: Bound the state.
+- `leakage_suppression::Bool = false`: Suppress leakage.
+- `R_leakage::Float64 = 1.0`: Leakage suppression parameter.
+- `free_phase_infidelity::Bool = false`: Free phase infidelity.
+- `phase_operators::Union{Nothing, AbstractVector{<:AbstractMatrix{<:Complex}}} = nothing`: Phase operators.
+- `phase_name::Symbol = :ϕ`: Name of the phase.
+"""
+@kwdef mutable struct PiccoloOptions
+    verbose::Bool = true
+    verbose_evaluator::Bool = false
+    free_time::Bool = true
+    timesteps_all_equal::Bool = true
+    integrator::Symbol = :pade
+    pade_order::Int = 4
+    rollout_integrator::Function = expv
+    eval_hessian::Bool = false
+    geodesic = true
+    blas_multithreading::Bool = true
+    build_trajectory_constraints::Bool = true
+    complex_control_norm_constraint_name::Union{Nothing, Symbol} = nothing
+    complex_control_norm_constraint_radius::Float64 = 1.0
+    bound_state::Bool = integrator == :exponential
+    leakage_suppression::Bool = false
+    R_leakage::Float64 = 1.0
+    free_phase_infidelity::Bool = false
+    phase_operators::Union{Nothing, AbstractVector{<:AbstractMatrix{<:Complex}}} = nothing
+    phase_name::Symbol = :ϕ
+end
+
+end

src/problem_templates/_problem_templates.jl

@@ -3,12 +3,14 @@ module ProblemTemplates
 using ..DirectSums
 using ..Rollouts
 using ..TrajectoryInitialization
-using ..Losses
+using ..QuantumObjectives
+using ..QuantumIntegrators
+using ..Options
 
 using Distributions
 using TrajectoryIndexingUtils
 using NamedTrajectories
-using QuantumCollocationCore
+using DirectTrajOpt
 using PiccoloQuantumObjects
 using LinearAlgebra
 using SparseArrays
@@ -58,16 +60,17 @@ function apply_piccolo_options!(
         if piccolo_options.timesteps_all_equal
             push!(
                 constraints,
-                TimeStepsAllEqualConstraint(timestep_name, traj)
+                TimeStepsAllEqualConstraint(traj)
             )
         end
     end
 
     if !isnothing(piccolo_options.complex_control_norm_constraint_name)
-        norm_con = ComplexModulusContraint(
+        norm_con = NonlinearKnotPointConstraint(
+            a -> -[norm(a)^2 - piccolo_options.complex_control_norm_constraint_radius^2],
             piccolo_options.complex_control_norm_constraint_name,
-            piccolo_options.complex_control_norm_constraint_radius,
             traj;
+            equality=false,
         )
         push!(constraints, norm_con)
     end

src/problem_templates/density_operator_smooth_pulse_problem.jl

@@ -111,7 +111,7 @@ function DensityOperatorSmoothPulseProblem(
         DerivativeIntegrator(:da, :dda, traj),
     ]
 
-    return QuantumControlProblem(
+    return DirectTrajOptProblem(
         traj,
         J,
         integrators;

src/problem_templates/quantum_state_minimum_time_problem.jl

@@ -5,7 +5,7 @@ export QuantumStateMinimumTimeProblem
     QuantumStateMinimumTimeProblem(traj, sys, obj, integrators, constraints; kwargs...)
     QuantumStateMinimumTimeProblem(prob; kwargs...)
 
-Construct a `QuantumControlProblem` for the minimum time problem of reaching a target state.
+Construct a `DirectTrajOptProblem` for the minimum time problem of reaching a target state.
 
 # Arguments
 - `traj::NamedTrajectory`: The initial trajectory.
@@ -14,15 +14,15 @@ Construct a `QuantumControlProblem` for the minimum time problem of reaching a t
 - `integrators::Vector{<:AbstractIntegrator}`: The integrators.
 - `constraints::Vector{<:AbstractConstraint}`: The constraints.
 or
-- `prob::QuantumControlProblem`: The quantum control problem.
+- `prob::DirectTrajOptProblem`: The quantum control problem.
 
 # Keyword Arguments
 - `state_name::Symbol=:ψ̃`: The symbol for the state variables.
 - `final_fidelity::Union{Real, Nothing}=nothing`: The final fidelity.
 - `D=1.0`: The cost weight on the time.
 - `ipopt_options::IpoptOptions=IpoptOptions()`: The Ipopt options.
 - `piccolo_options::PiccoloOptions=PiccoloOptions()`: The Piccolo options.
-- `kwargs...`: Additional keyword arguments, passed to `QuantumControlProblem`.
+- `kwargs...`: Additional keyword arguments, passed to `DirectTrajOptProblem`.
 
 """
 function QuantumStateMinimumTimeProblem end
@@ -62,7 +62,7 @@ function QuantumStateMinimumTimeProblem(
         push!(constraints, fidelity_constraint)
     end
 
-    return QuantumControlProblem(
+    return DirectTrajOptProblem(
         traj,
         obj,
         integrators;
@@ -75,7 +75,7 @@ function QuantumStateMinimumTimeProblem(
 end
 
 function QuantumStateMinimumTimeProblem(
-    prob::QuantumControlProblem;
+    prob::DirectTrajOptProblem;
     obj::Objective=get_objective(prob),
     constraints::AbstractVector{<:AbstractConstraint}=get_constraints(prob),
     ipopt_options::IpoptOptions=deepcopy(prob.ipopt_options),

src/problem_templates/quantum_state_sampling_problem.jl

@@ -121,7 +121,7 @@ function QuantumStateSamplingProblem(
         DerivativeIntegrator(control_names[2], control_names[3], traj),
     ]
 
-    return QuantumControlProblem(
+    return DirectTrajOptProblem(
         traj,
         J,
         integrators;

src/problem_templates/quantum_state_smooth_pulse_problem.jl

@@ -125,12 +125,12 @@ function QuantumStateSmoothPulseProblem(
     ]
 
     # Objective
-    J = QuadraticRegularizer(control_names[1], traj, R_a; timestep_name=timestep_name)
-    J += QuadraticRegularizer(control_names[2], traj, R_da; timestep_name=timestep_name)
-    J += QuadraticRegularizer(control_names[3], traj, R_dda; timestep_name=timestep_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)
 
     for name ∈ state_names
-        J += QuantumStateObjective(name, traj, Q)
+        J += KetInfidelityLoss(name, traj; Q=Q)
     end
 
     # Optional Piccolo constraints and objectives
@@ -139,71 +139,78 @@ function QuantumStateSmoothPulseProblem(
         state_leakage_indices=leakage_indices
     )
 
-    # Integrators
+    # # Integrators
+    # state_integrators = []
+    # if length(ψ_inits) == 1
+    #     if piccolo_options.integrator == :pade
+    #         state_integrators = [QuantumStatePadeIntegrator(
+    #             state_name,
+    #             control_name,
+    #             sys,
+    #             traj;
+    #             order=piccolo_options.pade_order
+    #         )]
+    #     elseif piccolo_options.integrator == :exponential
+    #         state_integrators = [QuantumStateExponentialIntegrator(
+    #             state_name,
+    #             control_name,
+    #             sys,
+    #             traj
+    #         )]
+    #     else
+    #         error("integrator must be one of (:pade, :exponential)")
+    #     end
+    # else
+    #     state_names = [
+    #         name for name ∈ traj.names
+    #             if startswith(string(name), string(state_name))
+    #     ]
+    #     state_integrators = []
+    #     for i = 1:length(ψ_inits)
+    #         if piccolo_options.integrator == :pade
+    #             state_integrator = QuantumStatePadeIntegrator(
+    #                 state_names[i],
+    #                 control_name,
+    #                 sys,
+    #                 traj;
+    #                 order=piccolo_options.pade_order
+    #             )
+    #         elseif piccolo_options.integrator == :exponential
+    #             state_integrator = QuantumStateExponentialIntegrator(
+    #                 state_names[i],
+    #                 control_name,
+    #                 sys,
+    #                 traj
+    #             )
+    #         else
+    #             error("integrator must be one of (:pade, :exponential)")
+    #         end
+    #         push!(state_integrators, state_integrator)
+    #     end
+    # end
+
+    state_names = [
+        name for name ∈ traj.names
+            if startswith(string(name), string(state_name))
+    ]
+
     state_integrators = []
-    if length(ψ_inits) == 1
-        if piccolo_options.integrator == :pade
-            state_integrators = [QuantumStatePadeIntegrator(
-                state_name,
-                control_name,
-                sys,
-                traj;
-                order=piccolo_options.pade_order
-            )]
-        elseif piccolo_options.integrator == :exponential
-            state_integrators = [QuantumStateExponentialIntegrator(
-                state_name,
-                control_name,
-                sys,
-                traj
-            )]
-        else
-            error("integrator must be one of (:pade, :exponential)")
-        end
-    else
-        state_names = [
-            name for name ∈ traj.names
-                if startswith(string(name), string(state_name))
-        ]
-        state_integrators = []
-        for i = 1:length(ψ_inits)
-            if piccolo_options.integrator == :pade
-                state_integrator = QuantumStatePadeIntegrator(
-                    state_names[i],
-                    control_name,
-                    sys,
-                    traj;
-                    order=piccolo_options.pade_order
-                )
-            elseif piccolo_options.integrator == :exponential
-                state_integrator = QuantumStateExponentialIntegrator(
-                    state_names[i],
-                    control_name,
-                    sys,
-                    traj
-                )
-            else
-                error("integrator must be one of (:pade, :exponential)")
-            end
-            push!(state_integrators, state_integrator)
-        end
+
+    for name ∈ state_names
+        push!(state_integrators, KetIntegrator(sys, traj, name, control_name))
     end
 
     integrators = [
         state_integrators...,
-        DerivativeIntegrator(control_name, control_names[2], traj),
-        DerivativeIntegrator(control_names[2], control_names[3], traj)
+        DerivativeIntegrator(traj, control_name, control_names[2]),
+        DerivativeIntegrator(traj, control_names[2], control_names[3])
     ]
 
-    return QuantumControlProblem(
+    return DirectTrajOptProblem(
         traj,
         J,
         integrators;
         constraints=constraints,
-        ipopt_options=ipopt_options,
-        piccolo_options=piccolo_options,
-        control_name=control_name,
-        kwargs...
     )
 end