Merge pull request #23 from aarontrowbridge/dev-aaron

Dev aaron

.gitignore

@@ -5,3 +5,4 @@
 /examples/**/plots/
 /examples/**/trajectories/
 pardiso.lic
+/.CondaPkg/

Project.toml

@@ -11,6 +11,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Einsum = "b7d42ee7-0b51-5a75-98ca-779d3107e4c0"
 ExponentialAction = "e24c0720-ea99-47e8-929e-571b494574d3"
+FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
@@ -27,12 +28,16 @@ Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
 NamedTrajectories = "538bc3a1-5ab9-4fc3-b776-35ca1e893e08"
 PGFPlotsX = "8314cec4-20b6-5062-9cdb-752b83310925"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+PythonPlot = "274fc56d-3b97-40fa-a1cd-1b4a50311bf9"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
+SignalAnalysis = "df1fea92-c066-49dd-8b36-eace3378ea47"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 TrajectoryIndexingUtils = "6dad8b7f-dd9a-4c28-9b70-85b9a079bfc8"
+Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]

examples/cavity/binomial_code.ipynb

@@ -2,17 +2,83 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
     "using NamedTrajectories\n",
-    "using QuantumCollocation"
+    "using QuantumCollocation\n",
+    "using HDF5"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "🗂️ HDF5.File: (read-only) /home/aaron/projects/JuliaPackages/QuantumCollocation.jl/examples/cavity/binomial_T_210_dt_12.5_R_0.1_iter_470ubound_0.0002_00000.h5\n",
+       "├─ 🔢 T\n",
+       "├─ 🔢 controls\n",
+       "├─ 🔢 delta_t\n",
+       "├─ 🔢 infidelity\n",
+       "└─ 🔢 total_time"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "controls_guess_path = joinpath(@__DIR__, \"binomial_T_210_dt_12.5_R_0.1_iter_470ubound_0.0002_00000.h5\")\n",
+    "\n",
+    "controls_file = h5open(controls_guess_path, \"r\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "210×4 Matrix{Float64}:\n",
+       "  0.0           0.0           0.0           0.0\n",
+       "  5.68517e-18   6.13106e-19   2.26775e-19  -1.21695e-18\n",
+       "  0.0150596     0.00811163    0.0421521     0.0558669\n",
+       " -0.00266762    0.00305039   -0.0154662    -0.0178372\n",
+       " -0.00849901   -0.00362743    0.0102498    -0.000682986\n",
+       "  0.00554702   -0.0116777    -0.0113973     0.008071\n",
+       "  0.00456502    0.0100983    -0.0184981    -0.0140855\n",
+       " -0.00209359    0.0229059     0.0450125     0.0131819\n",
+       " -0.00034722   -0.00827824   -0.011745      0.00381549\n",
+       "  0.0108724    -0.0136292    -0.0285457     0.0103408\n",
+       "  ⋮                                        \n",
+       "  0.0001775     0.00741911   -0.000840204  -0.00448363\n",
+       " -0.00198899    0.0020506    -0.00186831    0.000781397\n",
+       " -0.0116963    -0.00664443    0.00878012    0.00120831\n",
+       " -0.00637918   -0.00149828   -0.0186637     0.0308962\n",
+       "  0.0736114    -0.0243741    -0.0349503    -0.0120484\n",
+       " -0.0056758     0.0144296    -0.00049975   -0.00313207\n",
+       " -0.00638694   -0.0100058     0.0170821    -0.0089121\n",
+       " -0.0587661     0.00797249    0.0299866     0.018494\n",
+       "  0.0           0.0           0.0           0.0"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "controls_file[\"controls\"][]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
    "metadata": {},
    "outputs": [
     {
@@ -49,7 +115,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [
     {
@@ -72,6 +138,25 @@
     "integrators = data[\"integrators\"]"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "200"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "traj.T"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 4,

examples/cavity/binomial_code.jl

@@ -13,18 +13,18 @@ max_iter = 2000
 
 
 
-data_path = "/home/aditya/oc_projects/QuantumCollocation.jl/examples/cavity/data/binomial_code/transmon_3_T_200_dt_15.0_Q_200.0_R_L1_10.0_max_iter_5000_dda_bound_1.0e-5_00000.jld2"
+# data_path = "/home/aditya/oc_projects/QuantumCollocation.jl/examples/cavity/data/binomial_code/transmon_3_T_200_dt_15.0_Q_200.0_R_L1_10.0_max_iter_5000_dda_bound_1.0e-5_00000.jld2"
 
 controls_guess_path = joinpath(@__DIR__, "binomial_T_210_dt_12.5_R_0.1_iter_470ubound_0.0002_00000.h5")
 
 controls_file = h5open(controls_guess_path, "r")
 
 T = read(controls_file, "T")
 Δt = read(controls_file, "delta_t")
-A = read(controls_file, "controls")
+A = read(controls_file, "controls") |> transpose |> collect
 
-data = load_problem(data_path; return_data=true)
-init_traj = data["trajectory"]
+# data = load_problem(data_path; return_data=true)
+# init_traj = data["trajectory"]
 
 #data_path = "/home/aditya/oc_projects/QuantumCollocation.jl/examples/cavity/data/binomial_code/transmon_3_T_200_dt_15.0_Q_200.0_R_L1_10.0_max_iter_5000_dda_bound_1.0e-5_00000.jld2"
 

examples/qram/pi_gate.jl

@@ -0,0 +1,136 @@
+using QuantumCollocation
+import QuantumCollocation: lift
+using NamedTrajectories
+using LinearAlgebra
+
+levels = 2
+
+qubits = 4
+
+α = [225.78, 100.33, 189.32, 172.15] * 1e-3 # GHz
+
+χ = Symmetric([
+    0 -5.10982939 -0.18457118 -0.50235316;
+    0       0     -0.94914758 -1.07618574;
+    0       0           0     -0.44607489;
+    0       0           0           0
+]) * 1e-3 # GHz
+
+â_dag = create(levels)
+â = annihilate(levels)
+
+lift(op, i) = lift(op, i, qubits; l=levels)
+
+# drift hamiltonian for ith qubit
+H_q(i) = -α[i] / 2 * lift(â_dag, i)^2 * lift(â, i)^2
+
+# drift interaction hamiltonian for ith and jth qubit
+H_c_ij(i, j) = χ[i, j] * lift(â_dag, i) * lift(â, i) * lift(â_dag, j) * lift(â, j)
+
+# drive hamiltonian for ith qubit, real part
+H_d_real(i) = 1 / 2 * (lift(â_dag, i) + lift(â, i))
+
+# drive hamiltonian for ith qubit, imaginary part
+H_d_imag(i) = 1im / 2 * (lift(â_dag, i) - lift(â, i))
+
+# total drift hamiltonian
+H_drift =
+    sum(H_q(i) for i = 1:qubits) +
+    sum(H_c_ij(i, j) for i = 1:qubits, j = 1:qubits if j > i)
+
+H_drift *= 2π
+
+# make vector of drive hamiltonians: [H_d_real(1), H_d_imag(1), H_d_real(2), ...]
+# there's probably a cleaner way to do this lol
+# H_drives = collect.(vec(vcat(
+#     transpose(Matrix{ComplexF64}.([H_d_real(i) for i = 1:qubits])),
+#     transpose(Matrix{ComplexF64}.([H_d_imag(i) for i = 1:qubits]))
+# )))
+
+# H_drives = Matrix{ComplexF64}.([H_d_real(1), H_d_imag(1), H_d_real(2), H_d_imag(2)])
+H_drives = Matrix{ComplexF64}.([H_d_real(2), H_d_imag(2)])
+H_drives .*= 2π
+
+# make quantum system
+system = QuantumSystem(H_drift, H_drives)
+
+# create goal unitary
+Id = 1.0I(levels)
+g = cavity_state(0, levels)
+e = cavity_state(1, levels)
+eg = e * g'
+ge = g * e'
+U_goal = Id ⊗ (eg + ge) ⊗ Id ⊗ Id
+
+# time parameters
+duration = 100.0 # ns
+T = 100
+Δt = duration / T
+Δt_max = 1.2 * Δt
+Δt_min = 0.1 * Δt
+
+# drive constraint: 20 MHz (linear units)
+a_bound = 20 * 1e-3 # GHz
+
+# pulse acceleration (used to control smoothness)
+dda_bound = 2e-3
+
+# maximum number of iterations
+max_iter = 500
+
+# warm start
+warm_start = false
+
+if warm_start
+    data_path = joinpath(@__DIR__, "data/limited_drives_T_100_dt_1.0_dda_0.001_a_0.02_max_iter_500_00000.jld2")
+    data = load_problem(data_path; return_data=true)
+    init_traj = data["trajectory"]
+    init_drives = init_traj.a
+    init_Δt = init_traj.Δt[end]
+end
+
+prob = UnitarySmoothPulseProblem(
+    system,
+    U_goal,
+    T,
+    warm_start ? init_Δt : Δt;
+    Δt_max=Δt_max,
+    Δt_min=Δt_min,
+    a_bound=a_bound,
+    dda_bound=dda_bound,
+    max_iter=max_iter,
+    a_guess=warm_start ? init_drives : nothing,
+)
+
+save_dir = joinpath(@__DIR__, "data")
+plot_dir = joinpath(@__DIR__, "plots")
+
+experiment_name = "limited_drives_T_$(T)_dt_$(Δt)_dda_$(dda_bound)_a_$(a_bound)_max_iter_$(max_iter)"
+
+save_path = generate_file_path("jld2", experiment_name, save_dir)
+plot_path = generate_file_path("png", experiment_name, plot_dir)
+
+# plot the initial guess for the wavefunction and controls
+plot(plot_path, prob.trajectory, [:a])
+
+solve!(prob)
+
+# plot the final solution for the wavefunction and controls
+plot(plot_path, prob.trajectory, [:a])
+
+A = prob.trajectory.a
+Δt = prob.trajectory.Δt
+
+Ũ⃗_final = unitary_rollout(A, Δt, system; integrator=exp)[:, end]
+final_fidelity = unitary_fidelity(Ũ⃗_final, prob.trajectory.goal.Ũ⃗)
+println("Final fidelity: $final_fidelity")
+
+duration = times(prob.trajectory)[end]
+
+info = Dict(
+    "final_fidelity" => final_fidelity,
+    "duration" => duration,
+)
+
+# save the solution
+save_problem(save_path, prob, info)

src/evaluators.jl

@@ -138,6 +138,7 @@ end
     σ::T,
     μ::AbstractVector{T}
 ) where T
+
     σ∂²Ls = σ * evaluator.objective.∂²L(Z⃗, evaluator.trajectory)
 
     for (k, σ∂²Lₖ) in enumerate(σ∂²Ls)

src/losses.jl

@@ -15,12 +15,14 @@ export structure
 export infidelity
 export unitary_infidelity
 
-using TrajectoryIndexingUtils
 using ..QuantumUtils
 using ..QuantumSystems
+using ..StructureUtils
+
 
 using NamedTrajectories
 
+using TrajectoryIndexingUtils
 using LinearAlgebra
 using SparseArrays
 using ForwardDiff
@@ -161,33 +163,34 @@ struct UnitaryInfidelityLoss <: AbstractLoss
         name::Symbol,
         Ũ⃗_goal::AbstractVector
     )
-        l = Ũ⃗ -> unitary_infidelity(Ũ⃗, Ũ⃗_goal)
-        ∇l = Ũ⃗ -> ForwardDiff.gradient(l, Ũ⃗)
+        # l = Ũ⃗ -> unitary_infidelity(Ũ⃗, Ũ⃗_goal)
+        # ∇l = Ũ⃗ -> ForwardDiff.gradient(l, Ũ⃗)
 
-        Symbolics.@variables Ũ⃗[1:length(Ũ⃗_goal)]
-        Ũ⃗ = collect(Ũ⃗)
+        # Symbolics.@variables Ũ⃗[1:length(Ũ⃗_goal)]
+        # Ũ⃗ = collect(Ũ⃗)
 
-        ∇²l_symbolic = Symbolics.sparsehessian(l(Ũ⃗), Ũ⃗)
-        K, J, _ = findnz(∇²l_symbolic)
-        kjs = collect(zip(K, J))
-        filter!(((k, j),) -> k ≤ j, kjs)
-        ∇²l_structure = kjs
+        # ∇²l_symbolic = Symbolics.sparsehessian(l(Ũ⃗), Ũ⃗)
+        # K, J, _ = findnz(∇²l_symbolic)
+        # kjs = collect(zip(K, J))
+        # filter!(((k, j),) -> k ≤ j, kjs)
+        # ∇²l_structure = kjs
 
-        ∇²l_expression = Symbolics.build_function(∇²l_symbolic, Ũ⃗)
-        ∇²l = eval(∇²l_expression[1])
+        # ∇²l_expression = Symbolics.build_function(∇²l_symbolic, Ũ⃗)
+        # ∇²l = eval(∇²l_expression[1])
 
-        # l = Ũ⃗ -> unitary_infidelity(Ũ⃗, Ũ⃗_goal)
-        # ∇l = Ũ⃗ -> ForwardDiff.gradient(l, Ũ⃗)
-        # ∇²l = Ũ⃗ -> ForwardDiff.hessian(l, Ũ⃗)
-        # Ũ⃗_dim = length(Ũ⃗_goal)
+        l = Ũ⃗ -> unitary_infidelity(Ũ⃗, Ũ⃗_goal)
+        ∇l = Ũ⃗ -> ForwardDiff.gradient(l, Ũ⃗)
+        ∇²l = Ũ⃗ -> ForwardDiff.hessian(l, Ũ⃗)
+        Ũ⃗_dim = length(Ũ⃗_goal)
 
-        # ∇²l_structure = []
+        # ∇²l_structure = loss_hessian_structure(∇²l, Ũ⃗_dim)
 
-        # for (i, j) ∈ Iterators.product(1:Ũ⃗_dim, 1:Ũ⃗_dim)
-        #     if i ≤ j
-        #         push!(∇²l_structure, (i, j))
-        #     end
-        # end
+        ∇²l_structure = []
+        for (i, j) ∈ Iterators.product(1:Ũ⃗_dim, 1:Ũ⃗_dim)
+            if i ≤ j
+                push!(∇²l_structure, (i, j))
+            end
+        end
 
         return new(l, ∇l, ∇²l, ∇²l_structure, name)
     end