Merge #109 from branch J-a-avg-zero

Author: goerz

src/functionals.jl

@@ -1,7 +1,7 @@
 module Functionals
 
 export J_T_ss, J_T_sm, J_T_re
-export J_a_fluence
+export J_a_fluence, J_a_avg_zero
 export gate_functional
 export make_gate_chi
 export J_b
@@ -973,6 +973,60 @@ end
 make_analytic_grad_J_a(::typeof(J_a_fluence), tlist) = grad_J_a_fluence
 
 
+@doc raw"""Running cost penalizing a non-zero pulse area.
+
+```julia
+J_a = J_a_avg_zero(pulsevals, tlist)
+```
+
+calculates
+
+```math
+J_a = \sum_l \left(\int_0^T ϵ_l(t)\, dt\right)^2 \approx \sum_l A_l^2\,,
+\quad A_l = \sum_n ϵ_{nl}\, dt_n\,,
+```
+
+where ``ϵ_{nl}`` and ``dt_n`` are as in [`J_a_fluence`](@ref). The functional
+is zero when every control integrates to zero over the time grid, and positive
+otherwise.
+
+In conjunction with [`QuantumPropagators.Amplitudes.GuidedAmplitude`](@ref),
+this enables an optimization that preserves the pulse area of a given `guide`.
+
+# See also
+
+* [`grad_J_a_avg_zero`](@ref) — analytic (automatic) gradient
+"""
+function J_a_avg_zero(pulsevals, tlist)
+    N_T = length(tlist) - 1
+    dt = diff(tlist)                        # (N_T,)
+    pv = reshape(pulsevals, N_T, :)         # (N_T, N_L), no copy
+    A = vec(dt' * pv)                       # (N_L,): time integral per control
+    return sum(abs2.(A))
+end
+
+
+"""Analytic derivative for [`J_a_avg_zero`](@ref).
+
+```julia
+∇J_a = grad_J_a_avg_zero(pulsevals, tlist)
+```
+
+returns `∇J_a` with elements ``2 A_l\\, dt_n``, where
+``A_l = \\sum_n ϵ_{nl}\\, dt_n`` is the time integral of control `l`.
+"""
+function grad_J_a_avg_zero(pulsevals, tlist)
+    N_T = length(tlist) - 1
+    dt = diff(tlist)                        # (N_T,)
+    pv = reshape(pulsevals, N_T, :)         # (N_T, N_L), no copy
+    A = vec(dt' * pv)                       # (N_L,): time integral per control
+    return vec(2 .* reshape(dt, :, 1) .* reshape(A, 1, :))
+end
+
+
+make_analytic_grad_J_a(::typeof(J_a_avg_zero), tlist) = grad_J_a_avg_zero
+
+
 """
 Return a function that calculates ``|ξ_k(t_n)⟩ = -∂g_b/∂⟨Ψ_k(t_n)|``.
 

test/test_functionals.jl

@@ -7,6 +7,8 @@ using QuantumControl.Functionals:
     J_T_ss,
     J_a_fluence,
     grad_J_a_fluence,
+    J_a_avg_zero,
+    grad_J_a_avg_zero,
     make_grad_J_a,
     make_chi,
     make_xi,
@@ -148,6 +150,36 @@ end
 end
 
 
+@testset "J_a_avg_zero" begin
+
+    # Two controls on a non-uniform grid
+    tlist_nu = [0.0, 0.1, 0.3, 0.6, 1.0]
+    dt_nu = [0.1, 0.2, 0.3, 0.4]           # diff(tlist_nu)
+    pv1 = [1.0, -2.0, 3.0, -1.0]           # A₁ = 0.1 - 0.4 + 0.9 - 0.4 = 0.2
+    pv2 = [2.0, -1.0, -1.0, 0.0]           # A₂ = 0.2 - 0.2 - 0.3 + 0.0 = -0.3
+    pulsevals_nu = vcat(pv1, pv2)
+
+    A1 = dot(pv1, dt_nu)
+    A2 = dot(pv2, dt_nu)
+    @test J_a_avg_zero(pulsevals_nu, tlist_nu) ≈ A1^2 + A2^2
+
+    # Zero when each control integrates to zero:
+    # 1*0.1 + 0*0.2 + 0*0.3 + (-0.25)*0.4 = 0.1 - 0.1 = 0
+    pv_zeroA = [1.0, 0.0, 0.0, -0.25]
+    @test J_a_avg_zero(vcat(pv_zeroA, pv_zeroA), tlist_nu) ≈ 0.0 atol = 1e-15
+
+    # Analytic gradient matches manual calculation
+    G_expected = vcat(2 * A1 .* dt_nu, 2 * A2 .* dt_nu)
+    @test grad_J_a_avg_zero(pulsevals_nu, tlist_nu) ≈ G_expected
+
+    # Analytic gradient matches Zygote
+    grad_J_a_zygote =
+        make_grad_J_a(J_a_avg_zero, tlist_nu; mode = :automatic, automatic = Zygote)
+    @test norm(grad_J_a_zygote(pulsevals_nu, tlist_nu) - G_expected) < 1e-12
+
+end
+
+
 @testset "J_T without analytic derivative" begin
 
     QuantumControl.set_default_ad_framework(nothing; quiet = true)