From a3fdd2e82da3338217721cece13881867964d669 Mon Sep 17 00:00:00 2001
From: Michael Goerz <goerz@stanford.edu>
Date: Sun, 7 Nov 2021 01:17:02 -0400
Subject: [PATCH] =?UTF-8?q?Fix=20calculation=20of=20=E2=88=ABg=E2=82=90dt?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

See https://github.com/qucontrol/krotov/issues/96
---
 src/optimize.jl | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/src/optimize.jl b/src/optimize.jl
index b43bf45..ccb0d96 100644
--- a/src/optimize.jl
+++ b/src/optimize.jl
@@ -375,39 +375,40 @@ function krotov_iteration(wrk, ϵ⁽ⁱ⁾, ϵ⁽ⁱ⁺¹⁾)
         end
         ϵₙ⁽ⁱ⁺¹⁾ = [ϵ⁽ⁱ⁾[l][n]  for l ∈ 1:L]  # ϵₙ⁽ⁱ⁺¹⁾ ≈ ϵₙ⁽ⁱ⁾ for non-linear controls
         # TODO: we could add a self-consistent loop here for ϵₙ⁽ⁱ⁺¹⁾
-        Δuₙ = zeros(L)
+        Δuₙ′ = zeros(L)  # for step size 1
         for l = 1:L  # `l` is the index for the different controls
-            Sₗ = wrk.update_shapes[l]
-            λₐ = wrk.lambda_vals[l]
             ∂ϵₗ╱∂u :: Function = wrk.parametrization[l].de_du_derivative
             uₗ = wrk.parametrization[l].u_of_epsilon
             for k = 1:N  # k is the index over the objectives
                 ∂Hₖ╱∂ϵₗ :: Union{Function, Nothing} = wrk.control_derivs[k][l]
                 if !isnothing(∂Hₖ╱∂ϵₗ)
                     μₗₖₙ = (∂Hₖ╱∂ϵₗ)(ϵₙ⁽ⁱ⁺¹⁾[l])
-                    αₗ = (Sₗ[n]/λₐ)  # Krotov step size
                     if wrk.is_parametrized[l]
                         ∂ϵₗ╱∂uₗ = (∂ϵₗ╱∂u)(uₗ(ϵₙ⁽ⁱ⁺¹⁾[l]))
-                        Δuₙ[l] += αₗ * ∂ϵₗ╱∂uₗ * Im(dot(χ[k], μₗₖₙ, ϕ[k]))
+                        Δuₙ′[l] += ∂ϵₗ╱∂uₗ * Im(dot(χ[k], μₗₖₙ, ϕ[k]))
                     else
-                        Δuₙ[l] += αₗ * Im(dot(χ[k], μₗₖₙ, ϕ[k]))
+                        Δuₙ′[l] += Im(dot(χ[k], μₗₖₙ, ϕ[k]))
                     end
                 end
             end
         end
         # TODO: second order update
         for l = 1:L
+            Sₗ = wrk.update_shapes[l]
+            λₐ = wrk.lambda_vals[l]
+            αₗ = (Sₗ[n]/λₐ)  # Krotov step size
+            Δuₗₙ = αₗ * Δuₙ′[l]
             if wrk.is_parametrized[l]
                 uₗ = wrk.parametrization[l].u_of_epsilon
                 ϵₗ = wrk.parametrization[l].epsilon_of_u
-                ϵ⁽ⁱ⁺¹⁾[l][n] = ϵₗ(uₗ(ϵ⁽ⁱ⁾[l][n]) + Δuₙ[l])
+                ϵ⁽ⁱ⁺¹⁾[l][n] = ϵₗ(uₗ(ϵ⁽ⁱ⁾[l][n]) + Δuₗₙ)
             else
-                Δϵₗₙ = Δuₙ[l]
+                Δϵₗₙ = Δuₗₙ
                 ϵ⁽ⁱ⁺¹⁾[l][n] = ϵ⁽ⁱ⁾[l][n] + Δϵₗₙ
             end
+            (∫gₐdt)[l] += αₗ * abs(Δuₙ′[l])^2 * dt
         end
         # TODO: end of self-consistent loop
-        @. ∫gₐdt += abs(Δuₙ)^2 * dt
         @threadsif wrk.use_threads for k = 1:N
             local (G, dt) = _fw_gen(ϵ⁽ⁱ⁺¹⁾, k, n, wrk)
             propstep!(ϕ[k], G, dt, wrk.prop_wrk[k])