Merge pull request #136 from kestrelquantum/feat_rebased_robust_gateset

Feat rebased robust gateset

Manifest.toml

@@ -1,8 +1,8 @@
 # This file is machine-generated - editing it directly is not advised
 
-julia_version = "1.10.3"
+julia_version = "1.10.0"
 manifest_format = "2.0"
-project_hash = "0af13b9126caf21a7349df842fd536f87be35b2f"
+project_hash = "9419b1dd2f369832109faffc09e4c5a2ace40768"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "fc02d55798c1af91123d07915a990fbb9a10d146"
@@ -338,7 +338,7 @@ weakdeps = ["Dates", "LinearAlgebra"]
 [[deps.CompilerSupportLibraries_jll]]
 deps = ["Artifacts", "Libdl"]
 uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
-version = "1.1.1+0"
+version = "1.0.5+1"
 
 [[deps.CompositeTypes]]
 git-tree-sha1 = "bce26c3dab336582805503bed209faab1c279768"
@@ -1467,7 +1467,9 @@ version = "1.0.2"
 
 [[deps.NamedTrajectories]]
 deps = ["CairoMakie", "JLD2", "LaTeXStrings", "Latexify", "OrderedCollections", "Random", "Reexport", "Revise", "Unidecode"]
-git-tree-sha1 = "a109d8dee4055adb47bb435fd479dd67a50776af"
+git-tree-sha1 = "2be818959933a47c511d4f4ce7df035a256386e0"
+repo-rev = "main"
+repo-url = "https://github.com/aarontrowbridge/NamedTrajectories.jl.git"
 uuid = "538bc3a1-5ab9-4fc3-b776-35ca1e893e08"
 version = "0.1.5"
 
@@ -1542,7 +1544,7 @@ version = "0.3.24+0"
 [[deps.OpenBLAS_jll]]
 deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
 uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"
-version = "0.3.23+4"
+version = "0.3.23+2"
 
 [[deps.OpenEXR]]
 deps = ["Colors", "FileIO", "OpenEXR_jll"]

docs/src/contribution_guide.md

@@ -63,7 +63,7 @@ Tests are implemented using the [`TestItemRunner.jl`](https://github.com/julia-v
 
     prob = UnitarySmoothPulseProblem(
         H_drift, H_drives, U_goal, T, Δt,
-        ipopt_options=Options(print_level=1)
+        ipopt_options=IpoptOptions(print_level=1)
     )
 
     solve!(prob, max_iter=100)

src/QuantumCollocation.jl

@@ -38,8 +38,8 @@ include("dynamics.jl")
 include("evaluators.jl")
 @reexport using .Evaluators
 
-include("ipopt_options.jl")
-@reexport using .IpoptOptions
+include("options.jl")
+@reexport using .Options
 
 include("problems.jl")
 @reexport using .Problems

src/constraints.jl

@@ -122,15 +122,15 @@ minimum allowed fidelity.
 """
 
 function FinalFidelityConstraint(;
-    fidelity_function::Union{Function,Nothing}=nothing,
+    fidelity_function::Union{Symbol,Function,Nothing}=nothing,
     value::Union{Float64,Nothing}=nothing,
     comps::Union{AbstractVector{Int},Nothing}=nothing,
     goal::Union{AbstractVector{Float64},Nothing}=nothing,
     statedim::Union{Int,Nothing}=nothing,
     zdim::Union{Int,Nothing}=nothing,
     T::Union{Int,Nothing}=nothing,
     subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing,
-    hessian::Bool=true
+    eval_hessian::Bool=true
 )
     @assert !isnothing(fidelity_function) "must provide a fidelity function"
     @assert !isnothing(value) "must provide a fidelity value"
@@ -140,7 +140,12 @@ function FinalFidelityConstraint(;
     @assert !isnothing(zdim) "must provide a z dimension"
     @assert !isnothing(T) "must provide a T"
 
-    fidelity_function_symbol = Symbol(fidelity_function)
+    if fidelity_function isa Symbol
+        fidelity_function_symbol = fidelity_function
+        fidelity_function = eval(fidelity_function)
+    else
+        fidelity_function_symbol = Symbol(fidelity_function)
+    end
 
     if isnothing(subspace)
         fid = x -> fidelity_function(x, goal)
@@ -153,18 +158,20 @@ function FinalFidelityConstraint(;
     params = Dict{Symbol, Any}()
 
     if fidelity_function_symbol ∉ names(QuantumUtils)
-        @warn "fidelity function is not exported by QuantumUtils: will not be able to save this constraint"
+        @warn "Fidelity function :$(string(fidelity_function_symbol)) is not exported by QuantumUtils. Unable to save this constraint."
         params[:type] = :FinalFidelityConstraint
         params[:fidelity_function] = :not_saveable
     else
         params[:type] = :FinalFidelityConstraint
         params[:fidelity_function] = fidelity_function_symbol
         params[:value] = value
         params[:comps] = comps
+        params[:goal] = goal
         params[:statedim] = statedim
         params[:zdim] = zdim
         params[:T] = T
         params[:subspace] = subspace
+        params[:eval_hessian] = eval_hessian
     end
 
     state_slice = slice(T, comps, zdim)
@@ -194,7 +201,7 @@ function FinalFidelityConstraint(;
         end
     end
 
-    if hessian
+    if eval_hessian
         ∂²ℱ(x) = ForwardDiff.hessian(fid, x)
 
         ∂²ℱ_structure = hessian_of_lagrangian_structure(∂²ℱ, statedim, 1)
@@ -242,7 +249,7 @@ function FinalUnitaryFidelityConstraint(
     val::Float64,
     traj::NamedTrajectory;
     subspace::Union{AbstractVector{<:Integer}, Nothing}=nothing,
-    hessian::Bool=true
+    eval_hessian::Bool=true
 )
     @assert statesymb ∈ traj.names
     return FinalFidelityConstraint(;
@@ -254,7 +261,7 @@ function FinalUnitaryFidelityConstraint(
         zdim=traj.dim,
         T=traj.T,
         subspace=subspace,
-        hessian=hessian
+        eval_hessian=eval_hessian
     )
 end
 
@@ -268,7 +275,8 @@ is the NamedTrajectory symbol representing the unitary.
 function FinalQuantumStateFidelityConstraint(
     statesymb::Symbol,
     val::Float64,
-    traj::NamedTrajectory,
+    traj::NamedTrajectory;
+    kwargs...
 )
     @assert statesymb ∈ traj.names
     return FinalFidelityConstraint(;
@@ -278,26 +286,12 @@ function FinalQuantumStateFidelityConstraint(
         goal=traj.goal[statesymb],
         statedim=traj.dims[statesymb],
         zdim=traj.dim,
-        T=traj.T
+        T=traj.T,
+        kwargs...
     )
 end
 
 
-
-# function FinalStateFidelityConstraint(
-#     val::Float64,
-#     statesymb::Symbol,
-#     statedim::Int;
-#     fidelity_function::Function=fidelity
-# )
-#     return FinalFidelityConstraint(;
-#         fidelity_function=fidelity_function,
-#         value=val,
-#         statesymb=statesymb,
-#         statedim=statedim
-#     )
-# end
-
 """
     ComplexModulusContraint(<keyword arguments>)
 

src/direct_sums.jl

@@ -3,6 +3,7 @@ module DirectSums
 export add_suffix
 export get_suffix
 export direct_sum
+export merge_outer
 
 using ..Integrators
 using ..Problems
@@ -132,13 +133,6 @@ function direct_sum(
     )
 end
 
-
-function get_components(components::Union{Tuple, AbstractVector}, traj::NamedTrajectory)
-    symbs = Tuple(c for c in components)
-    vals = [traj[name] for name ∈ components]
-    return NamedTuple{symbs}(vals)
-end
-
 Base.endswith(symb::Symbol, suffix::AbstractString) = endswith(String(symb), suffix)
 Base.endswith(integrator::UnitaryPadeIntegrator, suffix::String) = endswith(integrator.unitary_symb, suffix)
 Base.endswith(integrator::DerivativeIntegrator, suffix::String) = endswith(integrator.variable, suffix)

src/dynamics.jl

@@ -199,7 +199,7 @@ end
 function QuantumDynamics(
     integrators::Vector{<:AbstractIntegrator},
     traj::NamedTrajectory;
-    hessian_approximation=false,
+    eval_hessian=true,
     jacobian_structure=true,
     verbose=false
 )
@@ -235,10 +235,10 @@ function QuantumDynamics(
 
     ∂f = dynamics_jacobian(integrators, traj)
 
-    if hessian_approximation
-        μ∂²f = nothing
-    else
+    if eval_hessian
         μ∂²f = dynamics_hessian_of_lagrangian(integrators, traj)
+    else
+        μ∂²f = nothing
     end
 
     if verbose
@@ -247,13 +247,7 @@ function QuantumDynamics(
 
     dynamics_dim = dim(integrators)
 
-    if hessian_approximation
-        ∂f_structure, ∂F_structure = dynamics_structure(∂f, traj, dynamics_dim;
-            verbose=verbose,
-            jacobian=jacobian_structure,
-        )
-        μ∂²F_structure = nothing
-    else
+    if eval_hessian
         ∂f_structure, ∂F_structure, μ∂²f_structure, μ∂²F_structure =
             dynamics_structure(∂f, μ∂²f, traj, dynamics_dim;
                 verbose=verbose,
@@ -263,6 +257,12 @@ function QuantumDynamics(
                 )
             )
         μ∂²f_nnz = length(μ∂²f_structure)
+    else
+        ∂f_structure, ∂F_structure = dynamics_structure(∂f, traj, dynamics_dim;
+            verbose=verbose,
+            jacobian=jacobian_structure,
+        )
+        μ∂²F_structure = nothing
     end
 
     ∂f_nnz = length(∂f_structure)
@@ -294,9 +294,7 @@ function QuantumDynamics(
         return ∂s
     end
 
-    if hessian_approximation
-        μ∂²F = nothing
-    else
+    if eval_hessian
         @views μ∂²F = (Z⃗::AbstractVector{<:Real}, μ⃗::AbstractVector{<:Real}) -> begin
             μ∂²s = Vector{eltype(Z⃗)}(undef, length(μ∂²F_structure))
             Threads.@threads for t = 1:traj.T-1
@@ -310,7 +308,10 @@ function QuantumDynamics(
             end
             return μ∂²s
         end
+    else
+        μ∂²F = nothing
     end
+
     return QuantumDynamics(
         integrators,
         F,

src/embedded_operators.jl

@@ -83,6 +83,16 @@ function Base.:*(
     )
 end
 
+function Base.kron(op1::EmbeddedOperator, op2::EmbeddedOperator)
+    levels = [size(op1, 1), size(op2, 2)]
+    indices = get_subspace_indices(
+        [op1.subspace_indices, op2.subspace_indices], levels
+    )
+    return EmbeddedOperator(unembed(op1) ⊗ unembed(op2), indices, levels)
+end
+
+QuantumUtils.:⊗(A::EmbeddedOperator, B::EmbeddedOperator) = kron(A, B)
+
 function EmbeddedOperator(
     op::AbstractMatrix{<:Number},
     system::QuantumSystem;
@@ -157,6 +167,8 @@ function EmbeddedOperator(
     return *(ops_embedded...)
 end
 
+# function LinearAlgebra
+
 # ====================
 
 basis_labels(subsystem_levels::AbstractVector{Int}; baseline=1) =

src/evaluators.jl

@@ -26,8 +26,8 @@ mutable struct PicoEvaluator <: MOI.AbstractNLPEvaluator
         trajectory::NamedTrajectory,
         objective::Objective,
         dynamics::QuantumDynamics,
-        nonlinear_constraints::Vector{<:NonlinearConstraint},
-        eval_hessian::Bool
+        nonlinear_constraints::Vector{<:NonlinearConstraint};
+        eval_hessian::Bool=true
     )
         n_dynamics_constraints = dynamics.dim * (trajectory.T - 1)
         n_nonlinear_constraints = sum(con.dim for con ∈ nonlinear_constraints; init=0)

src/integrators/pade_integrators.jl

@@ -96,8 +96,8 @@ struct UnitaryPadeIntegrator <: QuantumPadeIntegrator
         unitary_symb::Union{Symbol,Nothing}=nothing,
         drive_symb::Union{Symbol,Tuple{Vararg{Symbol}},Nothing}=nothing;
         order::Int=4,
-        autodiff::Bool=false,
-        G::Union{Function, Nothing}=nothing,
+        autodiff::Bool=order != 4,
+        G::Function=G_bilinear,
     )
         @assert order ∈ [4, 6, 8, 10] "order must be in [4, 6, 8, 10]"
         @assert !isnothing(unitary_symb) "must specify unitary symbol"
@@ -111,7 +111,6 @@ struct UnitaryPadeIntegrator <: QuantumPadeIntegrator
 
         G_drift = sys.G_drift
         G_drives = sys.G_drives
-        G = isnothing(G) ? G_bilinear : G
 
         drive_anticomms, drift_anticomms =
             order == 4 ? build_anticomms(G_drift, G_drives, n_drives) : (nothing, nothing)
@@ -189,7 +188,7 @@ struct QuantumStatePadeIntegrator <: QuantumPadeIntegrator
         drive_symb::Union{Symbol,Tuple{Vararg{Symbol}},Nothing}=nothing;
         order::Int=4,
         autodiff::Bool=order != 4,
-        G::Union{Function, Nothing}=nothing,
+        G::Function=G_bilinear,
     )
         @assert order ∈ [4, 6, 8, 10] "order must be in [4, 6, 8, 10]"
         @assert !isnothing(state_symb) "state_symb must be specified"
@@ -201,7 +200,6 @@ struct QuantumStatePadeIntegrator <: QuantumPadeIntegrator
 
         G_drift = sys.G_drift
         G_drives = sys.G_drives
-        G = isnothing(G) ? G_bilinear : G
 
         drive_anticomms, drift_anticomms =
             order == 4 ? build_anticomms(G_drift, G_drives, n_drives) : (nothing, nothing)