traj init for unitary

src/problem_templates/unitary_sampling_problem.jl

@@ -47,7 +47,6 @@ function UnitarySamplingProblem(
     operator::OperatorType,
     T::Int,
     Δt::Union{Float64,Vector{Float64}};
-    system_labels=string.(1:length(systems)),
     system_weights=fill(1.0, length(systems)),
     init_trajectory::Union{NamedTrajectory,Nothing}=nothing,
     ipopt_options::IpoptOptions=IpoptOptions(),
@@ -73,9 +72,6 @@ function UnitarySamplingProblem(
     R_dda::Union{Float64,Vector{Float64}}=R,
     kwargs...
 )
-    # Create keys for multiple systems
-    Ũ⃗_names = [add_suffix(state_name, ℓ) for ℓ ∈ system_labels]
-
     # Trajectory
     if !isnothing(init_trajectory)
         traj = init_trajectory
@@ -99,34 +95,38 @@ function UnitarySamplingProblem(
             a_guess=a_guess,
             system=systems,
             rollout_integrator=piccolo_options.rollout_integrator,
-            state_names=Ũ⃗_names
         )
     end
 
+    state_names = [
+        name for name ∈ traj.names
+        if startswith(string(name), string(state_name))
+    ]
+
     control_names = [
         name for name ∈ traj.names
         if endswith(string(name), string(control_name))
     ]
 
     # Objective
     J = NullObjective()
-    for (wᵢ, Ũ⃗_name) in zip(system_weights, Ũ⃗_names)
-        J += wᵢ * UnitaryInfidelityObjective(
-            Ũ⃗_name, traj, Q;
+    for (weight, name) in zip(system_weights, state_names)
+        J += weight * UnitaryInfidelityObjective(
+            name, traj, Q;
             subspace=operator isa EmbeddedOperator ? operator.subspace_indices : nothing
         )
     end
-    J += QuadraticRegularizer(control_names[1], traj, R_a; timestep_name=timestep_name)
-    J += QuadraticRegularizer(control_names[2], traj, R_da; timestep_name=timestep_name)
-    J += QuadraticRegularizer(control_names[3], traj, R_dda; timestep_name=timestep_name)
+    for (R_name, name) in zip([R_a, R_da, R_dda], control_names)
+        J += QuadraticRegularizer(name, traj, R_name; timestep_name=timestep_name)
+    end
 
     # Constraints
     if piccolo_options.leakage_suppression
         if operator isa EmbeddedOperator
             leakage_indices = get_iso_vec_leakage_indices(operator)
-            for Ũ⃗_name in Ũ⃗_names
+            for name in state_names
                 J += L1Regularizer!(
-                    constraints, Ũ⃗_name, traj,
+                    constraints, name, traj,
                     R_value=piccolo_options.R_leakage,
                     indices=leakage_indices,
                     eval_hessian=piccolo_options.eval_hessian
@@ -154,16 +154,16 @@ function UnitarySamplingProblem(
 
     # Integrators
     unitary_integrators = AbstractIntegrator[]
-    for (sys, Ũ⃗_name) in zip(systems, Ũ⃗_names)
+    for (sys, name) in zip(systems, state_names)
         if piccolo_options.integrator == :pade
             push!(
                 unitary_integrators,
-                UnitaryPadeIntegrator(sys, Ũ⃗_name, control_name, traj; order=piccolo_options.pade_order)
+                UnitaryPadeIntegrator(sys, name, control_name, traj; order=piccolo_options.pade_order)
             )
         elseif piccolo_options.integrator == :exponential
             push!(
                 unitary_integrators,
-                UnitaryExponentialIntegrator(sys, Ũ⃗_name, control_name, traj)
+                UnitaryExponentialIntegrator(sys, name, control_name, traj)
             )
         else
             error("integrator must be one of (:pade, :exponential)")
@@ -262,13 +262,13 @@ end
         piccolo_options=pi_ops
     )
 
-    @test g1_prob.trajectory.Ũ⃗1 ≈ g1_prob.trajectory.Ũ⃗2
+    @test g1_prob.trajectory.Ũ⃗_system_1 ≈ g1_prob.trajectory.Ũ⃗_system_2
 
     g2_prob = UnitarySamplingProblem(
         [systems(0), systems(0.05)], operator, T, Δt;
         a_guess=a_guess,
         piccolo_options=pi_ops
     )
 
-    @test ~(g2_prob.trajectory.Ũ⃗1 ≈ g2_prob.trajectory.Ũ⃗2)
+    @test ~(g2_prob.trajectory.Ũ⃗_system_1 ≈ g2_prob.trajectory.Ũ⃗_system_2)
 end

src/problem_templates/unitary_smooth_pulse_problem.jl

@@ -175,7 +175,9 @@ function UnitarySmoothPulseProblem(
     ]
 
     # Optional Piccolo constraints and objectives
-    apply_piccolo_options!(J, constraints, piccolo_options, traj, operator, state_name, timestep_name)
+    apply_piccolo_options!(
+        J, constraints, piccolo_options, traj, operator, state_name, timestep_name
+    )
 
     return QuantumControlProblem(
         system,

src/trajectory_initialization.jl

@@ -252,15 +252,15 @@ initialize_control_trajectory(a::AbstractMatrix, Δt::Real, n_derivatives::Int)
 # ----------------------------------------------------------------------------- #
 
 """
-    initialize_unitary_trajectory
+    initialize_trajectory
 
 
-Initialize a trajectory for a unitary control problem. The trajectory is initialized with
+Initialize a trajectory for a control problem. The trajectory is initialized with
 data that should be consistently the same type (in this case, Float64).
 
 """
 function initialize_trajectory(
-    state_data::Vector{Matrix{Float64}},
+    state_data::Vector{<:AbstractMatrix{Float64}},
     state_inits::Vector{<:AbstractVector{Float64}},
     state_goals::Vector{<:AbstractVector{Float64}},
     state_names::AbstractVector{Symbol},
@@ -277,21 +277,21 @@ function initialize_trajectory(
     drive_derivative_σ::Float64=0.1,
     a_guess::Union{AbstractMatrix{<:Float64}, Nothing}=nothing,
     global_data::Union{NamedTuple, Nothing}=nothing,
+    verbose=false,
 )
-    @assert !isnothing(state_inits) "state_init must be provided"
-    @assert !isnothing(state_goals) "state_goal must be provided"
     @assert length(state_data) == length(state_names) == length(state_inits) == length(state_goals) "state_data, state_names, state_inits, and state_goals must have the same length"
-
     @assert length(control_bounds) == n_control_derivatives + 1 "control_bounds must have $n_control_derivatives + 1 elements"
 
     # assert that state names are unique
     @assert length(state_names) == length(Set(state_names)) "state_names must be unique"
 
     # Control data
     control_derivative_names = [
-        Symbol("d"^i * string(control_name))
-            for i = 1:n_control_derivatives
+        Symbol("d"^i * string(control_name)) for i = 1:n_control_derivatives
     ]
+    if verbose
+        println("control_derivative_names: $control_derivative_names")
+    end
 
     control_names = (control_name, control_derivative_names...)
 
@@ -312,7 +312,6 @@ function initialize_trajectory(
         timestep = Δt
     end
 
-
     # Constraints
     initial = (;
         (state_names .=> state_inits)...,
@@ -380,12 +379,12 @@ function initialize_trajectory(
     Δt::Union{Real, AbstractVecOrMat{<:Real}},
     args...;
     state_name::Symbol=:Ũ⃗,
-    state_names::AbstractVector{<:Symbol}=[state_name],
     U_init::AbstractMatrix{<:Number}=Matrix{ComplexF64}(I(size(U_goal, 1))),
     a_guess::Union{AbstractMatrix{<:Float64}, Nothing}=nothing,
     system::Union{AbstractQuantumSystem, AbstractVector{<:AbstractQuantumSystem}, Nothing}=nothing,
     rollout_integrator::Function=expv,
     geodesic=true,
+    verbose=false,
     kwargs...
 )
     Ũ⃗_init = operator_to_iso_vec(U_init)
@@ -396,49 +395,54 @@ function initialize_trajectory(
         Ũ⃗_goal = operator_to_iso_vec(U_goal)
     end
 
+    # Construct state data
     if isnothing(a_guess)
         # No guess provided, initialize a geodesic and randomly sample controls
-
         Ũ⃗_traj = initialize_unitary_trajectory(U_init, U_goal, T; geodesic=geodesic)
-
-        state_data = repeat([Ũ⃗_traj], length(state_names))
-    else
-        if system isa AbstractQuantumSystem
-            @assert size(a_guess, 1) == length(system.H_drives) "a_guess must have the same number of drives as n_drives"
-        elseif system isa AbstractVector
-            @assert size(a_guess, 1) == length(system[1].H_drives) "a_guess must have the same number of drives as n_drives"
+        if system isa AbstractVector
+            state_data = repeat([Ũ⃗_traj], length(system))
+        else
+            state_data = [Ũ⃗_traj]
         end
-
-        @assert size(a_guess, 2) == T "a_guess must have the same number of timesteps as T"
-        @assert !isnothing(system) "system must be provided if a_guess is provided"
-
+    else
         if Δt isa AbstractMatrix
             timesteps = vec(Δt)
         elseif Δt isa Float64
             timesteps = fill(Δt, T)
         else
             timesteps = Δt
         end
-
-        if system isa AbstractVector
-            @assert length(system) == length(state_names) "systems must have the same length as state_names"
+
+        @assert size(a_guess, 2) == T "a_guess must have the same number of timesteps as T"
+        if system isa AbstractQuantumSystem
+            @assert size(a_guess, 1) == length(system.H_drives) "a_guess must have the same number of drives as n_drives"
+            state_data = [unitary_rollout(Ũ⃗_init, a_guess, timesteps, system; integrator=rollout_integrator)]
+        elseif system isa AbstractVector
             state_data = map(system) do sys
+                @assert size(a_guess, 1) == length(sys.H_drives) "a_guess must have the same number of drives as n_drives"
                 unitary_rollout(Ũ⃗_init, a_guess, timesteps, sys; integrator=rollout_integrator)
             end
         else
-            state = unitary_rollout(Ũ⃗_init, a_guess, timesteps, system; integrator=rollout_integrator)
-            state_data = [state]
+            error("System must be provided if a_guess is provided.")
         end
-        state_data = Matrix{Float64}.(state_data)
     end
 
-    if system isa AbstractVector && length(state_names) != length(system)
-        state_names = [string(state_name) * "_system_$i" for i = 1:length(system)]
-        @warn "length of state_names and number of systems ($(length(system))) are not equal, created state names for each system: " state_names
+    # Create a state name for each system
+    if system isa AbstractVector && length(system) > 1
+        state_names = Symbol.([string(state_name) * "_system_$i" for i in eachindex(system)])
+        if verbose
+            println("Created state names for ($(length(system))) systems: $state_names")
+        end
+        state_inits = repeat([Ũ⃗_init], length(state_names))
+        state_goals = repeat([Ũ⃗_goal], length(state_names))
+    else
+        state_names = [state_name]
+        state_inits = [Ũ⃗_init]
+        state_goals = [Ũ⃗_goal]
     end
 
-    state_inits = repeat([Ũ⃗_init], length(state_names))
-    state_goals = repeat([Ũ⃗_goal], length(state_names))
+    # Convert data to Float64
+    state_data = Matrix{Float64}.(state_data)
 
     return initialize_trajectory(
         state_data,
@@ -449,6 +453,7 @@ function initialize_trajectory(
         Δt,
         args...;
         a_guess=a_guess,
+        verbose=verbose,
         kwargs...
     )
 end
@@ -466,6 +471,7 @@ function initialize_trajectory(
     a_guess::Union{AbstractMatrix{<:Float64}, Nothing}=nothing,
     system::Union{AbstractQuantumSystem, AbstractVector{<:AbstractQuantumSystem}, Nothing}=nothing,
     rollout_integrator::Function=expv,
+    verbose=false,
     kwargs...
 )
     @assert length(ψ_inits) == length(ψ_goals) "ψ_inits and ψ_goals must have the same length"
@@ -525,7 +531,13 @@ function initialize_trajectory(
                 Symbol.(string.(state_names) .* "_system_$i")
                     for i = 1:length(system)
             ]...)
-            @warn "length of state_names and number of systems ($(length(system))) * number of states ($(length(ψ_goals))) are not equal, created state names for each system: " state_names
+            if verbose
+                println(
+                    "length of state_names and number of systems ($(length(system))) * ",
+                    "number of states ($(length(ψ_goals))) are not equal, created state ",
+                    "names for each system: $state_names"
+                )
+            end
         end
         ψ̃_inits = repeat(ψ̃_inits, length(system))
         ψ̃_goals = repeat(ψ̃_goals, length(system))
@@ -543,6 +555,7 @@ function initialize_trajectory(
         Δt,
         args...;
         a_guess=a_guess,
+        verbose=verbose,
         kwargs...
     )
 end