Refactor/direct collocation backend

Project.toml

@@ -1,9 +1,10 @@
 name = "QuantumCollocation"
 uuid = "0dc23a59-5ffb-49af-b6bd-932a8ae77adf"
 authors = ["Aaron Trowbridge <[email protected]> and contributors"]
-version = "0.6.0"
+version = "0.7.0"
 
 [deps]
+DirectTrajOpt = "c823fa1f-8872-4af5-b810-2b9b72bbbf56"
 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"
@@ -21,14 +21,14 @@ TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe"
 TrajectoryIndexingUtils = "6dad8b7f-dd9a-4c28-9b70-85b9a079bfc8"
 
 [compat]
+DirectTrajOpt = "0.1.0"
 Distributions = "0.25"
 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

@@ -3,9 +3,6 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 LiveServer = "16fef848-5104-11e9-1b77-fb7a48bbb589"
 NamedTrajectories = "538bc3a1-5ab9-4fc3-b776-35ca1e893e08"
+PiccoloQuantumObjects = "5a402ddf-f93c-42eb-975e-5582dcda653d"
 QuantumCollocation = "0dc23a59-5ffb-49af-b6bd-932a8ae77adf"
-QuantumCollocationCore = "2b384925-53cb-4042-a8d2-6faa627467e1"
 Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
-
-[compat]
-NamedTrajectories = "0.2"

docs/literate/man/ipopt_callbacks.jl → docs/dev/ipopt_callbacks.jl

@@ -1,9 +1,9 @@
 # ```@meta
 # CollapsedDocStrings = true
 # ```
-# # IpOpt Callbacks
+# # IPOPT Callbacks
 
-# This page describes the callback functions that can be used with the IpOpt solver (in the future, may describe more general callback behavior).
+# This page describes the callback functions that can be used with the IPOPT solver (in the future, may describe more general callback behavior).
 
 # ## Callbacks
 
@@ -13,7 +13,7 @@ using NamedTrajectories
 import ..QuantumStateSmoothPulseProblem
 import ..Callbacks
 
-# By default, IpOpt callbacks are called at each optimization step with the following signature:
+# By default, IPOPT callbacks are called at each optimization step with the following signature:
 function full_argument_list_callback(
     alg_mod::Cint,
     iter_count::Cint,

docs/literate/man/problem_templates.jl

@@ -1,19 +1,197 @@
 # ```@meta
 # CollapsedDocStrings = true
 # ```
+using NamedTrajectories
+using PiccoloQuantumObjects
+using QuantumCollocation
 
-# # Problem Templates
+#=
+# 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, which stores all the parts of the optimal control problem.
 
-# ## Unitary Smooth Pulse Problem
+This page provides a brief overview of each problem template, broken down by the state of
+the problem being solved.
 
-# ```@docs; canonical = false
-# UnitarySmoothPulseProblem
-# ```
+Ket Problem Templates:
+- [Quantum State Smooth Pulse Problem](#Quantum-State-Smooth-Pulse-Problem)
+- [Quantum State Minimum Time Problem](#Quantum-State-Minimum-Time-Problem)
+- [Quantum State Sampling Problem](#Quantum-State-Sampling-Problem)
+
+Unitary Problem Templates:
+- [Unitary Smooth Pulse Problem](#Unitary-Smooth-Pulse-Problem)
+- [Unitary Minimum Time Problem](#Unitary-Minimum-Time-Problem)
+- [Unitary Sampling Problem](#Unitary-Sampling-Problem)
+=#
+
+# -----
+# ## Ket Problem Templates
+
+#=
+#### Quantum State Smooth Pulse Problem
+```@docs; canonical = false
+QuantumStateSmoothPulseProblem
+```
+
+Each problem starts with a `QuantumSystem` object, which is used to define the system's
+Hamiltonian and control operators. The goal is to find a control pulse that drives the
+intial state, `ψ_init`, to a target state, `ψ_goal`.
+=#
+
+# _define the quantum system_
+system = QuantumSystem(0.1 * PAULIS.Z, [PAULIS.X, PAULIS.Y])
+ψ_init = Vector{ComplexF64}([1.0, 0.0])
+ψ_goal = Vector{ComplexF64}([0.0, 1.0])
+T = 51
+Δt = 0.2
+
+# _create the smooth pulse problem_
+state_prob = QuantumStateSmoothPulseProblem(system, ψ_init, ψ_goal, T, Δt);
+
+# _check the fidelity before solving_
+println("Before: ", rollout_fidelity(state_prob.trajectory, system))
+
+# _solve the problem_
+solve!(state_prob, max_iter=100, verbose=true, print_level=1);
+
+# _check the fidelity after solving_
+println("After: ", rollout_fidelity(state_prob.trajectory, system))
+
+# _extract the control pulses_
+state_prob.trajectory.a |> size
+
+#=
+#### Quantum State Minimum Time Problem 
+```@docs; canonical = false
+QuantumStateMinimumTimeProblem
+```
+=#
+
+# _create the minimum time problem_
+min_state_prob = QuantumStateMinimumTimeProblem(state_prob, ψ_goal);
+
+# _check the previous duration_
+println("Duration before: ", get_duration(state_prob.trajectory))
+
+# _solve the minimum time problem_
+solve!(min_state_prob, max_iter=100, verbose=true, print_level=1);
+
+# _check the new duration_
+println("Duration after: ", get_duration(min_state_prob.trajectory))
+
+# _the fidelity is preserved by a constraint_
+println("Fidelity after: ", rollout_fidelity(min_state_prob.trajectory, system))
+
+#=
+#### Quantum State Sampling Problem
+```@docs; canonical = false
+QuantumStateSamplingProblem
+```
+=#
+
+# _create a sampling problem_
+driftless_system = QuantumSystem([PAULIS.X, PAULIS.Y])
+sampling_state_prob = QuantumStateSamplingProblem([system, driftless_system], ψ_init, ψ_goal, T, Δt);
+
+# _new keys are added to the trajectory for the new states_
+println(sampling_state_prob.trajectory.state_names)
+
+# _solve the sampling problem for a few iterations_
+solve!(sampling_state_prob, max_iter=25, verbose=true, print_level=1);
+
+# _check the fidelity of the sampling problem (use the updated key to get the initial and goal)_
+println("After (original system): ", rollout_fidelity(sampling_state_prob.trajectory, system, state_name=:ψ̃1_system_1))
+println("After (new system): ", rollout_fidelity(sampling_state_prob.trajectory, driftless_system, state_name=:ψ̃1_system_1))
+
+# _compare this to using the original problem on the new system_
+println("After (new system, original `prob`): ", rollout_fidelity(state_prob.trajectory, driftless_system))
+
+
+# -----
+# ## Unitary Problem Templates
+
+#=
+#### Unitary Smooth Pulse Problem
+
+```@docs; canonical = false
+UnitarySmoothPulseProblem
+```
+
+The `UnitarySmoothPulseProblem` is similar to the `QuantumStateSmoothPulseProblem`, but
+instead of driving the system to a target state, the goal is to drive the system to a
+target unitary operator, `U_goal`.
+
+=#
+
+system = QuantumSystem(0.1 * PAULIS.Z, [PAULIS.X, PAULIS.Y])
+U_goal = GATES.H
+T = 51
+Δt = 0.2
+
+prob = UnitarySmoothPulseProblem(system, U_goal, T, Δt); 
+
+# _check the fidelity before solving_
+println("Before: ", unitary_rollout_fidelity(prob.trajectory, system))
+
+# _finding an optimal control is as simple as calling `solve!`_
+solve!(prob, max_iter=100, verbose=true, print_level=1);
+
+# _check the fidelity after solving_
+println("After: ", unitary_rollout_fidelity(prob.trajectory, system))
+
+#=
+The `NamedTrajectory` object stores the control pulse, state variables, and the time grid.
+=#
+
+# _extract the control pulses_
+prob.trajectory.a |> size
+
+#=
+#### Unitary Minimum Time Problem
+
+```@docs; canonical = false
+UnitaryMinimumTimeProblem
+```
+
+The goal of this problem is to find the shortest time it takes to drive the system to a
+target unitary operator, `U_goal`. The problem is solved by minimizing the sum of all of 
+the time steps. It is constructed from `prob` in the previous example.
+=#
+
+min_prob = UnitaryMinimumTimeProblem(prob, U_goal);
+
+# _check the previous duration_
+println("Duration before: ", get_duration(prob.trajectory))
+
+# _solve the minimum time problem_
+solve!(min_prob, max_iter=100, verbose=true, print_level=1);
+
+# _check the new duration_
+println("Duration after: ", get_duration(min_prob.trajectory))
+
+# _the fidelity is preserved by a constraint_
+println("Fidelity after: ", unitary_rollout_fidelity(min_prob.trajectory, system))
+
+#=
+#### Unitary Sampling Problem 
+
+```@docs; canonical = false
+UnitarySamplingProblem
+```
+
+A sampling problem is used to solve over multiple quantum systems with the same control.
+This can be useful for exploring robustness, for example.
+=#
+
+# _create a sampling problem_
+driftless_system = QuantumSystem([PAULIS.X, PAULIS.Y])
+sampling_prob = UnitarySamplingProblem([system, driftless_system], U_goal, T, Δt);
+
+# _new keys are addded to the trajectory for the new states_
+println(sampling_prob.trajectory.state_names)
+
+# _the `solve!` proceeds as in the [Quantum State Sampling Problem](#Quantum-State-Sampling-Problem)]_
 
-# ## Unitary Minimum Time Problem
 
-# ```@docs; canonical = false
-# UnitaryMinimumTimeProblem
-# ```

docs/make.jl

@@ -16,12 +16,11 @@ open(joinpath(@__DIR__, "src", "index.md"), write = true) do io
     end
 end
 
+# TODO: Callbacks are currently broken
 pages = [
     "Home" => "index.md",
     "Manual" => [
         "Problem Templates" => "generated/man/problem_templates.md",
-        "Rollouts" => "generated/man/rollouts.md",
-        "Callbacks" => "generated/man/ipopt_callbacks.md",
     ],
     "Library" => "lib.md",
 ]