qutip · ytdHuang · May 22, 2025 · May 22, 2025 · May 22, 2025 · May 22, 2025
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/qutip/QuantumToolbox.jl/tree/main)
 
+- Align statistical analysis methods for multi-trajectory solutions to `QuTiP`. ([#471])
+  - Breaking Changes:
+    - the origin field `expect` in multi-trajectory solutions now stores the expectation values for all trajectories at each time points.
+    - remove field `average_expect` from multi-trajectory solutions
+    - remove field `runs_expect` from multi-trajectory solutions
+    - add field `rng` (random number generator) to multi-trajectory solutions for reproducibility
+  - New statistical analysis functions:
+    - `average_states`
+    - `average_expect`
+    - `std_expect`
+
 ## [v0.31.1]
 Release date: 2025-05-16
 
@@ -228,3 +239,4 @@ Release date: 2024-11-13
 [#455]: https://github.com/qutip/QuantumToolbox.jl/issues/455
 [#456]: https://github.com/qutip/QuantumToolbox.jl/issues/456
 [#460]: https://github.com/qutip/QuantumToolbox.jl/issues/460
+[#471]: https://github.com/qutip/QuantumToolbox.jl/issues/471
diff --git a/Project.toml b/Project.toml
@@ -23,6 +23,7 @@ SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 
 [weakdeps]
@@ -63,6 +64,7 @@ SciMLOperators = "0.3, 0.4"
 SparseArrays = "1"
 SpecialFunctions = "2"
 StaticArraysCore = "1"
+Statistics = "1"
 StochasticDiffEq = "6"
 Test = "1"
 julia = "1.10"

diff --git a/docs/src/resources/api.md b/docs/src/resources/api.md
@@ -192,6 +192,9 @@ TimeEvolutionProblem
 TimeEvolutionSol
 TimeEvolutionMCSol
 TimeEvolutionStochasticSol
+average_states
+average_expect
+std_expect
 sesolveProblem
 mesolveProblem
 mcsolveProblem

diff --git a/docs/src/users_guide/steadystate.md b/docs/src/users_guide/steadystate.md
@@ -94,7 +94,7 @@ tlist = LinRange(0, 50, 100)
 
 # monte-carlo
 sol_mc = mcsolve(H, ψ0, tlist, c_ops, e_ops=e_ops, ntraj=100, progress_bar=false)
-exp_mc = real(sol_mc.expect[1, :])
+exp_mc = real(average_expect(sol_mc)[1, :])
 
 # master eq.
 sol_me = mesolve(H, ψ0, tlist, c_ops, e_ops=e_ops, progress_bar=false)

diff --git a/docs/src/users_guide/time_evolution/mcsolve.md b/docs/src/users_guide/time_evolution/mcsolve.md
@@ -70,6 +70,7 @@ c_ops = [sqrt(0.1) * a]
 e_ops = [a' * a, σm' * σm]
 
 sol_500 = mcsolve(H, ψ0, times, c_ops, e_ops=e_ops)
+expect_avg = average_expect(sol_500)
 
 # plot by CairoMakie.jl
 fig = Figure(size = (500, 350))
@@ -78,8 +79,8 @@ ax = Axis(fig[1, 1],
     ylabel = "Expectation values",
     title = "Monte Carlo time evolution (500 trajectories)",
 )
-lines!(ax, times, real(sol_500.expect[1,:]), label = "cavity photon number", linestyle = :solid)
-lines!(ax, times, real(sol_500.expect[2,:]), label = "atom excitation probability", linestyle = :dash)
+lines!(ax, times, real(expect_avg[1,:]), label = "cavity photon number", linestyle = :solid)
+lines!(ax, times, real(expect_avg[2,:]), label = "atom excitation probability", linestyle = :dash)
 
 axislegend(ax, position = :rt)
 
@@ -104,9 +105,9 @@ ax = Axis(fig[1, 1],
     ylabel = "Expectation values",
     title = "Monte Carlo time evolution",
 )
-lines!(ax, times, real(sol_1.expect[1,:]), label = "1 trajectory", linestyle = :dashdot)
-lines!(ax, times, real(sol_10.expect[1,:]), label = "10 trajectories", linestyle = :dash)
-lines!(ax, times, real(sol_100.expect[1,:]), label = "100 trajectories", linestyle = :solid)
+lines!(ax, times, real(average_expect(sol_1)[1,:]), label = "1 trajectory", linestyle = :dashdot)
+lines!(ax, times, real(average_expect(sol_10)[1,:]), label = "10 trajectories", linestyle = :dash)
+lines!(ax, times, real(average_expect(sol_100)[1,:]), label = "100 trajectories", linestyle = :solid)
 
 axislegend(ax, position = :rt)
 

diff --git a/docs/src/users_guide/time_evolution/solution.md b/docs/src/users_guide/time_evolution/solution.md
@@ -89,7 +89,23 @@ Some other solvers can have other output.
 
 ## [Multiple trajectories solution](@id doc-TE:Multiple-trajectories-solution)
 
-This part is still under construction, please read the docstrings for the following functions first:
+The solutions are different for solvers which compute multiple trajectories, such as the [`TimeEvolutionMCSol`](@ref) (Monte Carlo) or the [`TimeEvolutionStochasticSol`](@ref) (stochastic methods). The expectation values and states for all trajectories at each time points will be stored in these solutions. The index-order of the elements in fields `states` and `expect` are:
 
-- [`TimeEvolutionMCSol`](@ref)
-- [`TimeEvolutionStochasticSol`](@ref)
+```julia
+sol.states[trajectory][time]
+sol.expect[e_op,trajectory,time]
+```
+
+We also provide the following functions for statistical analysis of multi-trajectory `sol`utions:
+
+| **Functions** | **Description** |
+|:------------|:----------------|
+| [`average_states(sol)`](@ref average_states) | Return the trajectory-averaged result states (as density [`Operator`](@ref)) at each time point. |
+| [`average_expect(sol)`](@ref average_expect) | Return the trajectory-averaged expectation values at each time point. |
+| [`std_expect(sol)`](@ref std_expect) | Return the trajectory-wise standard deviation of the expectation values at each time point. |
+
+Multi-trajectory solutions also keep the initial random number generator (`rng`) to allow recomputing the results:
+
+```julia
+rng = sol.rng # this can be specified as a keyword argument (`rng = rng`) to the solvers
+```
diff --git a/docs/src/users_guide/time_evolution/stochastic.md b/docs/src/users_guide/time_evolution/stochastic.md
@@ -111,11 +111,13 @@ sse_sol = ssesolve(
 measurement_avg = sum(sse_sol.measurement, dims=2) / size(sse_sol.measurement, 2)
 measurement_avg = dropdims(measurement_avg, dims=2)
 
+expect_avg = average_expect(sse_sol)
+
 # plot by CairoMakie.jl
 fig = Figure(size = (500, 350))
 ax = Axis(fig[1, 1], xlabel = "Time")
 lines!(ax, tlist[1:end-1], real(measurement_avg[1,:]), label = L"J_x", color = :red, linestyle = :solid)
-lines!(ax, tlist, real(sse_sol.expect[1,:]),  label = L"\langle x \rangle", color = :black, linestyle = :solid)
+lines!(ax, tlist, real(expect_avg[1,:]),  label = L"\langle x \rangle", color = :black, linestyle = :solid)
 
 axislegend(ax, position = :rt)
 
@@ -141,14 +143,16 @@ sme_sol = smesolve(
     store_measurement = Val(true),
 )
 
+expect_avg = average_expect(sme_sol)
+
 measurement_avg = sum(sme_sol.measurement, dims=2) / size(sme_sol.measurement, 2)
 measurement_avg = dropdims(measurement_avg, dims=2)
 
 # plot by CairoMakie.jl
 fig = Figure(size = (500, 350))
 ax = Axis(fig[1, 1], xlabel = "Time")
 lines!(ax, tlist[1:end-1], real(measurement_avg[1,:]), label = L"J_x", color = :red, linestyle = :solid)
-lines!(ax, tlist, real(sme_sol.expect[1,:]),  label = L"\langle x \rangle", color = :black, linestyle = :solid)
+lines!(ax, tlist, real(expect_avg[1,:]),  label = L"\langle x \rangle", color = :black, linestyle = :solid)
 
 axislegend(ax, position = :rt)
 

diff --git a/docs/src/users_guide/time_evolution/time_dependent.md b/docs/src/users_guide/time_evolution/time_dependent.md
@@ -132,7 +132,7 @@ e_ops = [
 
 # solve dynamics 
 exp_me = mesolve(H_t, ψ0, tlist, c_ops; e_ops = e_ops, progress_bar = Val(false)).expect
-exp_mc = mcsolve(H_t, ψ0, tlist, c_ops; e_ops = e_ops, ntraj = 100, progress_bar = Val(false)).expect
+exp_mc = average_expect(mcsolve(H_t, ψ0, tlist, c_ops; e_ops = e_ops, ntraj = 100, progress_bar = Val(false)))
 
 # plot by CairoMakie.jl
 fig = Figure(size = (500, 350))

diff --git a/src/QuantumToolbox.jl b/src/QuantumToolbox.jl
@@ -5,6 +5,7 @@ using LinearAlgebra
 import LinearAlgebra: BlasInt, BlasFloat, checksquare
 import LinearAlgebra.LAPACK: hseqr!
 using SparseArrays
+import Statistics: mean, std
 
 # SciML packages (for QobjEvo, OrdinaryDiffEq, and LinearSolve)
 import SciMLBase:

diff --git a/src/qobj/functions.jl b/src/qobj/functions.jl
@@ -16,6 +16,9 @@ ket2dm(ψ::QuantumObject{Ket}) = ψ * ψ'
 
 ket2dm(ρ::QuantumObject{Operator}) = ρ
 
+# for special usage, e.g., average_states
+ket2dm(states::Vector{<:QuantumObject{T}}) where {T<:Union{Ket,Operator}} = map(s -> ket2dm(s), states)
+
 @doc raw"""
     expect(O::Union{AbstractQuantumObject,Vector{AbstractQuantumObject}}, ψ::Union{QuantumObject,Vector{QuantumObject}})
 

diff --git a/src/time_evolution/mcsolve.jl b/src/time_evolution/mcsolve.jl
@@ -235,7 +235,10 @@ function mcsolveEnsembleProblem(
     output_func::Union{Tuple,Nothing} = nothing,
     kwargs...,
 ) where {TJC<:LindbladJumpCallbackType}
-    _prob_func = isnothing(prob_func) ? _ensemble_dispatch_prob_func(rng, ntraj, tlist, _mcsolve_prob_func) : prob_func
+    copied_rng = copy(rng) # use the copied rng, and keep the initial one to pass it directly to solution later
+
+    _prob_func =
+        isnothing(prob_func) ? _ensemble_dispatch_prob_func(copied_rng, ntraj, tlist, _mcsolve_prob_func) : prob_func
     _output_func =
         output_func isa Nothing ?
         _ensemble_dispatch_output_func(ensemblealg, progress_bar, ntraj, _mcsolve_output_func) : output_func
@@ -247,7 +250,7 @@ function mcsolveEnsembleProblem(
         c_ops;
         e_ops = e_ops,
         params = params,
-        rng = rng,
+        rng = copied_rng,
         jump_callback = jump_callback,
         kwargs...,
     )
@@ -256,7 +259,7 @@ function mcsolveEnsembleProblem(
         EnsembleProblem(prob_mc.prob, prob_func = _prob_func, output_func = _output_func[1], safetycopy = false),
         prob_mc.times,
         prob_mc.dimensions,
-        (progr = _output_func[2], channel = _output_func[3]),
+        (progr = _output_func[2], channel = _output_func[3], rng = rng),
     )
 
     return ensemble_prob
@@ -407,20 +410,17 @@ function mcsolve(
     col_times = map(i -> _mc_get_jump_callback(sol[:, i]).affect!.col_times, eachindex(sol))
     col_which = map(i -> _mc_get_jump_callback(sol[:, i]).affect!.col_which, eachindex(sol))
 
-    expvals = _expvals_sol_1 isa Nothing ? nothing : dropdims(sum(expvals_all, dims = 2), dims = 2) ./ length(sol)
-
     return TimeEvolutionMCSol(
         ntraj,
         ens_prob_mc.times,
         states,
-        expvals,
-        expvals, # This is average_expect
         expvals_all,
+        ens_prob_mc.kwargs.rng,
         col_times,
         col_which,
         sol.converged,
         _sol_1.alg,
-        NamedTuple(_sol_1.prob.kwargs).abstol,
-        NamedTuple(_sol_1.prob.kwargs).reltol,
+        _sol_1.prob.kwargs[:abstol],
+        _sol_1.prob.kwargs[:reltol],
     )
 end
diff --git a/src/time_evolution/mesolve.jl b/src/time_evolution/mesolve.jl
@@ -215,7 +215,7 @@ function mesolve(prob::TimeEvolutionProblem, alg::OrdinaryDiffEqAlgorithm = Tsit
         _get_expvals(sol, SaveFuncMESolve),
         sol.retcode,
         sol.alg,
-        NamedTuple(sol.prob.kwargs).abstol,
-        NamedTuple(sol.prob.kwargs).reltol,
+        sol.prob.kwargs[:abstol],
+        sol.prob.kwargs[:reltol],
     )
 end
diff --git a/src/time_evolution/sesolve.jl b/src/time_evolution/sesolve.jl
@@ -165,7 +165,7 @@ function sesolve(prob::TimeEvolutionProblem, alg::OrdinaryDiffEqAlgorithm = Tsit
         _get_expvals(sol, SaveFuncSESolve),
         sol.retcode,
         sol.alg,
-        NamedTuple(sol.prob.kwargs).abstol,
-        NamedTuple(sol.prob.kwargs).reltol,
+        sol.prob.kwargs[:abstol],
+        sol.prob.kwargs[:reltol],
     )
 end
diff --git a/src/time_evolution/smesolve.jl b/src/time_evolution/smesolve.jl
@@ -239,10 +239,12 @@ function smesolveEnsembleProblem(
     store_measurement::Union{Val,Bool} = Val(false),
     kwargs...,
 ) where {StateOpType<:Union{Ket,Operator,OperatorKet}}
+    copied_rng = copy(rng) # use the copied rng, and keep the initial one to pass it directly to solution later
+
     _prob_func =
         isnothing(prob_func) ?
         _ensemble_dispatch_prob_func(
-            rng,
+            copied_rng,
             ntraj,
             tlist,
             _stochastic_prob_func;
@@ -261,7 +263,7 @@ function smesolveEnsembleProblem(
         sc_ops;
         e_ops = e_ops,
         params = params,
-        rng = rng,
+        rng = copied_rng,
         progress_bar = Val(false),
         store_measurement = makeVal(store_measurement),
         kwargs...,
@@ -271,7 +273,7 @@ function smesolveEnsembleProblem(
         EnsembleProblem(prob_sme, prob_func = _prob_func, output_func = _output_func[1], safetycopy = true),
         prob_sme.times,
         prob_sme.dimensions,
-        merge(prob_sme.kwargs, (progr = _output_func[2], channel = _output_func[3])),
+        merge(prob_sme.kwargs, (progr = _output_func[2], channel = _output_func[3], rng = rng)),
     )
 
     return ensemble_prob
@@ -418,17 +420,12 @@ function smesolve(
         _m_expvals_sol_1 isa Nothing ? nothing : map(i -> _get_m_expvals(sol[:, i], SaveFuncSMESolve), eachindex(sol))
     m_expvals = _m_expvals isa Nothing ? nothing : stack(_m_expvals, dims = 2)
 
-    expvals =
-        _get_expvals(_sol_1, SaveFuncMESolve) isa Nothing ? nothing :
-        dropdims(sum(expvals_all, dims = 2), dims = 2) ./ length(sol)
-
     return TimeEvolutionStochasticSol(
         ntraj,
         ens_prob.times,
         states,
-        expvals,
-        expvals, # This is average_expect
         expvals_all,
+        ens_prob.kwargs.rng,
         m_expvals, # Measurement expectation values
         sol.converged,
         _sol_1.alg,

diff --git a/src/time_evolution/ssesolve.jl b/src/time_evolution/ssesolve.jl
@@ -233,10 +233,12 @@ function ssesolveEnsembleProblem(
     store_measurement::Union{Val,Bool} = Val(false),
     kwargs...,
 )
+    copied_rng = copy(rng) # use the copied rng, and keep the initial one to pass it directly to solution later
+
     _prob_func =
         isnothing(prob_func) ?
         _ensemble_dispatch_prob_func(
-            rng,
+            copied_rng,
             ntraj,
             tlist,
             _stochastic_prob_func;
@@ -254,7 +256,7 @@ function ssesolveEnsembleProblem(
         sc_ops;
         e_ops = e_ops,
         params = params,
-        rng = rng,
+        rng = copied_rng,
         progress_bar = Val(false),
         store_measurement = makeVal(store_measurement),
         kwargs...,
@@ -264,7 +266,7 @@ function ssesolveEnsembleProblem(
         EnsembleProblem(prob_sme, prob_func = _prob_func, output_func = _output_func[1], safetycopy = true),
         prob_sme.times,
         prob_sme.dimensions,
-        (progr = _output_func[2], channel = _output_func[3]),
+        (progr = _output_func[2], channel = _output_func[3], rng = rng),
     )
 
     return ensemble_prob
@@ -412,17 +414,12 @@ function ssesolve(
         _m_expvals_sol_1 isa Nothing ? nothing : map(i -> _get_m_expvals(sol[:, i], SaveFuncSSESolve), eachindex(sol))
     m_expvals = _m_expvals isa Nothing ? nothing : stack(_m_expvals, dims = 2)
 
-    expvals =
-        _get_expvals(_sol_1, SaveFuncSSESolve) isa Nothing ? nothing :
-        dropdims(sum(expvals_all, dims = 2), dims = 2) ./ length(sol)
-
     return TimeEvolutionStochasticSol(
         ntraj,
         ens_prob.times,
         states,
-        expvals,
-        expvals, # This is average_expect
         expvals_all,
+        ens_prob.kwargs.rng,
         m_expvals, # Measurement expectation values
         sol.converged,
         _sol_1.alg,