Dev aaron

src/dynamics.jl

@@ -3,6 +3,11 @@ module Dynamics
 export AbstractDynamics
 export QuantumDynamics
 
+export dynamics
+export dynamics_jacobian
+export dynamics_hessian_of_lagrangian
+export dynamics_components
+
 using ..QuantumSystems
 using ..QuantumUtils
 using ..StructureUtils
@@ -31,74 +36,51 @@ struct QuantumDynamics <: AbstractDynamics
     dim::Int
 end
 
-function QuantumDynamics(
-    integrators::Vector{<:AbstractIntegrator},
-    traj::NamedTrajectory;
-    verbose=false
-)
-    if verbose
-        println("        constructing knot point dynamics functions...")
-    end
-
-    free_time = traj.timestep isa Symbol
-
-    if free_time
-        @assert all([
-            !isnothing(state(integrator)) &&
-            !isnothing(controls(integrator)) &&
-            !isnothing(timestep(integrator))
-                for integrator ∈ integrators
-        ])
-    else
-        @assert all([
-            !isnothing(state(integrator)) &&
-            !isnothing(controls(integrator))
-                for integrator ∈ integrators
-        ])
-    end
-
-    for integrator ∈ integrators
-        if integrator isa QuantumIntegrator && controls(integrator) isa Tuple
-            drive_comps = [traj.components[s] for s ∈ integrator.drive_symb]
-            number_of_drives = sum(length.(drive_comps))
-            @assert number_of_drives == integrator.n_drives "number of drives ($(number_of_drives)) does not match number of drive terms in Hamiltonian ($(integrator.n_drives))"
-        end
-    end
-
+function dynamics_components(integrators::Vector{<:AbstractIntegrator})
     dynamics_comps = []
-    let comp_mark = 0
-        for integrator ∈ integrators
-            integrator_comps = (comp_mark + 1):(comp_mark + dim(integrator))
-            push!(dynamics_comps, integrator_comps)
-            comp_mark += dim(integrator)
-        end
+    comp_mark = 0
+    for integrator ∈ integrators
+        integrator_comps = (comp_mark + 1):(comp_mark + dim(integrator))
+        push!(dynamics_comps, integrator_comps)
+        comp_mark += dim(integrator)
     end
+    return dynamics_comps
+end
 
+function dynamics(
+    integrators::Vector{<:AbstractIntegrator},
+    traj::NamedTrajectory
+)
+    dynamics_comps = dynamics_components(integrators)
     dynamics_dim = dim(integrators)
-
     function f(zₜ, zₜ₊₁)
         δ = Vector{eltype(zₜ)}(undef, dynamics_dim)
         for (integrator, integrator_comps) ∈ zip(integrators, dynamics_comps)
             δ[integrator_comps] = integrator(zₜ, zₜ₊₁, traj)
         end
         return δ
     end
+    return f
+end
 
-    function ∂f(zₜ, zₜ₊₁)
 
-        ∂ = zeros(eltype(zₜ), dynamics_dim, 2traj.dim)
 
+function dynamics_jacobian(
+    integrators::Vector{<:AbstractIntegrator},
+    traj::NamedTrajectory
+)
+    dynamics_comps = dynamics_components(integrators)
+    dynamics_dim = dim(integrators)
+    free_time = traj.timestep isa Symbol
+    function ∂f(zₜ, zₜ₊₁)
+        ∂ = zeros(eltype(zₜ), dynamics_dim, 2traj.dim)
         for (integrator, integrator_comps) ∈ zip(integrators, dynamics_comps)
-
             if integrator isa QuantumIntegrator
-
                 if integrator.autodiff
-
                     ∂P(z1, z2) = ForwardDiff.jacobian(
                         zz -> integrator(zz[1:traj.dim], zz[traj.dim+1:end], traj),
                         [z1; z2]
                     )
-
                     ∂[integrator_comps, 1:2traj.dim] = ∂P(zₜ, zₜ₊₁)
                 else
                     if free_time
@@ -111,22 +93,17 @@ function QuantumDynamics(
                         ∂xₜf, ∂xₜ₊₁f, ∂uₜf =
                             Integrators.jacobian(integrator, zₜ, zₜ₊₁, traj)
                     end
-
                     ∂[integrator_comps, x_comps] = ∂xₜf
                     ∂[integrator_comps, x_comps .+ traj.dim] = ∂xₜ₊₁f
-
                     if u_comps isa Tuple
                         for (uᵢ_comps, ∂uₜᵢf) ∈ zip(u_comps, ∂uₜf)
                             ∂[integrator_comps, uᵢ_comps] = ∂uₜᵢf
                         end
                     else
                         ∂[integrator_comps, u_comps] = ∂uₜf
                     end
-
                 end
-
             elseif integrator isa DerivativeIntegrator
-
                 if free_time
                     x_comps, dx_comps, Δt_comps = comps(integrator, traj)
                     ∂xₜf, ∂xₜ₊₁f, ∂dxₜf, ∂Δtₜf =
@@ -136,7 +113,6 @@ function QuantumDynamics(
                     ∂xₜf, ∂xₜ₊₁f, ∂dxₜf =
                         Integrators.jacobian(integrator, zₜ, zₜ₊₁, traj)
                 end
-
                 ∂[integrator_comps, x_comps] = ∂xₜf
                 ∂[integrator_comps, x_comps .+ traj.dim] = ∂xₜ₊₁f
                 ∂[integrator_comps, dx_comps] = ∂dxₜf
@@ -147,29 +123,28 @@ function QuantumDynamics(
                 error("integrator type not supported: $(typeof(integrator))")
             end
         end
-
         return sparse(∂)
     end
+    return ∂f
+end
 
+function dynamics_hessian_of_lagrangian(
+    integrators::Vector{<:AbstractIntegrator},
+    traj::NamedTrajectory
+)
+    dynamics_comps = dynamics_components(integrators)
+    free_time = traj.timestep isa Symbol
     function μ∂²f(zₜ, zₜ₊₁, μₜ)
-
         μ∂² = zeros(eltype(zₜ), 2traj.dim, 2traj.dim)
-
         for (integrator, integrator_comps) ∈ zip(integrators, dynamics_comps)
-
             if integrator isa QuantumIntegrator
-
                 if integrator.autodiff
-
                     μ∂²P(z1, z2, μ) = ForwardDiff.hessian(
                         zz -> μ' * integrator(zz[1:traj.dim], zz[traj.dim+1:end], traj),
                         [z1; z2]
                     )
-
                     μ∂²[1:2traj.dim, 1:2traj.dim] = sparse(μ∂²P(zₜ, zₜ₊₁, μₜ[integrator_comps]))
-
                 else
-
                     if free_time
                         x_comps, u_comps, Δt_comps = comps(integrator, traj)
                         μ∂uₜ∂xₜf, μ∂²uₜf, μ∂Δtₜ∂xₜf, μ∂Δtₜ∂uₜf, μ∂²Δtₜf, μ∂xₜ₊₁∂uₜf, μ∂xₜ₊₁∂Δtₜf =
@@ -179,7 +154,6 @@ function QuantumDynamics(
                         μ∂uₜ∂xₜf, μ∂²uₜf, μ∂xₜ₊₁∂uₜf =
                             hessian_of_the_lagrangian(integrator, zₜ, zₜ₊₁, μₜ[integrator_comps], traj)
                     end
-
                     if u_comps isa Tuple
                         for (uᵢ_comps, μ∂uₜᵢ∂xₜf) ∈ zip(u_comps, μ∂uₜ∂xₜf)
                             μ∂²[x_comps, uᵢ_comps] += μ∂uₜᵢ∂xₜf
@@ -209,25 +183,63 @@ function QuantumDynamics(
                         μ∂²[Δt_comps, Δt_comps] .+= μ∂²Δtₜf
                     end
                 end
-
             elseif integrator isa DerivativeIntegrator
                 if free_time
                     x_comps, dx_comps, Δt_comps = comps(integrator, traj)
-
                     μ∂dxₜ∂Δtₜf = -μₜ[integrator_comps]
-
                     μ∂²[dx_comps, Δt_comps] += μ∂dxₜ∂Δtₜf
                 end
             end
         end
-
         return sparse(μ∂²)
     end
+    return μ∂²f
+end
+
+
+function QuantumDynamics(
+    integrators::Vector{<:AbstractIntegrator},
+    traj::NamedTrajectory;
+    verbose=false
+)
+    if verbose
+        println("        constructing knot point dynamics functions...")
+    end
+
+    free_time = traj.timestep isa Symbol
+
+    if free_time
+        @assert all([
+            !isnothing(state(integrator)) &&
+            !isnothing(controls(integrator))
+                for integrator ∈ integrators
+        ])
+    else
+        @assert all([
+            !isnothing(state(integrator)) &&
+            !isnothing(controls(integrator))
+                for integrator ∈ integrators
+        ])
+    end
+
+    for integrator ∈ integrators
+        if integrator isa QuantumIntegrator && controls(integrator) isa Tuple
+            drive_comps = [traj.components[s] for s ∈ integrator.drive_symb]
+            number_of_drives = sum(length.(drive_comps))
+            @assert number_of_drives == integrator.n_drives "number of drives ($(number_of_drives)) does not match number of drive terms in Hamiltonian ($(integrator.n_drives))"
+        end
+    end
+
+    f = dynamics(integrators, traj)
+    ∂f = dynamics_jacobian(integrators, traj)
+    μ∂²f = dynamics_hessian_of_lagrangian(integrators, traj)
 
     if verbose
         println("        determining dynamics derivative structure...")
     end
 
+    dynamics_dim = dim(integrators)
+
     ∂f_structure, ∂F_structure, μ∂²f_structure, μ∂²F_structure =
         dynamics_structure(∂f, μ∂²f, traj, dynamics_dim)
 
@@ -275,7 +287,15 @@ function QuantumDynamics(
         return μ∂²s
     end
 
-    return QuantumDynamics(integrators, F, ∂F, ∂F_structure, μ∂²F, μ∂²F_structure, dynamics_dim)
+    return QuantumDynamics(
+        integrators,
+        F,
+        ∂F,
+        ∂F_structure,
+        μ∂²F,
+        μ∂²F_structure,
+        dynamics_dim
+    )
 end
 
 QuantumDynamics(P::AbstractIntegrator, traj::NamedTrajectory; kwargs...) =

src/losses.jl

@@ -43,16 +43,66 @@ abstract type AbstractLoss end
 function infidelity(ψ̃::AbstractVector, ψ̃goal::AbstractVector)
     ψ = iso_to_ket(ψ̃)
     ψgoal = iso_to_ket(ψ̃goal)
-    return abs(1 - abs2(ψ'ψgoal))
+    return abs(1 - abs2(ψgoal'ψ))
 end
 
-function unitary_infidelity(Ũ⃗::AbstractVector, Ũ⃗_goal::AbstractVector)
-    U = iso_vec_to_operator(Ũ⃗)
-    Ugoal = iso_vec_to_operator(Ũ⃗_goal)
-    N = size(U, 1)
-    return abs(1 - 1 / N * abs(tr(Ugoal'U)))
+@inline @views function isovec_unitary_fidelity(Ũ⃗::AbstractVector, Ũ⃗_goal::AbstractVector)
+    n = Int(sqrt(length(Ũ⃗) ÷ 2))
+    U⃗ᵣ = Ũ⃗[1:end ÷ 2]
+    U⃗ᵢ = Ũ⃗[end ÷ 2 + 1:end]
+    Ū⃗ᵣ = Ũ⃗_goal[1:end ÷ 2]
+    Ū⃗ᵢ = Ũ⃗_goal[end ÷ 2 + 1:end]
+    Tᵣ = Ū⃗ᵣ' * U⃗ᵣ + Ū⃗ᵢ' * U⃗ᵢ
+    Tᵢ = Ū⃗ᵣ' * U⃗ᵢ - Ū⃗ᵢ' * U⃗ᵣ
+    return 1 / n * sqrt(Tᵣ^2 + Tᵢ^2)
+end
+
+@views function unitary_infidelity(Ũ⃗::AbstractVector, Ũ⃗_goal::AbstractVector)
+    ℱ = isovec_unitary_fidelity(Ũ⃗, Ũ⃗_goal)
+    return abs(1 - ℱ)
+end
+
+@views function unitary_infidelity_gradient(Ũ⃗::AbstractVector, Ũ⃗_goal::AbstractVector)
+    n = Int(sqrt(length(Ũ⃗) ÷ 2))
+    U⃗ᵣ = Ũ⃗[1:end ÷ 2]
+    U⃗ᵢ = Ũ⃗[end ÷ 2 + 1:end]
+    Ū⃗ᵣ = Ũ⃗_goal[1:end ÷ 2]
+    Ū⃗ᵢ = Ũ⃗_goal[end ÷ 2 + 1:end]
+    Tᵣ = Ū⃗ᵣ' * U⃗ᵣ + Ū⃗ᵢ' * U⃗ᵢ
+    Tᵢ = Ū⃗ᵣ' * U⃗ᵢ - Ū⃗ᵢ' * U⃗ᵣ
+    ℱ = 1 / n * sqrt(Tᵣ^2 + Tᵢ^2)
+    ∇ᵣℱ = 1 / (n^2 * ℱ) * (Tᵣ * Ū⃗ᵣ - Tᵢ * Ū⃗ᵢ)
+    ∇ᵢℱ = 1 / (n^2 * ℱ) * (Tᵣ * Ū⃗ᵢ + Tᵢ * Ū⃗ᵣ)
+    ∇ℱ = [∇ᵣℱ; ∇ᵢℱ]
+    return -sign(1 - ℱ) * ∇ℱ
 end
 
+@views function unitary_infidelity_hessian(Ũ⃗::AbstractVector, Ũ⃗_goal::AbstractVector)
+    n = Int(sqrt(length(Ũ⃗) ÷ 2))
+    U⃗ᵣ = Ũ⃗[1:end ÷ 2]
+    U⃗ᵢ = Ũ⃗[end ÷ 2 + 1:end]
+    Ū⃗ᵣ = Ũ⃗_goal[1:end ÷ 2]
+    Ū⃗ᵢ = Ũ⃗_goal[end ÷ 2 + 1:end]
+    Tᵣ = Ū⃗ᵣ' * U⃗ᵣ + Ū⃗ᵢ' * U⃗ᵢ
+    Tᵢ = Ū⃗ᵣ' * U⃗ᵢ - Ū⃗ᵢ' * U⃗ᵣ
+    Wᵣᵣ = Ū⃗ᵣ * Ū⃗ᵣ'
+    Wᵣᵢ = Ū⃗ᵣ * Ū⃗ᵢ'
+    Wᵢᵢ = Ū⃗ᵢ * Ū⃗ᵢ'
+    ℱ = 1 / n * sqrt(Tᵣ^2 + Tᵢ^2)
+    ∇ᵣℱ = 1 / (n^2 * ℱ) * (Tᵣ * Ū⃗ᵣ - Tᵢ * Ū⃗ᵢ)
+    ∇ᵢℱ = 1 / (n^2 * ℱ) * (Tᵣ * Ū⃗ᵢ + Tᵢ * Ū⃗ᵣ)
+    ∇ℱ = [∇ᵣℱ; ∇ᵢℱ]
+    ∂ᵣ²ℱ = 1 / ℱ * (∇ᵣℱ * ∇ᵣℱ' + 1 / n^2 * (Wᵣᵣ + Wᵢᵢ))
+    ∂ᵢ²ℱ = 1 / ℱ * (∇ᵢℱ * ∇ᵢℱ' + 1 / n^2 * (Wᵣᵣ + Wᵢᵢ))
+    ∂ᵣ∂ᵢℱ = 1 / ℱ * (∇ᵣℱ * ∇ᵢℱ' + 1 / n^2 * (Wᵣᵢ - Wᵣᵢ'))
+    ∂²ℱ = [∂ᵣ²ℱ ∂ᵣ∂ᵢℱ; ∂ᵣ∂ᵢℱ' ∂ᵢ²ℱ]
+    return -sign(1 - ℱ) * (∇ℱ * ∇ℱ' + ∂²ℱ)
+end
+
+
+
+
+
 function unitary_trace_loss(Ũ⃗::AbstractVector, Ũ⃗_goal::AbstractVector)
     U = iso_vec_to_operator(Ũ⃗)
     Ugoal = iso_vec_to_operator(Ũ⃗_goal)
@@ -163,35 +213,16 @@ struct UnitaryInfidelityLoss <: AbstractLoss
         name::Symbol,
         Ũ⃗_goal::AbstractVector
     )
-        # l = Ũ⃗ -> unitary_infidelity(Ũ⃗, Ũ⃗_goal)
-        # ∇l = Ũ⃗ -> ForwardDiff.gradient(l, Ũ⃗)
-
-        # 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_expression = Symbolics.build_function(∇²l_symbolic, Ũ⃗)
-        # ∇²l = eval(∇²l_expression[1])
-
         l = Ũ⃗ -> unitary_infidelity(Ũ⃗, Ũ⃗_goal)
-        ∇l = Ũ⃗ -> ForwardDiff.gradient(l, Ũ⃗)
-        ∇²l = Ũ⃗ -> ForwardDiff.hessian(l, Ũ⃗)
+        ∇l = Ũ⃗ -> unitary_infidelity_gradient(Ũ⃗, Ũ⃗_goal)
+        ∇²l = Ũ⃗ -> unitary_infidelity_hessian(Ũ⃗, Ũ⃗_goal)
         Ũ⃗_dim = length(Ũ⃗_goal)
-
-        # ∇²l_structure = loss_hessian_structure(∇²l, Ũ⃗_dim)
-
         ∇²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
 end
@@ -202,7 +233,6 @@ function (loss::UnitaryInfidelityLoss)(
     hessian=false
 )
     @assert !(gradient && hessian)
-
     if !(gradient || hessian)
         return loss.l(Ũ⃗_end)
     elseif gradient