Robust control objectives

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
 

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

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
 # ------------------------------------------