From ddc4298d1f19141dfc02d186570d4077cd1a54a9 Mon Sep 17 00:00:00 2001 From: Michael Goerz Date: Fri, 19 Jan 2024 18:08:47 -0500 Subject: [PATCH] Avoid unnecessary evaluation of chi If an explicit `chi` is given, do not evaluate one with AD. Also, do not reconstruct a chi! routine in every Krotov iteration. Moving `make_chi` out of the iteration loop potentially avoids compilation overhead in each iteration. --- src/optimize.jl | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/optimize.jl b/src/optimize.jl index 02f1334..737dbcd 100644 --- a/src/optimize.jl +++ b/src/optimize.jl @@ -116,6 +116,14 @@ function optimize_krotov(problem) ϵ⁽ⁱ⁾ = wrk.pulses0 ϵ⁽ⁱ⁺¹⁾ = wrk.pulses1 + if haskey(wrk.kwargs, :chi) + chi! = wrk.kwargs[:chi] + else + # we only want to evaluate `make_chi` if `chi` is not a kwarg + J_T_func = wrk.kwargs[:J_T] + chi! = make_chi(J_T_func, wrk.objectives) + end + if skip_initial_forward_propagation @info "Skipping initial forward propagation" else @@ -145,7 +153,7 @@ function optimize_krotov(problem) try while !wrk.result.converged i = i + 1 - krotov_iteration(wrk, ϵ⁽ⁱ⁾, ϵ⁽ⁱ⁺¹⁾) + krotov_iteration(wrk, ϵ⁽ⁱ⁾, ϵ⁽ⁱ⁺¹⁾, chi!) update_result!(wrk, i) update_hook!(wrk, i, ϵ⁽ⁱ⁺¹⁾, ϵ⁽ⁱ⁾) info_tuple = info_hook(wrk, i, ϵ⁽ⁱ⁺¹⁾, ϵ⁽ⁱ⁾) @@ -207,11 +215,10 @@ _eval_mu(μ::Operator, _...) = μ _eval_mu(μ::AbstractMatrix, _...) = μ -function krotov_iteration(wrk, ϵ⁽ⁱ⁾, ϵ⁽ⁱ⁺¹⁾) +function krotov_iteration(wrk, ϵ⁽ⁱ⁾, ϵ⁽ⁱ⁺¹⁾, chi!) χ = [propagator.state for propagator in wrk.bw_propagators] J_T_func = wrk.kwargs[:J_T] - chi! = get(wrk.kwargs, :chi, make_chi(J_T_func, wrk.objectives)) tlist = wrk.result.tlist N_T = length(tlist) - 1 N = length(wrk.objectives)