From d469d80e26aa3bf749950390638d140804b540e1 Mon Sep 17 00:00:00 2001
From: Alberto Mercurio <61953577+albertomercurio@users.noreply.github.com>
Date: Wed, 28 Aug 2024 18:35:56 +0200
Subject: [PATCH 01/13] Improve OperatorSum (#209)

* Improve OperatorSum

* Include the file in the main file

* Add API docs for OperatorSum
---
 docs/src/api.md                      |  1 +
 src/QuantumToolbox.jl                |  1 +
 src/qobj/operator_sum.jl             | 49 ++++++++++++++++++++++++++++
 src/time_evolution/time_evolution.jl | 46 ++------------------------
 4 files changed, 53 insertions(+), 44 deletions(-)
 create mode 100644 src/qobj/operator_sum.jl

diff --git a/docs/src/api.md b/docs/src/api.md
index 3c27d4e45..fb66c49a8 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -26,6 +26,7 @@ OperatorKet
 SuperOperatorQuantumObject
 SuperOperator
 QuantumObject
+OperatorSum
 size
 eltype
 length
diff --git a/src/QuantumToolbox.jl b/src/QuantumToolbox.jl
index b7031fed0..bb99a645e 100644
--- a/src/QuantumToolbox.jl
+++ b/src/QuantumToolbox.jl
@@ -52,6 +52,7 @@ include("qobj/states.jl")
 include("qobj/operators.jl")
 include("qobj/superoperators.jl")
 include("qobj/synonyms.jl")
+include("qobj/operator_sum.jl")
 
 # time evolution
 include("time_evolution/time_evolution.jl")
diff --git a/src/qobj/operator_sum.jl b/src/qobj/operator_sum.jl
new file mode 100644
index 000000000..6f3642f5e
--- /dev/null
+++ b/src/qobj/operator_sum.jl
@@ -0,0 +1,49 @@
+export OperatorSum
+
+@doc raw"""
+    struct OperatorSum
+
+A structure to represent a sum of operators ``\sum_i c_i \hat{O}_i`` with a list of coefficients ``c_i`` and a list of operators ``\hat{O}_i``.
+
+This is very useful when we have to update only the coefficients, without allocating memory by performing the sum of the operators.
+"""
+struct OperatorSum{CT<:Vector{<:Number},OT<:AbstractVector} <: AbstractQuantumObject
+    coefficients::CT
+    operators::OT
+    function OperatorSum(coefficients::CT, operators::OT) where {CT<:Vector{<:Number},OT<:AbstractVector}
+        length(coefficients) == length(operators) ||
+            throw(DimensionMismatch("The number of coefficients must be the same as the number of operators."))
+        # Check if all the operators have the same dimensions
+        size_1 = size(operators[1])
+        mapreduce(x -> size(x) == size_1, &, operators) ||
+            throw(DimensionMismatch("All the operators must have the same dimensions."))
+        T = promote_type(
+            mapreduce(x -> eltype(x.data), promote_type, operators),
+            mapreduce(eltype, promote_type, coefficients),
+        )
+        coefficients2 = T.(coefficients)
+        return new{Vector{T},OT}(coefficients2, operators)
+    end
+end
+
+Base.size(A::OperatorSum) = size(A.operators[1])
+Base.size(A::OperatorSum, inds...) = size(A.operators[1], inds...)
+Base.length(A::OperatorSum) = length(A.operators[1])
+Base.copy(A::OperatorSum) = OperatorSum(copy(A.coefficients), copy(A.operators))
+Base.deepcopy(A::OperatorSum) = OperatorSum(deepcopy(A.coefficients), deepcopy(A.operators))
+
+function update_coefficients!(A::OperatorSum, coefficients)
+    length(A.coefficients) == length(coefficients) ||
+        throw(DimensionMismatch("The number of coefficients must be the same as the number of operators."))
+    return A.coefficients .= coefficients
+end
+
+@inline function LinearAlgebra.mul!(y::AbstractVector{T}, A::OperatorSum, x::AbstractVector, α, β) where {T}
+    # Note that β is applied only to the first term
+    mul!(y, A.operators[1], x, α * A.coefficients[1], β)
+    @inbounds for i in 2:length(A.operators)
+        A.coefficients[i] == 0 && continue
+        mul!(y, A.operators[i], x, α * A.coefficients[i], 1)
+    end
+    return y
+end
diff --git a/src/time_evolution/time_evolution.jl b/src/time_evolution/time_evolution.jl
index 42a9bd0f7..7c2fd89aa 100644
--- a/src/time_evolution/time_evolution.jl
+++ b/src/time_evolution/time_evolution.jl
@@ -1,4 +1,4 @@
-export OperatorSum, TimeDependentOperatorSum
+export TimeDependentOperatorSum
 export TimeEvolutionSol, TimeEvolutionMCSol
 
 export liouvillian, liouvillian_floquet, liouvillian_generalized
@@ -105,49 +105,7 @@ struct DiscreteLindbladJumpCallback <: LindbladJumpCallbackType end
 
 ContinuousLindbladJumpCallback(; interp_points::Int = 10) = ContinuousLindbladJumpCallback(interp_points)
 
-## Sum of operators
-
-struct OperatorSum{CT<:Vector{<:Number},OT<:Vector{<:QuantumObject}} <: AbstractQuantumObject
-    coefficients::CT
-    operators::OT
-    function OperatorSum(coefficients::CT, operators::OT) where {CT<:Vector{<:Number},OT<:Vector{<:QuantumObject}}
-        length(coefficients) == length(operators) ||
-            throw(DimensionMismatch("The number of coefficients must be the same as the number of operators."))
-        # Check if all the operators have the same dimensions
-        dims = operators[1].dims
-        optype = operators[1].type
-        mapreduce(x -> x.dims == dims && x.type == optype, &, operators) ||
-            throw(DimensionMismatch("All the operators must have the same dimensions."))
-        T = promote_type(
-            mapreduce(x -> eltype(x.data), promote_type, operators),
-            mapreduce(eltype, promote_type, coefficients),
-        )
-        coefficients2 = T.(coefficients)
-        return new{Vector{T},OT}(coefficients2, operators)
-    end
-end
-
-Base.size(A::OperatorSum) = size(A.operators[1])
-Base.size(A::OperatorSum, inds...) = size(A.operators[1], inds...)
-Base.length(A::OperatorSum) = length(A.operators[1])
-Base.copy(A::OperatorSum) = OperatorSum(copy(A.coefficients), copy(A.operators))
-Base.deepcopy(A::OperatorSum) = OperatorSum(deepcopy(A.coefficients), deepcopy(A.operators))
-
-function update_coefficients!(A::OperatorSum, coefficients)
-    length(A.coefficients) == length(coefficients) ||
-        throw(DimensionMismatch("The number of coefficients must be the same as the number of operators."))
-    return A.coefficients .= coefficients
-end
-
-@inline function LinearAlgebra.mul!(y::AbstractVector{T}, A::OperatorSum, x::AbstractVector, α, β) where {T}
-    # Note that β is applied only to the first term
-    mul!(y, A.operators[1], x, α * A.coefficients[1], β)
-    @inbounds for i in 2:length(A.operators)
-        A.coefficients[i] == 0 && continue
-        mul!(y, A.operators[i], x, α * A.coefficients[i], 1)
-    end
-    return y
-end
+## Time-dependent sum of operators
 
 struct TimeDependentOperatorSum{CFT,OST<:OperatorSum}
     coefficient_functions::CFT

From 4a1dbd559e2f61f58a2dec076c3389f2f0be5030 Mon Sep 17 00:00:00 2001
From: lgravina1997 <luca.gravina@epfl.ch>
Date: Wed, 28 Aug 2024 18:39:08 +0200
Subject: [PATCH 02/13] ssesolve problem

---
 src/QuantumToolbox.jl          |  1 +
 src/time_evolution/ssesolve.jl | 98 ++++++++++++++++++++++++++++++++++
 2 files changed, 99 insertions(+)
 create mode 100644 src/time_evolution/ssesolve.jl

diff --git a/src/QuantumToolbox.jl b/src/QuantumToolbox.jl
index b7031fed0..0eab8d48d 100644
--- a/src/QuantumToolbox.jl
+++ b/src/QuantumToolbox.jl
@@ -59,6 +59,7 @@ include("time_evolution/mesolve.jl")
 include("time_evolution/lr_mesolve.jl")
 include("time_evolution/sesolve.jl")
 include("time_evolution/mcsolve.jl")
+include("time_evolution/ssesolve.jl")
 include("time_evolution/time_evolution_dynamical.jl")
 
 # Others
diff --git a/src/time_evolution/ssesolve.jl b/src/time_evolution/ssesolve.jl
new file mode 100644
index 000000000..ee0a28ef0
--- /dev/null
+++ b/src/time_evolution/ssesolve.jl
@@ -0,0 +1,98 @@
+export ssesolveProblem
+
+#TODO: Check if works in GPU
+function _ssesolve_update_coefficients!(ψ, coefficients, c_ops)
+    _get_en = op -> real(dot(ψ, op, ψ))
+    @. coefficients[2:end-1] = _get_en(c_ops)
+    @. coefficients[end] = -sum(x->x^2, coefficients[2:end-1]) / 2
+    return nothing
+end
+
+function ssesolve_ti_drift!(du, u, p, t)
+    _ssesolve_update_coefficients!(u, p.K.coefficients, c_ops)
+
+    # mul!(view(du, :, 1), p.K, u)
+    # @views du[:, 2:end] .= du[:, 1] 
+    mul!(du, p.K, u)
+    return nothing
+end
+
+function ssesolve_ti_diffusion!(du, u, p, t)
+    D = p.D
+    en = p.K.coefficients[2:end-1]
+
+    du_reshaped = reshape(du, :)
+    mul!(du_reshaped, D, u)
+    du .-= en 
+end
+
+
+function ssesolveProblem(
+    H::QuantumObject{MT1,OperatorQuantumObject},
+    ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
+    tlist::AbstractVector,
+    c_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
+    alg::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm = Tsit5(),
+    e_ops::Union{Nothing,AbstractVector} = nothing,
+    H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
+    params::NamedTuple = NamedTuple(),
+    progress_bar::Union{Val,Bool} = Val(true),
+    kwargs...,
+) where {MT1<:AbstractMatrix,T2,Tc<:AbstractMatrix,TJC<:LindbladJumpCallbackType}
+
+    H.dims != ψ0.dims && throw(DimensionMismatch("The two quantum objects are not of the same Hilbert dimension."))
+
+    haskey(kwargs, :save_idxs) &&
+        throw(ArgumentError("The keyword argument \"save_idxs\" is not supported in QuantumToolbox."))
+
+    is_time_dependent = !(H_t isa Nothing)
+    progress_bar_val = makeVal(progress_bar)
+
+    t_l = convert(Vector{Float64}, tlist) # Convert it into Float64 to avoid type instabilities for OrdinaryDiffEq.jl
+
+    ϕ0 = get_data(ψ0)
+
+    H_eff = get_data(H - T2(0.5im) * mapreduce(op -> op' * op, +, c_ops))
+    c_ops2 = get_data.(c_ops)
+
+    coefficients = [1.0, fill(0.0, length(c_ops)+1)...]
+    operators = [H_eff, c_ops2..., I(prod(H.dims))]
+    K = OperatorSum(coefficients, operators)
+    _ssesolve_update_coefficients!(ϕ0, K.coefficients, c_ops2)
+
+    D = vcat(c_ops2...)
+
+    progr = ProgressBar(length(t_l), enable = getVal(progress_bar_val))
+
+    if e_ops isa Nothing
+        expvals = Array{ComplexF64}(undef, 0, length(t_l))
+        e_ops2 = MT1[]
+        is_empty_e_ops = true
+    else
+        expvals = Array{ComplexF64}(undef, length(e_ops), length(t_l))
+        e_ops2 = get_data.(e_ops)
+        is_empty_e_ops = isempty(e_ops)
+    end
+
+    p = (
+        K = K,
+        D = D,
+        e_ops = e_ops2,
+        c_ops = c_ops2,
+        expvals = expvals,
+        progr = progr,
+        Hdims = H.dims,
+        H_t = H_t,
+        is_empty_e_ops = is_empty_e_ops,
+        params...,
+    )
+
+    saveat = e_ops isa Nothing ? t_l : [t_l[end]]
+    default_values = (DEFAULT_ODE_SOLVER_OPTIONS..., saveat = saveat)
+    kwargs2 = merge(default_values, kwargs)
+    kwargs3 = _generate_sesolve_kwargs(e_ops, progress_bar_val, t_l, kwargs2)
+
+    tspan = (t_l[1], t_l[end])
+    A = similar(ϕ0, length(ϕ0), length(c_ops))
+    return SDEProblem(ssesolve_ti_drift!, ssesolve_ti_diffusion!, ϕ0, tspan, p; noise_rate_prototype=A, kwargs3...)
+end

From 0515244ddb43b14c5dac2752e842487c0ea71e75 Mon Sep 17 00:00:00 2001
From: lgravina1997 <luca.gravina@epfl.ch>
Date: Wed, 28 Aug 2024 18:41:04 +0200
Subject: [PATCH 03/13] fix operator_sum

---
 src/qobj/operator_sum.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/qobj/operator_sum.jl b/src/qobj/operator_sum.jl
index 6f3642f5e..6a32038c3 100644
--- a/src/qobj/operator_sum.jl
+++ b/src/qobj/operator_sum.jl
@@ -18,7 +18,7 @@ struct OperatorSum{CT<:Vector{<:Number},OT<:AbstractVector} <: AbstractQuantumOb
         mapreduce(x -> size(x) == size_1, &, operators) ||
             throw(DimensionMismatch("All the operators must have the same dimensions."))
         T = promote_type(
-            mapreduce(x -> eltype(x.data), promote_type, operators),
+            mapreduce(x -> eltype(x), promote_type, operators),
             mapreduce(eltype, promote_type, coefficients),
         )
         coefficients2 = T.(coefficients)

From eef4ce98850006f4a2e8740e1716bb326725d6e8 Mon Sep 17 00:00:00 2001
From: lgravina1997 <luca.gravina@epfl.ch>
Date: Wed, 28 Aug 2024 20:07:16 +0200
Subject: [PATCH 04/13] fixes

---
 Project.toml                   |  1 +
 src/QuantumToolbox.jl          |  1 +
 src/time_evolution/ssesolve.jl | 23 ++++++++++++++---------
 3 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/Project.toml b/Project.toml
index 7b33dab50..4962a98ea 100644
--- a/Project.toml
+++ b/Project.toml
@@ -18,6 +18,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
diff --git a/src/QuantumToolbox.jl b/src/QuantumToolbox.jl
index 948eaa7a7..38557f920 100644
--- a/src/QuantumToolbox.jl
+++ b/src/QuantumToolbox.jl
@@ -7,6 +7,7 @@ import Reexport: @reexport
 @reexport using LinearAlgebra
 @reexport using SparseArrays
 @reexport using OrdinaryDiffEq
+@reexport using StochasticDiffEq
 @reexport using LinearSolve
 
 # other functions in LinearAlgebra
diff --git a/src/time_evolution/ssesolve.jl b/src/time_evolution/ssesolve.jl
index ee0a28ef0..2af8a820f 100644
--- a/src/time_evolution/ssesolve.jl
+++ b/src/time_evolution/ssesolve.jl
@@ -4,26 +4,30 @@ export ssesolveProblem
 function _ssesolve_update_coefficients!(ψ, coefficients, c_ops)
     _get_en = op -> real(dot(ψ, op, ψ))
     @. coefficients[2:end-1] = _get_en(c_ops)
-    @. coefficients[end] = -sum(x->x^2, coefficients[2:end-1]) / 2
+    coefficients[end] = -sum(x->x^2, coefficients[2:end-1]) / 2
     return nothing
 end
 
 function ssesolve_ti_drift!(du, u, p, t)
-    _ssesolve_update_coefficients!(u, p.K.coefficients, c_ops)
-
-    # mul!(view(du, :, 1), p.K, u)
-    # @views du[:, 2:end] .= du[:, 1] 
+    _ssesolve_update_coefficients!(u, p.K.coefficients, p.c_ops)
+    
     mul!(du, p.K, u)
     return nothing
 end
 
 function ssesolve_ti_diffusion!(du, u, p, t)
     D = p.D
-    en = p.K.coefficients[2:end-1]
+    @views en = p.K.coefficients[2:end-1]
 
     du_reshaped = reshape(du, :)
     mul!(du_reshaped, D, u)
-    du .-= en 
+
+    #TODO: write something decent
+    for i in 1:length(en)
+        du[:,i] .-= en[i] .* u
+    end
+    
+    return nothing
 end
 
 
@@ -56,7 +60,7 @@ function ssesolveProblem(
     c_ops2 = get_data.(c_ops)
 
     coefficients = [1.0, fill(0.0, length(c_ops)+1)...]
-    operators = [H_eff, c_ops2..., I(prod(H.dims))]
+    operators = [-1im * H_eff, c_ops2..., I(prod(H.dims))]
     K = OperatorSum(coefficients, operators)
     _ssesolve_update_coefficients!(ϕ0, K.coefficients, c_ops2)
 
@@ -75,6 +79,7 @@ function ssesolveProblem(
     end
 
     p = (
+        U = -1im * get_data(H),
         K = K,
         D = D,
         e_ops = e_ops2,
@@ -88,7 +93,7 @@ function ssesolveProblem(
     )
 
     saveat = e_ops isa Nothing ? t_l : [t_l[end]]
-    default_values = (DEFAULT_ODE_SOLVER_OPTIONS..., saveat = saveat)
+    default_values = (saveat = saveat, )
     kwargs2 = merge(default_values, kwargs)
     kwargs3 = _generate_sesolve_kwargs(e_ops, progress_bar_val, t_l, kwargs2)
 

From 7b43c7ede7340f2c353fb5be3c8df844de0af21c Mon Sep 17 00:00:00 2001
From: lgravina1997 <luca.gravina@epfl.ch>
Date: Thu, 29 Aug 2024 09:43:33 +0200
Subject: [PATCH 05/13] improved diffusion and added comments

---
 src/time_evolution/ssesolve.jl | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/src/time_evolution/ssesolve.jl b/src/time_evolution/ssesolve.jl
index 2af8a820f..a68db4883 100644
--- a/src/time_evolution/ssesolve.jl
+++ b/src/time_evolution/ssesolve.jl
@@ -2,9 +2,9 @@ export ssesolveProblem
 
 #TODO: Check if works in GPU
 function _ssesolve_update_coefficients!(ψ, coefficients, c_ops)
-    _get_en = op -> real(dot(ψ, op, ψ))
-    @. coefficients[2:end-1] = _get_en(c_ops)
-    coefficients[end] = -sum(x->x^2, coefficients[2:end-1]) / 2
+    _get_en = op -> real(dot(ψ, op, ψ)) #this is en/2: <Sn + Sn'>/2 = Re<Sn>
+    @. coefficients[2:end-1] = _get_en(c_ops) #coefficients of the OperatorSum: Σ Sn * en/2
+    coefficients[end] = - sum(x->x^2, coefficients[2:end-1]) / 2 #this last coefficient is -Σen^2/8
     return nothing
 end
 
@@ -19,14 +19,12 @@ function ssesolve_ti_diffusion!(du, u, p, t)
     D = p.D
     @views en = p.K.coefficients[2:end-1]
 
+    # du:(H,W). du_reshaped:(H*W,). 
+    # H:Hilbert space dimension, W: number of c_ops
     du_reshaped = reshape(du, :)
-    mul!(du_reshaped, D, u)
+    mul!(du_reshaped, D, u) #du[:,i] = D[i] * u
 
-    #TODO: write something decent
-    for i in 1:length(en)
-        du[:,i] .-= en[i] .* u
-    end
-    
+    du .-= u .* reshape(en, 1, :) #du[:,i] -= en[i] * u
     return nothing
 end
 

From 06cd0677de6218641991dc5859d1345fac887f8e Mon Sep 17 00:00:00 2001
From: lgravina1997 <luca.gravina@epfl.ch>
Date: Thu, 29 Aug 2024 17:48:29 +0200
Subject: [PATCH 06/13] ensemble problem

---
 src/time_evolution/mcsolve.jl        |   4 +-
 src/time_evolution/ssesolve.jl       | 138 +++++++++++++++++++++++++--
 src/time_evolution/time_evolution.jl |  32 +++++++
 3 files changed, 162 insertions(+), 12 deletions(-)

diff --git a/src/time_evolution/mcsolve.jl b/src/time_evolution/mcsolve.jl
index 853588ada..c62c7783b 100644
--- a/src/time_evolution/mcsolve.jl
+++ b/src/time_evolution/mcsolve.jl
@@ -86,7 +86,7 @@ function _mcsolve_output_func(sol, i)
     return (sol, false)
 end
 
-function _mcsolve_generate_statistics(sol, i, times, states, expvals_all, jump_times, jump_which)
+function _mcsolve_generate_statistics!(sol, i, times, states, expvals_all, jump_times, jump_which)
     sol_i = sol[:, i]
     !isempty(sol_i.prob.kwargs[:saveat]) ?
     states[i] = [QuantumObject(sol_i.u[i], dims = sol_i.prob.p.Hdims) for i in 1:length(sol_i.u)] : nothing
@@ -524,7 +524,7 @@ function mcsolve(
     jump_which = Vector{Vector{Int16}}(undef, length(sol))
 
     foreach(
-        i -> _mcsolve_generate_statistics(sol, i, times, states, expvals_all, jump_times, jump_which),
+        i -> _mcsolve_generate_statistics!(sol, i, times, states, expvals_all, jump_times, jump_which),
         eachindex(sol),
     )
     expvals = dropdims(sum(expvals_all, dims = 1), dims = 1) ./ length(sol)
diff --git a/src/time_evolution/ssesolve.jl b/src/time_evolution/ssesolve.jl
index a68db4883..4f635942c 100644
--- a/src/time_evolution/ssesolve.jl
+++ b/src/time_evolution/ssesolve.jl
@@ -1,4 +1,4 @@
-export ssesolveProblem
+export ssesolveProblem, ssesolveEnsembleProblem, ssesolve
 
 #TODO: Check if works in GPU
 function _ssesolve_update_coefficients!(ψ, coefficients, c_ops)
@@ -8,46 +8,71 @@ function _ssesolve_update_coefficients!(ψ, coefficients, c_ops)
     return nothing
 end
 
-function ssesolve_ti_drift!(du, u, p, t)
+function ssesolve_drift!(du, u, p, t)
     _ssesolve_update_coefficients!(u, p.K.coefficients, p.c_ops)
     
     mul!(du, p.K, u)
     return nothing
 end
 
-function ssesolve_ti_diffusion!(du, u, p, t)
-    D = p.D
+function ssesolve_diffusion!(du, u, p, t)
     @views en = p.K.coefficients[2:end-1]
 
     # du:(H,W). du_reshaped:(H*W,). 
     # H:Hilbert space dimension, W: number of c_ops
     du_reshaped = reshape(du, :)
-    mul!(du_reshaped, D, u) #du[:,i] = D[i] * u
+    mul!(du_reshaped, p.D, u) #du[:,i] = D[i] * u
 
     du .-= u .* reshape(en, 1, :) #du[:,i] -= en[i] * u
     return nothing
 end
 
+function _ssesolve_prob_func(prob, i, repeat)
+    internal_params = prob.p
+
+    prm = merge(
+        internal_params,
+        (
+            expvals = similar(internal_params.expvals),
+            progr = ProgressBar(size(internal_params.expvals, 2), enable = false),
+        ),
+    )
+    return remake(prob, p=prm)
+end
+
+function _ssesolve_output_func(sol, i)
+    return (sol, false)
+end
+
+function _ssesolve_generate_statistics!(sol, i, states, expvals_all)
+    sol_i = sol[:, i]
+    !isempty(sol_i.prob.kwargs[:saveat]) ?
+    states[i] = [QuantumObject(sol_i.u[i], dims = sol_i.prob.p.Hdims) for i in 1:length(sol_i.u)] : nothing
+
+    copyto!(view(expvals_all, i, :, :), sol_i.prob.p.expvals)
+    return nothing
+end
 
 function ssesolveProblem(
     H::QuantumObject{MT1,OperatorQuantumObject},
     ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
     tlist::AbstractVector,
     c_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
-    alg::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm = Tsit5(),
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = SRA1(),
     e_ops::Union{Nothing,AbstractVector} = nothing,
     H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
     params::NamedTuple = NamedTuple(),
     progress_bar::Union{Val,Bool} = Val(true),
     kwargs...,
-) where {MT1<:AbstractMatrix,T2,Tc<:AbstractMatrix,TJC<:LindbladJumpCallbackType}
+) where {MT1<:AbstractMatrix,T2,Tc<:AbstractMatrix}
 
     H.dims != ψ0.dims && throw(DimensionMismatch("The two quantum objects are not of the same Hilbert dimension."))
 
     haskey(kwargs, :save_idxs) &&
         throw(ArgumentError("The keyword argument \"save_idxs\" is not supported in QuantumToolbox."))
 
-    is_time_dependent = !(H_t isa Nothing)
+    !(H_t isa Nothing) && throw(ArgumentError("Time-dependent Hamiltonians are not currently supported in ssesolve."))
+
     progress_bar_val = makeVal(progress_bar)
 
     t_l = convert(Vector{Float64}, tlist) # Convert it into Float64 to avoid type instabilities for OrdinaryDiffEq.jl
@@ -91,11 +116,104 @@ function ssesolveProblem(
     )
 
     saveat = e_ops isa Nothing ? t_l : [t_l[end]]
-    default_values = (saveat = saveat, )
+    default_values = (DEFAULT_SDE_SOLVER_OPTIONS..., saveat = saveat)
     kwargs2 = merge(default_values, kwargs)
     kwargs3 = _generate_sesolve_kwargs(e_ops, progress_bar_val, t_l, kwargs2)
 
     tspan = (t_l[1], t_l[end])
     A = similar(ϕ0, length(ϕ0), length(c_ops))
-    return SDEProblem(ssesolve_ti_drift!, ssesolve_ti_diffusion!, ϕ0, tspan, p; noise_rate_prototype=A, kwargs3...)
+    return SDEProblem(ssesolve_drift!, ssesolve_diffusion!, ϕ0, tspan, p; noise_rate_prototype=A, kwargs3...)
+end
+
+function ssesolveEnsembleProblem(
+    H::QuantumObject{MT1,OperatorQuantumObject},
+    ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
+    tlist::AbstractVector,
+    c_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = SRA1(),
+    e_ops::Union{Nothing,AbstractVector} = nothing,
+    H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
+    params::NamedTuple = NamedTuple(),
+    prob_func::Function = _ssesolve_prob_func,
+    output_func::Function = _ssesolve_output_func,
+    kwargs...,
+) where {MT1<:AbstractMatrix,T2,Tc<:AbstractMatrix}
+    prob_sse = ssesolveProblem(
+        H,
+        ψ0,
+        tlist,
+        c_ops;
+        alg = alg,
+        e_ops = e_ops,
+        H_t = H_t,
+        params = params,
+        kwargs...,
+    )
+
+    ensemble_prob = EnsembleProblem(prob_sse, prob_func = prob_func, output_func = output_func, safetycopy = false)
+
+    return ensemble_prob
+end
+
+function ssesolve(
+    H::QuantumObject{MT1,OperatorQuantumObject},
+    ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
+    tlist::AbstractVector,
+    c_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = SRA1(),
+    e_ops::Union{Nothing,AbstractVector} = nothing,
+    H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
+    params::NamedTuple = NamedTuple(),
+    n_traj::Int = 1,
+    ensemble_method = EnsembleThreads(),
+    prob_func::Function = _ssesolve_prob_func,
+    output_func::Function = _ssesolve_output_func,
+    kwargs...,
+) where {MT1<:AbstractMatrix,T2,Tc<:AbstractMatrix}
+
+    ens_prob = ssesolveEnsembleProblem(
+        H,
+        ψ0,
+        tlist,
+        c_ops;
+        alg = alg,
+        e_ops = e_ops,
+        H_t = H_t,
+        params = params,
+        prob_func = prob_func,
+        output_func = output_func,
+        kwargs...,
+    )
+
+    return ssesolve(ens_prob; alg = alg, n_traj = n_traj, ensemble_method = ensemble_method)
+end
+
+function ssesolve(
+    ens_prob::EnsembleProblem;
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = SRA1(),
+    n_traj::Int = 1,
+    ensemble_method = EnsembleThreads(),
+)
+    sol = solve(ens_prob, alg, ensemble_method, trajectories = n_traj)
+    _sol_1 = sol[:, 1]
+
+    expvals_all = Array{ComplexF64}(undef, length(sol), size(_sol_1.prob.p.expvals)...)
+    states =
+        isempty(_sol_1.prob.kwargs[:saveat]) ? fill(QuantumObject[], length(sol)) :
+        Vector{Vector{QuantumObject}}(undef, length(sol))
+
+    foreach(i -> _ssesolve_generate_statistics!(sol, i, states, expvals_all), eachindex(sol),)
+    expvals = dropdims(sum(expvals_all, dims = 1), dims = 1) ./ length(sol)
+
+    return TimeEvolutionSSESol(
+        n_traj,
+        _sol_1.t,
+        states,
+        expvals,
+        expvals_all,
+        sol.converged,
+        _sol_1.alg,
+        _sol_1.prob.kwargs[:abstol],
+        _sol_1.prob.kwargs[:reltol],
+    )
 end
diff --git a/src/time_evolution/time_evolution.jl b/src/time_evolution/time_evolution.jl
index 7c2fd89aa..9c5ec3179 100644
--- a/src/time_evolution/time_evolution.jl
+++ b/src/time_evolution/time_evolution.jl
@@ -4,6 +4,7 @@ export TimeEvolutionSol, TimeEvolutionMCSol
 export liouvillian, liouvillian_floquet, liouvillian_generalized
 
 const DEFAULT_ODE_SOLVER_OPTIONS = (abstol = 1e-8, reltol = 1e-6, save_everystep = false, save_end = true)
+const DEFAULT_SDE_SOLVER_OPTIONS = (abstol = 1e-2, reltol = 1e-2, save_everystep = false, save_end = true)
 
 @doc raw"""
     struct TimeEvolutionSol
@@ -95,6 +96,37 @@ function Base.show(io::IO, sol::TimeEvolutionMCSol)
     return nothing
 end
 
+struct TimeEvolutionSSESol{
+    TT<:Vector{<:Real},
+    TS<:AbstractVector,
+    TE<:Matrix{ComplexF64},
+    TEA<:Array{ComplexF64,3},
+    T<:Real,
+}
+    n_traj::Int
+    times::TT
+    states::TS
+    expect::TE
+    expect_all::TEA
+    converged::Bool
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm
+    abstol::T
+    reltol::T
+end
+
+function Base.show(io::IO, sol::TimeEvolutionSSESol)
+    print(io, "Solution of quantum trajectories\n")
+    print(io, "(converged: $(sol.converged))\n")
+    print(io, "--------------------------------\n")
+    print(io, "num_trajectories = $(sol.n_traj)\n")
+    # print(io, "num_states = $(length(sol.states[1]))\n")
+    print(io, "num_expect = $(size(sol.expect, 1))\n")
+    print(io, "SDE alg.: $(sol.alg)\n")
+    print(io, "abstol = $(sol.abstol)\n")
+    print(io, "reltol = $(sol.reltol)\n")
+    return nothing
+end
+
 abstract type LindbladJumpCallbackType end
 
 struct ContinuousLindbladJumpCallback <: LindbladJumpCallbackType

From 3e344ca94452edc7942e9d3c2dc58b382b5f9889 Mon Sep 17 00:00:00 2001
From: Alberto Mercurio <alberto.mercurio96@gmail.com>
Date: Thu, 29 Aug 2024 20:26:57 +0200
Subject: [PATCH 07/13] Fix instabilities

---
 Project.toml                   |  1 +
 src/QuantumToolbox.jl          |  1 +
 src/time_evolution/ssesolve.jl | 59 +++++++++++++++++-----------------
 3 files changed, 32 insertions(+), 29 deletions(-)

diff --git a/Project.toml b/Project.toml
index 4962a98ea..a4ff95900 100644
--- a/Project.toml
+++ b/Project.toml
@@ -6,6 +6,7 @@ version = "0.12.0"
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
+DiffEqNoiseProcess = "77a26b50-5914-5dd7-bc55-306e6241c503"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895"
diff --git a/src/QuantumToolbox.jl b/src/QuantumToolbox.jl
index 38557f920..46c35e56a 100644
--- a/src/QuantumToolbox.jl
+++ b/src/QuantumToolbox.jl
@@ -30,6 +30,7 @@ import Graphs: connected_components, DiGraph
 import IncompleteLU: ilu
 import LinearMaps: LinearMap
 import OrdinaryDiffEq: OrdinaryDiffEqAlgorithm
+import DiffEqNoiseProcess: RealWienerProcess, RealWienerProcess!
 import Pkg
 import Random
 import SpecialFunctions: loggamma
diff --git a/src/time_evolution/ssesolve.jl b/src/time_evolution/ssesolve.jl
index 4f635942c..ee5524690 100644
--- a/src/time_evolution/ssesolve.jl
+++ b/src/time_evolution/ssesolve.jl
@@ -4,19 +4,20 @@ export ssesolveProblem, ssesolveEnsembleProblem, ssesolve
 function _ssesolve_update_coefficients!(ψ, coefficients, c_ops)
     _get_en = op -> real(dot(ψ, op, ψ)) #this is en/2: <Sn + Sn'>/2 = Re<Sn>
     @. coefficients[2:end-1] = _get_en(c_ops) #coefficients of the OperatorSum: Σ Sn * en/2
-    coefficients[end] = - sum(x->x^2, coefficients[2:end-1]) / 2 #this last coefficient is -Σen^2/8
+    coefficients[end] = -sum(x -> x^2, coefficients[2:end-1]) / 2 #this last coefficient is -Σen^2/8
     return nothing
 end
 
 function ssesolve_drift!(du, u, p, t)
     _ssesolve_update_coefficients!(u, p.K.coefficients, p.c_ops)
-    
+
     mul!(du, p.K, u)
+
     return nothing
 end
 
 function ssesolve_diffusion!(du, u, p, t)
-    @views en = p.K.coefficients[2:end-1]
+    @inbounds en = @view(p.K.coefficients[2:end-1])
 
     # du:(H,W). du_reshaped:(H*W,). 
     # H:Hilbert space dimension, W: number of c_ops
@@ -24,12 +25,16 @@ function ssesolve_diffusion!(du, u, p, t)
     mul!(du_reshaped, p.D, u) #du[:,i] = D[i] * u
 
     du .-= u .* reshape(en, 1, :) #du[:,i] -= en[i] * u
+
     return nothing
 end
 
 function _ssesolve_prob_func(prob, i, repeat)
     internal_params = prob.p
 
+    noise_rate_prototype = similar(prob.u0, length(prob.u0), length(internal_params.c_ops))
+    noise = RealWienerProcess!(prob.tspan[1], zeros(length(internal_params.c_ops)))
+
     prm = merge(
         internal_params,
         (
@@ -37,13 +42,12 @@ function _ssesolve_prob_func(prob, i, repeat)
             progr = ProgressBar(size(internal_params.expvals, 2), enable = false),
         ),
     )
-    return remake(prob, p=prm)
-end
 
-function _ssesolve_output_func(sol, i)
-    return (sol, false)
+    return remake(prob, p = prm, noise_rate_prototype = noise_rate_prototype, noise = noise)
 end
 
+_ssesolve_output_func(sol, i) = (sol, false)
+
 function _ssesolve_generate_statistics!(sol, i, states, expvals_all)
     sol_i = sol[:, i]
     !isempty(sol_i.prob.kwargs[:saveat]) ?
@@ -58,14 +62,13 @@ function ssesolveProblem(
     ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
     tlist::AbstractVector,
     c_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
-    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = SRA1(),
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = EM(),
     e_ops::Union{Nothing,AbstractVector} = nothing,
     H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
     params::NamedTuple = NamedTuple(),
     progress_bar::Union{Val,Bool} = Val(true),
     kwargs...,
 ) where {MT1<:AbstractMatrix,T2,Tc<:AbstractMatrix}
-
     H.dims != ψ0.dims && throw(DimensionMismatch("The two quantum objects are not of the same Hilbert dimension."))
 
     haskey(kwargs, :save_idxs) &&
@@ -82,7 +85,7 @@ function ssesolveProblem(
     H_eff = get_data(H - T2(0.5im) * mapreduce(op -> op' * op, +, c_ops))
     c_ops2 = get_data.(c_ops)
 
-    coefficients = [1.0, fill(0.0, length(c_ops)+1)...]
+    coefficients = [1.0, fill(0.0, length(c_ops) + 1)...]
     operators = [-1im * H_eff, c_ops2..., I(prod(H.dims))]
     K = OperatorSum(coefficients, operators)
     _ssesolve_update_coefficients!(ϕ0, K.coefficients, c_ops2)
@@ -102,7 +105,6 @@ function ssesolveProblem(
     end
 
     p = (
-        U = -1im * get_data(H),
         K = K,
         D = D,
         e_ops = e_ops2,
@@ -121,8 +123,18 @@ function ssesolveProblem(
     kwargs3 = _generate_sesolve_kwargs(e_ops, progress_bar_val, t_l, kwargs2)
 
     tspan = (t_l[1], t_l[end])
-    A = similar(ϕ0, length(ϕ0), length(c_ops))
-    return SDEProblem(ssesolve_drift!, ssesolve_diffusion!, ϕ0, tspan, p; noise_rate_prototype=A, kwargs3...)
+    noise = RealWienerProcess!(t_l[1], zeros(length(c_ops)))
+    noise_rate_prototype = similar(ϕ0, length(ϕ0), length(c_ops))
+    return SDEProblem{true}(
+        ssesolve_drift!,
+        ssesolve_diffusion!,
+        ϕ0,
+        tspan,
+        p;
+        noise_rate_prototype = noise_rate_prototype,
+        noise = noise,
+        kwargs3...,
+    )
 end
 
 function ssesolveEnsembleProblem(
@@ -130,7 +142,7 @@ function ssesolveEnsembleProblem(
     ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
     tlist::AbstractVector,
     c_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
-    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = SRA1(),
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = EM(),
     e_ops::Union{Nothing,AbstractVector} = nothing,
     H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
     params::NamedTuple = NamedTuple(),
@@ -138,17 +150,7 @@ function ssesolveEnsembleProblem(
     output_func::Function = _ssesolve_output_func,
     kwargs...,
 ) where {MT1<:AbstractMatrix,T2,Tc<:AbstractMatrix}
-    prob_sse = ssesolveProblem(
-        H,
-        ψ0,
-        tlist,
-        c_ops;
-        alg = alg,
-        e_ops = e_ops,
-        H_t = H_t,
-        params = params,
-        kwargs...,
-    )
+    prob_sse = ssesolveProblem(H, ψ0, tlist, c_ops; alg = alg, e_ops = e_ops, H_t = H_t, params = params, kwargs...)
 
     ensemble_prob = EnsembleProblem(prob_sse, prob_func = prob_func, output_func = output_func, safetycopy = false)
 
@@ -160,7 +162,7 @@ function ssesolve(
     ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
     tlist::AbstractVector,
     c_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
-    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = SRA1(),
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = EM(),
     e_ops::Union{Nothing,AbstractVector} = nothing,
     H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
     params::NamedTuple = NamedTuple(),
@@ -170,7 +172,6 @@ function ssesolve(
     output_func::Function = _ssesolve_output_func,
     kwargs...,
 ) where {MT1<:AbstractMatrix,T2,Tc<:AbstractMatrix}
-
     ens_prob = ssesolveEnsembleProblem(
         H,
         ψ0,
@@ -190,7 +191,7 @@ end
 
 function ssesolve(
     ens_prob::EnsembleProblem;
-    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = SRA1(),
+    alg::StochasticDiffEq.StochasticDiffEqAlgorithm = EM(),
     n_traj::Int = 1,
     ensemble_method = EnsembleThreads(),
 )
@@ -202,7 +203,7 @@ function ssesolve(
         isempty(_sol_1.prob.kwargs[:saveat]) ? fill(QuantumObject[], length(sol)) :
         Vector{Vector{QuantumObject}}(undef, length(sol))
 
-    foreach(i -> _ssesolve_generate_statistics!(sol, i, states, expvals_all), eachindex(sol),)
+    foreach(i -> _ssesolve_generate_statistics!(sol, i, states, expvals_all), eachindex(sol))
     expvals = dropdims(sum(expvals_all, dims = 1), dims = 1) ./ length(sol)
 
     return TimeEvolutionSSESol(

From 90ad6a171ce61e2e42200ca6fcd23d4dc5df5a76 Mon Sep 17 00:00:00 2001
From: Alberto Mercurio <alberto.mercurio96@gmail.com>
Date: Thu, 29 Aug 2024 23:46:31 +0200
Subject: [PATCH 08/13] Improve WienerProcesses

---
 src/time_evolution/ssesolve.jl | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/time_evolution/ssesolve.jl b/src/time_evolution/ssesolve.jl
index ee5524690..5e14b64d3 100644
--- a/src/time_evolution/ssesolve.jl
+++ b/src/time_evolution/ssesolve.jl
@@ -32,8 +32,7 @@ end
 function _ssesolve_prob_func(prob, i, repeat)
     internal_params = prob.p
 
-    noise_rate_prototype = similar(prob.u0, length(prob.u0), length(internal_params.c_ops))
-    noise = RealWienerProcess!(prob.tspan[1], zeros(length(internal_params.c_ops)))
+    noise = RealWienerProcess!(prob.tspan[1], zeros(length(internal_params.c_ops)), zeros(length(internal_params.c_ops)), save_everystep = false)
 
     prm = merge(
         internal_params,
@@ -43,7 +42,7 @@ function _ssesolve_prob_func(prob, i, repeat)
         ),
     )
 
-    return remake(prob, p = prm, noise_rate_prototype = noise_rate_prototype, noise = noise)
+    return remake(prob, p = prm, noise = noise)
 end
 
 _ssesolve_output_func(sol, i) = (sol, false)
@@ -123,7 +122,7 @@ function ssesolveProblem(
     kwargs3 = _generate_sesolve_kwargs(e_ops, progress_bar_val, t_l, kwargs2)
 
     tspan = (t_l[1], t_l[end])
-    noise = RealWienerProcess!(t_l[1], zeros(length(c_ops)))
+    noise = RealWienerProcess!(t_l[1], zeros(length(c_ops)), zeros(length(c_ops)), save_everystep = false)
     noise_rate_prototype = similar(ϕ0, length(ϕ0), length(c_ops))
     return SDEProblem{true}(
         ssesolve_drift!,
@@ -152,7 +151,7 @@ function ssesolveEnsembleProblem(
 ) where {MT1<:AbstractMatrix,T2,Tc<:AbstractMatrix}
     prob_sse = ssesolveProblem(H, ψ0, tlist, c_ops; alg = alg, e_ops = e_ops, H_t = H_t, params = params, kwargs...)
 
-    ensemble_prob = EnsembleProblem(prob_sse, prob_func = prob_func, output_func = output_func, safetycopy = false)
+    ensemble_prob = EnsembleProblem(prob_sse, prob_func = prob_func, output_func = output_func)
 
     return ensemble_prob
 end

From e95ea5dd6ee44045e63a3f8f9427a36d13501c3f Mon Sep 17 00:00:00 2001
From: Alberto Mercurio <alberto.mercurio96@gmail.com>
Date: Fri, 30 Aug 2024 19:27:01 +0200
Subject: [PATCH 09/13] Fix type instabilities

---
 src/QuantumToolbox.jl                |  2 +-
 src/time_evolution/ssesolve.jl       | 11 ++++++++---
 src/time_evolution/time_evolution.jl |  7 ++++---
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/src/QuantumToolbox.jl b/src/QuantumToolbox.jl
index 46c35e56a..03c17776b 100644
--- a/src/QuantumToolbox.jl
+++ b/src/QuantumToolbox.jl
@@ -25,12 +25,12 @@ end
 # other dependencies (in alphabetical order)
 import ArrayInterface: allowed_getindex, allowed_setindex!
 import DiffEqCallbacks: DiscreteCallback, PeriodicCallback, PresetTimeCallback, TerminateSteadyState
+import DiffEqNoiseProcess: RealWienerProcess, RealWienerProcess!
 import FFTW: fft, fftshift
 import Graphs: connected_components, DiGraph
 import IncompleteLU: ilu
 import LinearMaps: LinearMap
 import OrdinaryDiffEq: OrdinaryDiffEqAlgorithm
-import DiffEqNoiseProcess: RealWienerProcess, RealWienerProcess!
 import Pkg
 import Random
 import SpecialFunctions: loggamma
diff --git a/src/time_evolution/ssesolve.jl b/src/time_evolution/ssesolve.jl
index 5e14b64d3..528b154e9 100644
--- a/src/time_evolution/ssesolve.jl
+++ b/src/time_evolution/ssesolve.jl
@@ -32,7 +32,12 @@ end
 function _ssesolve_prob_func(prob, i, repeat)
     internal_params = prob.p
 
-    noise = RealWienerProcess!(prob.tspan[1], zeros(length(internal_params.c_ops)), zeros(length(internal_params.c_ops)), save_everystep = false)
+    noise = RealWienerProcess!(
+        prob.tspan[1],
+        zeros(length(internal_params.c_ops)),
+        zeros(length(internal_params.c_ops)),
+        save_everystep = false,
+    )
 
     prm = merge(
         internal_params,
@@ -85,11 +90,11 @@ function ssesolveProblem(
     c_ops2 = get_data.(c_ops)
 
     coefficients = [1.0, fill(0.0, length(c_ops) + 1)...]
-    operators = [-1im * H_eff, c_ops2..., I(prod(H.dims))]
+    operators = [-1im * H_eff, c_ops2..., MT1(I(prod(H.dims)))]
     K = OperatorSum(coefficients, operators)
     _ssesolve_update_coefficients!(ϕ0, K.coefficients, c_ops2)
 
-    D = vcat(c_ops2...)
+    D = reduce(vcat, c_ops2)
 
     progr = ProgressBar(length(t_l), enable = getVal(progress_bar_val))
 
diff --git a/src/time_evolution/time_evolution.jl b/src/time_evolution/time_evolution.jl
index 9c5ec3179..985e0afae 100644
--- a/src/time_evolution/time_evolution.jl
+++ b/src/time_evolution/time_evolution.jl
@@ -101,7 +101,8 @@ struct TimeEvolutionSSESol{
     TS<:AbstractVector,
     TE<:Matrix{ComplexF64},
     TEA<:Array{ComplexF64,3},
-    T<:Real,
+    T1<:Real,
+    T2<:Real,
 }
     n_traj::Int
     times::TT
@@ -110,8 +111,8 @@ struct TimeEvolutionSSESol{
     expect_all::TEA
     converged::Bool
     alg::StochasticDiffEq.StochasticDiffEqAlgorithm
-    abstol::T
-    reltol::T
+    abstol::T1
+    reltol::T2
 end
 
 function Base.show(io::IO, sol::TimeEvolutionSSESol)

From 34327c0cf735fe7cb4e36b516792207fbe57e36a Mon Sep 17 00:00:00 2001
From: lgravina1997 <luca.gravina@epfl.ch>
Date: Fri, 30 Aug 2024 22:24:19 +0200
Subject: [PATCH 10/13] docs

---
 src/time_evolution/ssesolve.jl | 224 ++++++++++++++++++++++++++++++---
 1 file changed, 204 insertions(+), 20 deletions(-)

diff --git a/src/time_evolution/ssesolve.jl b/src/time_evolution/ssesolve.jl
index 528b154e9..38a5cb6f8 100644
--- a/src/time_evolution/ssesolve.jl
+++ b/src/time_evolution/ssesolve.jl
@@ -1,15 +1,15 @@
 export ssesolveProblem, ssesolveEnsembleProblem, ssesolve
 
 #TODO: Check if works in GPU
-function _ssesolve_update_coefficients!(ψ, coefficients, c_ops)
+function _ssesolve_update_coefficients!(ψ, coefficients, sc_ops)
     _get_en = op -> real(dot(ψ, op, ψ)) #this is en/2: <Sn + Sn'>/2 = Re<Sn>
-    @. coefficients[2:end-1] = _get_en(c_ops) #coefficients of the OperatorSum: Σ Sn * en/2
+    @. coefficients[2:end-1] = _get_en(sc_ops) #coefficients of the OperatorSum: Σ Sn * en/2
     coefficients[end] = -sum(x -> x^2, coefficients[2:end-1]) / 2 #this last coefficient is -Σen^2/8
     return nothing
 end
 
 function ssesolve_drift!(du, u, p, t)
-    _ssesolve_update_coefficients!(u, p.K.coefficients, p.c_ops)
+    _ssesolve_update_coefficients!(u, p.K.coefficients, p.sc_ops)
 
     mul!(du, p.K, u)
 
@@ -20,7 +20,7 @@ function ssesolve_diffusion!(du, u, p, t)
     @inbounds en = @view(p.K.coefficients[2:end-1])
 
     # du:(H,W). du_reshaped:(H*W,). 
-    # H:Hilbert space dimension, W: number of c_ops
+    # H:Hilbert space dimension, W: number of sc_ops
     du_reshaped = reshape(du, :)
     mul!(du_reshaped, p.D, u) #du[:,i] = D[i] * u
 
@@ -34,8 +34,8 @@ function _ssesolve_prob_func(prob, i, repeat)
 
     noise = RealWienerProcess!(
         prob.tspan[1],
-        zeros(length(internal_params.c_ops)),
-        zeros(length(internal_params.c_ops)),
+        zeros(length(internal_params.sc_ops)),
+        zeros(length(internal_params.sc_ops)),
         save_everystep = false,
     )
 
@@ -61,11 +61,67 @@ function _ssesolve_generate_statistics!(sol, i, states, expvals_all)
     return nothing
 end
 
+@doc raw"""
+    ssesolveProblem(H::QuantumObject,
+        ψ0::QuantumObject,
+        tlist::AbstractVector;
+        sc_ops::Vector{QuantumObject{Tc, OperatorQuantumObject}}=QuantumObject{Matrix, OperatorQuantumObject}[];
+        alg::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm=Tsit5()
+        e_ops::Union{Nothing,AbstractVector} = nothing,
+        H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing,
+        params::NamedTuple=NamedTuple(),
+        progress_bar::Union{Val,Bool}=Val(true),
+        kwargs...)
+
+Generates 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)
+    ```
+
+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),
+    ```
+    ```math
+    M_n = C_n - \frac{e_n}{2},
+    ```
+    and
+    ```math
+    e_n = \langle C_n + 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`.
+
+# Arguments
+
+- `H::QuantumObject`: The Hamiltonian of the system ``\hat{H}``.
+- `ψ0::QuantumObject`: The initial state of the system ``|\psi(0)\rangle``.
+- `tlist::AbstractVector`: The time list of the evolution.
+- `sc_ops::Vector`: List of stochastic collapse operators ``\{\hat{C}_n\}_n``.
+- `alg::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm`: The algorithm used for the time evolution.
+- `e_ops::Union{Nothing,AbstractVector}`: The list of operators to be evaluated during the evolution.
+- `H_t::Union{Nothing,Function,TimeDependentOperatorSum}`: The time-dependent Hamiltonian of the system. If `nothing`, the Hamiltonian is time-independent.
+- `params::NamedTuple`: The parameters of the system.
+- `progress_bar::Union{Val,Bool}`: Whether to show the progress bar. Using non-`Val` types might lead to type instabilities.
+- `kwargs...`: The keyword arguments passed to the `ODEProblem` constructor.
+
+# Notes
+
+- The states will be saved depend on the keyword argument `saveat` in `kwargs`.
+- If `e_ops` is specified, the default value of `saveat=[tlist[end]]` (only save the final state), otherwise, `saveat=tlist` (saving the states corresponding to `tlist`). You can also specify `e_ops` and `saveat` separately.
+- For more details about `alg` and extra `kwargs`, please refer to [`DifferentialEquations.jl`](https://diffeq.sciml.ai/stable/)
+
+# Returns
+
+- `prob`: The `SDEProblem` for the Stochastic Schrödinger time evolution of the system.
+"""
 function ssesolveProblem(
     H::QuantumObject{MT1,OperatorQuantumObject},
     ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
     tlist::AbstractVector,
-    c_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
+    sc_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
     alg::StochasticDiffEq.StochasticDiffEqAlgorithm = EM(),
     e_ops::Union{Nothing,AbstractVector} = nothing,
     H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
@@ -86,15 +142,15 @@ function ssesolveProblem(
 
     ϕ0 = get_data(ψ0)
 
-    H_eff = get_data(H - T2(0.5im) * mapreduce(op -> op' * op, +, c_ops))
-    c_ops2 = get_data.(c_ops)
+    H_eff = get_data(H - T2(0.5im) * mapreduce(op -> op' * op, +, sc_ops))
+    sc_ops2 = get_data.(sc_ops)
 
-    coefficients = [1.0, fill(0.0, length(c_ops) + 1)...]
-    operators = [-1im * H_eff, c_ops2..., MT1(I(prod(H.dims)))]
+    coefficients = [1.0, fill(0.0, length(sc_ops) + 1)...]
+    operators = [-1im * H_eff, sc_ops2..., MT1(I(prod(H.dims)))]
     K = OperatorSum(coefficients, operators)
-    _ssesolve_update_coefficients!(ϕ0, K.coefficients, c_ops2)
+    _ssesolve_update_coefficients!(ϕ0, K.coefficients, sc_ops2)
 
-    D = reduce(vcat, c_ops2)
+    D = reduce(vcat, sc_ops2)
 
     progr = ProgressBar(length(t_l), enable = getVal(progress_bar_val))
 
@@ -112,7 +168,7 @@ function ssesolveProblem(
         K = K,
         D = D,
         e_ops = e_ops2,
-        c_ops = c_ops2,
+        sc_ops = sc_ops2,
         expvals = expvals,
         progr = progr,
         Hdims = H.dims,
@@ -127,8 +183,8 @@ function ssesolveProblem(
     kwargs3 = _generate_sesolve_kwargs(e_ops, progress_bar_val, t_l, kwargs2)
 
     tspan = (t_l[1], t_l[end])
-    noise = RealWienerProcess!(t_l[1], zeros(length(c_ops)), zeros(length(c_ops)), save_everystep = false)
-    noise_rate_prototype = similar(ϕ0, length(ϕ0), length(c_ops))
+    noise = RealWienerProcess!(t_l[1], zeros(length(sc_ops)), zeros(length(sc_ops)), save_everystep = false)
+    noise_rate_prototype = similar(ϕ0, length(ϕ0), length(sc_ops))
     return SDEProblem{true}(
         ssesolve_drift!,
         ssesolve_diffusion!,
@@ -141,11 +197,69 @@ function ssesolveProblem(
     )
 end
 
+@doc raw"""
+    ssesolveEnsembleProblem(H::QuantumObject,
+        ψ0::QuantumObject,
+        tlist::AbstractVector;
+        sc_ops::Vector{QuantumObject{Tc, OperatorQuantumObject}}=QuantumObject{Matrix, OperatorQuantumObject}[];
+        alg::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm=Tsit5()
+        e_ops::Union{Nothing,AbstractVector} = nothing,
+        H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing,
+        params::NamedTuple=NamedTuple(),
+        prob_func::Function=_mcsolve_prob_func,
+        output_func::Function=_mcsolve_output_func,
+        kwargs...)
+
+Generates 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)
+    ```
+
+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),
+    ```
+    ```math
+    M_n = C_n - \frac{e_n}{2},
+    ```
+    and
+    ```math
+    e_n = \langle C_n + 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`.
+
+# Arguments
+
+- `H::QuantumObject`: The Hamiltonian of the system ``\hat{H}``.
+- `ψ0::QuantumObject`: The initial state of the system ``|\psi(0)\rangle``.
+- `tlist::AbstractVector`: The time list of the evolution.
+- `sc_ops::Vector`: List of stochastic collapse operators ``\{\hat{C}_n\}_n``.
+- `alg::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm`: The algorithm used for the time evolution.
+- `e_ops::Union{Nothing,AbstractVector}`: The list of operators to be evaluated during the evolution.
+- `H_t::Union{Nothing,Function,TimeDependentOperatorSum}`: The time-dependent Hamiltonian of the system. If `nothing`, the Hamiltonian is time-independent.
+- `params::NamedTuple`: The parameters of the system.
+- `prob_func::Function`: Function to use for generating the SDEProblem.
+- `output_func::Function`: Function to use for generating the output of a single trajectory.
+- `kwargs...`: The keyword arguments passed to the `ODEProblem` constructor.
+
+# Notes
+
+- The states will be saved depend on the keyword argument `saveat` in `kwargs`.
+- If `e_ops` is specified, the default value of `saveat=[tlist[end]]` (only save the final state), otherwise, `saveat=tlist` (saving the states corresponding to `tlist`). You can also specify `e_ops` and `saveat` separately.
+- For more details about `alg` and extra `kwargs`, please refer to [`DifferentialEquations.jl`](https://diffeq.sciml.ai/stable/)
+
+# Returns
+
+- `prob::EnsembleProblem with SDEProblem`: The Ensemble SDEProblem for the Stochastic Shrödinger time evolution.
+"""
 function ssesolveEnsembleProblem(
     H::QuantumObject{MT1,OperatorQuantumObject},
     ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
     tlist::AbstractVector,
-    c_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
+    sc_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
     alg::StochasticDiffEq.StochasticDiffEqAlgorithm = EM(),
     e_ops::Union{Nothing,AbstractVector} = nothing,
     H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
@@ -154,18 +268,88 @@ function ssesolveEnsembleProblem(
     output_func::Function = _ssesolve_output_func,
     kwargs...,
 ) where {MT1<:AbstractMatrix,T2,Tc<:AbstractMatrix}
-    prob_sse = ssesolveProblem(H, ψ0, tlist, c_ops; alg = alg, e_ops = e_ops, H_t = H_t, params = params, kwargs...)
+    prob_sse = ssesolveProblem(H, ψ0, tlist, sc_ops; alg = alg, e_ops = e_ops, H_t = H_t, params = params, kwargs...)
 
     ensemble_prob = EnsembleProblem(prob_sse, prob_func = prob_func, output_func = output_func)
 
     return ensemble_prob
 end
 
+
+@doc raw"""
+    ssesolve(H::QuantumObject,
+        ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
+        tlist::AbstractVector,
+        sc_ops::Vector{QuantumObject{Tc, OperatorQuantumObject}}=QuantumObject{Matrix, OperatorQuantumObject}[];
+        alg::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm=Tsit5(),
+        e_ops::Vector{QuantumObject{Te, OperatorQuantumObject}}=QuantumObject{Matrix, OperatorQuantumObject}[],
+        H_t::Union{Nothing,Function,TimeDependentOperatorSum}=nothing,
+        params::NamedTuple=NamedTuple(),
+        n_traj::Int=1,
+        ensemble_method=EnsembleThreads(),
+        prob_func::Function=_mcsolve_prob_func,
+        output_func::Function=_mcsolve_output_func,
+        kwargs...)
+
+
+Stochastic Schrödinger equation evolution of a quantum system given the system Hamiltonian ``\hat{H}`` and a list of stochadtic collapse (jump) operators ``\{\hat{C}_n\}_n``.
+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)
+    ```
+
+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),
+    ```
+    ```math
+    M_n = C_n - \frac{e_n}{2},
+    ```
+    and
+    ```math
+    e_n = \langle C_n + 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`.
+
+
+# Arguments
+
+- `H::QuantumObject`: Hamiltonian of the system ``\hat{H}``.
+- `ψ0::QuantumObject`: Initial state of the system ``|\psi(0)\rangle``.
+- `tlist::AbstractVector`: List of times at which to save the state of the system.
+- `sc_ops::Vector`: List of stochastic collapse operators ``\{\hat{C}_n\}_n``.
+- `alg::OrdinaryDiffEq.OrdinaryDiffEqAlgorithm`: Algorithm to use for the time evolution.
+- `e_ops::Vector`: List of operators for which to calculate expectation values.
+- `H_t::Union{Nothing,Function,TimeDependentOperatorSum}`: Time-dependent part of the Hamiltonian.
+- `params::NamedTuple`: Dictionary of parameters to pass to the solver.
+- `seeds::Union{Nothing, Vector{Int}}`: List of seeds for the random number generator. Length must be equal to the number of trajectories provided.
+- `n_traj::Int`: Number of trajectories to use.
+- `ensemble_method`: Ensemble method to use.
+- `prob_func::Function`: Function to use for generating the SDEProblem.
+- `output_func::Function`: Function to use for generating the output of a single trajectory.
+- `kwargs...`: Additional keyword arguments to pass to the solver.
+
+# Notes
+
+- `ensemble_method` can be one of `EnsembleThreads()`, `EnsembleSerial()`, `EnsembleDistributed()`
+- The states will be saved depend on the keyword argument `saveat` in `kwargs`.
+- If `e_ops` is specified, the default value of `saveat=[tlist[end]]` (only save the final state), otherwise, `saveat=tlist` (saving the states corresponding to `tlist`). You can also specify `e_ops` and `saveat` separately.
+- The default tolerances in `kwargs` are given as `reltol=1e-6` and `abstol=1e-8`.
+- For more details about `alg` and extra `kwargs`, please refer to [`DifferentialEquations.jl`](https://diffeq.sciml.ai/stable/)
+
+# Returns
+
+- `sol::TimeEvolutionSSESol`: The solution of the time evolution. See also [`TimeEvolutionSSESol`](@ref)
+"""
+
 function ssesolve(
     H::QuantumObject{MT1,OperatorQuantumObject},
     ψ0::QuantumObject{<:AbstractArray{T2},KetQuantumObject},
     tlist::AbstractVector,
-    c_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
+    sc_ops::Vector{QuantumObject{Tc,OperatorQuantumObject}} = QuantumObject{MT1,OperatorQuantumObject}[];
     alg::StochasticDiffEq.StochasticDiffEqAlgorithm = EM(),
     e_ops::Union{Nothing,AbstractVector} = nothing,
     H_t::Union{Nothing,Function,TimeDependentOperatorSum} = nothing,
@@ -180,7 +364,7 @@ function ssesolve(
         H,
         ψ0,
         tlist,
-        c_ops;
+        sc_ops;
         alg = alg,
         e_ops = e_ops,
         H_t = H_t,

From 139c017b4762b7a7029b0a234a55ea37115270f7 Mon Sep 17 00:00:00 2001
From: lgravina1997 <luca.gravina@epfl.ch>
Date: Sat, 31 Aug 2024 09:58:50 +0200
Subject: [PATCH 11/13] api

---
 docs/src/api.md | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/docs/src/api.md b/docs/src/api.md
index fb66c49a8..da5de2f3f 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -179,11 +179,14 @@ sesolveProblem
 mesolveProblem
 lr_mesolveProblem
 mcsolveProblem
+ssesolveProblem
 mcsolveEnsembleProblem
+ssesolveEnsembleProblem
 sesolve
 mesolve
 lr_mesolve
 mcsolve
+ssesolve
 dfd_mesolve
 dsf_mesolve
 dsf_mcsolve

From 5730f8c31fe6149e9ad2b31689b022aa8269ad6e Mon Sep 17 00:00:00 2001
From: lgravina1997 <luca.gravina@epfl.ch>
Date: Sat, 31 Aug 2024 14:41:16 +0200
Subject: [PATCH 12/13] update project

---
 Project.toml | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/Project.toml b/Project.toml
index a4ff95900..453ad23d6 100644
--- a/Project.toml
+++ b/Project.toml
@@ -31,6 +31,7 @@ QuantumToolboxCUDAExt = "CUDA"
 ArrayInterface = "6, 7"
 CUDA = "5"
 DiffEqCallbacks = "2 - 3.1, 3.8"
+DiffEqNoiseProcess = "5"
 FFTW = "1.5"
 Graphs = "1.7"
 IncompleteLU = "0.2"
@@ -43,6 +44,7 @@ Random = "<0.0.1, 1"
 Reexport = "1"
 SparseArrays = "<0.0.1, 1"
 SpecialFunctions = "2"
+StochasticDiffEq = "6"
 Test = "<0.0.1, 1"
 julia = "1.7"
 

From ef253ab10792c8c4f0f24849eeac59de06a65edf Mon Sep 17 00:00:00 2001
From: lgravina1997 <luca.gravina@epfl.ch>
Date: Sat, 31 Aug 2024 18:23:55 +0200
Subject: [PATCH 13/13] add docs time_ev_v1

---
 docs/Project.toml                             |   1 +
 docs/make.jl                                  |   6 +-
 docs/src/bibliography.md                      |   3 +
 docs/src/refs.bib                             | 117 ++++++++++++++++++
 .../users_guide/time_evolution/callbacks.md   |   0
 docs/src/users_guide/time_evolution/intro.md  |  42 ++++++-
 .../src/users_guide/time_evolution/mcsolve.md |   0
 .../src/users_guide/time_evolution/mesolve.md |   3 +
 .../time_evolution/parallel_computing.md      |   0
 .../src/users_guide/time_evolution/sesolve.md |  22 ++++
 .../users_guide/time_evolution/ssesolve.md    |   0
 11 files changed, 191 insertions(+), 3 deletions(-)
 create mode 100644 docs/src/bibliography.md
 create mode 100644 docs/src/refs.bib
 create mode 100644 docs/src/users_guide/time_evolution/callbacks.md
 create mode 100644 docs/src/users_guide/time_evolution/mcsolve.md
 create mode 100644 docs/src/users_guide/time_evolution/mesolve.md
 create mode 100644 docs/src/users_guide/time_evolution/parallel_computing.md
 create mode 100644 docs/src/users_guide/time_evolution/sesolve.md
 create mode 100644 docs/src/users_guide/time_evolution/ssesolve.md

diff --git a/docs/Project.toml b/docs/Project.toml
index 08fe32855..5fe17437a 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -1,4 +1,5 @@
 [deps]
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244"
 QuantumToolbox = "6c2fb7c5-b903-41d2-bc5e-5a7c320b9fab"
diff --git a/docs/make.jl b/docs/make.jl
index 01d5a01d5..cc5320fcb 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -3,6 +3,7 @@
 
 using QuantumToolbox
 using Documenter
+using DocumenterCitations
 
 DocMeta.setdocmeta!(QuantumToolbox, :DocTestSetup, :(using QuantumToolbox); recursive = true)
 
@@ -31,7 +32,7 @@ const PAGES = [
         "Manipulating States and Operators" => "users_guide/states_and_operators.md",
         "Tensor Products and Partial Traces" => "users_guide/tensor.md",
         "Time Evolution and Dynamics" => [
-            "Introduction" => "users_guide/time_evolution/intro.md",
+            # "Introduction" => "users_guide/time_evolution/intro.md",
         ],
         "Solving for Steady-State Solutions" => [],
         "Symmetries" => [],
@@ -52,12 +53,15 @@ const PAGES = [
     # "Change Log" => "changelog.md",
 ]
 
+bib = CitationBibliography(joinpath(@__DIR__, "src", "refs.bib"))
+
 makedocs(;
     modules = [QuantumToolbox],
     authors = "Alberto Mercurio, Luca Gravina and Yi-Te Huang",
     repo = Remotes.GitHub("qutip", "QuantumToolbox.jl"),
     sitename = "QuantumToolbox.jl",
     pages = PAGES,
+    plugins=[bib],
     format = Documenter.HTML(;
         prettyurls = get(ENV, "CI", "false") == "true",
         canonical = "https://qutip.github.io/QuantumToolbox.jl",
diff --git a/docs/src/bibliography.md b/docs/src/bibliography.md
new file mode 100644
index 000000000..05f60f46f
--- /dev/null
+++ b/docs/src/bibliography.md
@@ -0,0 +1,3 @@
+# References
+
+@bibliography
\ No newline at end of file
diff --git a/docs/src/refs.bib b/docs/src/refs.bib
new file mode 100644
index 000000000..fe9ed2133
--- /dev/null
+++ b/docs/src/refs.bib
@@ -0,0 +1,117 @@
+@article{Sinibaldi2023,
+  title = {Unbiasing time-dependent Variational Monte Carlo by projected quantum evolution},
+  volume = {7},
+  ISSN = {2521-327X},
+  url = {http://dx.doi.org/10.22331/q-2023-10-10-1131},
+  DOI = {10.22331/q-2023-10-10-1131},
+  journal = {Quantum},
+  publisher = {Verein zur Forderung des Open Access Publizierens in den Quantenwissenschaften},
+  author = {Sinibaldi,  Alessandro and Giuliani,  Clemens and Carleo,  Giuseppe and Vicentini,  Filippo},
+  year = {2023},
+  month = oct,
+  pages = {1131}
+}
+
+@article{Donatella2023,
+  title = {Dynamics with autoregressive neural quantum states: Application to critical quench dynamics},
+  author = {Donatella, Kaelan and Denis, Zakari and Le Boit\'e, Alexandre and Ciuti, Cristiano},
+  journal = {Phys. Rev. A},
+  volume = {108},
+  issue = {2},
+  pages = {022210},
+  numpages = {11},
+  year = {2023},
+  month = {Aug},
+  publisher = {American Physical Society},
+  doi = {10.1103/PhysRevA.108.022210},
+  url = {https://link.aps.org/doi/10.1103/PhysRevA.108.022210}
+}
+
+@misc{Lange2024,
+  doi = {10.48550/ARXIV.2402.09402},
+  url = {https://arxiv.org/abs/2402.09402},
+  author = {Lange,  Hannah and Van de Walle,  Anka and Abedinnia,  Atiye and Bohrdt,  Annabelle},
+  keywords = {Disordered Systems and Neural Networks (cond-mat.dis-nn),  Quantum Gases (cond-mat.quant-gas),  Strongly Correlated Electrons (cond-mat.str-el),  Quantum Physics (quant-ph),  FOS: Physical sciences,  FOS: Physical sciences},
+  title = {From Architectures to Applications: A Review of Neural Quantum States},
+  publisher = {arXiv},
+  year = {2024},
+  copyright = {arXiv.org perpetual,  non-exclusive license}
+}
+
+@article{Yuan2019,
+  title = {Theory of variational quantum simulation},
+  volume = {3},
+  ISSN = {2521-327X},
+  url = {http://dx.doi.org/10.22331/q-2019-10-07-191},
+  DOI = {10.22331/q-2019-10-07-191},
+  journal = {Quantum},
+  publisher = {Verein zur Forderung des Open Access Publizierens in den Quantenwissenschaften},
+  author = {Yuan,  Xiao and Endo,  Suguru and Zhao,  Qi and Li,  Ying and Benjamin,  Simon C.},
+  year = {2019},
+  month = oct,
+  pages = {191}
+}
+
+@article{Gutirrez2022,
+  title = {Real time evolution with neural-network quantum states},
+  volume = {6},
+  ISSN = {2521-327X},
+  url = {http://dx.doi.org/10.22331/q-2022-01-20-627},
+  DOI = {10.22331/q-2022-01-20-627},
+  journal = {Quantum},
+  publisher = {Verein zur Forderung des Open Access Publizierens in den Quantenwissenschaften},
+  author = {Gutiérrez,  Irene López and Mendl,  Christian B.},
+  year = {2022},
+  month = jan,
+  pages = {627}
+}
+
+@misc{Poletti2024,
+  doi = {10.48550/ARXIV.2406.03381},
+  url = {https://arxiv.org/abs/2406.03381},
+  author = {Zhang,  Wenxuan and Xing,  Bo and Xu,  Xiansong and Poletti,  Dario},
+  keywords = {Quantum Physics (quant-ph),  Disordered Systems and Neural Networks (cond-mat.dis-nn),  Quantum Gases (cond-mat.quant-gas),  Statistical Mechanics (cond-mat.stat-mech),  FOS: Physical sciences,  FOS: Physical sciences},
+  title = {Paths towards time evolution with larger neural-network quantum states},
+  publisher = {arXiv},
+  year = {2024},
+  copyright = {arXiv.org perpetual,  non-exclusive license}
+}
+
+
+
+@article{Havlicek2023,
+  title = {Amplitude Ratios and Neural Network Quantum States},
+  volume = {7},
+  ISSN = {2521-327X},
+  url = {http://dx.doi.org/10.22331/q-2023-03-02-938},
+  DOI = {10.22331/q-2023-03-02-938},
+  journal = {Quantum},
+  publisher = {Verein zur Forderung des Open Access Publizierens in den Quantenwissenschaften},
+  author = {Havlicek,  Vojtech},
+  year = {2023},
+  month = mar,
+  pages = {938}
+}
+
+@article{Medvidovi2021,
+  title = {Classical variational simulation of the Quantum Approximate Optimization Algorithm},
+  volume = {7},
+  ISSN = {2056-6387},
+  url = {http://dx.doi.org/10.1038/s41534-021-00440-z},
+  DOI = {10.1038/s41534-021-00440-z},
+  number = {1},
+  journal = {npj Quantum Information},
+  publisher = {Springer Science and Business Media LLC},
+  author = {Medvidović,  Matija and Carleo,  Giuseppe},
+  year = {2021},
+  month = jun 
+}
+
+@book{stanley1999,
+  title={Enumerative Combinatorics, Volume 2},
+  author={Stanley, Richard P.},
+  volume={2},
+  year={1999},
+  publisher={Cambridge University Press},
+  series={Cambridge Studies in Advanced Mathematics},
+}
\ No newline at end of file
diff --git a/docs/src/users_guide/time_evolution/callbacks.md b/docs/src/users_guide/time_evolution/callbacks.md
new file mode 100644
index 000000000..e69de29bb
diff --git a/docs/src/users_guide/time_evolution/intro.md b/docs/src/users_guide/time_evolution/intro.md
index fe0097ff0..4f340f9e9 100644
--- a/docs/src/users_guide/time_evolution/intro.md
+++ b/docs/src/users_guide/time_evolution/intro.md
@@ -1,3 +1,41 @@
-# [Time Evolution and Quantum System Dynamics](@id doc:Time-Evolution-and-Quantum-System-Dynamics)
+# [Quantum System Dynamics](@id doc:Quantum-System-Dynamics)
 
-This page is still under construction, please visit [API](@ref doc-API) first.
\ No newline at end of file
+The time evolution of quantum systems lies at the heart of understanding and predicting the behavior of physical phenomena at the quantum level. Whether in theoretical research or practical applications such as quantum computing, quantum chemistry, and quantum optics, simulating the dynamics of quantum states is a fundamental task. 
+
+Quantum systems can broadly be classified into two types: closed systems and open systems. 
+
+- **Closed systems** are isolated from their environment and do not exchange energy or information with it. The dynamics of closed systems are unitary and reversible. The time evolution of a closed quantum system is governed by the Schrödinger equation.
+
+- **Open systems** interact with their environment and exchange energy and information with it. The dynamics of open systems are non-unitary and irreversible. While the average dynamics of an open quantum system is described by the Lindblad master equation, a more complete description is provided by quantum trajectory approach whereby system follows a deterministic evolution conditioned on the measurement outcomes of the environment. 
+
+The following table lists of the solvers QuantumToolbox.jl provides for dynamic quantum systems and indicates the type of solver that is returned by each function:
+
+| Solver | Problem | Description | Return Type |
+| --- | --- | --- | --- |
+| [`sesolve`](@ref) | [`sesolveProblem`](@ref) | Schrödinger equation solver | [`TimeEvolutionSol`](@ref) |
+| [`mesolve`](@ref) | [`mesolveProblem`](@ref) | Master equation solver | [`TimeEvolutionSol`](@ref) |
+| [`lr_mesolve`](@ref) | [`lr_mesolveProblem`](@ref) | Low-rank master equation solver | [`LRTimeEvolutionSol`](@ref) |
+| [`mcsolve`](@ref) | [`mcsolveProblem`](@ref) | Monte Carlo wave function solver | [`TimeEvolutionMCSol`](@ref) |
+| [`ssesolve`](@ref) | [`ssesolveProblem`](@ref) | Stochastic Schrödinger equation solver | [`TimeEvolutionSSESol`](@ref) |
+
+
+
+
+The following sections provide an overview of the different solvers and how to use them to simulate the time evolution of quantum systems. 
+
+## Main differences from QuTiP
+
+QuTip is a widely used Python library and offers grat flexibility for simulating the dynamics of quantum systems. QuantumToolbox.jl is inspired by QuTiP and aims to provide similar functionalities in Julia with close-to identical syntax. However, there are some key differences between the two libraries:
+
+- **Performance**: QuantumToolbox.jl is built on top of the Julia programming language, which is known for its high performance and efficiency. QuantumToolbox.jl leverages Julia's just-in-time (JIT) compilation and parallel computing capabilities to vastly outperform QuTip in terms of speed and scalability. We refer to the [Performance](@ref) section for a detailed comparison of the performance of QuantumToolbox.jl and QuTip.
+
+- **DifferentialEquations.jl**: QuantumToolbox.jl is based on DifferentialEquations.jl: a state of the art library for solving ordinary, partial, and stochastic differential equations. This library is the most efficient and flexible library for solving these problems constituting the backbone of most solvers in Julia and Python alike. The seamless integration with DifferentialEquations.jl allows QuantumToolbox.jl to take advantage of the latest advancements in numerical methods and scientific computing.
+
+- **GPU Support**: QuantumToolbox.jl provides built-in support for GPU acceleration, allowing users to take advantage of the computational power of GPUs for simulating quantum systems.
+
+- **Distributed Computing**: QuantumToolbox.jl supports distributed computing, enabling users with to distribute large scale problems across multiple processors or nodes for parallel execution. The structure of the solvers in QuantumToolbox.jl is designed around this feature, as discussed in the [Parallel Computing](@ref) section.
+
+- **Callbacks**: QuantumToolbox.jl allows users to define custom callbacks that are triggered during the simulation, providing flexibility in monitoring and controlling the simulation. We encourage users to use a prefixed callback structure via the `f_ops` argument in the solvers to compute complex functions of the evolving state while avoiding the need to store the full state vector at all desired times. We refer to the [Callbacks](@ref) section for more details.
+
+- **Additional Methods**: QuantumToolbox.jl provides additional methods and functionalities for the investigation of closed and open quantum systems. 
+A few examples are: the low-rank solver [`lr_mesolve`](@ref) for the dynamics of low-entropy open quantum systems, the Arnoldi-Lindblad method [`eigsolve_al`](@ref) for faster-than-the-clock estimation of the spectrum of time-independent and Floquet open quantum systems, the breath-first seach algorithm [`bdf`](@ref) for the block-diagonalization of aribitrary Hamiltonians/Liouvillians with unknown symmetries, and the U(1) symmetric Monte-Carlo solver [`mcsolve`](@ref).
\ No newline at end of file
diff --git a/docs/src/users_guide/time_evolution/mcsolve.md b/docs/src/users_guide/time_evolution/mcsolve.md
new file mode 100644
index 000000000..e69de29bb
diff --git a/docs/src/users_guide/time_evolution/mesolve.md b/docs/src/users_guide/time_evolution/mesolve.md
new file mode 100644
index 000000000..7ba526757
--- /dev/null
+++ b/docs/src/users_guide/time_evolution/mesolve.md
@@ -0,0 +1,3 @@
+
+
+## Low rank approximation
\ No newline at end of file
diff --git a/docs/src/users_guide/time_evolution/parallel_computing.md b/docs/src/users_guide/time_evolution/parallel_computing.md
new file mode 100644
index 000000000..e69de29bb
diff --git a/docs/src/users_guide/time_evolution/sesolve.md b/docs/src/users_guide/time_evolution/sesolve.md
new file mode 100644
index 000000000..c42ccbbd7
--- /dev/null
+++ b/docs/src/users_guide/time_evolution/sesolve.md
@@ -0,0 +1,22 @@
+# Unitary dynamics
+
+The evolution of a closed quantum system is governed by the Schrödinger equation
+```math
+i \frac{d}{dt} \left| \psi(t) \right\rangle = \hat{H} \left| \psi(t) \right\rangle,
+```
+where $\hat{H}$ is the Hamiltonian of the system, and $\left| \psi(t) \right\rangle$ is the state vector. 
+Numerically, $\left| \psi(t) \right\rangle$ is represented as a column vector of size $N\times 1$, while $\hat{H}$ as a matrix of size $N\times N$. $N$ is the dimension of the Hilbert space scaling exponentially with the number of degrees of freedom in the system (e.g., the number of qubits in a quantum circuit). 
+ The formal solution to the Schrödinger equation is given by
+```math 
+|\psi(t)\rangle = e^{-i\hat{H}t}|\psi(0)\rangle.
+```
+The Hermiticity of $\hat{H}$, makes the evolution operation unitary and the dynamics reversible. The norm of the state vector, as well as the expectation value of $\hat{H}$ are preserved under unitary evolution.
+
+Analytic solutions to the Schrödinger equation are rare and limited to very small systems making numerical methods the only practical approach for larger systems. A direct approach is to directly compute the matrix exponential of $\hat{H}$. This scales as $O(N^3)$, similar to matrix diagonalization, making it impractical for large systems. In contrast, matrix-vector multiplication scales as $O(N^2)$.
+Integration schemes for the Schrödinger equation involve solving the equation directly using matrix-vector products with adaptive time steps and error control. The Euler method is a simple example of a numerical integration scheme that approximates the time evolution over an infinitesimal time step $\Delta t$ as
+```math
+|\psi(t+\Delta t)\rangle = (1-i\hat{H}\Delta t)|\psi(t)\rangle
+```
+More sophisticated schemes like Runge-Kutta or Crank-Nicolson methods are often used in practice to achieve better accuracy and stability.
+
+QuantumToolbox.jl provides the `sesolve` function to simulate the unitary dynamics of a quantum system. The function takes a `sesolveProblem` object as input and returns a `TimeEvolutionSol` object that contains the solution to the Schrödinger equation. The following example demonstrates how to use the `sesolve` function to simulate the time evolution of a simple quantum system.
\ No newline at end of file
diff --git a/docs/src/users_guide/time_evolution/ssesolve.md b/docs/src/users_guide/time_evolution/ssesolve.md
new file mode 100644
index 000000000..e69de29bb