Change order of expect_all

Author: albertomercurio

src/time_evolution/mcsolve.jl

@@ -401,12 +401,12 @@ 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))
     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,

src/time_evolution/smesolve.jl

@@ -354,12 +354,12 @@ function smesolve(
     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 -> _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,

src/time_evolution/ssesolve.jl

@@ -444,12 +444,12 @@ function ssesolve(
         dims = _sol_1.prob.p.Hdims
         _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,

test/core-test/time_evolution.jl

@@ -664,17 +664,17 @@
             @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.expect_all ≈ sol_mc3.expect_all[:, 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.expect_all ≈ sol_sse3.expect_all[:, :, 1:50] atol = 1e-10
+            @test sol_sse1.expect_all ≈ sol_sse3.expect_all[:, 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.expect_all ≈ sol_sme3.expect_all[:, :, 1:50] atol = 1e-10
+            @test sol_sme1.expect_all ≈ sol_sme3.expect_all[:, 1:50, :] atol = 1e-10
         end
     end