From 203846f9a39afc7d4fa0f707e26ca8d01ee519ab Mon Sep 17 00:00:00 2001
From: AmitRotem <66641519+AmitRotem@users.noreply.github.com>
Date: Sun, 22 Jan 2023 10:15:37 -0800
Subject: [PATCH 1/3] Improve coverage

---
 test/test_ForwardDiff.jl | 26 +++++++++++---------------
 1 file changed, 11 insertions(+), 15 deletions(-)

diff --git a/test/test_ForwardDiff.jl b/test/test_ForwardDiff.jl
index a34f84d7..2002d348 100644
--- a/test/test_ForwardDiff.jl
+++ b/test/test_ForwardDiff.jl
@@ -17,7 +17,6 @@ import ForwardDiff
 # system
 ba0 = FockBasis(2)
 psi = basisstate(ba0, 1)
-target0 = basisstate(ba0, 2)
 function getHt(p)
     op = [create(ba0)+destroy(ba0)]
     f(t) = sin(p*t)
@@ -30,27 +29,24 @@ function getHt(p)
 end
 
 # cost function
-function cost(par; kwargs...)
+function cost(par, ψ0; kwargs...)
     opti = (;dtmax=exp2(-4), dt=exp2(-4))
     Ht = getHt(par)
-    # this will rebuild the Bra with Dual elements
-    _, ψT = timeevolution.schroedinger_dynamic((0.0, 0.2), psi'     , Ht; opti..., kwargs...)
-    # this will not rebuild the Bra
-    _, ψT = timeevolution.schroedinger_dynamic((0.2, 0.4), last(ψT) , Ht; opti..., kwargs...) 
-    # this will not rebuild the Ket
-    # also tests static schroedinger
-    _, ψT = timeevolution.schroedinger((0.4, 0.6), last(ψT)', Ht(1.0, ψT); opti..., kwargs...)
-    # this will not rebuild the Ket
-    _, ψT = timeevolution.schroedinger_dynamic((0.6, 0.8), last(ψT)⊗last(ψT)', Ht; opti..., kwargs...)
-    abs2(target0'*last(ψT)*target0)
+    # this will rebuild the state with Dual elements
+    _, ψT = timeevolution.schroedinger_dynamic((0.0, 1.0), ψ0      , Ht; opti..., kwargs...)
+    # this will not rebuild the state
+    _, ψT = timeevolution.schroedinger((1.0, 2.0), last(ψT), Ht(0.5); opti..., kwargs...)
+    abs2(tr(ψ0'*last(ψT)))
 end
 
 # setup
 p0 = rand()
 δp = √eps()
 # test
-finite_diff_derivative = ( cost(p0+δp) - cost(p0) ) / δp
-Auto_diff_derivative = ForwardDiff.derivative(cost, p0)
-@test isapprox(Auto_diff_derivative, finite_diff_derivative; atol=1e-5)
+for u0 = (psi, psi', psi⊗psi') # test all methods of `rebuild`
+    finite_diff_derivative = ( cost(p0+δp, u0) - cost(p0, u0) ) / δp
+    Auto_diff_derivative = ForwardDiff.derivative(Base.Fix2(cost, u0), p0)
+    @test isapprox(Auto_diff_derivative, finite_diff_derivative; atol=1e-5)
+end
 
 end # testset

From f3e59fc2287393a35ffa9ddb2faf1beeacac562e Mon Sep 17 00:00:00 2001
From: AmitRotem <66641519+AmitRotem@users.noreply.github.com>
Date: Sun, 22 Jan 2023 11:25:36 -0800
Subject: [PATCH 2/3] Update test_ForwardDiff.jl

---
 test/test_ForwardDiff.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/test_ForwardDiff.jl b/test/test_ForwardDiff.jl
index 2002d348..dfa79007 100644
--- a/test/test_ForwardDiff.jl
+++ b/test/test_ForwardDiff.jl
@@ -33,9 +33,9 @@ function cost(par, ψ0; kwargs...)
     opti = (;dtmax=exp2(-4), dt=exp2(-4))
     Ht = getHt(par)
     # this will rebuild the state with Dual elements
-    _, ψT = timeevolution.schroedinger_dynamic((0.0, 1.0), ψ0      , Ht; opti..., kwargs...)
+    _, ψT = timeevolution.schroedinger_dynamic((0.0, 1.0), ψ0, Ht; opti..., kwargs...)
     # this will not rebuild the state
-    _, ψT = timeevolution.schroedinger((1.0, 2.0), last(ψT), Ht(0.5); opti..., kwargs...)
+    _, ψT = timeevolution.schroedinger((1.0, 2.0), last(ψT), Ht(0.5, ψ0); opti..., kwargs...)
     abs2(tr(ψ0'*last(ψT)))
 end
 

From 9020a118931ec14a2cda0c5c7320b7b51f71310e Mon Sep 17 00:00:00 2001
From: Amit Rotem <AmiRotem@amazon.com>
Date: Sun, 22 Jan 2023 12:19:40 -0800
Subject: [PATCH 3/3] small fix

---
 test/test_ForwardDiff.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/test_ForwardDiff.jl b/test/test_ForwardDiff.jl
index dfa79007..01d8135c 100644
--- a/test/test_ForwardDiff.jl
+++ b/test/test_ForwardDiff.jl
@@ -20,7 +20,7 @@ psi = basisstate(ba0, 1)
 function getHt(p)
     op = [create(ba0)+destroy(ba0)]
     f(t) = sin(p*t)
-    H_at_t = LazySum([f(0)], op)
+    H_at_t = LazySum([f(0.0)], op)
     function Ht(t,_)
         H_at_t.factors .= (f(t),)
         return H_at_t
@@ -36,7 +36,7 @@ function cost(par, ψ0; kwargs...)
     _, ψT = timeevolution.schroedinger_dynamic((0.0, 1.0), ψ0, Ht; opti..., kwargs...)
     # this will not rebuild the state
     _, ψT = timeevolution.schroedinger((1.0, 2.0), last(ψT), Ht(0.5, ψ0); opti..., kwargs...)
-    abs2(tr(ψ0'*last(ψT)))
+    (abs2∘tr)( ψ0.data' * last(ψT).data ) # getting the data so this will work with also Bra states
 end
 
 # setup