[no ci] Add state normalization during evolution in ssesolve

Author: albertomercurio

src/QuantumToolbox.jl

@@ -99,6 +99,7 @@ include("time_evolution/time_evolution.jl")
 include("time_evolution/callback_helpers/sesolve_callback_helpers.jl")
 include("time_evolution/callback_helpers/mesolve_callback_helpers.jl")
 include("time_evolution/callback_helpers/mcsolve_callback_helpers.jl")
+include("time_evolution/callback_helpers/ssesolve_callback_helpers.jl")
 include("time_evolution/callback_helpers/callback_helpers.jl")
 include("time_evolution/mesolve.jl")
 include("time_evolution/lr_mesolve.jl")

src/time_evolution/callback_helpers/callback_helpers.jl

@@ -32,9 +32,26 @@ end
 
 _get_e_ops_data(e_ops, ::Type{SaveFuncSESolve}) = get_data.(e_ops)
 _get_e_ops_data(e_ops, ::Type{SaveFuncMESolve}) = [_generate_mesolve_e_op(op) for op in e_ops] # Broadcasting generates type instabilities on Julia v1.10
+_get_e_ops_data(e_ops, ::Type{SaveFuncSSESolve}) = get_data.(e_ops)
 
 _generate_mesolve_e_op(op) = mat2vec(adjoint(get_data(op)))
 
+#=
+This function add the normalization callback to the kwargs. It is needed to stabilize the integration when using the ssesolve method.
+=#
+function _ssesolve_add_normalize_cb(kwargs)
+    _condition = (u, t, integrator) -> true
+    _affect! = (integrator) -> normalize!(integrator.u)
+    cb = DiscreteCallback(_condition, _affect!; save_positions = (false, false))
+    # return merge(kwargs, (callback = CallbackSet(kwargs[:callback], cb),))
+
+    cb_set = haskey(kwargs, :callback) ? CallbackSet(kwargs[:callback], cb) : cb
+
+    kwargs2 = merge(kwargs, (callback = cb_set,))
+
+    return kwargs2
+end
+
 ##
 
 # When e_ops is Nothing. Common for both mesolve and sesolve
@@ -80,10 +97,10 @@ function _se_me_sse_get_save_callback(cb::CallbackSet)
         return nothing
     end
 end
-_se_me_sse_get_save_callback(cb::DiscreteCallback) =
-    if (cb.affect! isa SaveFuncSESolve) || (cb.affect! isa SaveFuncMESolve)
+function _se_me_sse_get_save_callback(cb::DiscreteCallback)
+    if typeof(cb.affect!) <: Union{SaveFuncSESolve,SaveFuncMESolve,SaveFuncSSESolve}
         return cb
-    else
-        return nothing
     end
+    return nothing
+end
 _se_me_sse_get_save_callback(cb::ContinuousCallback) = nothing

src/time_evolution/callback_helpers/ssesolve_callback_helpers.jl

@@ -0,0 +1,25 @@
+#=
+Helper functions for the ssesolve callbacks. Equal to the sesolve case, but with an additional normalization before saving the expectation values.
+=#
+
+struct SaveFuncSSESolve{TE,PT<:Union{Nothing,ProgressBar},IT,TEXPV<:Union{Nothing,AbstractMatrix}}
+    e_ops::TE
+    progr::PT
+    iter::IT
+    expvals::TEXPV
+end
+
+(f::SaveFuncSSESolve)(integrator) = _save_func_ssesolve(integrator, f.e_ops, f.progr, f.iter, f.expvals)
+(f::SaveFuncSSESolve{Nothing})(integrator) = _save_func(integrator, f.progr) # Common for both mesolve and sesolve
+
+##
+
+# When e_ops is a list of operators
+function _save_func_ssesolve(integrator, e_ops, progr, iter, expvals)
+    ψ = normalize!(integrator.u)
+    _expect = op -> dot(ψ, op, ψ)
+    @. expvals[:, iter[]] = _expect(e_ops)
+    iter[] += 1
+
+    return _save_func(integrator, progr)
+end

src/time_evolution/ssesolve.jl

@@ -35,13 +35,14 @@ end
         Base.@nexprs $N i -> begin
             mul!(@view(v[:, i]), L.ops[i], u)
         end
-        v
+        return v
     end
 end
 
 # TODO: Implement the three-argument dot function for SciMLOperators.jl
 # Currently, we are assuming a time-independent MatrixOperator
 function _ssesolve_update_coeff(u, p, t, op)
+    normalize!(u)
     return real(dot(u, op.A, u)) #this is en/2: <Sn + Sn'>/2 = Re<Sn>
 end
 
@@ -104,23 +105,23 @@ _ScalarOperator_e2_2(op, f = +) =
 Generate the SDEProblem for the Stochastic Schrödinger time evolution of a quantum system. This is defined by the following stochastic differential equation:
     
 ```math
-d|\psi(t)\rangle = -i K |\psi(t)\rangle dt + \sum_n M_n |\psi(t)\rangle dW_n(t)
+d|\psi(t)\rangle = -i \hat{K} |\psi(t)\rangle dt + \sum_n \hat{M}_n |\psi(t)\rangle dW_n(t)
 ```
 
 where 
     
 ```math
-K = \hat{H} + i \sum_n \left(\frac{e_j} C_n - \frac{1}{2} \sum_{j} C_n^\dagger C_n - \frac{e_j^2}{8}\right),
+\hat{K} = \hat{H} + i \sum_n \left(\frac{e_n}{2} \hat{C}_n - \frac{1}{2} \hat{C}_n^\dagger \hat{C}_n - \frac{e_n^2}{8}\right),
 ```
 ```math
-M_n = C_n - \frac{e_n}{2},
+\hat{M}_n = \hat{C}_n - \frac{e_n}{2},
 ```
 and
 ```math
-e_n = \langle C_n + C_n^\dagger \rangle.
+e_n = \langle \hat{C}_n + \hat{C}_n^\dagger \rangle.
 ```
 
-Above, `C_n` is the `n`-th collapse operator and `dW_j(t)` is the real Wiener increment associated to `C_n`.
+Above, `\hat{C}_n` is the `n`-th collapse operator and `dW_j(t)` is the real Wiener increment associated to `\hat{C}_n`.
 
 # Arguments
 
@@ -193,13 +194,14 @@ function ssesolveProblem(
     saveat = is_empty_e_ops ? tlist : [tlist[end]]
     default_values = (DEFAULT_SDE_SOLVER_OPTIONS..., saveat = saveat)
     kwargs2 = merge(default_values, kwargs)
-    kwargs3 = _generate_se_me_kwargs(e_ops, makeVal(progress_bar), tlist, kwargs2, SaveFuncSESolve)
+    kwargs3 = _generate_se_me_kwargs(e_ops, makeVal(progress_bar), tlist, kwargs2, SaveFuncSSESolve)
+    kwargs4 = _ssesolve_add_normalize_cb(kwargs3)
 
     tspan = (tlist[1], tlist[end])
     noise =
         RealWienerProcess!(tlist[1], zeros(length(sc_ops)), zeros(length(sc_ops)), save_everystep = false, rng = rng)
     noise_rate_prototype = similar(ψ0, length(ψ0), length(sc_ops))
-    return SDEProblem{true}(K, D, ψ0, tspan, p; noise_rate_prototype = noise_rate_prototype, noise = noise, kwargs3...)
+    return SDEProblem{true}(K, D, ψ0, tspan, p; noise_rate_prototype = noise_rate_prototype, noise = noise, kwargs4...)
 end
 
 @doc raw"""
@@ -222,23 +224,23 @@ end
 Generate the SDE EnsembleProblem for the Stochastic Schrödinger time evolution of a quantum system. This is defined by the following stochastic differential equation:
     
 ```math
-d|\psi(t)\rangle = -i K |\psi(t)\rangle dt + \sum_n M_n |\psi(t)\rangle dW_n(t)
+d|\psi(t)\rangle = -i \hat{K} |\psi(t)\rangle dt + \sum_n \hat{M}_n |\psi(t)\rangle dW_n(t)
 ```
 
 where 
     
 ```math
-K = \hat{H} + i \sum_n \left(\frac{e_j} C_n - \frac{1}{2} \sum_{j} C_n^\dagger C_n - \frac{e_j^2}{8}\right),
+\hat{K} = \hat{H} + i \sum_n \left(\frac{e_n}{2} \hat{C}_n - \frac{1}{2} \hat{C}_n^\dagger \hat{C}_n - \frac{e_n^2}{8}\right),
 ```
 ```math
-M_n = C_n - \frac{e_n}{2},
+\hat{M}_n = \hat{C}_n - \frac{e_n}{2},
 ```
 and
 ```math
-e_n = \langle C_n + C_n^\dagger \rangle.
+e_n = \langle \hat{C}_n + \hat{C}_n^\dagger \rangle.
 ```
 
-Above, `C_n` is the `n`-th collapse operator and  `dW_j(t)` is the real Wiener increment associated to `C_n`.
+Above, `\hat{C}_n` is the `n`-th collapse operator and  `dW_j(t)` is the real Wiener increment associated to `\hat{C}_n`.
 
 # Arguments
 
@@ -345,23 +347,23 @@ Stochastic Schrödinger equation evolution of a quantum system given the system
 The stochastic evolution of the state ``|\psi(t)\rangle`` is defined by:
     
 ```math
-d|\psi(t)\rangle = -i K |\psi(t)\rangle dt + \sum_n M_n |\psi(t)\rangle dW_n(t)
+d|\psi(t)\rangle = -i \hat{K} |\psi(t)\rangle dt + \sum_n \hat{M}_n |\psi(t)\rangle dW_n(t)
 ```
 
 where 
     
 ```math
-K = \hat{H} + i \sum_n \left(\frac{e_j} C_n - \frac{1}{2} \sum_{j} C_n^\dagger C_n - \frac{e_j^2}{8}\right),
+\hat{K} = \hat{H} + i \sum_n \left(\frac{e_n}{2} \hat{C}_n - \frac{1}{2} \hat{C}_n^\dagger \hat{C}_n - \frac{e_n^2}{8}\right),
 ```
 ```math
-M_n = C_n - \frac{e_n}{2},
+\hat{M}_n = \hat{C}_n - \frac{e_n}{2},
 ```
 and
 ```math
-e_n = \langle C_n + C_n^\dagger \rangle.
+e_n = \langle \hat{C}_n + \hat{C}_n^\dagger \rangle.
 ```
 
-Above, `C_n` is the `n`-th collapse operator and `dW_j(t)` is the real Wiener increment associated to `C_n`.
+Above, `\hat{C}_n` is the `n`-th collapse operator and `dW_j(t)` is the real Wiener increment associated to `\hat{C}_n`.
 
 
 # Arguments

test/core-test/time_evolution.jl

@@ -18,6 +18,8 @@
     e_ops = [a' * a, σz]
     c_ops = [sqrt(γ * (1 + nth)) * a, sqrt(γ * nth) * a', sqrt(γ * (1 + nth)) * σm, sqrt(γ * nth) * σm']
 
+    ψ0_int = Qobj(round.(Int, ψ0.data), dims = ψ0.dims) # Used for testing the type inference
+
     @testset "sesolve" begin
         tlist = range(0, 20 * 2π / g, 1000)
 
@@ -83,7 +85,7 @@
         @testset "Type Inference sesolve" begin
             @inferred sesolveProblem(H, ψ0, tlist, progress_bar = Val(false))
             @inferred sesolveProblem(H, ψ0, [0, 10], progress_bar = Val(false))
-            @inferred sesolveProblem(H, Qobj(zeros(Int64, N * 2); dims = (N, 2)), tlist, progress_bar = Val(false))
+            @inferred sesolveProblem(H, ψ0_int, tlist, progress_bar = Val(false))
             @inferred sesolve(H, ψ0, tlist, e_ops = e_ops, progress_bar = Val(false))
             @inferred sesolve(H, ψ0, tlist, progress_bar = Val(false))
             @inferred sesolve(H, ψ0, tlist, e_ops = e_ops, saveat = tlist, progress_bar = Val(false))
@@ -367,14 +369,7 @@
             ad_t = QobjEvo(a', coef)
             @inferred mesolveProblem(H, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false))
             @inferred mesolveProblem(H, ψ0, [0, 10], c_ops, e_ops = e_ops, progress_bar = Val(false))
-            @inferred mesolveProblem(
-                H,
-                tensor(Qobj(zeros(Int64, N)), Qobj([0, 1])),
-                tlist,
-                c_ops,
-                e_ops = e_ops,
-                progress_bar = Val(false),
-            )
+            @inferred mesolveProblem(H, ψ0_int, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false))
             @inferred mesolve(H, ψ0, tlist, c_ops, e_ops = e_ops, progress_bar = Val(false))
             @inferred mesolve(H, ψ0, tlist, c_ops, progress_bar = Val(false))
             @inferred mesolve(H, ψ0, tlist, c_ops, e_ops = e_ops, saveat = tlist, progress_bar = Val(false))
@@ -398,15 +393,7 @@
             @inferred mcsolve(H, ψ0, tlist, c_ops, ntraj = 5, e_ops = e_ops, progress_bar = Val(false), rng = rng)
             @inferred mcsolve(H, ψ0, tlist, c_ops, ntraj = 5, progress_bar = Val(true), rng = rng)
             @inferred mcsolve(H, ψ0, [0, 10], c_ops, ntraj = 5, progress_bar = Val(false), rng = rng)
-            @inferred mcsolve(
-                H,
-                tensor(Qobj(zeros(Int64, N)), Qobj([0, 1])),
-                tlist,
-                c_ops,
-                ntraj = 5,
-                progress_bar = Val(false),
-                rng = rng,
-            )
+            @inferred mcsolve(H, ψ0_int, tlist, c_ops, ntraj = 5, progress_bar = Val(false), rng = rng)
             @inferred mcsolve(
                 H,
                 ψ0,
@@ -454,15 +441,7 @@
             )
             @inferred ssesolve(H, ψ0, tlist, c_ops_tuple, ntraj = 5, progress_bar = Val(true), rng = rng)
             @inferred ssesolve(H, ψ0, [0, 10], c_ops_tuple, ntraj = 5, progress_bar = Val(false), rng = rng)
-            @inferred ssesolve(
-                H,
-                tensor(Qobj(zeros(Int64, N)), Qobj([0, 1])),
-                tlist,
-                c_ops_tuple,
-                ntraj = 5,
-                progress_bar = Val(false),
-                rng = rng,
-            )
+            @inferred ssesolve(H, ψ0_int, tlist, c_ops_tuple, ntraj = 5, progress_bar = Val(false), rng = rng)
             @inferred ssesolve(
                 H,
                 ψ0,