diff --git a/src/problem_templates/unitary_smooth_pulse_problem.jl b/src/problem_templates/unitary_smooth_pulse_problem.jl index 2a560f44..4062353b 100644 --- a/src/problem_templates/unitary_smooth_pulse_problem.jl +++ b/src/problem_templates/unitary_smooth_pulse_problem.jl @@ -73,6 +73,7 @@ function UnitarySmoothPulseProblem( unitary_integrator=UnitaryIntegrator, state_name::Symbol = :Ũ⃗, control_name::Symbol = :a, + time_name::Symbol = :t, timestep_name::Symbol = :Δt, init_trajectory::Union{NamedTrajectory, Nothing}=nothing, a_guess::Union{Matrix{Float64}, Nothing}=nothing, @@ -121,6 +122,11 @@ function UnitarySmoothPulseProblem( ) end + if typeof(system) == TimeDependentQuantumSystem + add_component!(traj, time_name, get_times(traj)) + update_bound!(traj, time_name, (0.0, Δt_max*T)) + end + # Objective J = UnitaryInfidelityObjective(goal, state_name, traj; Q=Q) @@ -138,12 +144,23 @@ function UnitarySmoothPulseProblem( 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, control_name), DerivativeIntegrator(traj, control_name, control_names[2]), DerivativeIntegrator(traj, control_names[2], control_names[3]), ] + if typeof(system) == TimeDependentQuantumSystem + integrators = [ + TimeDependentUnitaryIntegrator(system, traj, state_name, control_name, time_name), + integrators..., + TimeIntegrator(traj, time_name) + ] + else + integrators = [ + unitary_integrator(system, traj, state_name, control_name), + integrators... + ] + end return DirectTrajOptProblem( traj, diff --git a/src/quantum_integrators.jl b/src/quantum_integrators.jl index 8108a1c8..919c47f9 100644 --- a/src/quantum_integrators.jl +++ b/src/quantum_integrators.jl @@ -4,6 +4,8 @@ export KetIntegrator export UnitaryIntegrator export DensityMatrixIntegrator export VariationalUnitaryIntegrator +export TimeDependentKetIntegrator +export TimeDependentUnitaryIntegrator using LinearAlgebra using NamedTrajectories @@ -80,5 +82,30 @@ function VariationalUnitaryIntegrator( return BilinearIntegrator(Ĝ, traj, var_Ũ⃗, a) end +# ----------------------------------------------------------------------------- # +# Default Integrators +# ----------------------------------------------------------------------------- # + +function TimeDependentKetIntegrator( + sys::TimeDependentQuantumSystem, + traj::NamedTrajectory, + ψ̃::Symbol, + a::Symbol , + t::Symbol +) + return TimeDependentBilinearIntegrator(sys.G, traj, ψ̃, a, t) +end + +function TimeDependentUnitaryIntegrator( + sys::TimeDependentQuantumSystem, + traj::NamedTrajectory, + Ũ⃗::Symbol, + a::Symbol, + t::Symbol +) + Ĝ = (a_, t_) -> I(sys.levels) ⊗ sys.G(a_,t_) + return TimeDependentBilinearIntegrator(Ĝ, traj, Ũ⃗, a, t) +end + end \ No newline at end of file