in progress variational problem template

Project.toml

@@ -1,7 +1,7 @@
 name = "QuantumCollocation"
 uuid = "0dc23a59-5ffb-49af-b6bd-932a8ae77adf"
 authors = ["Aaron Trowbridge <[email protected]> and contributors"]
-version = "0.7.0"
+version = "0.7.1"
 
 [deps]
 DirectTrajOpt = "c823fa1f-8872-4af5-b810-2b9b72bbbf56"
@@ -21,7 +21,7 @@ TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe"
 TrajectoryIndexingUtils = "6dad8b7f-dd9a-4c28-9b70-85b9a079bfc8"
 
 [compat]
-DirectTrajOpt = "0.1.0"
+DirectTrajOpt = "0.1"
 Distributions = "0.25"
 ExponentialAction = "0.2"
 Interpolations = "0.15"

src/problem_templates/_problem_templates.jl

@@ -19,8 +19,7 @@ using JLD2
 using TestItems
 
 include("unitary_smooth_pulse_problem.jl")
-include("adjoint_smooth_pulse_problem.jl")
-
+include("unitary_variational_problem.jl")
 include("unitary_minimum_time_problem.jl")
 include("unitary_sampling_problem.jl")
 

src/problem_templates/adjoint_smooth_pulse_problem.jl → src/problem_templates/unitary_variational_problem.jl

@@ -1,11 +1,13 @@
-export AdjointUnitarySmoothPulseProblem
-function AdjointUnitarySmoothPulseProblem(
-    system::ParameterizedQuantumSystem,
+export UnitaryVariationalProblem
+
+
+function UnitaryVariationalProblem(
+    system::VariationalQuantumSystem,
     goal::AbstractPiccoloOperator,
     T::Int,
     Δt::Union{Float64, <:AbstractVector{Float64}};
-    times::Vector{<:Union{AbstractVector,Nothing}}=[nothing for s∈system.Gₐ],
-    down_times::Vector{<:Union{AbstractVector,Nothing}}=[nothing for s∈system.Gₐ],
+    times::Vector{<:Union{AbstractVector,Nothing}}=[nothing for s∈system.G_vars],
+    down_times::Vector{<:Union{AbstractVector,Nothing}}=[nothing for s∈system.G_vars],
     unitary_integrator=AdjointUnitaryIntegrator,
     state_name::Symbol = :Ũ⃗,
     state_adjoint_name::Symbol = :Ũ⃗ₐ,
@@ -31,7 +33,7 @@ function AdjointUnitarySmoothPulseProblem(
     final_fidelity::Float64 = 1.0,
 )
     if piccolo_options.verbose
-        println("    constructing AdjointUnitarySmoothPulseProblem...")
+        println("    constructing UnitaryVariationalProblem...")
         println("\tusing integrator: $(typeof(unitary_integrator))")
     end
 
@@ -57,30 +59,26 @@ function AdjointUnitarySmoothPulseProblem(
         )
     end
 
+    state_vars = [traj[state_name] for _ ∈ system.G_vars]
 
+    full_rollout =  variational_unitary_rollout(
+        traj, 
+        system;
+        unitary_name=state_name,
+        drive_name=control_name
+    )[2]
 
-    adj = [traj[state_name] for s ∈ system.Gₐ]
-
-
-    # for a∈adj
-    #     a[:,1] *= 0 
-    # end
-
-
-    full_rollout =  adjoint_unitary_rollout(traj, system ;unitary_name=state_name,drive_name = control_name)[2]
-
-    for i in 1:length(system.Gₐ)
-        adj[i] = full_rollout[i]
+    for i in eachindex(system.G_vars)
+        state_vars[i] = full_rollout[i]
     end
 
+    state_var_names = [
+        add_suffix(state_adjoint_name, string(i)) for i in eachindex(system.G_vars)
+    ]
 
-    state_adjoint_names = [add_suffix(state_adjoint_name,string(i)) for i in 1:length(system.Gₐ)]
-
-
-    for (name,a) in zip(state_adjoint_names,adj)
-        add_component!(traj,name,a;type=:state)
-        traj.initial = merge(traj.initial, (name=> a[:,1], ))
-
+    for (name, var) in zip(state_var_names, state_vars)
+        add_component!(traj, name, var; type=:state)
+        traj.initial = merge(traj.initial, (name => var[:,1], ))
     end
 
     fidelity_constraint = FinalUnitaryFidelityConstraint(
@@ -98,13 +96,13 @@ function AdjointUnitarySmoothPulseProblem(
     J += QuadraticRegularizer(control_names[2], traj, R_da)
     J += QuadraticRegularizer(control_names[3], traj, R_dda)
 
-    for (name,t) ∈ zip(state_adjoint_names,times)
+    for (name, t) ∈ zip(state_var_names,times)
         if(!isnothing(t))
             J += UnitaryNormLoss(name, traj, t; Q=Q)
         end
     end
 
-    for (name,down_t) ∈ zip(state_adjoint_names,down_times)
+    for (name, down_t) ∈ zip(state_var_names,down_times)
         if(!isnothing(down_t))
             J += UnitaryNormLoss(name, traj, down_t; Q=Qd, rep = false)
         end
@@ -114,19 +112,18 @@ function AdjointUnitarySmoothPulseProblem(
     if(!isnothing(times) || !isnothing(down_times))
         Q_reg = Qd/100
     end 
-    for name ∈ state_adjoint_names
+    for name ∈ state_var_names
         J += UnitaryNormLoss(name, traj, 1:T; Q=Q_reg, rep = false) ###small regularization 
     end
 
-
     # Optional Piccolo constraints and objectives
     apply_piccolo_options!(
         J, constraints, piccolo_options, traj, state_name, timestep_name;
         state_leakage_indices=goal isa EmbeddedOperator ? get_leakage_indices(goal) : nothing
     )
 
     integrators = [
-        unitary_integrator(system, traj, state_name,state_adjoint_names,control_name),
+        unitary_integrator(system, traj, state_name,state_var_names,control_name),
         DerivativeIntegrator(traj, control_name, control_names[2]),
         DerivativeIntegrator(traj, control_names[2], control_names[3]),
     ]

src/quantum_integrators.jl

@@ -3,7 +3,7 @@ module QuantumIntegrators
 export KetIntegrator
 export UnitaryIntegrator
 export DensityMatrixIntegrator
-export AdjointUnitaryIntegrator
+export VariationalUnitaryIntegrator
 
 using LinearAlgebra
 using NamedTrajectories
@@ -13,6 +13,10 @@ using SparseArrays
 
 const ⊗ = kron
 
+# ----------------------------------------------------------------------------- #
+# Default Integrators
+# ----------------------------------------------------------------------------- #
+
 function KetIntegrator(
     sys::QuantumSystem,
     traj::NamedTrajectory, 
@@ -41,37 +45,34 @@ function DensityMatrixIntegrator(
     return BilinearIntegrator(sys.𝒢, traj, ρ̃, a)
 end
 
+# ----------------------------------------------------------------------------- #
+# Variational Integrators
+# ----------------------------------------------------------------------------- #
 
-function AdjointUnitaryIntegrator(
-    sys::ParameterizedQuantumSystem,
+function VariationalKetIntegrator(
+    sys::VariationalQuantumSystem,
     traj::NamedTrajectory, 
-    Ũ⃗::Symbol, 
-    Ũ⃗ₐ::Vector{Symbol},
+    ψ̃::Symbol, 
+    ψ̃_variations::AbstractVector{Symbol},
     a::Symbol
 ) 
-    n_sys = length(sys.Gₐ)
-
-    G = a_ -> I(sys.levels) ⊗ sys.G(a_)
-
-    Gai = (i,a_) -> I(sys.levels) ⊗ sys.Gₐ[i](a_)
+    var_ψ̃ = hcat(ψ̃, ψ̃_variations...)
+    G = a -> Isomorphisms.var_G(sys.G(a), [G(a) for G in sys.G_vars])
+    return BilinearIntegrator(G, traj, var_ψ̃, a)
+end
 
-    function Ĝ(a_)
-        G_eval = G(a_)
-        dim = size(G_eval)[1]    
-        Gx_index, Gy_index, G_val = findnz(G_eval)
-        G_full = spzeros((n_sys+1).*size(G_eval))
-
-        for i ∈ 0:n_sys
-            G_full +=    sparse((i*dim) .+ Gx_index, (i*dim) .+ Gy_index, G_val, size(G_full)...)
-            if(i<n_sys)
-                Ga_x_index, Ga_y_index, Ga_val = findnz(Gai(i+1,a_))
-                G_full +=    sparse((i*dim) .+ Ga_x_index, (n_sys*dim) .+ Ga_y_index, Ga_val, size(G_full)...)
-            end
-        end 
-        return G_full
-    end
-
-    return AdjointBilinearIntegrator(Ĝ, traj, Ũ⃗, Ũ⃗ₐ, a)
+function VariationalUnitaryIntegrator(
+    sys::VariationalQuantumSystem,
+    traj::NamedTrajectory, 
+    Ũ⃗::Symbol, 
+    Ũ⃗_variations::AbstractVector{Symbol},
+    a::Symbol
+) 
+    var_Ũ⃗ = hcat(Ũ⃗, Ũ⃗_variations...)
+    Ĝ = a -> Isomorphisms.var_G(
+        I(sys.levels) ⊗ sys.G(a), [I(sys.levels) ⊗ G(a) for G in sys.G_vars]
+    )
+    return BilinearIntegrator(Ĝ, traj, var_Ũ⃗, a)
 end