Fix free phase objectives

Project.toml

@@ -1,7 +1,7 @@
 name = "QuantumCollocation"
 uuid = "0dc23a59-5ffb-49af-b6bd-932a8ae77adf"
 authors = ["Aaron Trowbridge <[email protected]> and contributors"]
-version = "0.3.1"
+version = "0.3.2"
 
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"

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/isomorphisms.jl

@@ -59,11 +59,24 @@ Convert a real vector `Ũ⃗` into a complex matrix representing an operator.
 Must be differentiable.
 """
 function iso_vec_to_operator(Ũ⃗::AbstractVector{R}) where R
+    if !(R <: Real)
+        try
+            Ũ⃗ = convert(Vector{Real}, Ũ⃗)
+        catch 
+            throw(TypeError(
+                :iso_vec_to_operator,
+                "unable to convert input to a Real vector",
+                Real,
+                R
+            ))
+        end
+    end
+
     Ũ⃗_dim = div(length(Ũ⃗), 2)
     N = Int(sqrt(Ũ⃗_dim))
-    U = Matrix{complex(R)}(undef, N, N)
+    U = Matrix{complex(eltype(Ũ⃗))}(undef, N, N)
     for i=0:N-1
-        U[:, i+1] .= @view(Ũ⃗[i * 2N .+ (1:N)]) + one(R) * im * @view(Ũ⃗[i * 2N .+ (N+1:2N)])
+        U[:, i+1] .= @view(Ũ⃗[i * 2N .+ (1:N)]) + one(eltype(Ũ⃗)) * im * @view(Ũ⃗[i * 2N .+ (N+1:2N)])
     end
     return U
 end
@@ -99,8 +112,21 @@ Convert a complex matrix `U` representing an operator into a real vector.
 Must be differentiable.
 """
 function operator_to_iso_vec(U::AbstractMatrix{R}) where R
+    if !(R <: Number)
+        try
+            U = convert(Matrix{Complex}, U)
+        catch 
+            throw(TypeError(
+                :operator_to_iso_vec,
+                "unable to convert input to a Complex matrix",
+                Number,
+                R
+            ))
+        end
+    end
+
     N = size(U,1)
-    Ũ⃗ = Vector{real(R)}(undef, N^2 * 2)
+    Ũ⃗ = Vector{real(eltype(U))}(undef, N^2 * 2)
     for i=0:N-1
         Ũ⃗[i*2N .+ (1:N)] .= real(@view(U[:, i+1]))
         Ũ⃗[i*2N .+ (N+1:2N)] .= imag(@view(U[:, i+1]))
@@ -202,4 +228,32 @@ iso_dm(ρ::AbstractMatrix) = ket_to_iso(vec(ρ))
     @test iso_operator_to_iso_vec(iso_vec_to_iso_operator(iso_vec)) == iso_vec
 end
 
+@testitem "Test isomorphism type promotion" begin
+    # Check that generic types are converted to Complex
+    X = Any[0 1; 1 0]
+    @test operator_to_iso_vec(X) == [0; 1; 0; 0; 1; 0; 0; 0]
+
+    X = Number[0 1; 1 0]
+    @test operator_to_iso_vec(X) == [0; 1; 0; 0; 1; 0; 0; 0]
+
+
+    # Check that Any types are converted to Real
+    x = Any[0; 1; 0; 0; 1; 0; 0; 0]
+    @test iso_vec_to_operator(x) == [0 1; 1 0]
+
+    # Check that valid Complex can be converted to Real
+    x = Complex[0; 1; 0; 0; 1; 0; 0; 0]
+    @test iso_vec_to_operator(x) == [0 1; 1 0]
+
+    # Check that actual Complex throws an error
+    x = Complex[0; im; 0; 0; 1; 0; 0; 0]
+    @test_throws TypeError iso_vec_to_operator(x)
+
+    # Check the non-converting Isomorphisms
+    x = Any[0; 1; 0; 0; 1; 0; 0; 0]
+    X = Any[0 1 0 0; 1 0 0 0; 0 0 0 1; 0 0 1 0]
+    @test iso_vec_to_iso_operator(x) == X
+    @test iso_operator_to_iso_vec(X) == x
+end
+
 end

src/losses/_losses.jl

@@ -16,6 +16,9 @@ using ForwardDiff
 using Symbolics
 using TestItemRunner
 
+# For differentiable exponential
+using ExponentialAction
+
 # TODO:
 # - [ ] Do not reference the Z object in the loss (components only / remove "name")
 

src/losses/unitary_infidelity_loss.jl

@@ -30,23 +30,15 @@ where $n$ is the dimension of the unitary operators.
     U_goal::Matrix;
     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
 
+
 @doc raw"""
     iso_vec_unitary_fidelity(Ũ⃗::AbstractVector, Ũ⃗_goal::AbstractVector)
 
@@ -78,7 +70,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 +79,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 +100,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 +123,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 +141,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 +162,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 +238,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(
@@ -302,7 +297,11 @@ end
     phase_operators::AbstractVector{<:AbstractMatrix};
     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(
@@ -324,7 +323,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 +349,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)
@@ -379,7 +384,6 @@ struct UnitaryFreePhaseInfidelityLoss <: AbstractLoss
 
             # phase slice
             ∇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 +419,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
@@ -509,5 +518,36 @@ end
 
 
 @testitem "Isovec Unitary Fidelity Gradient" begin
+    @test_skip true
+end
 
+@testitem "Free phase fidelity" begin
+    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

@@ -119,8 +119,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,
@@ -151,10 +151,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
 
@@ -163,3 +168,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

src/problem_templates/_problem_templates.jl

@@ -1,5 +1,6 @@
 module ProblemTemplates
 
+using ..QuantumObjectUtils
 using ..QuantumSystems
 using ..Isomorphisms
 using ..EmbeddedOperators