Align some attributes of mcsolve, ssesolve and smesolve results with QuTiP (#402)

Author: albertomercurio

CHANGELOG.md

@@ -34,6 +34,7 @@ Release date: 2025-01-29
 - Fix Dynamical Fock Dimension states saving due to wrong saving of dimensions. ([#375])
 - Support a list of observables for `expect`. ([#374], [#376])
 - Add checks for `tlist` in time evolution solvers. The checks are to ensure that `tlist` is not empty, the elements are in increasing order, and the elements are unique. ([#378])
+- Change the definition of jump_times and jump_which into col_times and col_which, respectively. ([#402])
 
 ## [v0.25.0]
 Release date: 2025-01-20
@@ -129,4 +130,5 @@ Release date: 2024-11-13
 [#395]: https://github.com/qutip/QuantumToolbox.jl/issues/395
 [#396]: https://github.com/qutip/QuantumToolbox.jl/issues/396
 [#398]: https://github.com/qutip/QuantumToolbox.jl/issues/398
+[#402]: https://github.com/qutip/QuantumToolbox.jl/issues/402
 [#403]: https://github.com/qutip/QuantumToolbox.jl/issues/403

src/time_evolution/callback_helpers/mcsolve_callback_helpers.jl

@@ -28,9 +28,9 @@ struct LindbladJump{
     cache_mc::CT
     weights_mc::WT
     cumsum_weights_mc::WT
-    jump_times::JTT
-    jump_which::JWT
-    jump_times_which_idx::JTWIT
+    col_times::JTT
+    col_which::JWT
+    col_times_which_idx::JTWIT
 end
 
 (f::LindbladJump)(integrator) = _lindblad_jump_affect!(
@@ -42,9 +42,9 @@ end
     f.cache_mc,
     f.weights_mc,
     f.cumsum_weights_mc,
-    f.jump_times,
-    f.jump_which,
-    f.jump_times_which_idx,
+    f.col_times,
+    f.col_which,
+    f.col_times_which_idx,
 )
 
 ##
@@ -71,9 +71,9 @@ function _generate_mcsolve_kwargs(ψ0, T, e_ops, tlist, c_ops, jump_callback, rn
     weights_mc = Vector{Float64}(undef, length(c_ops))
     cumsum_weights_mc = similar(weights_mc)
 
-    jump_times = Vector{Float64}(undef, JUMP_TIMES_WHICH_INIT_SIZE)
-    jump_which = Vector{Int}(undef, JUMP_TIMES_WHICH_INIT_SIZE)
-    jump_times_which_idx = Ref(1)
+    col_times = Vector{Float64}(undef, COL_TIMES_WHICH_INIT_SIZE)
+    col_which = Vector{Int}(undef, COL_TIMES_WHICH_INIT_SIZE)
+    col_times_which_idx = Ref(1)
 
     random_n = Ref(rand(rng))
 
@@ -85,9 +85,9 @@ function _generate_mcsolve_kwargs(ψ0, T, e_ops, tlist, c_ops, jump_callback, rn
         cache_mc,
         weights_mc,
         cumsum_weights_mc,
-        jump_times,
-        jump_which,
-        jump_times_which_idx,
+        col_times,
+        col_which,
+        col_times_which_idx,
     )
 
     if jump_callback isa DiscreteLindbladJumpCallback
@@ -129,9 +129,9 @@ function _lindblad_jump_affect!(
     cache_mc,
     weights_mc,
     cumsum_weights_mc,
-    jump_times,
-    jump_which,
-    jump_times_which_idx,
+    col_times,
+    col_which,
+    col_times_which_idx,
 )
     ψ = integrator.u
 
@@ -147,13 +147,13 @@ function _lindblad_jump_affect!(
 
     random_n[] = rand(traj_rng)
 
-    idx = jump_times_which_idx[]
-    @inbounds jump_times[idx] = integrator.t
-    @inbounds jump_which[idx] = collapse_idx
-    jump_times_which_idx[] += 1
-    if jump_times_which_idx[] > length(jump_times)
-        resize!(jump_times, length(jump_times) + JUMP_TIMES_WHICH_INIT_SIZE)
-        resize!(jump_which, length(jump_which) + JUMP_TIMES_WHICH_INIT_SIZE)
+    idx = col_times_which_idx[]
+    @inbounds col_times[idx] = integrator.t
+    @inbounds col_which[idx] = collapse_idx
+    col_times_which_idx[] += 1
+    if col_times_which_idx[] > length(col_times)
+        resize!(col_times, length(col_times) + COL_TIMES_WHICH_INIT_SIZE)
+        resize!(col_which, length(col_which) + COL_TIMES_WHICH_INIT_SIZE)
     end
     u_modified!(integrator, true)
     return nothing
@@ -309,9 +309,9 @@ function _similar_affect!(affect::LindbladJump, traj_rng)
     cache_mc = similar(affect.cache_mc)
     weights_mc = similar(affect.weights_mc)
     cumsum_weights_mc = similar(affect.cumsum_weights_mc)
-    jump_times = similar(affect.jump_times)
-    jump_which = similar(affect.jump_which)
-    jump_times_which_idx = Ref(1)
+    col_times = similar(affect.col_times)
+    col_which = similar(affect.col_which)
+    col_times_which_idx = Ref(1)
 
     return LindbladJump(
         affect.c_ops,
@@ -321,9 +321,9 @@ function _similar_affect!(affect::LindbladJump, traj_rng)
         cache_mc,
         weights_mc,
         cumsum_weights_mc,
-        jump_times,
-        jump_which,
-        jump_times_which_idx,
+        col_times,
+        col_which,
+        col_times_which_idx,
     )
 end
 

src/time_evolution/mcsolve.jl

@@ -14,9 +14,9 @@ end
 
 # Standard output function
 function _mcsolve_output_func(sol, i)
-    idx = _mc_get_jump_callback(sol).affect!.jump_times_which_idx[]
-    resize!(_mc_get_jump_callback(sol).affect!.jump_times, idx - 1)
-    resize!(_mc_get_jump_callback(sol).affect!.jump_which, idx - 1)
+    idx = _mc_get_jump_callback(sol).affect!.col_times_which_idx[]
+    resize!(_mc_get_jump_callback(sol).affect!.col_times, idx - 1)
+    resize!(_mc_get_jump_callback(sol).affect!.col_which, idx - 1)
     return (sol, false)
 end
 
@@ -401,21 +401,22 @@ function mcsolve(
     _expvals_sol_1 = _mcsolve_get_expvals(_sol_1)
 
     _expvals_all = _expvals_sol_1 isa Nothing ? nothing : map(i -> _mcsolve_get_expvals(sol[:, i]), eachindex(sol))
-    expvals_all = _expvals_all isa Nothing ? nothing : stack(_expvals_all)
+    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))
-    jump_times = map(i -> _mc_get_jump_callback(sol[:, i]).affect!.jump_times, eachindex(sol))
-    jump_which = map(i -> _mc_get_jump_callback(sol[:, i]).affect!.jump_which, eachindex(sol))
+    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 = 3), dims = 3) ./ length(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,
-        jump_times,
-        jump_which,
+        col_times,
+        col_which,
         sol.converged,
         _sol_1.alg,
         NamedTuple(_sol_1.prob.kwargs).abstol,

src/time_evolution/smesolve.jl

@@ -361,18 +361,19 @@ function smesolve(
 
     dims = ens_prob.dimensions
     _expvals_all = _expvals_sol_1 isa Nothing ? nothing : map(i -> _se_me_sse_get_expvals(sol[:, i]), eachindex(sol))
-    expvals_all = _expvals_all isa Nothing ? nothing : stack(_expvals_all)
+    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)), eachindex(sol))
 
     expvals =
         _se_me_sse_get_expvals(_sol_1) isa Nothing ? nothing :
-        dropdims(sum(expvals_all, dims = 3), dims = 3) ./ length(sol)
+        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,
         sol.converged,
         _sol_1.alg,

src/time_evolution/ssesolve.jl

@@ -362,18 +362,19 @@ function ssesolve(
     normalize_states = Val(false)
     dims = ens_prob.dimensions
     _expvals_all = _expvals_sol_1 isa Nothing ? nothing : map(i -> _se_me_sse_get_expvals(sol[:, i]), eachindex(sol))
-    expvals_all = _expvals_all isa Nothing ? nothing : stack(_expvals_all)
+    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))
 
     expvals =
         _se_me_sse_get_expvals(_sol_1) isa Nothing ? nothing :
-        dropdims(sum(expvals_all, dims = 3), dims = 3) ./ length(sol)
+        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,
         sol.converged,
         _sol_1.alg,

src/time_evolution/time_evolution.jl

@@ -4,7 +4,7 @@ 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-2, reltol = 1e-2, save_everystep = false, save_end = true)
-const JUMP_TIMES_WHICH_INIT_SIZE = 200
+const COL_TIMES_WHICH_INIT_SIZE = 200
 
 @doc raw"""
     struct TimeEvolutionProblem
@@ -99,9 +99,10 @@ A structure storing the results and some information from solving quantum trajec
 - `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`.
-- `expect_all::Union{AbstractMatrix,Nothing}`: The expectation values corresponding to each trajectory and each time point in `times`
-- `jump_times::Vector{Vector{Real}}`: The time records of every quantum jump occurred in each trajectory.
-- `jump_which::Vector{Vector{Int}}`: The indices of the jump operators in `c_ops` that describe the corresponding quantum jumps occurred in each trajectory.
+- `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`
+- `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.
@@ -122,9 +123,10 @@ struct TimeEvolutionMCSol{
     times::TT
     states::TS
     expect::TE
-    expect_all::TEA
-    jump_times::TJT
-    jump_which::TJW
+    average_expect::TE # Currently just a synonym for `expect`
+    runs_expect::TEA
+    col_times::TJT
+    col_which::TJW
     converged::Bool
     alg::AlgT
     abstol::AT
@@ -140,7 +142,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.expect, 1))\n")
+        print(io, "num_expect = $(size(sol.average_expect, 1))\n")
     end
     print(io, "ODE alg.: $(sol.alg)\n")
     print(io, "abstol = $(sol.abstol)\n")
@@ -159,7 +161,8 @@ A structure storing the results and some information from solving trajectories o
 - `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`.
-- `expect_all::Union{AbstractArray,Nothing}`: The expectation values corresponding to each trajectory and 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`
 - `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.
@@ -178,7 +181,8 @@ struct TimeEvolutionStochasticSol{
     times::TT
     states::TS
     expect::TE
-    expect_all::TEA
+    average_expect::TE # Currently just a synonym for `expect`
+    runs_expect::TEA
     converged::Bool
     alg::AlgT
     abstol::AT
@@ -194,7 +198,7 @@ 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.expect, 1))\n")
+        print(io, "num_expect = $(size(sol.average_expect, 1))\n")
     end
     print(io, "SDE alg.: $(sol.alg)\n")
     print(io, "abstol = $(sol.abstol)\n")

test/core-test/time_evolution.jl

@@ -660,21 +660,21 @@
             )
 
             @test sol_mc1.expect ≈ sol_mc2.expect atol = 1e-10
-            @test sol_mc1.expect_all ≈ sol_mc2.expect_all atol = 1e-10
-            @test sol_mc1.jump_times ≈ sol_mc2.jump_times atol = 1e-10
-            @test sol_mc1.jump_which ≈ sol_mc2.jump_which atol = 1e-10
+            @test sol_mc1.runs_expect ≈ sol_mc2.runs_expect atol = 1e-10
+            @test sol_mc1.col_times ≈ sol_mc2.col_times atol = 1e-10
+            @test sol_mc1.col_which ≈ sol_mc2.col_which atol = 1e-10
 
-            @test sol_mc1.expect_all ≈ sol_mc3.expect_all[:, :, 1:500] atol = 1e-10
+            @test sol_mc1.runs_expect ≈ sol_mc3.runs_expect[:, 1:500, :] atol = 1e-10
 
             @test sol_sse1.expect ≈ sol_sse2.expect atol = 1e-10
-            @test sol_sse1.expect_all ≈ sol_sse2.expect_all atol = 1e-10
+            @test sol_sse1.runs_expect ≈ sol_sse2.runs_expect atol = 1e-10
 
-            @test sol_sse1.expect_all ≈ sol_sse3.expect_all[:, :, 1:50] atol = 1e-10
+            @test sol_sse1.runs_expect ≈ sol_sse3.runs_expect[:, 1:50, :] atol = 1e-10
 
             @test sol_sme1.expect ≈ sol_sme2.expect atol = 1e-10
-            @test sol_sme1.expect_all ≈ sol_sme2.expect_all atol = 1e-10
+            @test sol_sme1.runs_expect ≈ sol_sme2.runs_expect atol = 1e-10
 
-            @test sol_sme1.expect_all ≈ sol_sme3.expect_all[:, :, 1:50] atol = 1e-10
+            @test sol_sme1.runs_expect ≈ sol_sme3.runs_expect[:, 1:50, :] atol = 1e-10
         end
     end