cleaning and renaming

src/constraints/fidelity_constraint.jl

@@ -286,8 +286,8 @@ function FinalUnitaryFreePhaseFidelityConstraint(;
 end
 
 function FinalUnitaryFreePhaseFidelityConstraint(
-    state_symbol::Symbol,
-    global_symbol::Symbol,
+    state_name::Symbol,
+    phase_name::Symbol,
     phase_operators::AbstractVector{<:AbstractMatrix{<:Complex}},
     val::Float64,
     traj::NamedTrajectory;
@@ -296,9 +296,9 @@ function FinalUnitaryFreePhaseFidelityConstraint(
 )
     return FinalUnitaryFreePhaseFidelityConstraint(;
         value=val,
-        state_slice=slice(traj.T, traj.components[state_symbol], traj.dim),
-        phase_slice=trajectory.global_components[global_symbol],
-        goal=traj.goal[state_symbol],
+        state_slice=slice(traj.T, traj.components[state_name], traj.dim),
+        phase_slice=traj.global_components[phase_name],
+        goal=traj.goal[state_name],
         phase_operators=phase_operators,
         zdim=length(traj),
         subspace=subspace,

src/integrators/exponential_integrators.jl

@@ -149,18 +149,18 @@ end
 
     Id = I(ℰ.ketdim)
 
-    expĜₜ = Id ⊗ exp_eigen(Δtₜ * Gₜ)
+    expGₜ = exp_eigen(Δtₜ * Gₜ)
 
     ∂Ũ⃗ₜ₊₁ℰ = sparse(I, ℰ.dim, ℰ.dim)
-    ∂Ũ⃗ₜℰ = -expĜₜ
+    ∂Ũ⃗ₜℰ = -Id ⊗ expGₜ
 
     ∂aₜℰ = ForwardDiff.jacobian(
         a -> -expv(Δtₜ, Id ⊗ ℰ.G(a), Ũ⃗ₜ),
         aₜ
     )
 
     if ℰ.freetime
-        ∂Δtₜℰ = -(Id ⊗ Gₜ) * (expĜₜ * Ũ⃗ₜ)
+        ∂Δtₜℰ = -(Id ⊗ (Gₜ * expGₜ)) * Ũ⃗ₜ
         return ∂Ũ⃗ₜℰ, ∂Ũ⃗ₜ₊₁ℰ, ∂aₜℰ, ∂Δtₜℰ
     else
         return ∂Ũ⃗ₜℰ, ∂Ũ⃗ₜ₊₁ℰ, ∂aₜℰ

src/losses/_losses.jl

@@ -13,6 +13,7 @@ using SparseArrays
 using ForwardDiff
 using Symbolics
 using TestItemRunner
+using ExponentialAction
 
 # TODO:
 # - [ ] Do not reference the Z object in the loss (components only / remove "name")

src/losses/unitary_infidelity_loss.jl

@@ -26,24 +26,15 @@ where $n$ is the dimension of the unitary operators.
 - `subspace::AbstractVector{Int}`: The subspace to calculate the fidelity over.
 """
 @views function unitary_fidelity(
-    U::Matrix,
-    U_goal::Matrix;
+    U::AbstractMatrix,
+    U_goal::AbstractMatrix;
     subspace::AbstractVector{Int}=axes(U_goal, 1)
 )
-    U_goal = U_goal[subspace, subspace]
-    U = U[subspace, subspace]
-    return 1 / size(U_goal, 1) * abs(tr(U_goal'U))
-end
-
-@views function unitary_fidelity(
-    Ũ⃗::AbstractVector,
-    Ũ⃗_goal::AbstractVector;
-    kwargs...
-)
-    return unitary_fidelity(
-        iso_vec_to_operator(Ũ⃗),
-        iso_vec_to_operator(Ũ⃗_goal);
-        kwargs...
+    # avoids type coercion neceessary because tr(Matrix{Any}) fails
+    return iso_vec_unitary_fidelity(
+        operator_to_iso_vec(U),
+        operator_to_iso_vec(U_goal),
+        subspace=subspace
     )
 end
 
@@ -78,7 +69,7 @@ where $T_R = \langle \vec{\widetilde{U}}_{\text{goal}, R}, \vec{\widetilde{U}}_R
     return 1 / n * sqrt(Tᵣ^2 + Tᵢ^2)
 end
 
-@views function unitary_infidelity(
+@views function iso_vec_unitary_infidelity(
     Ũ⃗::AbstractVector,
     Ũ⃗_goal::AbstractVector;
     subspace::AbstractVector{Int}=axes(iso_vec_to_operator(Ũ⃗_goal), 1)
@@ -87,7 +78,7 @@ end
     return abs(1 - ℱ)
 end
 
-@views function unitary_infidelity_gradient(
+@views function iso_vec_unitary_infidelity_gradient(
     Ũ⃗::AbstractVector,
     Ũ⃗_goal::AbstractVector;
     subspace::AbstractVector{Int}=axes(iso_vec_to_operator(Ũ⃗_goal), 1)
@@ -108,7 +99,7 @@ end
     return -sign(1 - ℱ) * ∇ℱ
 end
 
-@views function unitary_infidelity_hessian(
+@views function iso_vec_unitary_infidelity_hessian(
     Ũ⃗::AbstractVector,
     Ũ⃗_goal::AbstractVector;
     subspace::AbstractVector{Int}=axes(iso_vec_to_operator(Ũ⃗_goal), 1)
@@ -131,6 +122,7 @@ end
     ∂ᵢ²ℱ = 1 / ℱ * (-∇ᵢℱ * ∇ᵢℱ' + 1 / n^2 * (Wᵣᵣ + Wᵢᵢ))
     ∂ᵣ∂ᵢℱ = 1 / ℱ * (-∇ᵢℱ * ∇ᵣℱ' + 1 / n^2 * (Wᵢᵣ - Wᵣᵢ))
     ∂²ℱ = [∂ᵣ²ℱ ∂ᵣ∂ᵢℱ; ∂ᵣ∂ᵢℱ' ∂ᵢ²ℱ]
+    # TODO: This should be moved to Isomorphisms.jl
     permutation = vcat(vcat([[slice(j, n), slice(j, n) .+ n^2] for j = 1:n]...)...)
     ∂²ℱ = ∂²ℱ[permutation, permutation]
     return -sign(1 - ℱ) * ∂²ℱ
@@ -148,13 +140,13 @@ struct UnitaryInfidelityLoss <: AbstractLoss
         Ũ⃗_goal::AbstractVector;
         subspace::AbstractVector{Int}=axes(iso_vec_to_operator(Ũ⃗_goal), 1)
     )
-        l = Ũ⃗ -> unitary_infidelity(Ũ⃗, Ũ⃗_goal, subspace=subspace)
+        l = Ũ⃗ -> iso_vec_unitary_infidelity(Ũ⃗, Ũ⃗_goal, subspace=subspace)
 
         @views ∇l = Ũ⃗ -> begin
             subspace_rows, subspace_cols = (subspace, subspace)
             n_subspace = length(subspace_rows)
             n_full = Int(sqrt(length(Ũ⃗) ÷ 2))
-            ∇l_subspace = unitary_infidelity_gradient(Ũ⃗, Ũ⃗_goal, subspace=subspace)
+            ∇l_subspace = iso_vec_unitary_infidelity_gradient(Ũ⃗, Ũ⃗_goal, subspace=subspace)
             ∇l_full = zeros(2 * n_full^2)
             for j ∈ eachindex(subspace_cols)
                 ∇l_full[slice(2subspace_cols[j] - 1, subspace_rows, n_full)] =
@@ -169,7 +161,7 @@ struct UnitaryInfidelityLoss <: AbstractLoss
             subspace_rows, subspace_cols = (subspace, subspace)
             n_subspace = length(subspace_rows)
             n_full = Int(sqrt(length(Ũ⃗) ÷ 2))
-            ∇²l_subspace = unitary_infidelity_hessian(Ũ⃗, Ũ⃗_goal, subspace=subspace)
+            ∇²l_subspace = iso_vec_unitary_infidelity_hessian(Ũ⃗, Ũ⃗_goal, subspace=subspace)
             ∇²l_full = zeros(2 * n_full^2, 2 * n_full^2)
             # NOTE: Assumes subspace_rows = subspace_cols
             for k ∈ eachindex(subspace_cols)
@@ -245,7 +237,9 @@ function free_phase(
     Hs::AbstractVector{<:AbstractMatrix}
 )
     # NOTE: switch to expv for ForwardDiff
-    return reduce(kron, [exp(im * ϕ * H) for (ϕ, H) ∈ zip(ϕs, Hs)])
+    # return reduce(kron, [exp(im * ϕ * H) for (ϕ, H) ∈ zip(ϕs, Hs)])
+    Id = Matrix{eltype(Hs[1])}(I, size(Hs[1]))
+    return reduce(kron, [expv(im * ϕ, H, Id) for (ϕ, H) ∈ zip(ϕs, Hs)])
 end
 
 function free_phase_gradient(
@@ -284,10 +278,10 @@ end
 end
 
 @views function iso_vec_unitary_free_phase_fidelity(
-    Ũ⃗::AbstractVector,
-    Ũ⃗_goal::AbstractVector,
-    phases::AbstractVector,
-    phase_operators::AbstractVector{<:AbstractMatrix};
+    Ũ⃗::AbstractVector{<:Real},
+    Ũ⃗_goal::AbstractVector{<:Real},
+    phases::AbstractVector{<:Real},
+    phase_operators::AbstractVector{<:AbstractMatrix{ComplexF64}};
     subspace::AbstractVector{Int}=axes(iso_vec_to_operator(Ũ⃗_goal), 1)
 )
     U = iso_vec_to_operator(Ũ⃗)
@@ -296,20 +290,24 @@ end
 end
 
 @views function iso_vec_unitary_free_phase_infidelity(
-    Ũ⃗::AbstractVector,
-    Ũ⃗_goal::AbstractVector,
-    phases::AbstractVector,
-    phase_operators::AbstractVector{<:AbstractMatrix};
+    Ũ⃗::AbstractVector{<:Real},
+    Ũ⃗_goal::AbstractVector{<:Real},
+    phases::AbstractVector{<:Real},
+    phase_operators::AbstractVector{<:AbstractMatrix{ComplexF64}};
     subspace::AbstractVector{Int}=axes(iso_vec_to_operator(Ũ⃗_goal), 1)
 )
-    return 1 - iso_vec_unitary_free_phase_fidelity(Ũ⃗, Ũ⃗_goal, phases, phase_operators, subspace=subspace)
+    ℱ = iso_vec_unitary_free_phase_fidelity(
+        Ũ⃗, Ũ⃗_goal, phases, phase_operators,
+        subspace=subspace
+    )
+    return abs(1 - ℱ)
 end
 
 @views function iso_vec_unitary_free_phase_infidelity_gradient(
-    Ũ⃗::AbstractVector,
-    Ũ⃗_goal::AbstractVector,
-    phases::AbstractVector,
-    phase_operators::AbstractVector{<:AbstractMatrix};
+    Ũ⃗::AbstractVector{<:Real},
+    Ũ⃗_goal::AbstractVector{<:Real},
+    phases::AbstractVector{<:Real},
+    phase_operators::AbstractVector{<:AbstractMatrix{ComplexF64}};
     subspace::AbstractVector{Int}=axes(iso_vec_to_operator(Ũ⃗_goal), 1)
 )
     n_phases = length(phases)
@@ -324,7 +322,9 @@ end
     R[subspace, subspace] = free_phase(phases, phase_operators)
 
     # loss gradient in subspace
-    ∂ℱ_∂Ũ⃗_subspace = unitary_infidelity_gradient(operator_to_iso_vec(R * U), Ũ⃗_goal, subspace=subspace)
+    ∂ℱ_∂Ũ⃗_subspace = iso_vec_unitary_infidelity_gradient(
+        operator_to_iso_vec(R * U), Ũ⃗_goal, subspace=subspace
+    )
 
     # state slice in subspace
     ∂[1:n_subspace] = operator_to_iso_vec(R[subspace, subspace]'iso_vec_to_operator(∂ℱ_∂Ũ⃗_subspace))
@@ -348,16 +348,20 @@ struct UnitaryFreePhaseInfidelityLoss <: AbstractLoss
     function UnitaryFreePhaseInfidelityLoss(
         Ũ⃗_goal::AbstractVector,
         phase_operators::AbstractVector{<:AbstractMatrix};
-        subspace::AbstractVector{Int}=axes(iso_vec_to_operator(Ũ⃗_goal), 1),
+        subspace::Union{AbstractVector{Int}, Nothing}=nothing,
     )
-        @assert reduce(.*, size.(phase_operators)) == length.(subspace) "phase operators must span the subspace"
+        if isnothing(subspace)
+            subspace = axes(iso_vec_to_operator(Ũ⃗_goal), 1)
+        end
+
+        @assert reduce(*, size.(phase_operators, 1)) == length(subspace) "phase operators must span the subspace"
 
         @views function l(Ũ⃗::AbstractVector, ϕ⃗::AbstractVector)
             return iso_vec_unitary_free_phase_infidelity(Ũ⃗, Ũ⃗_goal, ϕ⃗, phase_operators, subspace=subspace)
         end
 
         @views function ∇l(Ũ⃗::AbstractVector, ϕ⃗::AbstractVector)
-            subspace_rows, subspace_cols = subspace
+            subspace_rows = subspace_cols = subspace
             n_phase = length(ϕ⃗)
             n_rows = length(subspace_rows)
             n_cols = length(subspace_cols)
@@ -378,7 +382,8 @@ struct UnitaryFreePhaseInfidelityLoss <: AbstractLoss
             end
 
             # phase slice
-            ∇l_full[2 * n_full^2 .+ (1:n_phase)] = ∇l_subspace[2 * n_rows * n_cols .+ (1:n_phase)]
+            ∇l_full[2 * n_full^2 .+ (1:n_phase)] =
+                ∇l_subspace[2 * n_rows * n_cols .+ (1:n_phase)]
 
             return ∇l_full
         end
@@ -415,6 +420,11 @@ end
     Y = [0 -im; im 0]
     @test unitary_fidelity(X, X) ≈ 1
     @test unitary_fidelity(X, Y) ≈ 0
+
+    # Tr undefined on Type{Any}
+    # X = Any[0 1; 1 0]
+    # Y = Complex[0 -im; im 0]
+    # @test unitary_fidelity(X, Y) ≈ 0
 end
 
 @testitem "Isovec Unitary Fidelity" begin
@@ -511,5 +521,37 @@ end
 
 
 @testitem "Isovec Unitary Fidelity Gradient" begin
+    @test_skip true
+end
 
+@testitem "Free phase fidelity" begin
+    using PiccoloQuantumObjects
+    n_levels = 3
+    phase_data = [1.9, 2.7]
+    phase_operators = [PAULIS[:Z], PAULIS[:Z]]
+    subspace = get_subspace_indices([1:2, 1:2], [n_levels, n_levels])
+
+    R = Losses.free_phase(phase_data, phase_operators)
+    @test R'R ≈ [1 0 0 0; 0 1 0 0; 0 0 1 0; 0 0 0 1]
+    @test size(R) == (2^2, 2^2)
+
+    U_goal = EmbeddedOperator(GATES[:CZ], subspace, [n_levels, n_levels]).operator
+    U_final = EmbeddedOperator(R'GATES[:CZ], subspace, [n_levels, n_levels]).operator
+    # Value is ~0.3 without phases
+    @test unitary_fidelity(U_final, U_goal, subspace=subspace) < 0.5
+    @test unitary_free_phase_fidelity(
+        U_final,
+        U_goal,
+        phase_data,
+        phase_operators,
+        subspace=subspace
+    ) ≈ 1
+
+    # Forgot subspace
+    @test_throws DimensionMismatch iso_vec_unitary_free_phase_fidelity(
+        operator_to_iso_vec(U_final),
+        operator_to_iso_vec(U_goal),
+        phase_data,
+        phase_operators,
+    )
 end

src/objectives/unitary_infidelity_objective.jl

@@ -120,8 +120,8 @@ Fields:
 """
 function UnitaryFreePhaseInfidelityObjective(;
     name::Union{Nothing,Symbol}=nothing,
-    phase_name::Union{Nothing,Symbol}=nothing,
     goal::Union{Nothing,AbstractVector{<:R}}=nothing,
+    phase_name::Union{Nothing,Symbol}=nothing,
     phase_operators::Union{Nothing,AbstractVector{<:AbstractMatrix{<:Complex{R}}}}=nothing,
 	Q::R=1.0,
 	eval_hessian::Bool=false,
@@ -152,10 +152,15 @@ function UnitaryFreePhaseInfidelityObjective(;
     end
 
     @views function ∇L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
-        # TODO: implement analytic
-        ∇ = ForwardDiff.gradient(Z⃗ -> L(Z⃗, Z), Z⃗)
-        # println(nnz(sparse(∇)))
-        # println(∇[Z.global_components[global_name]])
+        ∇ = zeros(Z.dim * Z.T + Z.global_dim)
+        Ũ⃗_slice = slice(Z.T, Z.components[name], Z.dim)
+        Ũ⃗ = Z⃗[Ũ⃗_slice]
+        ϕ⃗_slice = Z.global_components[phase_name]
+        ϕ⃗ = Z⃗[ϕ⃗_slice]
+        ∇l = l(Ũ⃗, ϕ⃗; gradient=true)
+        ∇[Ũ⃗_slice] = Q * ∇l[1:length(Ũ⃗)]
+        # WARNING: 2π periodic; using Q≠1 is not recommended
+        ∇[ϕ⃗_slice] = Q * ∇l[length(Ũ⃗) .+ (1:length(ϕ⃗))]
         return ∇
     end
 
@@ -164,3 +169,23 @@ function UnitaryFreePhaseInfidelityObjective(;
 
 	return Objective(L, ∇L, ∂²L, ∂²L_structure, Dict[params])
 end
+
+function UnitaryFreePhaseInfidelityObjective(
+    name::Symbol,
+    phase_name::Symbol,
+    phase_operators::AbstractVector{<:AbstractMatrix{<:Complex}},
+    traj::NamedTrajectory,
+    Q::Float64;
+    subspace=nothing,
+	eval_hessian::Bool=true
+)
+    return UnitaryFreePhaseInfidelityObjective(
+        name=name,
+        goal=traj.goal[name],
+        phase_name=phase_name,
+        phase_operators=phase_operators,
+        Q=Q,
+        subspace=subspace,
+        eval_hessian=eval_hessian,
+    )
+end