Fix Taylor gradient for static operators

src/optimize.jl

@@ -355,7 +355,11 @@ function optimize_grape(problem)
                             control => val for (control, val) ∈ zip(wrk.controls, ϵₙ⁽ⁱ⁾)
                         )
                         local μₗₖₙ = evaluate(μₖₗ, tlist, n; vals_dict)
-                        evaluate!(Hₖₙ⁺, Hₖ⁺, tlist, n; vals_dict)
+                        if supports_inplace(Hₖₙ⁺)
+                            evaluate!(Hₖₙ⁺, Hₖ⁺, tlist, n; vals_dict)
+                        else
+                            Hₖₙ⁺ = evaluate(Hₖ⁺, tlist, n; vals_dict)
+                        end
                         local χ̃ₗₖ = wrk.taylor_grad_states[l, k][1]
                         local ϕ_temp = wrk.taylor_grad_states[l, k][2:5]
                         local dt = tlist[n] - tlist[n+1]
@@ -627,6 +631,7 @@ end
 #   |ϕ_n⟩ &= μ̂ Ĥⁿ⁻¹ |Ψ⟩ + Ĥ |Φₙ₋₁⟩
 # \end{align}
 # ```
+# TODO: this should probably be adapted to static states (avoiding in-place)
 function taylor_grad_step!(
     Ψ̃,
     Ψ,

test/test_tls_optimization.jl

@@ -2,6 +2,7 @@ using Test
 using QuantumControl
 using QuantumPropagators: ExpProp
 using QuantumControl.Functionals: J_T_sm
+using QuantumControl.Interfaces: check_generator
 using GRAPE
 import Krotov
 using LinearAlgebra
@@ -43,10 +44,61 @@ function tls_hamiltonian_static(Ω=1.0, ϵ=ϵ)
     ]
     Ĥ₀ = -0.5 * Ω * σ̂_z
     Ĥ₁ = σ̂_x
-    return hamiltonian(Ĥ₀, (Ĥ₁, ϵ))
+    OT = typeof(Ĥ₀)
+    PT = typeof(ϵ)
+    return TLSHamiltonianStatic{OT,PT}(Ĥ₀, Ĥ₁, ϵ)
 end;
 
 
+struct TLSHamiltonianStatic{OT,PT<:Function}
+    H₀::OT
+    H₁::OT
+    ϵ::PT
+end
+
+import QuantumControl.Controls: get_controls, evaluate, get_control_deriv
+import QuantumControl.Interfaces: supports_inplace
+
+
+Base.adjoint(H::TLSHamiltonianStatic) = H
+
+
+function Base.copy(H::TLSHamiltonianStatic{OT,PT}) where {OT,PT}
+    return TLSHamiltonianStatic{OT,PT}(H.H₀, H.H₁, H.ϵ)
+end
+
+
+function get_controls(H::TLSHamiltonianStatic)
+    return (H.ϵ,)
+end
+
+
+function evaluate(H::TLSHamiltonianStatic, args...; kwargs...)
+    ϵ_val = evaluate(H.ϵ, args...; kwargs...)
+    return H.H₀ + ϵ_val * H.H₁
+end
+
+function get_control_deriv(H::TLSHamiltonianStatic, control)
+    if control ≡ H.ϵ
+        return H.H₁
+    else
+        return nothing
+    end
+end
+
+supports_inplace(::TLSHamiltonianStatic) = false
+
+
+@testset "TLS Static Hamiltonian" begin
+
+    H = tls_hamiltonian_static()
+    tlist = collect(range(0, 5, length=501))
+    Ψ₀ = @SVector ComplexF64[1, 0]
+    @test check_generator(H; state=Ψ₀, tlist)
+
+end
+
+
 function ls_info_hook(wrk, iter)
     g = gradient(wrk)
     s = search_direction(wrk)
@@ -268,6 +320,34 @@ end
 end
 
 
+@testset "TLS (static Taylor)" begin
+
+    println("\n============= TLS (static Taylor) ===================\n")
+    H = tls_hamiltonian_static()
+    tlist = collect(range(0, 5, length=501))
+    Ψ₀ = @SVector ComplexF64[1, 0]
+    Ψtgt = @SVector ComplexF64[0, 1]
+    problem = ControlProblem(
+        [Trajectory(Ψ₀, H, target_state=Ψtgt)],
+        tlist;
+        iter_stop=5,
+        prop_method=ExpProp,
+        J_T=J_T_sm,
+        rethrow_exceptions=true,
+        gradient_method=:taylor,
+        check_convergence=res -> begin
+            ((res.J_T < 1e-10) && (res.converged = true) && (res.message = "J_T < 10⁻¹⁰"))
+        end,
+    )
+    res = optimize(problem; method=GRAPE)
+    display(res)
+    @test res.J_T < 1e-3
+    @test 0.75 < maximum(abs.(res.optimized_controls[1])) < 0.85
+    println("===================================================\n")
+
+end
+
+
 @testset "TLS (continue from Krotov)" begin
 
     println("\n============ TLS (Krotov continuation) ============\n")