Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Robust control objectives #54

Merged
merged 9 commits into from
Nov 29, 2023
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Robust template test
  • Loading branch information
andgoldschmidt committed Nov 27, 2023
commit 0aa4f26e32fdc05545d8cc462a402f475c2a14ce
2 changes: 1 addition & 1 deletion src/constraints.jl
Original file line number Diff line number Diff line change
Expand Up @@ -233,7 +233,7 @@ is the NamedTrajectory symbol representing the unitary.
function FinalUnitaryFidelityConstraint(
statesymb::Symbol,
val::Float64,
traj::NamedTrajectory,
traj::NamedTrajectory;
subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing
)
@assert statesymb ∈ traj.names
Expand Down
14 changes: 8 additions & 6 deletions src/objectives.jl
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ export QuadraticSmoothnessRegularizer
export L1Regularizer

using TrajectoryIndexingUtils
using ..QuantumUtils
using ..QuantumSystems
using ..Losses
using ..Constraints
Expand Down Expand Up @@ -665,11 +666,11 @@ end

@doc raw"""
InfidelityRobustnessObjective(
Hₑ::AbstractMatrix{<:Number},
Z::NamedTrajectory;
eval_hessian::Bool=false,
subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing
)
Hₑ::AbstractMatrix{<:Number},
Z::NamedTrajectory;
eval_hessian::Bool=false,
subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing
)

Create a control objective which penalizes the sensitivity of the infidelity
to the provided operator defined in the subspace of the control dynamics,
Expand Down Expand Up @@ -707,7 +708,6 @@ function InfidelityRobustnessObjective(
return [j for (s, j) ∈ zip(operator_to_iso_vec(A), Z.components[:Ũ⃗]) if convert(Bool, s)]
end
ivs = iso_vec_subspace(isnothing(subspace) ? collect(1:size(Hₑ, 1)) : subspace, Z)
T = sum(NT.timesteps(Z))

@views function timesteps(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
return map(1:Z.T) do t
Expand All @@ -722,6 +722,7 @@ function InfidelityRobustnessObjective(
# Control frame
@views function toggle(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory)
Δts = timesteps(Z⃗, Z)
T = sum(Δts)
R = sum(
map(1:Z.T) do t
Uₜ = iso_vec_to_operator(Z⃗[slice(t, ivs, Z.dim)])
Expand All @@ -740,6 +741,7 @@ function InfidelityRobustnessObjective(
∇ = zeros(Z.dim * Z.T)
R = toggle(Z⃗, Z)
Δts = timesteps(Z⃗, Z)
T = sum(Δts)
Threads.@threads for t ∈ 1:Z.T
# State gradients
Uₜ_slice = slice(t, ivs, Z.dim)
Expand Down
3 changes: 2 additions & 1 deletion src/problem_templates.jl
Original file line number Diff line number Diff line change
Expand Up @@ -412,6 +412,7 @@ function UnitaryRobustnessProblem(
constraints::Vector{<:AbstractConstraint};
unitary_symbol::Symbol=:Ũ⃗,
final_fidelity::Float64=unitary_fidelity(trajectory[end][unitary_symbol], trajectory.goal[unitary_symbol]),
subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing,
eval_hessian::Bool=false,
verbose::Bool=false,
ipopt_options::Options=Options(),
Expand All @@ -433,7 +434,7 @@ function UnitaryRobustnessProblem(
fidelity_constraint = FinalUnitaryFidelityConstraint(
unitary_symbol,
final_fidelity,
trajectory,
trajectory;
subspace=subspace
)

Expand Down
74 changes: 9 additions & 65 deletions src/quantum_utils.jl
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
module QuantumUtils

export GATES
export gate
export ⊗
export apply
export qubit_system_state
Expand Down Expand Up @@ -35,30 +34,16 @@ using LinearAlgebra


"""
⊗(A::AbstractVecOrMat, B::AbstractVecOrMat)

Kronecker product of two vectors or matrices.
kronecker product utility
"""
⊗(A::AbstractVecOrMat, B::AbstractVecOrMat) = kron(A, B)


@doc raw"""
GATES::Dict{Symbol, Matrix{ComplexF64}}

Dictionary of common quantum gates.

# Elements
- `:I` - identity, $I$
- `:X` - Pauli X, $\sigma_x$
- `:Y` - Pauli Y, $\sigma_y$
- `:Z` - Pauli Z, $\sigma_z$
- `:H` - Hadamard, $H$
- `:CX` - controlled-X, $C_X$
- `:XI` - $-i X \otimes I$
- `:sqrtiSWAP` - $\sqrt{i \text{SWAP}}$
⊗(A::AbstractVecOrMat, B::AbstractVecOrMat) = kron(A, B)


"""
quantum gates
"""

const GATES = Dict(
:I => [1 0;
0 1],
Expand Down Expand Up @@ -91,18 +76,8 @@ const GATES = Dict(
0 0 0 1]
)

"""
gate(U::Symbol)

Get a gate from the `GATES` dictionary.
"""
gate(U::Symbol) = GATES[U]

"""
apply(gate::Symbol, ψ::Vector{<:Number})

Apply a gate from the GATES dictionary to a quantum state.
"""
function apply(gate::Symbol, ψ::Vector{<:Number})
@assert norm(ψ) ≈ 1.0
@assert gate in keys(GATES) "gate not found"
Expand All @@ -111,11 +86,6 @@ function apply(gate::Symbol, ψ::Vector{<:Number})
return ComplexF64.(normalize(Û * ψ))
end

@doc raw"""
qubit_system_state(ket::String)

Get a quantum state from a string of qubit states, e.g. `"01"` -> $\ket{01}$.
"""
function qubit_system_state(ket::String)
cs = [c for c ∈ ket]
@assert all(c ∈ "01" for c ∈ cs)
Expand All @@ -125,43 +95,17 @@ function qubit_system_state(ket::String)
return ψ
end


@doc raw"""
lift(
U::AbstractMatrix{<:Number},
i::Int,
n::Int;
levels::Int=size(U, 1)
)

Lift a matrix to a larger Hilbert space by placing it on the `i`-th qubit of a `n`-qubit system. I.e.,
```math
U \mapsto I_1 \otimes \cdots \otimes I_{i-1} \otimes U \otimes I_{i+1} \otimes \cdots \otimes I_n
```
"""
function lift(
U::AbstractMatrix{<:Number},
i::Int,
n::Int;
qubit_index::Int,
n_qubits::Int;
levels::Int=size(U, 1)
)::Matrix{ComplexF64}
Is = Matrix{Complex}[I(levels) for _ = 1:n]
Is[i] = U
Is = Matrix{Complex}[I(levels) for _ = 1:n_qubits]
Is[qubit_index] = U
return foldr(⊗, Is)
end

@doc raw"""
lift(
op::AbstractMatrix{<:Number},
i::Int,
levels::Vector{Int}
)

Lift a matrix to a larger Hilbert space by placing it on the `i`-th qubit of a `n`-qubit system. I.e.,
```math
U \mapsto I_1 \otimes \cdots \otimes I_{i-1} \otimes U \otimes I_{i+1} \otimes \cdots \otimes I_n
```
"""
function lift(
op::AbstractMatrix{<:Number},
i::Int,
Expand Down
88 changes: 88 additions & 0 deletions test/problem_templates_test.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
#
# 1. test UnitarySmoothPulseProblem
# 2. test UnitaryMinimumTimeProblem
# 3. test UnitaryRobustnessProblem
# ------------------------------------------------


Expand Down Expand Up @@ -38,5 +39,92 @@

@test times(mintime_prob.trajectory)[end] < times(prob.trajectory)[end]

end



@testset "Robust and Subspace Templates" begin
# --------------------------------------------
# Initialize with UnitarySmoothPulseProblem
# --------------------------------------------
H_drift = zeros(2, 2)
H_drives = [GATES[:X], GATES[:Y]]
U_goal = GATES[:X]
T = 50
Δt = .2

probs = Dict()

probs["qubit"] = UnitarySmoothPulseProblem(
H_drift, H_drives, U_goal, T, Δt;
verbose=false
)

solve!(probs["qubit"]; max_iter=100)

@test unitary_fidelity(probs["qubit"]) > 0.99

# --------------------------------------------
# 1. test UnitarySmoothPulseProblem with subspace
# --------------------------------------------
H_error = GATES[:Z]
H_drift = zeros(3, 3)
H_drives = [create(3) + annihilate(3), im * (create(3) - annihilate(3))]
U_goal = [0 1 0; 1 0 0; 0 0 0]
T = probs["qubit"].trajectory.T
Δt = probs["qubit"].trajectory[:Δt][1]
a_guess = probs["qubit"].trajectory[:a]

subspace = subspace_indices([3])
probs["transmon"] = UnitarySmoothPulseProblem(
H_drift, H_drives, U_goal, T, Δt;
a_guess=a_guess, subspace=subspace, geodesic=false, verbose=false
)
solve!(probs["transmon"]; max_iter=100)

# Subspace gate success
@test unitary_fidelity(probs["transmon"]; subspace=subspace) > 0.99


# --------------------------------------------
# 3. test UnitaryRobustnessProblem from previous problem
# --------------------------------------------
probs["robust"] = UnitaryRobustnessProblem(
H_error, probs["transmon"];
final_fidelity=0.99, subspace=subspace, verbose=false
)
solve!(probs["robust"]; max_iter=100)

EvalLoss(problem, Loss) = Loss(vec(problem.trajectory.data), problem.trajectory)
Loss = InfidelityRobustnessObjective(H_error, probs["transmon"].trajectory).L

# Robustness improvement over default
@test EvalLoss(probs["robust"], Loss) < EvalLoss(probs["transmon"], Loss)

# Fidelity constraint approximately satisfied
@test isapprox(unitary_fidelity(probs["robust"]; subspace=subspace), 0.99, atol=0.025)

# --------------------------------------------
# 4. test UnitaryRobustnessProblem from default struct
# --------------------------------------------
params = deepcopy(probs["transmon"].params)
trajectory = copy(probs["transmon"].trajectory)
system = probs["transmon"].system
objective = QuadraticRegularizer(:dda, trajectory, 1e-4)
integrators = probs["transmon"].integrators
constraints = AbstractConstraint[]

probs["unconstrained"] = UnitaryRobustnessProblem(
H_error, trajectory, system, objective, integrators, constraints;
final_fidelity=0.99, subspace=subspace, verbose=false
)
solve!(probs["unconstrained"]; max_iter=100)

# Additonal robustness improvement after relaxed objective
@test EvalLoss(probs["unconstrained"], Loss) < EvalLoss(probs["transmon"], Loss)

# Fidelity constraint approximately satisfied
@test isapprox(unitary_fidelity(probs["unconstrained"]; subspace=subspace), 0.99, atol=0.025)


end
Loading