diff --git a/docs/make.jl b/docs/make.jl index ac9eb97..73cb732 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -28,6 +28,7 @@ links = InterLinks( joinpath(@__DIR__, "src", "inventories", "TimerOutputs.toml") ), "QuantumPropagators" => "https://juliaquantumcontrol.github.io/QuantumPropagators.jl/$DEV_OR_STABLE", + "QuantumGradientGenerators" => "https://juliaquantumcontrol.github.io/QuantumGradientGenerators.jl/$DEV_OR_STABLE", "QuantumControl" => "https://juliaquantumcontrol.github.io/QuantumControl.jl/$DEV_OR_STABLE", "GRAPE" => "https://juliaquantumcontrol.github.io/GRAPE.jl/$DEV_OR_STABLE", "Examples" => "https://juliaquantumcontrol.github.io/QuantumControlExamples.jl/$DEV_OR_STABLE", diff --git a/docs/src/refs.bib b/docs/src/refs.bib index fff45f8..511d0ac 100644 --- a/docs/src/refs.bib +++ b/docs/src/refs.bib @@ -231,3 +231,13 @@ @article{MachnesPRL2018 Pages = {150401}, Volume = {120}, } + +@article{GoerzNJP2014, + Author = {Goerz, Michael H. and Reich, Daniel M. and Koch, Christiane P.}, + Title = {Optimal control theory for a unitary operation under dissipative evolution}, + Journal = njp, + Year = {2014}, + Doi = {10.1088/1367-2630/16/5/055012}, + Pages = {055012}, + Volume = {16}, +} diff --git a/src/optimize.jl b/src/optimize.jl index 323e997..19383e1 100644 --- a/src/optimize.jl +++ b/src/optimize.jl @@ -1,10 +1,10 @@ using QuantumControlBase.QuantumPropagators.Generators: Operator using QuantumControlBase.QuantumPropagators.Controls: evaluate, evaluate! -using QuantumControlBase.QuantumPropagators: prop_step!, set_state!, reinit_prop! +using QuantumControlBase.QuantumPropagators: prop_step!, set_state!, reinit_prop!, propagate using QuantumControlBase.QuantumPropagators.Storage: write_to_storage!, get_from_storage! using QuantumGradientGenerators: resetgradvec! using QuantumControlBase: make_chi, make_grad_J_a, set_atexit_save_optimization -using QuantumControlBase: @threadsif +using QuantumControlBase: @threadsif, Trajectory using LinearAlgebra using Printf @@ -15,9 +15,8 @@ import QuantumControlBase: optimize result = optimize(problem; method=:GRAPE, kwargs...) ``` -optimizes the given -control [`problem`](@ref QuantumControlBase.ControlProblem) via the GRAPE -method, by minimizing the functional +optimizes the given control [`problem`](@ref QuantumControlBase.ControlProblem) +via the GRAPE method, by minimizing the functional ```math J(\{ϵ_{ln}\}) = J_T(\{|ϕ_k(T)⟩\}) + λ_a J_a(\{ϵ_{ln}\}) @@ -31,19 +30,21 @@ the time grid. Returns a [`GrapeResult`](@ref). Keyword arguments that control the optimization are taken from the keyword -arguments used in the instantiation of `problem`. +arguments used in the instantiation of `problem`; any of these can be overridden +with explicit keyword arguments to `optimize`. + # Required problem keyword arguments -* `J_T`: A function `J_T(ϕ, objectives; τ=τ)` that evaluates the final time +* `J_T`: A function `J_T(ϕ, trajectories; τ=τ)` that evaluates the final time functional from a vector `ϕ` of forward-propagated states and - `problem.objectives`. For all `objectives` that define a `target_state`, the - element `τₖ` of the vector `τ` will contain the overlap of the state `ϕₖ` - with the `target_state` of the `k`'th objective, or `NaN` otherwise. + `problem.trajectories`. For all `trajectories` that define a `target_state`, + the element `τₖ` of the vector `τ` will contain the overlap of the state `ϕₖ` + with the `target_state` of the `k`'th trajectory, or `NaN` otherwise. # Optional problem keyword arguments -* `chi`: A function `chi!(χ, ϕ, objectives)` what receives a list `ϕ` +* `chi`: A function `chi!(χ, ϕ, trajectories)` what receives a list `ϕ` of the forward propagated states and must set ``|χₖ⟩ = -∂J_T/∂⟨ϕₖ|``. If not given, it will be automatically determined from `J_T` via [`make_chi`](@ref) with the default parameters. @@ -81,7 +82,7 @@ arguments used in the instantiation of `problem`. * `pulse_options`: A dictionary that maps every control (as obtained by [`get_controls`](@ref QuantumControlBase.QuantumPropagators.Controls.get_controls) from the - `problem.objectives`) to a dict with the following possible keys: + `problem.trajectories`) to a dict with the following possible keys: - `:upper_bounds`: A vector of upper bound values, one for each intervals of the time grid. Values of `Inf` indicate an unconstrained upper bound for @@ -112,26 +113,34 @@ arguments used in the instantiation of `problem`. * `optimizer`: An optional Optim.jl optimizer (`Optim.AbstractOptimizer` instance). If not given, an [L-BFGS-B](https://github.com/Gnimuc/LBFGSB.jl) optimizer will be used. -* `prop_method`/`fw_prop_method`/`bw_prop_method`: The propagation method to - use for each objective, see below. -* `prop_method`/`fw_prop_method`/`grad_prop_method`: The propagation method to - use for the extended gradient vector for each objective, see below. +* `prop_method`: The propagation method to use for each trajectory, see below. * `verbose=false`: If `true`, print information during initialization -The propagation method for the forward propagation of each objective is -determined by the first available item of the following: +# Trajectory propagation + +GRAPE may involve three types of propagation: + +* A forward propagation for every [`Trajectory`](@ref) in the `problem` +* A backward propagation for every trajectory +* A backward propagation of a + [gradient generator](@extref QuantumGradientGenerators.GradGenerator) + for every trajectory. + +The keyword arguments for each propagation (see [`propagate`](@ref)) are +determined from any properties of each [`Trajectory`](@ref) that have a `prop_` +prefix, cf. [`init_prop_trajectory`](@ref). -* a `fw_prop_method` keyword argument -* a `prop_method` keyword argument -* a property `fw_prop_method` of the objective -* a property `prop_method` of the objective -* the value `:auto` +In situations where different parameters are required for the forward and +backward propagation, instead of the `prop_` prefix, the `fw_prop_` and +`bw_prop_` prefix can be used, respectively. These override any setting with +the `prop_` prefix. Similarly, properties for the backward propagation of the +gradient generators can be set with properties that have a `grad_prop_` prefix. +These prefixes apply both to the properties of each [`Trajectory`](@ref) and +the problem keyword arguments. -The propagation method for the backward propagation is determined similarly, -but with `bw_prop_method` instead of `fw_prop_method`. The propagation method -for the backward propagation of the extended gradient vector for each objective -is determined from `grad_prop_method`, `fw_prop_method`, `prop_method` in order -of precedence. +Note that the propagation method for each propagation must be specified. In +most cases, it is sufficient (and recommended) to pass a global `prop_method` +problem keyword argument. """ optimize(problem, method::Val{:GRAPE}) = optimize_grape(problem) optimize(problem, method::Val{:grape}) = optimize_grape(problem) @@ -156,10 +165,10 @@ function optimize_grape(problem) wrk = GrapeWrk(problem; verbose) χ = wrk.chi_states - Ψ₀ = [obj.initial_state for obj ∈ wrk.objectives] + Ψ₀ = [traj.initial_state for traj ∈ wrk.trajectories] Ψtgt = Union{eltype(Ψ₀),Nothing}[ - (hasproperty(obj, :target_state) ? obj.target_state : nothing) for - obj ∈ wrk.objectives + (hasproperty(traj, :target_state) ? traj.target_state : nothing) for + traj ∈ wrk.trajectories ] J = wrk.J_parts @@ -173,7 +182,7 @@ function optimize_grape(problem) chi! = wrk.kwargs[:chi] else # we only want to evaluate `make_chi` if `chi` is not a kwarg - chi! = make_chi(J_T_func, wrk.objectives) + chi! = make_chi(J_T_func, wrk.trajectories) end grad_J_a! = nothing if !isnothing(J_a_func) @@ -183,7 +192,7 @@ function optimize_grape(problem) τ = wrk.result.tau_vals ∇τ = wrk.tau_grads N_T = length(tlist) - 1 - N = length(wrk.objectives) + N = length(wrk.trajectories) L = length(wrk.controls) Φ = wrk.fw_storage @@ -215,7 +224,7 @@ function optimize_grape(problem) τ[k] = isnothing(Ψtgt[k]) ? NaN : (Ψtgt[k] ⋅ Ψₖ) end Ψ = [p.state for p ∈ wrk.fw_propagators] - J[1] = J_T_func(Ψ, wrk.objectives; τ=τ) + J[1] = J_T_func(Ψ, wrk.trajectories; τ=τ) if !isnothing(J_a_func) J[2] = λₐ * J_a_func(pulsevals, tlist) end @@ -239,7 +248,7 @@ function optimize_grape(problem) # backward propagation of combined χ-state and gradient Ψ = [p.state for p ∈ wrk.fw_propagators] - chi!(χ, Ψ, wrk.objectives; τ=τ) # τ from f(...) + chi!(χ, Ψ, wrk.trajectories; τ=τ) # τ from f(...) @threadsif wrk.use_threads for k = 1:N local Ψₖ = wrk.fw_propagators[k].state local χ̃ₖ = wrk.bw_grad_propagators[k].state @@ -283,12 +292,12 @@ function optimize_grape(problem) # backward propagation of χ-state Ψ = [p.state for p ∈ wrk.fw_propagators] - chi!(χ, Ψ, wrk.objectives; τ=τ) # τ from f(...) + chi!(χ, Ψ, wrk.trajectories; τ=τ) # τ from f(...) @threadsif wrk.use_threads for k = 1:N local Ψₖ = wrk.fw_propagators[k].state reinit_prop!(wrk.bw_propagators[k], χ[k]; transform_control_ranges) local χₖ = wrk.bw_propagators[k].state - local Hₖ⁺ = wrk.adjoint_objectives[k].generator + local Hₖ⁺ = wrk.adjoint_trajectories[k].generator local Hₖₙ⁺ = wrk.taylor_genops[k] for n = N_T:-1:1 # N_T is the number of time slices # TODO: It would be cleaner to encapsulate this in a @@ -491,7 +500,7 @@ end # minus sign in front of the derivative, compensated by the minus sign in the # factor ``(-2)`` of the final ``(∇J_T)_{ln}``. function _grad_J_T_via_chi!(∇J_T, τ, ∇τ) - N = length(τ) # number of objectives + N = length(τ) # number of trajectories L, N_T = size(∇τ[1]) # number of controls/time intervals ∇J_T′ = reshape(∇J_T, L, N_T) # writing to ∇J_T′ modifies ∇J_T for l = 1:L diff --git a/src/result.jl b/src/result.jl index e873f53..a33139b 100644 --- a/src/result.jl +++ b/src/result.jl @@ -25,17 +25,17 @@ mutable struct GrapeResult{STST} function GrapeResult(problem) tlist = problem.tlist - controls = get_controls(problem.objectives) + controls = get_controls(problem.trajectories) iter_start = get(problem.kwargs, :iter_start, 0) iter_stop = get(problem.kwargs, :iter_stop, 5000) iter = iter_start secs = 0 - tau_vals = zeros(ComplexF64, length(problem.objectives)) + tau_vals = zeros(ComplexF64, length(problem.trajectories)) guess_controls = [discretize(control, tlist) for control in controls] J_T = 0.0 J_T_prev = 0.0 optimized_controls = [copy(guess) for guess in guess_controls] - states = [similar(obj.initial_state) for obj in problem.objectives] + states = [similar(traj.initial_state) for traj in problem.trajectories] start_local_time = now() end_local_time = now() records = Vector{Tuple}() @@ -74,7 +74,7 @@ Base.show(io::IO, ::MIME"text/plain", r::GrapeResult) = print( GRAPE Optimization Result ------------------------- - Started at $(r.start_local_time) -- Number of objectives: $(length(r.states)) +- Number of trajectories: $(length(r.states)) - Number of iterations: $(max(r.iter - r.iter_start, 0)) - Number of pure func evals: $(r.f_calls) - Number of func/grad evals: $(r.fg_calls) diff --git a/src/workspace.jl b/src/workspace.jl index dedb126..aee0d1e 100644 --- a/src/workspace.jl +++ b/src/workspace.jl @@ -2,19 +2,22 @@ import QuantumControlBase using QuantumControlBase.QuantumPropagators: init_prop using QuantumControlBase.QuantumPropagators.Storage: init_storage using QuantumControlBase.QuantumPropagators.Controls: get_controls, discretize_on_midpoints -using QuantumControlBase: get_control_derivs +using QuantumControlBase: Trajectory, get_control_derivs, init_prop_trajectory using QuantumGradientGenerators: GradVector, GradGenerator using ConcreteStructs # GRAPE workspace (for internal use) @concrete terse mutable struct GrapeWrk - # a copy of the objectives - objectives + # a copy of the trajectories + trajectories - # the adjoint objectives, containing the adjoint generators for the + # the adjoint trajectories, containing the adjoint generators for the # backward propagation - adjoint_objectives + adjoint_trajectories + + # trajectories for bw-prop of gradients + grad_trajectories # The kwargs from the control problem kwargs @@ -58,7 +61,7 @@ using ConcreteStructs result ################################# - # scratch objects, per objective: + # scratch objects, per trajectory: # backward-propagated states chi_states @@ -66,11 +69,7 @@ using ConcreteStructs # gradients ∂τₖ/ϵₗ(tₙ) tau_grads::Vector{Matrix{ComplexF64}} - # dynamical generator for grad-bw-propagation, time-dependent - # gradient_method=:gradgen only - gradgen - - # backward storage array (per objective) + # backward storage array fw_storage # for normal forward propagation @@ -92,7 +91,7 @@ using ConcreteStructs # gradient_method=:taylor only control_derivs - # 5 temporary states for each objective and each control, for evaluating + # 5 temporary states for each trajectory and each control, for evaluating # gradients via Taylor expansions # gradient_method=:taylor only taylor_grad_states @@ -104,11 +103,11 @@ end function GrapeWrk(problem::QuantumControlBase.ControlProblem; verbose=false) use_threads = get(problem.kwargs, :use_threads, false) gradient_method = get(problem.kwargs, :gradient_method, :gradgen) - objectives = [obj for obj in problem.objectives] - adjoint_objectives = [adjoint(obj) for obj in problem.objectives] - controls = get_controls(objectives) + trajectories = [traj for traj in problem.trajectories] + adjoint_trajectories = [adjoint(traj) for traj in problem.trajectories] + controls = get_controls(trajectories) if length(controls) == 0 - error("no controls in objectives: cannot optimize") + error("no controls in trajectories: cannot optimize") end tlist = problem.tlist # interleave the pulse values as [ϵ₁(t̃₁), ϵ₂(t̃₁), ..., ϵ₁(t̃₂), ϵ₂(t̃₂), ...] @@ -182,103 +181,82 @@ function GrapeWrk(problem::QuantumControlBase.ControlProblem; verbose=false) end alpha = 0.0 dummy_vals = IdDict(control => 1.0 for (i, control) in enumerate(controls)) - fw_storage = [init_storage(obj.initial_state, tlist) for obj in objectives] + fw_storage = [init_storage(traj.initial_state, tlist) for traj in trajectories] kwargs[:piecewise] = true # only accept piecewise propagators - fw_prop_method = [ - QuantumControlBase.get_objective_prop_method( - obj, - :fw_prop_method, - :prop_method; - kwargs... - ) for obj in objectives - ] - bw_prop_method = [ - QuantumControlBase.get_objective_prop_method( - obj, - :bw_prop_method, - :prop_method; - kwargs... - ) for obj in objectives - ] - grad_prop_method = [ - QuantumControlBase.get_objective_prop_method( - obj, - :grad_prop_method, - :bw_prop_method, - :prop_method; - kwargs... - ) for obj in objectives - ] - + _prefixes = ["prop_", "fw_prop_"] fw_propagators = [ - begin - verbose && - @info "Initializing fw-prop of objective $k with method $(fw_prop_method[k])" - init_prop( - obj.initial_state, - obj.generator, - tlist; - method=fw_prop_method[k], - parameters=parameters, - kwargs... - ) - end for (k, obj) in enumerate(objectives) + init_prop_trajectory( + traj, + tlist; + verbose, + _msg="Initializing fw-prop of trajectory $k", + _prefixes, + _filter_kwargs=true, + fw_prop_parameters=parameters, # will filter to `parameters` + kwargs... + ) for (k, traj) in enumerate(trajectories) ] - chi_states = [similar(obj.initial_state) for obj in objectives] + chi_states = [similar(traj.initial_state) for traj in trajectories] tau_grads::Vector{Matrix{ComplexF64}} = - [zeros(ComplexF64, length(controls), length(tlist) - 1) for _ in objectives] + [zeros(ComplexF64, length(controls), length(tlist) - 1) for _ in trajectories] if gradient_method == :gradgen - gradgen = [GradGenerator(obj.generator) for obj in adjoint_objectives] - bw_grad_propagators = [ + grad_trajectories = [ begin - verbose && - @info "Initializing gradient bw-prop of objective $k with method $(grad_prop_method[k])" χ̃ₖ = GradVector(chi_states[k], length(controls)) - init_prop( - χ̃ₖ, - gradgen[k], - tlist; - method=grad_prop_method[k], - backward=true, - parameters=parameters, - kwargs... - ) - end for k ∈ eachindex(objectives) + G̃ₖ = GradGenerator(traj.generator) + Trajectory(χ̃ₖ, G̃ₖ, getfield(traj, :kwargs)...) + end for (k, traj) in enumerate(adjoint_trajectories) + ] + _prefixes = ["prop_", "bw_prop_", "grad_prop_"] + bw_grad_propagators = [ + init_prop_trajectory( + traj, + tlist; + verbose, + _msg="Initializing bw-gradient-prop of trajectory $k", + _prefixes, + _filter_kwargs=true, + grad_prop_backward=true, # will filter to `backward=true` + grad_prop_parameters=parameters, # will filter to `parameters` + kwargs... + ) for (k, traj) in enumerate(grad_trajectories) ] bw_propagators = [] taylor_genops = [] control_derivs = [] taylor_grad_states = [] elseif gradient_method == :taylor - gradgen = [] + grad_trajectories = [] bw_grad_propagators = [] + _prefixes = ["prop_", "bw_prop_"] bw_propagators = [ - begin - verbose && - @info "Initializing bw-prop of objective $k with method $(bw_prop_method[k])" - init_prop( - obj.initial_state, - obj.generator, - tlist; - method=bw_prop_method[k], - backward=true, - parameters=parameters, - kwargs... - ) - end for (k, obj) in enumerate(adjoint_objectives) + init_prop_trajectory( + traj, + tlist; + verbose, + _msg="Initializing bw-prop of trajectory $k", + _prefixes, + _filter_kwargs=true, + bw_prop_backward=true, # will filter to `backward=true` + bw_prop_parameters=parameters, # will filter to `parameters` + kwargs... + ) for (k, traj) in enumerate(adjoint_trajectories) ] - taylor_genops = [evaluate(obj.generator, tlist, 1) for obj in adjoint_objectives] - control_derivs = [get_control_derivs(obj.generator, controls) for obj in objectives] + taylor_genops = + [evaluate(traj.generator, tlist, 1) for traj in adjoint_trajectories] + control_derivs = + [get_control_derivs(traj.generator, controls) for traj in trajectories] taylor_grad_states = [ - Tuple(similar(objectives[k].initial_state) for _ = 1:5) for - l = 1:length(controls), k = 1:length(objectives) + Tuple(similar(trajectories[k].initial_state) for _ = 1:5) for + l = 1:length(controls), k = 1:length(trajectories) ] else error("Invalid gradient_method=$(repr(gradient_method)) ∉ (:gradgen, :taylor)") end GrapeWrk( - objectives, - adjoint_objectives, + trajectories, + adjoint_trajectories, + grad_trajectories, kwargs, controls, pulsevals_guess, @@ -296,7 +274,6 @@ function GrapeWrk(problem::QuantumControlBase.ControlProblem; verbose=false) result, chi_states, tau_grads, - gradgen, fw_storage, fw_propagators, bw_grad_propagators, diff --git a/test/test_empty_optimization.jl b/test/test_empty_optimization.jl index d413263..8fd9cb0 100644 --- a/test/test_empty_optimization.jl +++ b/test/test_empty_optimization.jl @@ -1,6 +1,6 @@ using Test using StableRNGs -using QuantumControl: hamiltonian, optimize, ControlProblem, Objective +using QuantumControl: hamiltonian, optimize, ControlProblem, Trajectory using QuantumControl.Controls: get_controls using QuantumControl.Functionals: J_T_re using QuantumControlTestUtils.RandomObjects: random_matrix, random_state_vector @@ -14,21 +14,21 @@ using QuantumControlTestUtils.RandomObjects: random_matrix, random_state_vector N = 10 H = random_matrix(N; rng) - objectives = [ - Objective(; + trajectories = [ + Trajectory(; initial_state=random_state_vector(N; rng), generator=H, target_state=random_state_vector(N; rng) ) ] - @test length(get_controls(objectives)) == 0 + @test length(get_controls(trajectories)) == 0 tlist = collect(range(0; length=1001, step=1.0)) - problem = ControlProblem(; objectives, tlist, J_T=J_T_re) + problem = ControlProblem(trajectories, tlist; J_T=J_T_re) - msg = "no controls in objectives: cannot optimize" + msg = "no controls in trajectories: cannot optimize" @test_throws ErrorException(msg) optimize(problem; method=:GRAPE) end diff --git a/test/test_pulse_optimization.jl b/test/test_pulse_optimization.jl index 50565f8..51458b6 100644 --- a/test/test_pulse_optimization.jl +++ b/test/test_pulse_optimization.jl @@ -10,7 +10,7 @@ using QuantumControl.Functionals: J_T_re # Test the resolution of # https://github.com/JuliaQuantumControl/Krotov.jl/issues/28 - # + # # While this hasn't been a problem for GRAPE, we'd want to make sure that # any future changes won't result in the optimization mutating the guess # controls @@ -19,16 +19,16 @@ using QuantumControl.Functionals: J_T_re problem = dummy_control_problem(; pulses_as_controls=true) nt = length(problem.tlist) - guess_pulse = QuantumControl.Controls.get_controls(problem.objectives)[1] + guess_pulse = QuantumControl.Controls.get_controls(problem.trajectories)[1] @test length(guess_pulse) == nt - 1 - guess_pulse_copy = copy(QuantumControl.Controls.get_controls(problem.objectives)[1]) + guess_pulse_copy = copy(QuantumControl.Controls.get_controls(problem.trajectories)[1]) # Optimizing this should not modify the original generator in any way res = optimize(problem; method=:GRAPE, J_T=J_T_re, iter_stop=2) opt_control = res.optimized_controls[1] @test length(opt_control) == nt # optimized_controls are always *on* tlist opt_pulse = discretize_on_midpoints(opt_control, problem.tlist) - post_pulse = QuantumControl.Controls.get_controls(problem.objectives)[1] + post_pulse = QuantumControl.Controls.get_controls(problem.trajectories)[1] # * The generator should still have the exact same objects as controls @test guess_pulse ≡ post_pulse