align statistical analysis methods for multi-trajectory with qutip

Author: ytdHuang

docs/src/resources/api.md

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

src/time_evolution/mcsolve.jl

@@ -407,14 +407,10 @@ 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,
         col_times,
         col_which,

src/time_evolution/smesolve.jl

@@ -418,16 +418,10 @@ 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,
         m_expvals, # Measurement expectation values
         sol.converged,

src/time_evolution/ssesolve.jl

@@ -412,16 +412,10 @@ 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,
         m_expvals, # Measurement expectation values
         sol.converged,

src/time_evolution/time_evolution.jl

@@ -1,11 +1,15 @@
-export TimeEvolutionSol, TimeEvolutionMCSol, TimeEvolutionStochasticSol
+export TimeEvolutionSol
+export TimeEvolutionMultiTrajSol, TimeEvolutionMCSol, TimeEvolutionStochasticSol
+export average_states, average_expect, std_expect
 
 export 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-3, reltol = 2e-3, save_everystep = false, save_end = true)
 const COL_TIMES_WHICH_INIT_SIZE = 200
 
+abstract type TimeEvolutionMultiTrajSol{Texpect} end
+
 @doc raw"""
     struct TimeEvolutionProblem
 
@@ -89,42 +93,45 @@ function Base.show(io::IO, sol::TimeEvolutionSol)
 end
 
 @doc raw"""
-    struct TimeEvolutionMCSol
+    struct TimeEvolutionMCSol <: TimeEvolutionMultiTrajSol
 
 A structure storing the results and some information from solving quantum trajectories of the Monte Carlo wave function time evolution.
 
 # Fields (Attributes)
 
 - `ntraj::Int`: Number of trajectories
 - `times::AbstractVector`: The time list of the evolution.
-- `states::Vector{Vector{QuantumObject}}`: The list of result states in each trajectory.
-- `expect::Union{AbstractMatrix,Nothing}`: The expectation values (averaging all trajectories) corresponding to each time point in `times`.
-- `average_expect::Union{AbstractMatrix,Nothing}`: The expectation values (averaging all trajectories) corresponding to each time point in `times`.
-- `runs_expect::Union{AbstractArray,Nothing}`: The expectation values corresponding to each trajectory and each time point in `times`
+- `states::Vector{Vector{QuantumObject}}`: The list of result states in each trajectory and each time point in `times`.
+- `expect::Union{AbstractArray,Nothing}`: The expectation values corresponding to each trajectory and each time point in `times`.
 - `col_times::Vector{Vector{Real}}`: The time records of every quantum jump occurred in each trajectory.
 - `col_which::Vector{Vector{Int}}`: The indices of which collapse operator was responsible for each quantum jump in `col_times`.
 - `converged::Bool`: Whether the solution is converged or not.
 - `alg`: The algorithm which is used during the solving process.
 - `abstol::Real`: The absolute tolerance which is used during the solving process.
 - `reltol::Real`: The relative tolerance which is used during the solving process.
+
+# Methods
+
+We also provide the following functions for analyzing statistics from multi-trajectory solutions.
+
+- [`average_states`](@ref)
+- [`average_expect`](@ref)
+- [`std_expect`](@ref)
 """
 struct TimeEvolutionMCSol{
     TT<:AbstractVector{<:Real},
     TS<:AbstractVector,
-    TE<:Union{AbstractMatrix,Nothing},
-    TEA<:Union{AbstractArray,Nothing},
+    TE<:Union{AbstractArray,Nothing},
     TJT<:Vector{<:Vector{<:Real}},
     TJW<:Vector{<:Vector{<:Integer}},
     AlgT<:OrdinaryDiffEqAlgorithm,
     AT<:Real,
     RT<:Real,
-}
+} <: TimeEvolutionMultiTrajSol{TE}
     ntraj::Int
     times::TT
     states::TS
     expect::TE
-    average_expect::TE # Currently just a synonym for `expect`
-    runs_expect::TEA
     col_times::TJT
     col_which::TJW
     converged::Bool
@@ -142,7 +149,7 @@ function Base.show(io::IO, sol::TimeEvolutionMCSol)
     if sol.expect isa Nothing
         print(io, "num_expect = 0\n")
     else
-        print(io, "num_expect = $(size(sol.average_expect, 1))\n")
+        print(io, "num_expect = $(size(sol.expect, 1))\n")
     end
     print(io, "ODE alg.: $(sol.alg)\n")
     print(io, "abstol = $(sol.abstol)\n")
@@ -151,39 +158,42 @@ function Base.show(io::IO, sol::TimeEvolutionMCSol)
 end
 
 @doc raw"""
-    struct TimeEvolutionStochasticSol
+    struct TimeEvolutionStochasticSol <: TimeEvolutionMultiTrajSol
 
 A structure storing the results and some information from solving trajectories of the Stochastic time evolution.
 
 # Fields (Attributes)
 
 - `ntraj::Int`: Number of trajectories
 - `times::AbstractVector`: The time list of the evolution.
-- `states::Vector{Vector{QuantumObject}}`: The list of result states in each trajectory.
-- `expect::Union{AbstractMatrix,Nothing}`: The expectation values (averaging all trajectories) corresponding to each time point in `times`.
-- `average_expect::Union{AbstractMatrix,Nothing}`: The expectation values (averaging all trajectories) corresponding to each time point in `times`.
-- `runs_expect::Union{AbstractArray,Nothing}`: The expectation values corresponding to each trajectory and each time point in `times`
+- `states::Vector{Vector{QuantumObject}}`: The list of result states in each trajectory and each time point in `times`.
+- `expect::Union{AbstractArray,Nothing}`: The expectation values corresponding to each trajectory and each time point in `times`.
 - `converged::Bool`: Whether the solution is converged or not.
 - `alg`: The algorithm which is used during the solving process.
 - `abstol::Real`: The absolute tolerance which is used during the solving process.
 - `reltol::Real`: The relative tolerance which is used during the solving process.
+
+# Methods
+
+We also provide the following functions for analyzing statistics from multi-trajectory solutions.
+
+- [`average_states`](@ref)
+- [`average_expect`](@ref)
+- [`std_expect`](@ref)
 """
 struct TimeEvolutionStochasticSol{
     TT<:AbstractVector{<:Real},
     TS<:AbstractVector,
-    TE<:Union{AbstractMatrix,Nothing},
-    TEA<:Union{AbstractArray,Nothing},
+    TE<:Union{AbstractArray,Nothing},
     TEM<:Union{AbstractArray,Nothing},
     AlgT<:StochasticDiffEqAlgorithm,
     AT<:Real,
     RT<:Real,
-}
+} <: TimeEvolutionMultiTrajSol{TE}
     ntraj::Int
     times::TT
     states::TS
     expect::TE
-    average_expect::TE # Currently just a synonym for `expect`
-    runs_expect::TEA
     measurement::TEM
     converged::Bool
     alg::AlgT
@@ -200,14 +210,37 @@ function Base.show(io::IO, sol::TimeEvolutionStochasticSol)
     if sol.expect isa Nothing
         print(io, "num_expect = 0\n")
     else
-        print(io, "num_expect = $(size(sol.average_expect, 1))\n")
+        print(io, "num_expect = $(size(sol.expect, 1))\n")
     end
     print(io, "SDE alg.: $(sol.alg)\n")
     print(io, "abstol = $(sol.abstol)\n")
     print(io, "reltol = $(sol.reltol)\n")
     return nothing
 end
 
+@doc raw"""
+    average_states(sol::TimeEvolutionMultiTrajSol)
+
+Return the trajectory-averaged result states at each time point.
+"""
+average_states(sol::TimeEvolutionMultiTrajSol) = mean(sol.states)
+
+@doc raw"""
+    average_expect(sol::TimeEvolutionMultiTrajSol)
+
+Return the trajectory-averaged expectation values at each time point.
+"""
+average_expect(sol::TimeEvolutionMultiTrajSol{TE}) where {TE<:AbstractArray} = dropdims(mean(sol.runs_expect, dims = 2), dims = 2)
+average_expect(sol::TimeEvolutionMultiTrajSol{Nothing}) = nothing
+
+@doc raw"""
+    std_expect(sol::TimeEvolutionMultiTrajSol)
+
+Return the trajectory-wise standard deviation of the expectation values at each time point.
+"""
+std_expect(sol::TimeEvolutionMultiTrajSol{TE}) where {TE<:AbstractArray} = dropdims(std(sol.runs_expect, dims = 2), dims = 2)
+std_expect(sol::TimeEvolutionMultiTrajSol{Nothing}) = nothing
+
 #######################################
 #=
     Callbacks for Monte Carlo quantum trajectories