Use trapezoidal rule to evaluate J_b

Author: goerz

src/functionals.jl

@@ -1101,37 +1101,50 @@ J_b_val = J_b(storage, trajectories, tlist; g_b)
 returns
 
 ```math
-J_b = ∑_k ∑_n g_b(|Ψ_k(t_n)⟩) Δt_n
+J_b = ∑_k ∑_{n=0}^{N_T}{}^′ g_b(|Ψ_k(t_n)⟩) Δt_n
 ```
 
-as the piecewise-constant discretization of
+as the [trapezoid discretization](https://en.wikipedia.org/wiki/Trapezoidal_rule)
+of the integral
 
 ```math
-J_b = ∑_k \\int g_b(|Ψ_k(t)⟩) dt
+J_b = ∑_k ∫_0^T g_b(|Ψ_k(t)⟩) dt ,
 ```
 
 where `storage[k]` contains the forward-propagated states for trajectory `k`
-at each point of `tlist`; for example, the `fw_storage` component of the
-[GRAPE workspace](@extref `GRAPE.GrapeWrk`). The ``Δt_n`` is the
-[time step around the time grid point ``t_n``](@extref GRAPE Overview-Running-Costs).
+at each time grid point in `tlist`, ``t_0 = 0, t_1, … t_{N_T} = T``;
+for example, the `fw_storage` component of the
+[GRAPE workspace](@extref `GRAPE.GrapeWrk`). The ``Δt_n`` are defined at
+the endpoints as ``Δt_0 = t_1 - t_0`` and ``Δt_{N_T} = (t_{N_T} - t_{N_T-1})``,
+and at the interior points as ``Δt_n = \\frac{(t_{n+1} - t_{n-1})}{2}``.
+For uniform time grids, ``Δt_n ≡ dt``. The prime in the sum over ``n``
+indicates that the terms at the endpoints ``n=0`` and ``n=N_T`` are halved.
 
 Note that `g_b` is a mandatory keyword argument and must be a function
-`g_b(Ψ, trajectory, tlist, n) -> Float64`.
+`g_b(Ψ, trajectory, tlist, n) -> Float64` where `Ψ` is the state ``|Ψ(t)⟩``
+at ``t`` corresponding to the the (one-based) time grid point `tlist[n]`, and
+`trajectory` is a [`QuantumControl.Trajectory`](@ref) that may hold additional
+data in a custom property that is relevant to the calculation of ``g_b``.
+
+The definition of `J_b` here is compatible with the
+[treatment of running costs in GRAPE](@extref GRAPE Overview-Running-Costs).
 """
 function J_b(storage, trajectories, tlist; g_b)
     N = length(trajectories)
     result = 0.0
     for k = 1:N
         Ψ₁ = get_from_storage(storage[k], 1)
-        result += g_b(Ψ₁, trajectories[k], tlist, 1) * (tlist[2] - tlist[1])
+        dt = tlist[2] - tlist[1]
+        result += g_b(Ψ₁, trajectories[k], tlist, 1) * (dt/2)
         for n_tl = 2:length(tlist)
             Ψₙ = get_from_storage(storage[k], n_tl)
             if n_tl < length(tlist)
                 dt = 0.5 * (tlist[n_tl+1] - tlist[n_tl-1])
+                result += g_b(Ψₙ, trajectories[k], tlist, n_tl) * dt
             else
                 dt = tlist[end] - tlist[end-1]
+                result += g_b(Ψₙ, trajectories[k], tlist, n_tl) * (dt/2)
             end
-            result += g_b(Ψₙ, trajectories[k], tlist, n_tl) * dt
         end
     end
     return result

test/test_functionals.jl

@@ -771,32 +771,27 @@ end
 @testset "J_b" begin
 
     # Test that J_b correctly integrates g_b over all trajectories and time.
-    # Uses simple storage (vector of vectors) and a constant g_b = 1 whose
-    # integral is fully determined by the time grid and number of trajectories.
+    # Uses simple storage (vector of vectors) and a constant g_b = |Ψ| = 1
+    # whose integral is fully determined by the time grid and number of
+    # trajectories.
 
     tlist = PROBLEM.tlist
+    T = tlist[end]
     trajectories = PROBLEM.trajectories
+    N = length(trajectories)
     N_tl = length(tlist)
 
     # Build fake storage: each trajectory gets a vector of random states
     storage = [[random_state_vector(N_HILBERT; rng = RNG) for _ = 1:N_tl] for _ = 1:N]
 
     function g_b_const(Ψ, traj, tlist, n)
-        return 1.0
+        return norm(Ψ)
     end
 
     J_b_val = J_b(storage, trajectories, tlist; g_b = g_b_const)
 
-    # Manually reproduce J_b's integration weights for a constant g_b = 1
-    expected = 0.0
-    for _ = 1:N
-        expected += tlist[2] - tlist[1]
-        for n_tl = 2:(N_tl-1)
-            expected += 0.5 * (tlist[n_tl+1] - tlist[n_tl-1])
-        end
-        expected += tlist[end] - tlist[end-1]
-    end
-    @test J_b_val ≈ expected
+    # Trapezoidal rule for integral should return the exact result
+    @test J_b_val ≈ T * N
 
     # With g_b ∝ n, J_b should reflect the weighted sum over time indices
     function g_b_index(Ψ, traj, tlist, n)
@@ -805,14 +800,7 @@ end
 
     J_b_idx = J_b(storage, trajectories, tlist; g_b = g_b_index)
 
-    expected_idx = 0.0
-    for _ = 1:N
-        expected_idx += 1.0 * (tlist[2] - tlist[1])
-        for n_tl = 2:(N_tl-1)
-            expected_idx += n_tl * 0.5 * (tlist[n_tl+1] - tlist[n_tl-1])
-        end
-        expected_idx += N_tl * (tlist[end] - tlist[end-1])
-    end
-    @test J_b_idx ≈ expected_idx
+    # Again, the integral should be exact
+    @test J_b_idx ≈ ((N_tl - 1) / 2 + 1) * T * N
 
 end