diff --git a/src/problem_templates/_problem_templates.jl b/src/problem_templates/_problem_templates.jl index 72f0b851..65aa9801 100644 --- a/src/problem_templates/_problem_templates.jl +++ b/src/problem_templates/_problem_templates.jl @@ -5,7 +5,6 @@ export UnitaryMinimumTimeProblem export UnitaryRobustnessProblem export UnitaryDirectSumProblem export UnitarySamplingProblem -export UnitaryBangBangProblem export QuantumStateSmoothPulseProblem export QuantumStateMinimumTimeProblem @@ -35,7 +34,6 @@ include("unitary_minimum_time_problem.jl") include("unitary_robustness_problem.jl") include("unitary_direct_sum_problem.jl") include("unitary_sampling_problem.jl") -include("unitary_bang_bang_problem.jl") include("quantum_state_smooth_pulse_problem.jl") include("quantum_state_minimum_time_problem.jl") diff --git a/src/problem_templates/unitary_bang_bang_problem.jl b/src/problem_templates/unitary_bang_bang_problem.jl deleted file mode 100644 index 9805b5ed..00000000 --- a/src/problem_templates/unitary_bang_bang_problem.jl +++ /dev/null @@ -1,175 +0,0 @@ -@doc raw""" - UnitaryBangBangProblem(system::QuantumSystem, operator, T, Δt; kwargs...) - -TODO: Docstring -""" -function UnitaryBangBangProblem( - system::AbstractQuantumSystem, - operator::Union{EmbeddedOperator, AbstractMatrix{<:Number}}, - T::Int, - Δt::Union{Float64, Vector{Float64}}; - free_time=true, - init_trajectory::Union{NamedTrajectory, Nothing}=nothing, - a_bound::Float64=1.0, - a_bounds=fill(a_bound, length(system.G_drives)), - a_guess::Union{Matrix{Float64}, Nothing}=nothing, - da_bound::Float64=1.0, - da_bounds=fill(da_bound, length(system.G_drives)), - Δt_min::Float64=0.5 * Δt, - Δt_max::Float64=1.5 * Δt, - drive_derivative_σ::Float64=0.01, - Q::Float64=100.0, - R=1e-2, - R_a::Union{Float64, Vector{Float64}}=R, - R_da::Union{Float64, Vector{Float64}}=R, - leakage_suppression=false, - R_leakage=1e-1, - max_iter::Int=1000, - linear_solver::String="mumps", - ipopt_options::Options=Options(), - constraints::Vector{<:AbstractConstraint}=AbstractConstraint[], - timesteps_all_equal::Bool=true, - verbose::Bool=false, - integrator::Symbol=:pade, - rollout_integrator=exp, - bound_unitary=integrator == :exponential, - # TODO: control modulus norm, advanced feature, needs documentation - control_norm_constraint=false, - control_norm_constraint_components=nothing, - control_norm_R=nothing, - geodesic=true, - pade_order=4, - autodiff=pade_order != 4, - jacobian_structure=true, - hessian_approximation=false, - blas_multithreading=true, - kwargs... -) - if !blas_multithreading - BLAS.set_num_threads(1) - end - - if hessian_approximation - ipopt_options.hessian_approximation = "limited-memory" - end - - # Trajectory - if !isnothing(init_trajectory) - traj = init_trajectory - else - n_drives = length(system.G_drives) - traj = initialize_unitary_trajectory( - operator, - T, - Δt, - n_drives, - (a = a_bounds, da = da_bounds); - n_derivatives=1, - free_time=free_time, - Δt_bounds=(Δt_min, Δt_max), - geodesic=geodesic, - bound_unitary=bound_unitary, - drive_derivative_σ=drive_derivative_σ, - a_guess=a_guess, - system=system, - rollout_integrator=rollout_integrator, - ) - end - - # Objective - J = UnitaryInfidelityObjective(:Ũ⃗, traj, Q; - subspace=operator isa EmbeddedOperator ? operator.subspace_indices : nothing, - ) - J += QuadraticRegularizer(:a, traj, R_a) - J += QuadraticRegularizer(:da, traj, R_da) - - # Constraints - if leakage_suppression - if operator isa EmbeddedOperator - leakage_indices = get_unitary_isomorphism_leakage_indices(operator) - J += L1Regularizer!( - constraints, :Ũ⃗, traj, - R_value=R_leakage, - indices=leakage_indices, - eval_hessian=!hessian_approximation - ) - else - @warn "leakage_suppression is not supported for non-embedded operators, ignoring." - end - end - - if free_time - if timesteps_all_equal - push!(constraints, TimeStepsAllEqualConstraint(:Δt, traj)) - end - end - - if control_norm_constraint - @assert !isnothing(control_norm_constraint_components) "control_norm_constraint_components must be provided" - @assert !isnothing(control_norm_R) "control_norm_R must be provided" - norm_con = ComplexModulusContraint( - :a, - control_norm_R, - traj; - name_comps=control_norm_constraint_components, - ) - push!(constraints, norm_con) - end - - # Integrators - if integrator == :pade - unitary_integrator = - UnitaryPadeIntegrator(system, :Ũ⃗, :a; order=pade_order, autodiff=autodiff) - elseif integrator == :exponential - unitary_integrator = - UnitaryExponentialIntegrator(system, :Ũ⃗, :a) - else - error("integrator must be one of (:pade, :exponential)") - end - - integrators = [ - unitary_integrator, - DerivativeIntegrator(:a, :da, traj), - ] - - return QuantumControlProblem( - system, - traj, - J, - integrators; - constraints=constraints, - max_iter=max_iter, - linear_solver=linear_solver, - verbose=verbose, - ipopt_options=ipopt_options, - jacobian_structure=jacobian_structure, - hessian_approximation=hessian_approximation, - eval_hessian=!hessian_approximation, - kwargs... - ) -end - - -""" - UnitarySmoothPulseProblem( - H_drift::AbstractMatrix{<:Number}, - H_drives::Vector{<:AbstractMatrix{<:Number}}, - operator, - T, - Δt; - kwargs... - ) - -Constructor for a `UnitarySmoothPulseProblem` from a drift Hamiltonian and a set of control Hamiltonians. -""" -function UnitarySmoothPulseProblem( - H_drift::AbstractMatrix{<:Number}, - H_drives::Vector{<:AbstractMatrix{<:Number}}, - args...; - kwargs... -) - system = QuantumSystem(H_drift, H_drives) - return UnitarySmoothPulseProblem(system, args...; kwargs...) -end - -# *************************************************************************** # diff --git a/src/trajectory_initialization.jl b/src/trajectory_initialization.jl index 9b5c65d1..57e866e2 100644 --- a/src/trajectory_initialization.jl +++ b/src/trajectory_initialization.jl @@ -163,49 +163,37 @@ end function initialize_controls( n_drives::Int, - n_derivatives::Int, T::Int, - bounds::VectorBound, - drive_derivative_σ::Float64, + a_bounds::VectorBound, + drive_derivative_σ::Float64 ) - if bounds isa AbstractVector - a_dists = [Uniform(-bounds[i], bounds[i]) for i = 1:n_drives] - elseif bounds isa Tuple - a_dists = [Uniform(aᵢ_lb, aᵢ_ub) for (aᵢ_lb, aᵢ_ub) ∈ zip(bounds...)] + if a_bounds isa AbstractVector + a_dists = [Uniform(-a_bounds[i], a_bounds[i]) for i = 1:n_drives] + elseif a_bounds isa Tuple + a_dists = [Uniform(aᵢ_lb, aᵢ_ub) for (aᵢ_lb, aᵢ_ub) ∈ zip(a_bounds...)] else - error("bounds must be a Vector or Tuple") + error("a_bounds must be a Vector or Tuple") end - controls = Matrix{Float64}[] - a = hcat([ zeros(n_drives), vcat([rand(a_dists[i], 1, T - 2) for i = 1:n_drives]...), zeros(n_drives) ]...) - push!(controls, a) - - for _ in 1:n_derivatives - push!(controls, randn(n_drives, T) * drive_derivative_σ) - end - return controls + da = randn(n_drives, T) * drive_derivative_σ + dda = randn(n_drives, T) * drive_derivative_σ + return a, da, dda end -function initialize_controls(a::AbstractMatrix, Δt::AbstractVecOrMat, n_derivatives::Int) - controls = Matrix{Float64}[a] - push!(controls, a) +function initialize_controls(a::AbstractMatrix, Δt::AbstractVecOrMat) + da = derivative(a, Δt) + dda = derivative(da, Δt) - for n in 1:n_derivatives - # next derivative - push!(controls, derivative(controls[end], Δt)) + # to avoid constraint violation error at initial iteration + da[:, end] = da[:, end-1] + Δt[end-1] * dda[:, end-1] - # to avoid constraint violation error at initial iteration for da, dda, ... - if n > 1 - controls[end-1][:, end] = controls[end-1][:, end-1] + Δt[end-1] * controls[end][:, end-1] - end - end - return controls + return a, da, dda end function initialize_unitary_trajectory( @@ -213,9 +201,9 @@ function initialize_unitary_trajectory( T::Int, Δt::Real, n_drives::Int, - all_a_bounds::NamedTuple{anames, <:Tuple{Vararg{VectorBound}}} where anames; + a_bounds::VectorBound, + dda_bounds::VectorBound; U_init::AbstractMatrix{<:Number}=Matrix{ComplexF64}(I(size(U_goal, 1))), - n_derivatives::Int=2, geodesic=true, bound_unitary=false, free_time=false, @@ -225,10 +213,7 @@ function initialize_unitary_trajectory( system::Union{AbstractQuantumSystem, AbstractVector{<:AbstractQuantumSystem}, Nothing}=nothing, rollout_integrator::Function=exp, Ũ⃗_keys::AbstractVector{<:Symbol}=[:Ũ⃗], - a_keys::AbstractVector{<:Symbol}=[Symbol("d"^i * "a") for i in 0:n_derivatives] ) - @assert length(a_keys) == n_derivatives + 1 "a_keys must have the same length as n_derivatives + 1" - if free_time if Δt isa Float64 Δt = fill(Δt, 1, T) @@ -257,25 +242,19 @@ function initialize_unitary_trajectory( goal = (; (Ũ⃗_keys .=> Ũ⃗_goals)...) # Bounds - bounds = all_a_bounds + bounds = (a = a_bounds, dda = dda_bounds,) if bound_unitary Ũ⃗_dim = length(Ũ⃗_init) Ũ⃗_bounds = repeat([(-ones(Ũ⃗_dim), ones(Ũ⃗_dim))], length(Ũ⃗_keys)) - bounds = merge(bounds, (; (Ũ⃗_keys .=> Ũ⃗_bounds)...)) + merge!(bounds, (; (Ũ⃗_keys .=> Ũ⃗_bounds)...)) end # Initial state and control values if isnothing(a_guess) Ũ⃗ = initialize_unitaries(U_init, U_goal, T, geodesic=geodesic) Ũ⃗_values = repeat([Ũ⃗], length(Ũ⃗_keys)) - a_values = initialize_controls( - n_drives, - n_derivatives, - T, - bounds[a_keys[1]], - drive_derivative_σ - ) + a, da, dda = initialize_controls(n_drives, T, a_bounds, drive_derivative_σ) else @assert !isnothing(system) "system must be provided if a_guess is provided" if system isa AbstractVector @@ -287,21 +266,21 @@ function initialize_unitary_trajectory( unitary_rollout(Ũ⃗_init, a_guess, Δt, system; integrator=rollout_integrator) Ũ⃗_values = repeat([Ũ⃗], length(Ũ⃗_keys)) end - a_values = initialize_controls(a_guess, Δt, n_derivatives) + a, da, dda = initialize_controls(a_guess, Δt) end # Trajectory - keys = [Ũ⃗_keys..., a_keys...] - values = [Ũ⃗_values..., a_values...] + keys = [Ũ⃗_keys..., :a, :da, :dda] + values = [Ũ⃗_values..., a, da, dda] if free_time push!(keys, :Δt) push!(values, Δt) - controls = (a_keys[end], :Δt) + controls = (:dda, :Δt) timestep = :Δt bounds = merge(bounds, (Δt = Δt_bounds,)) else - controls = (a_keys[end],) + controls = (:dda,) timestep = Δt end @@ -322,16 +301,15 @@ function initialize_state_trajectory( T::Int, Δt::Real, n_drives::Int, - all_a_bounds::NamedTuple{anames, <:Tuple{Vararg{VectorBound}}} where anames; - n_derivatives::Int=2, + a_bounds::VectorBound, + dda_bounds::VectorBound; free_time=false, Δt_bounds::ScalarBound=(0.5 * Δt, 1.5 * Δt), drive_derivative_σ::Float64=0.1, a_guess::Union{AbstractMatrix{<:Float64}, Nothing}=nothing, system::Union{AbstractQuantumSystem, AbstractVector{<:AbstractQuantumSystem}, Nothing}=nothing, rollout_integrator::Function=exp, - ψ̃_keys::AbstractVector{<:Symbol}=[Symbol("ψ̃$i") for i = 1:length(ψ̃_goals)], - a_keys::AbstractVector{<:Symbol}=[Symbol("d"^i * "a") for i in 0:n_derivatives] + ψ̃_keys::AbstractVector{<:Symbol}=[Symbol("ψ̃$i") for i = 1:length(ψ̃_goals)] ) @assert length(ψ̃_inits) == length(ψ̃_goals) "ψ̃_inits and ψ̃_goals must have the same length" @assert length(ψ̃_keys) == length(ψ̃_goals) "ψ̃_keys and ψ̃_goals must have the same length" @@ -352,41 +330,35 @@ function initialize_state_trajectory( goal = (; (ψ̃_keys .=> ψ̃_goals)...) # Bounds - bounds = all_a_bounds + bounds = (a = a_bounds, dda = dda_bounds,) # Initial state and control values if isnothing(a_guess) - ψ̃_values = NamedTuple([ + ψ̃s = NamedTuple([ k => linear_interpolation(ψ̃_init, ψ̃_goal, T) for (k, ψ̃_init, ψ̃_goal) in zip(ψ̃_keys, ψ̃_inits, ψ̃_goals) ]) - a_values = initialize_controls( - n_drives, - n_derivatives, - T, - bounds[a_keys[1]], - drive_derivative_σ - ) + a, da, dda = initialize_controls(n_drives, T, a_bounds, drive_derivative_σ) else - ψ̃_values = NamedTuple([ + ψ̃s = NamedTuple([ k => rollout(ψ̃_init, a_guess, Δt, system, integrator=rollout_integrator) for (i, ψ̃_init) in zip(ψ̃_keys, ψ̃_inits) ]) - a_values = initialize_controls(a_guess, Δt, n_derivatives) + a, da, dda = initialize_controls(a_guess, Δt) end # Trajectory - keys = [ψ̃_keys..., a_keys...] - values = [ψ̃_values..., a_values...] + keys = [ψ̃_keys..., :a, :da, :dda] + values = [ψ̃s..., a, da, dda] if free_time push!(keys, :Δt) push!(values, Δt) - controls = (a_keys[end], :Δt) + controls = (:dda, :Δt) timestep = :Δt bounds = merge(bounds, (Δt = Δt_bounds,)) else - controls = (a_keys[end],) + controls = (:dda,) timestep = Δt end