Add save positions (#71)

* ✨added save_positions as a keyword argument to build_ensembles interface

* ✅ added save_positions test for build_ensemble interface

* 🔖 bump version 0.6.2

Project.toml

@@ -1,7 +1,7 @@
 name = "OpenQuantumTools"
 uuid = "e429f160-8886-11e9-20cb-0dbe84e78965"
 authors = ["Huo Chen <[email protected]>"]
-version = "0.6.1"
+version = "0.6.2"
 
 [deps]
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"

src/QControl/callback_lib.jl

@@ -49,7 +49,7 @@ function InstPulseCallback(tstops, pulse_update)
     PresetTimeCallback(tstops, affect!, initialize=initialize)
 end
 
-function FluctuatorCallback(F::OpenQuantumBase.FluctuatorLiouvillian, initialize)
+function FluctuatorCallback(F::OpenQuantumBase.FluctuatorLiouvillian, initialize, save_positions)
     time_choice = function (integrator)
         next_t = integrator.t + F.next_τ
         if next_t > integrator.sol.prob.tspan[2]
@@ -74,11 +74,11 @@ function FluctuatorCallback(F::OpenQuantumBase.FluctuatorLiouvillian, initialize
         time_choice,
         affect!,
         initialize=initialize_fluc,
-        save_positions=(false, false),
+        save_positions=save_positions,
     )
 end
 
-function LindbladJumpCallback()
+function LindbladJumpCallback(save_positions)
     r = Ref{Float64}(rand(Float64))
     condition = function (u, t, integrator)
         real(u' * u) - r[]
@@ -103,7 +103,7 @@ function LindbladJumpCallback()
         u_modified!(integrator, false)
     end
 
-    ContinuousCallback(condition, affect!, save_positions=(true, true), initialize=initialize)
+    ContinuousCallback(condition, affect!, save_positions=save_positions, initialize=initialize)
 end
 
 """

src/QSolver/ame_solver.jl

@@ -55,14 +55,15 @@ function build_ensemble_ame(
     lambshift_S=nothing,
     lvl::Int=size(A.H, 1),
     initializer=initializer,
+    save_positions=(false, false),
     kwargs...,
 )
     u0 = build_u0(A.u0, :v)
 
     flist = OpenQuantumBase.fluctuator_from_interactions(A.interactions)
     dlist = OpenQuantumBase.davies_from_interactions(A.interactions, ω_hint, lambshift, lambshift_S)
     L = DiffEqLiouvillian(A.H, dlist, flist, lvl)
-    cb = build_jump_callback(dlist, flist, initializer)
+    cb = build_jump_callback(dlist, flist, initializer, save_positions)
     p = ODEParams(L, tf, A.annealing_parameter)
     update_func = function (cache, u, p, t)
         update_cache!(cache, p.L, p, t)
@@ -83,18 +84,18 @@ function build_ensemble_ame(
     ensemble_prob
 end
 
-function build_jump_callback(dlist, flist, initializer)
+function build_jump_callback(dlist, flist, initializer, save_positions)
     initializer =
         initializer == DEFAULT_INITIALIZER ? (x, y) -> rand([-1, 1], x, y) :
         initializer
     if isempty(flist)
-        cb = LindbladJumpCallback()
+        cb = LindbladJumpCallback(save_positions)
     elseif isempty(dlist)
         error("No interactions support AME. Use other ensemble instead.")
     else
         cb = CallbackSet(
-            LindbladJumpCallback(),
-            [FluctuatorCallback(f, initializer) for f in flist]...,
+            LindbladJumpCallback(save_positions),
+            [FluctuatorCallback(f, initializer, save_positions) for f in flist]...,
         )
     end
     cb

src/QSolver/ensemble_builder.jl

@@ -10,6 +10,7 @@ Build `EnsembleProblem` object for different open-system models. For `:lindblad`
 - `type::Symbol`: type of the ensemble to build. Available options are `:lindblad`, `:ame`, `:stochastic` and `:redfield`.
 - `tspan = (0, tf)`: time interval to solve the dynamics.
 - `initializer = DEFAULT_INITIALIZER`: initializer for the ensemble problem. Currently it is only supported by the `:stochastic` ensemble.
+- `save_positions = (false, false)`: Boolean tuple for whether to save before and after the callbacks. This saving will occur just before and after the event, only at event times, and does not depend on options like saveat, save_everystep, etc. (i.e. if saveat=[1.0,2.0,3.0], this can still add a save point at 2.1 if true).
 - `output_func`: The function determines what is saved from the solution to the output array. Defaults to saving the solution itself. The output is (out,rerun) where out is the output and rerun is a boolean which designates whether to rerun. It is part of the `DifferentialEquations`'s parallel interface.
 - `reduction`: This function determines how to reduce the data in each batch. Defaults to appending the data from the batches. The second part of the output determines whether the simulation has converged. If true, the simulation will exit early. By default, this is always false. It is part of the `DifferentialEquations`'s parallel interface.
 - `kwargs`: other keyword arguments supported by the specific solver or by `DifferentialEquations`.
@@ -23,6 +24,7 @@ function build_ensembles(
     output_func = (sol, i) -> (sol, false),
     reduction = (u, data, I) -> (append!(u, data), false),
     initializer = DEFAULT_INITIALIZER,
+    save_positions = (false, false),
     kwargs...,
 )
     if type == :stochastic
@@ -33,6 +35,7 @@ function build_ensembles(
             reduction;
             tspan = tspan,
             initializer = initializer,
+            save_positions = save_positions,
             kwargs...,
         )
     elseif type == :ame
@@ -43,6 +46,7 @@ function build_ensembles(
             reduction;
             tspan = tspan,
             initializer = initializer,
+            save_positions = save_positions,
             kwargs...,
         )
     elseif type == :lindblad
@@ -52,6 +56,7 @@ function build_ensembles(
             output_func,
             reduction;
             tspan = tspan,
+            save_positions = save_positions,
             kwargs...,
         )
     else
@@ -72,6 +77,7 @@ function build_ensembles(
     int_rtol = 1e-6,
     Ta = tf,
     initializer = DEFAULT_INITIALIZER,
+    save_positions = (false, false),
     kwargs...,
 )
     if type == :redfield
@@ -85,7 +91,8 @@ function build_ensembles(
             int_atol = int_atol,
             int_rtol = int_rtol,
             Ta = Ta,
-            initializer = DEFAULT_INITIALIZER,
+            initializer = initializer,
+            save_positions = save_positions,
             kwargs...,
         )
     else

src/QSolver/lindblad_solver.jl

@@ -36,11 +36,12 @@ function build_ensemble_lindblad(
     output_func,
     reduction;
     tspan=(0.0, tf),
+    save_positions=(false, false),
     kwargs...,
 )
     u0 = build_u0(A.u0, :v)
     L = DiffEqLiouvillian(A.H, [], OpenQuantumBase.lindblad_from_interactions(A.interactions), size(A.H, 1))
-    cb = LindbladJumpCallback()
+    cb = LindbladJumpCallback(save_positions)
     p = ODEParams(L, tf, A.annealing_parameter)
     update_func = function (cache, u, p, t)
         update_cache!(cache, p.L, p, t)

src/QSolver/redfield_solver.jl

@@ -160,6 +160,7 @@ function build_ensemble_redfield(
     int_rtol=1e-6,
     Ta=tf,
     initializer=DEFAULT_INITIALIZER,
+    save_positions=(false, false),
     kwargs...,
 )
     u0 = build_u0(A.u0, :m, vectorize=vectorize)
@@ -174,9 +175,9 @@ function build_ensemble_redfield(
         initializer
 
     if length(stocs) == 1
-        cb = FluctuatorCallback(stocs[1], initializer)
+        cb = FluctuatorCallback(stocs[1], initializer, save_positions)
     else
-        cb = CallbackSet([FluctuatorCallback(f, initializer) for f in stocs]...)
+        cb = CallbackSet([FluctuatorCallback(f, initializer, save_positions) for f in stocs]...)
     end
     R = DiffEqLiouvillian(A.H, [], [reds; stocs], size(A.H, 1))
     p = ODEParams(R, float(tf), A.annealing_parameter)

src/QSolver/stochastic_schrodinger_solver.jl

@@ -5,6 +5,7 @@ function build_ensemble_stochastic(
     reduction;
     tspan=(0.0, tf),
     initializer=DEFAULT_INITIALIZER,
+    save_positions=(false, false),
     kwargs...
 )
     u0 = build_u0(A.u0, :v)
@@ -14,11 +15,11 @@ function build_ensemble_stochastic(
 
     flist = OpenQuantumBase.fluctuator_from_interactions(A.interactions)
     if length(flist) == 1
-        cb = FluctuatorCallback(flist[1], initializer)
+        cb = FluctuatorCallback(flist[1], initializer, save_positions)
     else
-        cb = CallbackSet([FluctuatorCallback(f, initializer) for f in flist]...)
+        cb = CallbackSet([FluctuatorCallback(f, initializer, save_positions) for f in flist]...)
     end
-    ff = DiffEqLiouvillian(A.H, [], flist, size(A.H,1))
+    ff = DiffEqLiouvillian(A.H, [], flist, size(A.H, 1))
     p = ODEParams(ff, float(tf), A.annealing_parameter)
 
     update_func = function (cache, u, p, t)

test/QSolvers/ame_solver_test.jl

@@ -1,4 +1,4 @@
-using OpenQuantumTools, Test, OrdinaryDiffEq
+using OpenQuantumTools, Random, Test, OrdinaryDiffEq
 
 
 H = DenseHamiltonian([(s) -> 1.0], [σi], unit = :ħ)
@@ -31,4 +31,33 @@ sol = solve_ame(
     abstol = 1e-8,
     reltol = 1e-8,
 )
-@test sol(tf)[1, 2] ≈ exp(-2 * γ * tf) * 0.5 atol = 1e-5 rtol = 1e-5
+@test sol(tf)[1, 2] ≈ exp(-2 * γ * tf) * 0.5 atol = 1e-5 rtol = 1e-5
+
+tf = 1000
+prob = build_ensembles(annealing, tf, :ame, save_positions=(true, true))
+Random.seed!(1234)
+sol1 = solve(prob, Tsit5(), EnsembleSerial(), trajectories=1, save_everystep=false)[1]
+
+prob = build_ensembles(annealing, tf, :ame)
+Random.seed!(1234)
+sol2 = solve(prob, Tsit5(), EnsembleSerial(), trajectories=1, save_everystep=false)[1]
+
+@test !(length(sol1) == 2)
+@test sol1[end] ≈ sol2[end]
+
+# test suite for hybrid spin-fluctuator and ame
+fbath = EnsembleFluctuator([0.1], [0.1])
+interactions = InteractionSet(Interaction(coupling, bath), Interaction(coupling, fbath))
+annealing = Annealing(H, u0; interactions=interactions)
+
+tf = 1000
+prob = build_ensembles(annealing, tf, :ame, save_positions=(true, true))
+Random.seed!(1234)
+sol1 = solve(prob, Tsit5(), EnsembleSerial(), trajectories=1, save_everystep=false)[1]
+
+prob = build_ensembles(annealing, tf, :ame)
+Random.seed!(1234)
+sol2 = solve(prob, Tsit5(), EnsembleSerial(), trajectories=1, save_everystep=false)[1]
+
+@test !(length(sol1) == 2)
+@test length(sol2) == 2

test/QSolvers/redfield_solver_test.jl

@@ -1,4 +1,4 @@
-using OpenQuantumTools, Test, OrdinaryDiffEq, QuadGK
+using OpenQuantumTools, Test, OrdinaryDiffEq, QuadGK, Random
 
 H = DenseHamiltonian([(s) -> 1.0], [σi], unit=:ħ)
 u0 = PauliVec[1][1]
@@ -20,6 +20,21 @@ sol = solve_redfield(annealing, tf, InplaceUnitary(U), vectorize=true,
     alg=TRBDF2(), reltol=1e-6)
 @test sol(10)[2] ≈ exp(-4 * γ) * 0.5 atol = 1e-5 rtol = 1e-5
 
+# test for hybrid spin-fluctuator and Redfield equation
+fbath = EnsembleFluctuator([0.1], [1.0])
+interactions = InteractionSet(Interaction(coupling, bath), Interaction(coupling, fbath))
+annealing = Annealing(H, u0; interactions=interactions)
+tf = 10
+U = solve_unitary(annealing, tf, alg=Tsit5(), abstol=1e-8, reltol=1e-8)
+prob1 = build_ensembles(annealing, tf, U, :redfield)
+prob2 = build_ensembles(annealing, tf, U, :redfield, save_positions=(true, true))
+Random.seed!(1234)
+sol1 = solve(prob1, Tsit5(), EnsembleSerial(), trajectories=1, save_everystep=false)[1]
+Random.seed!(1234)
+sol2 = solve(prob2, Tsit5(), EnsembleSerial(), trajectories=1, save_everystep=false)[1]
+@test length(sol1) == 2
+@test !(length(sol2) == 2)
+
 # test for CustomBath
 H = DenseHamiltonian([(s) -> 0.0], [σi], unit=:ħ)
 u0 = PauliVec[1][1]
@@ -82,20 +97,20 @@ isol = solve_redfield(iannealing, tf, (t) -> σi, alg=Tsit5(), reltol=1e-6)
 β = 4
 T = β_2_temperature(β)
 η = 0.1
-fc= 10/(2π)
+fc = 10 / (2π)
 bath = Ohmic(η, fc, T)
 
-Hp = 0.5*σz⊗σi - 0.7*σi⊗σz + 0.3*σz⊗σz
+Hp = 0.5 * σz ⊗ σi - 0.7 * σi ⊗ σz + 0.3 * σz ⊗ σz
 Hd = standard_driver(2)
-H = DenseHamiltonian([(s)->1-s, (s)->s], [-Hd, Hp], unit=:ħ)
+H = DenseHamiltonian([(s) -> 1 - s, (s) -> s], [-Hd, Hp], unit=:ħ)
 
 tf = 20
-ρ0 = (σi+σx)⊗(σi+σx)/4
-coupling = ConstantCouplings([σz⊗σi, σi⊗σz], unit=:ħ)
+ρ0 = (σi + σx) ⊗ (σi + σx) / 4
+coupling = ConstantCouplings([σz ⊗ σi, σi ⊗ σz], unit=:ħ)
 annealing = Annealing(H, ρ0, bath=bath, coupling=coupling)
 
-U = solve_unitary(annealing, tf, alg = Tsit5(), abstol=1e-7, reltol=1e-7)
+U = solve_unitary(annealing, tf, alg=Tsit5(), abstol=1e-7, reltol=1e-7)
 
-redfield_sol = solve_redfield(annealing, tf, U, alg = Tsit5(), abstol=1e-7, reltol=1e-7, callback=PositivityCheckCallback())
+redfield_sol = solve_redfield(annealing, tf, U, alg=Tsit5(), abstol=1e-7, reltol=1e-7, callback=PositivityCheckCallback())
 
 @test redfield_sol.t[end] ≈ 3.45278167 atol = 1e-5 rtol = 1e-5

test/QSolvers/stochastic_schrodinger.jl

@@ -12,8 +12,9 @@ prob = build_ensembles(annealing, tf, :stochastic)
 Random.seed!(1234)
 sol1 = solve(prob, Tsit5(), EnsembleSerial(), trajectories=1, save_everystep=false)
 
+prob2 = build_ensembles(annealing, tf, :stochastic, save_positions=(true, true))
 Random.seed!(1234)
-sol2 = solve(prob, Tsit5(), EnsembleSerial(), trajectories=1, save_everystep=false)
+sol2 = solve(prob2, Tsit5(), EnsembleSerial(), trajectories=1, save_everystep=false)
 
 @test sol1[1][end] ≈ sol2[1][end] atol = 1e-6 rtol = 1e-6