Add keep_runs_results option for multi-trajectory solvers to align with qutip (#512)

Co-authored-by: Alberto Mercurio <[email protected]>

Author: ytdHuang

CHANGELOG.md

@@ -10,6 +10,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Improve efficiency of `bloch_redfield_tensor` by avoiding unnecessary conversions. ([#509])
 - Support `SciMLOperators v1.4+`. ([#470])
 - Fix compatibility with `Makie v0.24+`. ([#513])
+- Add `keep_runs_results` option for multi-trajectory solvers to align with `QuTiP`. ([#512])
+  - Breaking changes for multi-trajectory solutions:
+    - the original fields `expect` and `states` now store the results depend on keyword argument `keep_runs_results` (decide whether to store all trajectories results or not).
+    - remove field `average_expect`
+    - remove field `runs_expect`
+  - New statistical analysis functions:
+    - `average_states`
+    - `average_expect`
+    - `std_expect`
 
 ## [v0.33.0]
 Release date: 2025-07-22
@@ -273,4 +282,5 @@ Release date: 2024-11-13
 [#506]: https://github.com/qutip/QuantumToolbox.jl/issues/506
 [#507]: https://github.com/qutip/QuantumToolbox.jl/issues/507
 [#509]: https://github.com/qutip/QuantumToolbox.jl/issues/509
+[#512]: https://github.com/qutip/QuantumToolbox.jl/issues/512
 [#513]: https://github.com/qutip/QuantumToolbox.jl/issues/513

Project.toml

@@ -24,6 +24,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]
@@ -65,5 +66,6 @@ SciMLOperators = "1.4"
 SparseArrays = "1"
 SpecialFunctions = "2"
 StaticArraysCore = "1"
+Statistics = "1"
 StochasticDiffEq = "6"
 julia = "1.10"

docs/src/resources/api.md

@@ -197,6 +197,9 @@ TimeEvolutionProblem
 TimeEvolutionSol
 TimeEvolutionMCSol
 TimeEvolutionStochasticSol
+average_states
+average_expect
+std_expect
 sesolveProblem
 mesolveProblem
 mcsolveProblem

docs/src/users_guide/time_evolution/solution.md

@@ -97,7 +97,47 @@ 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 storage of expectation values and states depends on the keyword argument `keep_runs_results`, which determines whether the results of all trajectories are stored in the solution.
 
-- [`TimeEvolutionMCSol`](@ref)
-- [`TimeEvolutionStochasticSol`](@ref)
+When the keyword argument `keep_runs_results` is passed as `Val(false)` to a multi-trajectory solver, the `states` and `expect` fields store only the average results (averaged over all trajectories). The results can be accessed by the following index-order:
+
+```julia
+sol.states[time_idx]
+sol.expect[e_op,time_idx]
+```
+
+For example:
+
+```@example TE-solution
+tlist = LinRange(0, 1, 11)
+c_ops = (destroy(2),)
+e_ops = (num(2),)
+
+sol_mc1 = mcsolve(H, ψ0, tlist, c_ops, e_ops=e_ops, ntraj=25, keep_runs_results=Val(false), progress_bar=Val(false))
+
+size(sol_mc1.expect)
+```
+
+If the keyword argument `keep_runs_results = Val(true)`, the results for each trajectory and each time are stored, and the index-order of the elements in fields `states` and `expect` are:
+
+
+```julia
+sol.states[trajectory,time_idx]
+sol.expect[e_op,trajectory,time_idx]
+```
+
+For example:
+
+```@example TE-solution
+sol_mc2 = mcsolve(H, ψ0, tlist, c_ops, e_ops=e_ops, ntraj=25, keep_runs_results=Val(true), progress_bar=Val(false))
+
+size(sol_mc2.expect)
+```
+
+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. |

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:

src/time_evolution/mcsolve.jl

@@ -22,7 +22,7 @@ end
 
 function _normalize_state!(u, dims, normalize_states)
     getVal(normalize_states) && normalize!(u)
-    return QuantumObject(u, type = Ket(), dims = dims)
+    return QuantumObject(u, Ket(), dims)
 end
 
 function _mcsolve_make_Heff_QobjEvo(H::QuantumObject, c_ops)
@@ -278,6 +278,7 @@ end
         progress_bar::Union{Val,Bool} = Val(true),
         prob_func::Union{Function, Nothing} = nothing,
         output_func::Union{Tuple,Nothing} = nothing,
+        keep_runs_results::Union{Val,Bool} = Val(false),
         normalize_states::Union{Val,Bool} = Val(true),
         kwargs...,
     )
@@ -332,6 +333,7 @@ If the environmental measurements register a quantum jump, the wave function und
 - `progress_bar`: Whether to show the progress bar. Using non-`Val` types might lead to type instabilities.
 - `prob_func`: Function to use for generating the ODEProblem.
 - `output_func`: a `Tuple` containing the `Function` to use for generating the output of a single trajectory, the (optional) `ProgressBar` object, and the (optional) `RemoteChannel` object.
+- `keep_runs_results`: Whether to save the results of each trajectory. Default to `Val(false)`.
 - `normalize_states`: Whether to normalize the states. Default to `Val(true)`.
 - `kwargs`: The keyword arguments for the ODEProblem.
 
@@ -363,6 +365,7 @@ function mcsolve(
     progress_bar::Union{Val,Bool} = Val(true),
     prob_func::Union{Function,Nothing} = nothing,
     output_func::Union{Tuple,Nothing} = nothing,
+    keep_runs_results::Union{Val,Bool} = Val(false),
     normalize_states::Union{Val,Bool} = Val(true),
     kwargs...,
 ) where {TJC<:LindbladJumpCallbackType}
@@ -384,14 +387,15 @@ function mcsolve(
         kwargs...,
     )
 
-    return mcsolve(ens_prob_mc, alg, ntraj, ensemblealg, normalize_states)
+    return mcsolve(ens_prob_mc, alg, ntraj, ensemblealg, makeVal(keep_runs_results), normalize_states)
 end
 
 function mcsolve(
     ens_prob_mc::TimeEvolutionProblem,
     alg::OrdinaryDiffEqAlgorithm = Tsit5(),
     ntraj::Int = 500,
     ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
+    keep_runs_results = Val(false),
     normalize_states = Val(true),
 )
     sol = _ensemble_dispatch_solve(ens_prob_mc, alg, ensemblealg, ntraj)
@@ -403,25 +407,24 @@ function mcsolve(
     _expvals_all =
         _expvals_sol_1 isa Nothing ? nothing : map(i -> _get_expvals(sol[:, i], SaveFuncMCSolve), eachindex(sol))
     expvals_all = _expvals_all isa Nothing ? nothing : stack(_expvals_all, dims = 2) # Stack on dimension 2 to align with QuTiP
-    states = map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states), eachindex(sol))
+
+    # stack to transform Vector{Vector{QuantumObject}} -> Matrix{QuantumObject}
+    states_all = stack(map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states), eachindex(sol)), dims = 1)
+
     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,
         _sol_1.t,
-        states,
-        expvals,
-        expvals, # This is average_expect
-        expvals_all,
+        _store_multitraj_states(states_all, keep_runs_results),
+        _store_multitraj_expect(expvals_all, keep_runs_results),
         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

src/time_evolution/mesolve.jl

@@ -211,7 +211,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

src/time_evolution/sesolve.jl

@@ -161,7 +161,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

src/time_evolution/smesolve.jl

@@ -293,6 +293,7 @@ end
         prob_func::Union{Function, Nothing} = nothing,
         output_func::Union{Tuple,Nothing} = nothing,
         progress_bar::Union{Val,Bool} = Val(true),
+        keep_runs_results::Union{Val,Bool} = Val(false),
         store_measurement::Union{Val,Bool} = Val(false),
         kwargs...,
     )
@@ -333,6 +334,7 @@ Above, ``\hat{C}_i`` represent the collapse operators related to pure dissipatio
 - `prob_func`: Function to use for generating the SDEProblem.
 - `output_func`: a `Tuple` containing the `Function` to use for generating the output of a single trajectory, the (optional) `ProgressBar` object, and the (optional) `RemoteChannel` object.
 - `progress_bar`: Whether to show the progress bar. Using non-`Val` types might lead to type instabilities.
+- `keep_runs_results`: Whether to save the results of each trajectory. Default to `Val(false)`.
 - `store_measurement`: Whether to store the measurement expectation values. Default is `Val(false)`.
 - `kwargs`: The keyword arguments for the ODEProblem.
 
@@ -365,6 +367,7 @@ function smesolve(
     prob_func::Union{Function,Nothing} = nothing,
     output_func::Union{Tuple,Nothing} = nothing,
     progress_bar::Union{Val,Bool} = Val(true),
+    keep_runs_results::Union{Val,Bool} = Val(false),
     store_measurement::Union{Val,Bool} = Val(false),
     kwargs...,
 ) where {StateOpType<:Union{Ket,Operator,OperatorKet}}
@@ -392,14 +395,15 @@ function smesolve(
         alg = sc_ops_isa_Qobj ? SRIW1() : SRA2()
     end
 
-    return smesolve(ensemble_prob, alg, ntraj, ensemblealg)
+    return smesolve(ensemble_prob, alg, ntraj, ensemblealg, makeVal(keep_runs_results))
 end
 
 function smesolve(
     ens_prob::TimeEvolutionProblem,
     alg::StochasticDiffEqAlgorithm = SRA2(),
     ntraj::Int = 500,
     ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
+    keep_runs_results = Val(false),
 )
     sol = _ensemble_dispatch_solve(ens_prob, alg, ensemblealg, ntraj)
 
@@ -412,24 +416,22 @@ function smesolve(
         _expvals_sol_1 isa Nothing ? nothing : map(i -> _get_expvals(sol[:, i], SaveFuncMESolve), eachindex(sol))
     expvals_all = _expvals_all isa Nothing ? nothing : stack(_expvals_all, dims = 2) # Stack on dimension 2 to align with QuTiP
 
-    states = map(i -> _smesolve_generate_state.(sol[:, i].u, Ref(dims), ens_prob.kwargs.isoperket), eachindex(sol))
+    # stack to transform Vector{Vector{QuantumObject}} -> Matrix{QuantumObject}
+    states_all = stack(
+        map(i -> _smesolve_generate_state.(sol[:, i].u, Ref(dims), ens_prob.kwargs.isoperket), eachindex(sol)),
+        dims = 1,
+    )
 
     _m_expvals =
         _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)
+    m_expvals = _m_expvals isa Nothing ? nothing : stack(_m_expvals, dims = 2) # Stack on dimension 2 to align with QuTiP
 
     return TimeEvolutionStochasticSol(
         ntraj,
         ens_prob.times,
         _sol_1.t,
-        states,
-        expvals,
-        expvals, # This is average_expect
-        expvals_all,
+        _store_multitraj_states(states_all, keep_runs_results),
+        _store_multitraj_expect(expvals_all, keep_runs_results),
         m_expvals, # Measurement expectation values
         sol.converged,
         _sol_1.alg,

src/time_evolution/ssesolve.jl

@@ -285,6 +285,7 @@ end
         prob_func::Union{Function, Nothing} = nothing,
         output_func::Union{Tuple,Nothing} = nothing,
         progress_bar::Union{Val,Bool} = Val(true),
+        keep_runs_results::Union{Val,Bool} = Val(false),
         store_measurement::Union{Val,Bool} = Val(false),
         kwargs...,
     )
@@ -328,6 +329,7 @@ Above, ``\hat{S}_n`` are the stochastic collapse operators and ``dW_n(t)`` is th
 - `prob_func`: Function to use for generating the SDEProblem.
 - `output_func`: a `Tuple` containing the `Function` to use for generating the output of a single trajectory, the (optional) `ProgressBar` object, and the (optional) `RemoteChannel` object.
 - `progress_bar`: Whether to show the progress bar. Using non-`Val` types might lead to type instabilities.
+- `keep_runs_results`: Whether to save the results of each trajectory. Default to `Val(false)`.
 - `store_measurement`: Whether to store the measurement results. Default is `Val(false)`.
 - `kwargs`: The keyword arguments for the ODEProblem.
 
@@ -360,6 +362,7 @@ function ssesolve(
     prob_func::Union{Function,Nothing} = nothing,
     output_func::Union{Tuple,Nothing} = nothing,
     progress_bar::Union{Val,Bool} = Val(true),
+    keep_runs_results::Union{Val,Bool} = Val(false),
     store_measurement::Union{Val,Bool} = Val(false),
     kwargs...,
 )
@@ -386,14 +389,15 @@ function ssesolve(
         alg = sc_ops_isa_Qobj ? SRIW1() : SRA2()
     end
 
-    return ssesolve(ens_prob, alg, ntraj, ensemblealg)
+    return ssesolve(ens_prob, alg, ntraj, ensemblealg, makeVal(keep_runs_results))
 end
 
 function ssesolve(
     ens_prob::TimeEvolutionProblem,
     alg::StochasticDiffEqAlgorithm = SRA2(),
     ntraj::Int = 500,
     ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
+    keep_runs_results = Val(false),
 )
     sol = _ensemble_dispatch_solve(ens_prob, alg, ensemblealg, ntraj)
 
@@ -406,24 +410,20 @@ function ssesolve(
     _expvals_all =
         _expvals_sol_1 isa Nothing ? nothing : map(i -> _get_expvals(sol[:, i], SaveFuncSSESolve), eachindex(sol))
     expvals_all = _expvals_all isa Nothing ? nothing : stack(_expvals_all, dims = 2) # Stack on dimension 2 to align with QuTiP
-    states = map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states), eachindex(sol))
+
+    # stack to transform Vector{Vector{QuantumObject}} -> Matrix{QuantumObject}
+    states_all = stack(map(i -> _normalize_state!.(sol[:, i].u, Ref(dims), normalize_states), eachindex(sol)), dims = 1)
 
     _m_expvals =
         _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,
         _sol_1.t,
-        states,
-        expvals,
-        expvals, # This is average_expect
-        expvals_all,
+        _store_multitraj_states(states_all, keep_runs_results),
+        _store_multitraj_expect(expvals_all, keep_runs_results),
         m_expvals, # Measurement expectation values
         sol.converged,
         _sol_1.alg,