diff --git a/src/constraints.jl b/src/constraints.jl index 72c09e56..0043e9e0 100644 --- a/src/constraints.jl +++ b/src/constraints.jl @@ -117,6 +117,7 @@ minimum allowed fidelity. - `statedim::Union{Int,Nothing}=nothing`: the dimension of the state - `zdim::Union{Int,Nothing}=nothing`: the dimension of a single time step of the trajectory - `T::Union{Int,Nothing}=nothing`: the number of time steps +- `subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing`: the subspace indices of the fidelity """ @@ -128,6 +129,7 @@ function FinalFidelityConstraint(; statedim::Union{Int,Nothing}=nothing, zdim::Union{Int,Nothing}=nothing, T::Union{Int,Nothing}=nothing, + subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing ) @assert !isnothing(fidelity_function) "must provide a fidelity function" @assert !isnothing(value) "must provide a fidelity value" @@ -139,7 +141,11 @@ function FinalFidelityConstraint(; fidelity_function_symbol = Symbol(fidelity_function) - fid = x -> fidelity_function(x, goal) + if isnothing(subspace) + fid = x -> fidelity_function(x, goal) + else + fid = x -> fidelity_function(x, goal; subspace=subspace) + end @assert fid(randn(statedim)) isa Float64 "fidelity function must return a scalar" @@ -157,6 +163,7 @@ function FinalFidelityConstraint(; params[:statedim] = statedim params[:zdim] = zdim params[:T] = T + params[:subspace] = subspace end state_slice = slice(T, comps, zdim) @@ -226,7 +233,8 @@ 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 return FinalFidelityConstraint(; @@ -236,7 +244,8 @@ function FinalUnitaryFidelityConstraint( goal=traj.goal[statesymb], statedim=traj.dims[statesymb], zdim=traj.dim, - T=traj.T + T=traj.T, + subspace=subspace ) end diff --git a/src/objectives.jl b/src/objectives.jl index 6d972049..1de467da 100644 --- a/src/objectives.jl +++ b/src/objectives.jl @@ -8,12 +8,14 @@ export QuantumUnitaryObjective export UnitaryInfidelityObjective export MinimumTimeObjective +export InfidelityRobustnessObjective export QuadraticRegularizer export QuadraticSmoothnessRegularizer export L1Regularizer using TrajectoryIndexingUtils +using ..QuantumUtils using ..QuantumSystems using ..Losses using ..Constraints @@ -662,4 +664,119 @@ function MinimumTimeObjective(traj::NamedTrajectory; D=1.0) return MinimumTimeObjective(; D=D, Δt_indices=Δt_indices) end +@doc raw""" +InfidelityRobustnessObjective( + 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, +thereby realizing robust control. + +The control dynamics are +```math +U_C(a)= \prod_t \exp{-i H_C(a_t)} +``` + +In the control frame, the Hₑ operator is (proportional to) +```math +R_{Robust}(a) = \frac{1}{T \norm{H_e}_2} \sum_t U_C(a_t)^\dag H_e U_C(a_t) \Delta t +``` +where we have adjusted to a unitless expression of the operator. + +The robustness objective is +```math +R_{Robust}(a) = \frac{1}{N} \norm{R}^2_F +``` +where N is the dimension of the Hilbert space. +""" +function InfidelityRobustnessObjective( + Hₑ::AbstractMatrix{<:Number}, + Z::NamedTrajectory; + eval_hessian::Bool=false, + subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing +) + # Indices of all non-zero subspace components for iso_vec_operators + function iso_vec_subspace(subspace::AbstractVector{<:Integer}, Z::NamedTrajectory) + d = isqrt(Z.dims[:Ũ⃗] ÷ 2) + A = zeros(Complex, d, d) + A[subspace, subspace] .= 1 + im + # Return any index where there is a 1. + 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) + + @views function timesteps(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory) + return map(1:Z.T) do t + if Z.timestep isa Symbol + Z⃗[slice(t, Z.components[Z.timestep], Z.dim)][1] + else + Z.timestep + end + end + end + + # 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)]) + Uₜ'Hₑ*Uₜ .* Δts[t] + end + ) / norm(Hₑ) / T + return R + end + + function L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory) + R = toggle(Z⃗, Z) + return real(tr(R'R)) / size(R, 1) + end + + @views function ∇L(Z⃗::AbstractVector{<:Real}, Z::NamedTrajectory) + ∇ = 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) + Uₜ = iso_vec_to_operator(Z⃗[Uₜ_slice]) + ∇[Uₜ_slice] .= 2 .* operator_to_iso_vec( + Hₑ * Uₜ * R .* Δts[t] + ) / norm(Hₑ) / T + # Time gradients + if Z.timestep isa Symbol + t_slice = slice(t, Z.components[Z.timestep], Z.dim) + ∂R = Uₜ'Hₑ*Uₜ + ∇[t_slice] .= tr(∂R*R + R*∂R) / norm(Hₑ) / T + end + end + return ∇ / size(R, 1) + end + + # Hessian is dense (Control frame R contains sum over all unitaries). + if eval_hessian + # TODO + ∂²L = (Z⃗, Z) -> [] + ∂²L_structure = Z -> [] + else + ∂²L = nothing + ∂²L_structure = nothing + end + + params = Dict( + :type => :QuantumRobustnessObjective, + :error => Hₑ, + :eval_hessian => eval_hessian + ) + + return Objective(L, ∇L, ∂²L, ∂²L_structure, Dict[params]) +end + + end diff --git a/src/problem_templates.jl b/src/problem_templates.jl index 9442c254..5a9fa2ef 100644 --- a/src/problem_templates.jl +++ b/src/problem_templates.jl @@ -2,6 +2,7 @@ module ProblemTemplates export UnitarySmoothPulseProblem export UnitaryMinimumTimeProblem +export UnitaryRobustnessProblem export QuantumStateSmoothPulseProblem export QuantumStateMinimumTimeProblem @@ -359,7 +360,7 @@ function UnitaryMinimumTimeProblem( kwargs... ) params = deepcopy(prob.params) - traj = copy(prob.trajectory) + trajectory = copy(prob.trajectory) system = prob.system objective = Objective(params[:objective_terms]) integrators = prob.integrators @@ -368,7 +369,7 @@ function UnitaryMinimumTimeProblem( NonlinearConstraint.(params[:nonlinear_constraints])... ] return UnitaryMinimumTimeProblem( - traj, + trajectory, system, objective, integrators, @@ -378,7 +379,6 @@ function UnitaryMinimumTimeProblem( ) end - function UnitaryMinimumTimeProblem( data_path::String; kwargs... @@ -403,6 +403,84 @@ function UnitaryMinimumTimeProblem( ) end +function UnitaryRobustnessProblem( + Hₑ::AbstractMatrix{<:Number}, + trajectory::NamedTrajectory, + system::QuantumSystem, + objective::Objective, + integrators::Vector{<:AbstractIntegrator}, + 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(), + kwargs... +) + @assert unitary_symbol ∈ trajectory.names + + if !eval_hessian + ipopt_options.hessian_approximation = "limited-memory" + end + + objective += InfidelityRobustnessObjective( + Hₑ, + trajectory, + eval_hessian=eval_hessian, + subspace=subspace + ) + + fidelity_constraint = FinalUnitaryFidelityConstraint( + unitary_symbol, + final_fidelity, + trajectory; + subspace=subspace + ) + + constraints = AbstractConstraint[constraints..., fidelity_constraint] + + return QuantumControlProblem( + system, + trajectory, + objective, + integrators; + constraints=constraints, + verbose=verbose, + ipopt_options=ipopt_options, + eval_hessian=eval_hessian, + kwargs... + ) +end + +function UnitaryRobustnessProblem( + Hₑ::AbstractMatrix{<:Number}, + prob::QuantumControlProblem; + kwargs... +) + params = deepcopy(prob.params) + trajectory = copy(prob.trajectory) + system = prob.system + objective = Objective(params[:objective_terms]) + integrators = prob.integrators + constraints = [ + params[:linear_constraints]..., + NonlinearConstraint.(params[:nonlinear_constraints])... + ] + + return UnitaryRobustnessProblem( + Hₑ, + trajectory, + system, + objective, + integrators, + constraints; + build_trajectory_constraints=false, + kwargs... + ) +end + + # ------------------------------------------ # Quantum State Problem Templates # ------------------------------------------ diff --git a/src/quantum_utils.jl b/src/quantum_utils.jl index 9f8cf271..50e2f07c 100644 --- a/src/quantum_utils.jl +++ b/src/quantum_utils.jl @@ -1,7 +1,6 @@ module QuantumUtils export GATES -export gate export ⊗ export apply export qubit_system_state @@ -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], @@ -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" @@ -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) @@ -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, diff --git a/test/problem_templates_test.jl b/test/problem_templates_test.jl index eb6af881..ea1e6915 100644 --- a/test/problem_templates_test.jl +++ b/test/problem_templates_test.jl @@ -3,6 +3,7 @@ # # 1. test UnitarySmoothPulseProblem # 2. test UnitaryMinimumTimeProblem +# 3. test UnitaryRobustnessProblem # ------------------------------------------------ @@ -38,5 +39,76 @@ @test times(mintime_prob.trajectory)[end] < times(prob.trajectory)[end] +end + + + +@testset "Robust and Subspace Templates" begin + # -------------------------------------------- + # Initialize with UnitarySmoothPulseProblem + # -------------------------------------------- + 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] + subspace = subspace_indices([3]) + T = 50 + Δt = .2 + probs = Dict() + + # -------------------------------------------- + # 1. test UnitarySmoothPulseProblem with subspace + # - rely on random initialization + # -------------------------------------------- + probs["transmon"] = UnitarySmoothPulseProblem( + H_drift, H_drives, U_goal, T, Δt; + subspace=subspace, geodesic=false, verbose=false + ) + solve!(probs["transmon"]; max_iter=200) + + # Subspace gate success + @test unitary_fidelity(probs["transmon"]; subspace=subspace) > 0.99 + + + # -------------------------------------------- + # 2. 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) + + eval_loss(problem, Loss) = Loss(vec(problem.trajectory.data), problem.trajectory) + loss = InfidelityRobustnessObjective(H_error, probs["transmon"].trajectory).L + + # Robustness improvement over default + @test eval_loss(probs["robust"], loss) < eval_loss(probs["transmon"], loss) + + # Fidelity constraint approximately satisfied + @test isapprox(unitary_fidelity(probs["robust"]; subspace=subspace), 0.99, atol=0.025) + + # -------------------------------------------- + # 3. 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 eval_loss(probs["unconstrained"], loss) < eval_loss(probs["transmon"], loss) + + # Fidelity constraint approximately satisfied + @test isapprox(unitary_fidelity(probs["unconstrained"]; subspace=subspace), 0.99, atol=0.025) + end