Feature/variational

README.md

@@ -67,30 +67,7 @@
 ```
 where $\mathbf{Z}$ is a trajectory  containing states and controls, from [NamedTrajectories.jl](https://github.com/harmoniqs/NamedTrajectories.jl).
 
-### 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,
-```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 \\
-    \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 \\
-        \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 collocation*. For details of our implementation please see our award-winning IEEE QCE 2023 paper, [Direct Collocation for Quantum Optimal Control](https://arxiv.org/abs/2305.03261). If you use QuantumCollocation.jl in your work, please cite :raised_hands:!
-
-Problem templates give the user the ability to add other constraints and objective functions to this problem and solve it efficiently using [Ipopt.jl](https://github.com/jump-dev/Ipopt.jl) and [MathOptInterface.jl](https://github.com/jump-dev/MathOptInterface.jl) under the hood.
+For details of our implementation please see our IEEE QCE 2023 paper, [Direct Collocation for Quantum Optimal Control](https://arxiv.org/abs/2305.03261). If you use QuantumCollocation.jl in your work, please cite :raised_hands:!
 
 ## Installation
 

docs/Project.toml

@@ -0,0 +1,8 @@
+[deps]
+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"
+Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"

docs/dev/ipopt_callbacks.jl

@@ -0,0 +1,88 @@
+# ```@meta
+# CollapsedDocStrings = true
+# ```
+# # 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).
+
+# ## Callbacks
+
+using QuantumCollocation
+using NamedTrajectories
+
+import ..QuantumStateSmoothPulseProblem
+import ..Callbacks
+
+# 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,
+    obj_value::Float64,
+    inf_pr::Float64,
+    inf_du::Float64,
+    mu::Float64,
+    d_norm::Float64,
+    regularization_size::Float64,
+    alpha_du::Float64,
+    alpha_pr::Float64,
+    ls_trials::Cint,
+)
+    return true
+end
+
+# This gives the user access to some of the optimization state internals at each iteration.
+# A callback function with any subset of these arguments can be passed into the `solve!` function via the `callback` keyword argument see below.
+
+# The callback function can be used to stop the optimization early by returning `false`. The following callback when passed to `solve!` will stop the optimization after the first iteration:
+my_callback = (kwargs...) -> false
+
+# Single initial and target states
+# --------------------------------
+T = 50
+Δt = 0.2
+sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]])
+ψ_init =  Vector{ComplexF64}([1.0, 0.0])
+ψ_target =  Vector{ComplexF64}([0.0, 1.0])
+
+prob = QuantumStateSmoothPulseProblem(
+    sys, ψ_init, ψ_target, T, Δt;
+    ipopt_options=IpoptOptions(print_level=1), 
+    piccolo_options=PiccoloOptions(verbose=false)
+)
+
+
+# The callback function can be used to monitor the optimization progress, save intermediate results, or modify the optimization process.
+# For example, the following callback function saves the optimization trajectory at each iteration - this can be useful for debugging or plotting the optimization progress.
+# `trajectory_history_callback` from the `Callbacks` module
+callback, trajectory_history = QuantumCollocation.Callbacks.trajectory_history_callback(prob)
+solve!(prob, max_iter=20, callback=callback)
+
+# Save trajectory images into files which can be used to create a gif like the following: 
+for (iter, traj) in enumerate(trajectory_history)
+    str_index = lpad(iter, length(string(length(trajectory_history))), "0")
+    plot("./iteration-$str_index-trajectory.png", traj, [:ψ̃, :a], xlims=(-Δt, (T+5)*Δt), ylims=(ψ̃1 = (-2, 2), a = (-1.1, 1.1)))
+end
+
+# ![pulse optimization animation](../../assets/animation.gif)
+
+# Using a callback to get the best trajectory from all the optimization iterations
+sys2 = QuantumSystem(0.15 * GATES[:Z], [GATES[:X], GATES[:Y]])
+ψ_init2 =  Vector{ComplexF64}([0.0, 1.0])
+ψ_target2 =  Vector{ComplexF64}([1.0, 0.0])
+
+# Using other callbacks from the callback library 
+# --------------------------------
+# Callback used here is `best_rollout_fidelity_callback` which appends the best trajectories based on monotonically increasing fidelity of the rollout
+prob2 = QuantumStateSmoothPulseProblem(
+    sys2, ψ_init2, ψ_target2, T, Δt;
+    ipopt_options=IpoptOptions(print_level=1), 
+    piccolo_options=PiccoloOptions(verbose=false)
+)
+
+best_trajectory_callback, best_trajectory_list = best_rollout_fidelity_callback(prob2)
+solve!(prob2, max_iter=20, callback=best_trajectory_callback)
+# fidelity of the last iterate
+@show Losses.fidelity(prob2)
+
+# fidelity of the best iterate
+@show QuantumCollocation.fidelity(best_trajectory_list[end], prob2.system)

docs/literate/man/ket_problem_templates.jl

@@ -0,0 +1,97 @@
+# ```@meta
+# CollapsedDocStrings = true
+# ```
+using NamedTrajectories
+using PiccoloQuantumObjects
+using QuantumCollocation
+
+# -----
+
+#=
+## 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))
+
+# -----

docs/literate/man/unitary_problem_templates.jl

@@ -0,0 +1,96 @@
+# ```@meta
+# CollapsedDocStrings = true
+# ```
+using NamedTrajectories
+using PiccoloQuantumObjects
+using QuantumCollocation
+
+# -----
+
+#=
+## 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)]_
+
+# -----

docs/make.jl

@@ -0,0 +1,61 @@
+using QuantumCollocation
+using Documenter
+using Literate
+
+push!(LOAD_PATH, joinpath(@__DIR__, "..", "src"))
+
+@info "Building Documenter site for QuantumCollocation.jl"
+
+pages = [
+    "Home" => "index.md",
+    "Library" => [
+        "Ket Problem Templates" => "generated/man/ket_problem_templates.md",
+        "Unitary Problem Templates" => "generated/man/unitary_problem_templates.md",
+    ],
+]
+
+format = Documenter.HTML(;
+    prettyurls=get(ENV, "CI", "false") == "true",
+    canonical="https://docs.harmoniqs.co/QuantumCollocation.jl",
+    edit_link="main",
+    assets=String[],
+    mathengine = MathJax3(Dict(
+        :loader => Dict("load" => ["[tex]/physics"]),
+        :tex => Dict(
+            "inlineMath" => [["\$","\$"], ["\\(","\\)"]],
+            "tags" => "ams",
+            "packages" => [
+                "base",
+                "ams",
+                "autoload",
+                "physics"
+            ],
+        ),
+    )),
+)
+
+src = joinpath(@__DIR__, "src")
+lit = joinpath(@__DIR__, "literate")
+
+lit_output = joinpath(src, "generated")
+
+for (root, _, files) ∈ walkdir(lit), file ∈ files
+    splitext(file)[2] == ".jl" || continue
+    ipath = joinpath(root, file)
+    opath = splitdir(replace(ipath, lit=>lit_output))[1]
+    Literate.markdown(ipath, opath)
+end
+
+makedocs(;
+    modules=[QuantumCollocation],
+    authors="Aaron Trowbridge <[email protected]> and contributors",
+    sitename="QuantumCollocation.jl",
+    format=format,
+    pages=pages,
+    warnonly=true,
+)
+
+deploydocs(;
+    repo="github.com/harmoniqs/QuantumCollocation.jl.git",
+    devbranch="main",
+)