tests passing for feat

src/problem_templates/quantum_state_smooth_pulse_problem.jl

@@ -54,6 +54,7 @@ function QuantumStateSmoothPulseProblem(
     L1_regularized_names=Symbol[],
     L1_regularized_indices::NamedTuple=NamedTuple(),
     verbose=false,
+    kwargs...
 )
     @assert all(name ∈ L1_regularized_names for name in keys(L1_regularized_indices) if !isempty(L1_regularized_indices[name]))
 
@@ -239,6 +240,7 @@ function QuantumStateSmoothPulseProblem(
         linear_solver=linear_solver,
         verbose=verbose,
         ipopt_options=ipopt_options,
+        kwargs...
     )
 end
 
@@ -250,4 +252,4 @@ function QuantumStateSmoothPulseProblem(
 )
     system = QuantumSystem(H_drift, H_drives)
     return QuantumStateSmoothPulseProblem(system, args...; kwargs...)
-end
+end

src/problem_templates/unitary_smooth_pulse_problem.jl

@@ -104,6 +104,7 @@ function UnitarySmoothPulseProblem(
     jacobian_structure=true,
     hessian_approximation=false,
     blas_multithreading=true,
+    kwargs...
 )
     if operator isa EmbeddedOperator
         U_goal = operator.operator
@@ -274,7 +275,8 @@ function UnitarySmoothPulseProblem(
         ipopt_options=ipopt_options,
         jacobian_structure=jacobian_structure,
         hessian_approximation=hessian_approximation,
-        eval_hessian=!hessian_approximation
+        eval_hessian=!hessian_approximation,
+        kwargs...
     )
 end
 

src/problems.jl

@@ -8,6 +8,7 @@ export initialize_trajectory!
 export update_trajectory!
 export get_traj_data
 export get_datavec
+export get_objective
 
 using ..QuantumSystems
 using ..Integrators
@@ -50,6 +51,7 @@ function QuantumControlProblem(
     traj::NamedTrajectory,
     obj::Objective,
     dynamics::QuantumDynamics;
+    additional_objective::Union{Nothing, Objective}=nothing,
     eval_hessian::Bool=true,
     options::Options=Options(),
     constraints::Vector{<:AbstractConstraint}=AbstractConstraint[],
@@ -65,6 +67,10 @@ function QuantumControlProblem(
 
     nonlinear_constraints = NonlinearConstraint[con for con ∈ constraints if con isa NonlinearConstraint]
 
+    if !isnothing(additional_objective)
+        obj += additional_objective
+    end
+
     if verbose
         println("    building evaluator...")
     end
@@ -286,5 +292,14 @@ end
     prob.trajectory = NamedTrajectory(Z⃗, prob.trajectory)
 end
 
+"""
+    get_objective(prob::QuantumControlProblem)
+
+Return the objective function of the `prob::QuantumControlProblem`.
+"""
+function get_objective(prob::QuantumControlProblem)
+    return Objective(prob.params[:objective_terms])
+end
+
 
 end

test/problem_templates_test.jl

@@ -7,7 +7,7 @@
 # ------------------------------------------------
 
 
-@testset "Problem Templates" begin
+@testitem "Problem Templates" begin
     H_drift = GATES[:Z]
     H_drives = [GATES[:X], GATES[:Y]]
     U_goal = GATES[:H]
@@ -18,7 +18,7 @@
     # 1. test UnitarySmoothPulseProblem
     # --------------------------------------------
     prob = UnitarySmoothPulseProblem(
-        H_drift, H_drives, U_goal, T, Δt, 
+        H_drift, H_drives, U_goal, T, Δt,
         ipopt_options=Options(print_level=4)
     )
 
@@ -33,8 +33,8 @@
     final_fidelity = 0.99
 
     mintime_prob = UnitaryMinimumTimeProblem(
-        prob, 
-        final_fidelity=final_fidelity, 
+        prob,
+        final_fidelity=final_fidelity,
         ipopt_options=Options(print_level=4)
     )
 
@@ -43,12 +43,11 @@
     @test unitary_fidelity(mintime_prob) > final_fidelity
 
     @test sum(mintime_prob.trajectory[:Δt]) < sum(prob.trajectory[:Δt])
-
 end
 
 
 
-@testset "Robust and Subspace Templates" begin
+@testitem "Robust and Subspace Templates" begin
     # TODO: Improve these tests.
     # --------------------------------------------
     # Initialize with UnitarySmoothPulseProblem
@@ -112,8 +111,8 @@ end
 
     probs["unconstrained"] = UnitaryRobustnessProblem(
         H_error, trajectory, system, objective, integrators, constraints,
-        final_fidelity=0.99, 
-        subspace=subspace, 
+        final_fidelity=0.99,
+        subspace=subspace,
         ipopt_options=Options(recalc_y="yes", recalc_y_feas_tol=1e-1, print_level=4)
     )
     solve!(probs["unconstrained"]; max_iter=100)

test/problems_test.jl

@@ -20,3 +20,35 @@
 
 
 end
+
+@testitem "Additional Objective" begin
+    H_drift = GATES[:Z]
+    H_drives = [GATES[:X], GATES[:Y]]
+    U_goal = GATES[:H]
+    T = 50
+    Δt = 0.2
+
+    prob_vanilla = UnitarySmoothPulseProblem(
+        H_drift, H_drives, U_goal, T, Δt,
+        ipopt_options=Options(print_level=4)
+    )
+
+    J_extra = QuadraticSmoothnessRegularizer(:dda, prob_vanilla.trajectory, 10.0)
+
+    prob_additional = UnitarySmoothPulseProblem(
+        H_drift, H_drives, U_goal, T, Δt,
+        ipopt_options=Options(print_level=4),
+        additional_objective=J_extra,
+    )
+
+    J_prob_vanilla = Problems.get_objective(prob_vanilla)
+
+    J_additional = Problems.get_objective(prob_additional)
+
+    Z = prob_vanilla.trajectory
+    Z⃗ = prob_vanilla.trajectory.datavec
+
+    @test J_prob_vanilla.L(Z⃗, Z) + J_extra.L(Z⃗, Z) ≈ J_additional.L(Z⃗, Z)
+
+
+end