phase rollout fidelity

src/losses/unitary_infidelity_loss.jl

@@ -30,9 +30,12 @@ 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))
+    # 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
 
 
@@ -240,7 +243,6 @@ function free_phase(
     return reduce(kron, [expv(im * ϕ, H, Id) for (ϕ, H) ∈ zip(ϕs, Hs)])
 end
 
-# TODO: in-place
 function free_phase_gradient(
     ϕs::AbstractVector,
     Hs::AbstractVector{<:AbstractMatrix}

src/rollouts.jl

@@ -82,6 +82,7 @@ function rollout(
 )
     T = size(controls, 2)
 
+    # Real type enables ForwardDiff
     Ψ̃ = zeros(Real, length(ψ̃_init), T)
 
     Ψ̃[:, 1] .= ψ̃_init
@@ -179,6 +180,7 @@ function open_rollout(
 )
     T = size(controls, 2)
 
+    # Real type enables ForwardDiff
     ρ⃗̃ = zeros(Real, 2length(ρ⃗₁), T)
 
     ρ⃗̃[:, 1] .= ket_to_iso(ρ⃗₁)
@@ -222,6 +224,7 @@ function unitary_rollout(
 )
     T = size(controls, 2)
 
+    # Real type enables ForwardDiff
     Ũ⃗ = zeros(Real, length(Ũ⃗_init), T)
 
     Ũ⃗[:, 1] .= Ũ⃗_init
@@ -272,17 +275,26 @@ function unitary_rollout(
     )
 end
 
+"""
+Compute the rollout fidelity.
+"""
 function Losses.iso_vec_unitary_fidelity(
     Ũ⃗_init::AbstractVector{<:Real},
     Ũ⃗_goal::AbstractVector{<:Real},
     controls::AbstractMatrix,
     Δt::AbstractVector,
     system::AbstractQuantumSystem;
     subspace::AbstractVector{Int}=axes(iso_vec_to_operator(Ũ⃗_goal), 1),
+    phases::Union{Nothing, AbstractVector{<:Real}}=nothing,
+    phase_operators::Union{Nothing, AbstractVector{<:AbstractMatrix{<:Complex}}}=nothing,
     kwargs...
 )
-    Ũ⃗ = unitary_rollout(Ũ⃗_init, controls, Δt, system; kwargs...)
-    return iso_vec_unitary_fidelity(Ũ⃗[:, end], Ũ⃗_goal; subspace=subspace)
+    Ũ⃗_T = unitary_rollout(Ũ⃗_init, controls, Δt, system; kwargs...)[:, end]
+    if !isnothing(phases)
+        return iso_vec_unitary_free_phase_fidelity(Ũ⃗_T, Ũ⃗_goal, phases, phase_operators; subspace=subspace)
+    else
+        return iso_vec_unitary_fidelity(Ũ⃗_T, Ũ⃗_goal; subspace=subspace)
+    end
 end
 
 function Losses.iso_vec_unitary_fidelity(