Add spre and spost methods for ComposedOperator and cache propagator in every time evolution solver (#596)

Author: albertomercurio

CHANGELOG.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/qutip/QuantumToolbox.jl/tree/main)
 
+- Add `spre` and `spost` methods for `ComposedOperator` and cache propagator in every time evolution solver. ([#596])
+
 ## [v0.39.0]
 Release date: 2025-11-17
 
@@ -374,3 +376,4 @@ Release date: 2024-11-13
 [#588]: https://github.com/qutip/QuantumToolbox.jl/issues/588
 [#589]: https://github.com/qutip/QuantumToolbox.jl/issues/589
 [#591]: https://github.com/qutip/QuantumToolbox.jl/issues/591
+[#596]: https://github.com/qutip/QuantumToolbox.jl/issues/596

src/QuantumToolbox.jl

@@ -50,6 +50,7 @@ import SciMLOperators:
     ScalarOperator,
     ScaledOperator,
     AddedOperator,
+    ComposedOperator,
     IdentityOperator,
     update_coefficients!,
     concretize

src/qobj/superoperators.jl

@@ -22,6 +22,7 @@ _sprepost(A, B) = _spre(A) * _spost(B) # for any other input types
 _spre(A::MatrixOperator) = MatrixOperator(_spre(A.A))
 _spre(A::ScaledOperator) = ScaledOperator(A.λ, _spre(A.L))
 _spre(A::AddedOperator) = AddedOperator(map(op -> _spre(op), A.ops))
+_spre(A::ComposedOperator) = ComposedOperator(map(_spre, A.ops), nothing)
 function _spre(A::AbstractSciMLOperator)
     Id = Eye(size(A, 1))
     _lazy_tensor_warning(Id, A)
@@ -31,6 +32,7 @@ end
 _spost(B::MatrixOperator) = MatrixOperator(_spost(B.A))
 _spost(B::ScaledOperator) = ScaledOperator(B.λ, _spost(B.L))
 _spost(B::AddedOperator) = AddedOperator(map(op -> _spost(op), B.ops))
+_spost(B::ComposedOperator) = ComposedOperator(map(_spost, reverse(B.ops)), nothing)
 function _spost(B::AbstractSciMLOperator)
     B_T = transpose(B)
     Id = Eye(size(B, 1))

src/time_evolution/mesolve.jl

@@ -112,7 +112,7 @@ function mesolveProblem(
     else
         to_dense(_complex_float_type(T), mat2vec(ket2dm(ψ0).data))
     end
-    L = L_evo.data
+    L = cache_operator(L_evo.data, ρ0)
 
     kwargs2 = _merge_saveat(tlist, e_ops, DEFAULT_ODE_SOLVER_OPTIONS; kwargs...)
     kwargs3 = _merge_tstops(kwargs2, isconstant(L), tlist)

src/time_evolution/sesolve.jl

@@ -80,7 +80,7 @@ function sesolveProblem(
 
     T = Base.promote_eltype(H_evo, ψ0)
     ψ0 = to_dense(_complex_float_type(T), get_data(ψ0)) # Convert it to dense vector with complex element type
-    U = H_evo.data
+    U = cache_operator(H_evo.data, ψ0)
 
     kwargs2 = _merge_saveat(tlist, e_ops, DEFAULT_ODE_SOLVER_OPTIONS; kwargs...)
     kwargs3 = _merge_tstops(kwargs2, isconstant(U), tlist)

src/time_evolution/smesolve.jl

@@ -109,7 +109,7 @@ function smesolveProblem(
 
     sc_ops_evo_data = Tuple(map(get_data ∘ QobjEvo, sc_ops_list))
 
-    K = get_data(L_evo)
+    K = cache_operator(get_data(L_evo), ρ0)
 
     Id_op = IdentityOperator(prod(dims)^2)
     D_l = map(sc_ops_evo_data) do op

src/time_evolution/ssesolve.jl

@@ -111,7 +111,7 @@ function ssesolveProblem(
         sc_ops_evo_data,
     )
 
-    K = get_data(-1im * QuantumObjectEvolution(H_eff_evo)) + K_l
+    K = cache_operator(get_data(-1im * QuantumObjectEvolution(H_eff_evo)) + K_l, ψ0)
 
     D_l = map(op -> op + _ScalarOperator_e(op, -) * IdentityOperator(prod(dims)), sc_ops_evo_data)
     D = DiffusionOperator(D_l)

test/core-test/quantum_objects_evo.jl

@@ -216,7 +216,7 @@
             coef2(p, t) * conj(coef3(p, t)) * (spre(a) * spost(X') - 0.5 * spre(X' * a) - 0.5 * spost(X' * a)) + # cross terms
             conj(coef2(p, t)) * coef3(p, t) * (spre(X) * spost(a') - 0.5 * spre(a' * X) - 0.5 * spost(a' * X))   # cross terms
         L_ti = liouvillian(H_ti) + D1_ti + D2_ti
-        L_td = @test_logs (:warn,) (:warn,) liouvillian(H_td, c_ops) # warnings from lazy tensor in `lindblad_dissipator(c_op2)`
+        L_td = @test_logs liouvillian(H_td, c_ops) # Test there are no warnings from lazy tensor
         ρvec = mat2vec(rand_dm(N))
         @test L_td(p, t) ≈ L_ti
         @test iscached(L_td) == false
@@ -237,7 +237,7 @@
 
         coef_wrong1(t) = nothing
         coef_wrong2(p, t::ComplexF64) = nothing
-        @test_logs (:warn,) (:warn,) liouvillian(H_td * H_td) # warnings from lazy tensor
+        @test_logs liouvillian(H_td * H_td) # Check there are no warnings from lazy tensor
         @test_throws ArgumentError QobjEvo(a, coef_wrong1)
         @test_throws ArgumentError QobjEvo(a, coef_wrong2)
         @test_throws MethodError QobjEvo([[a, coef1], a' * a, [a', coef2]])

test/core-test/time_evolution.jl

@@ -1256,3 +1256,26 @@ end
     @test all(isapprox.(Z_analytic, sol_se.expect[2, :]; atol = 1e-6))
     @test all(isapprox.(Z_analytic, sol_me.expect[2, :]; atol = 1e-6))
 end
+
+@testitem "ComposedOperator support in time evolution" begin
+    N = 10  # number of basis states
+    a = destroy(N)
+    H = a' * a
+
+    ψ0 = basis(N, 9)  # initial state
+
+    γ0 = 0.5
+    γ1(p, t) = sqrt(p[1] * exp(-t))
+    c_ops_1 = [QobjEvo(2a, γ1)] # This generates a ScaledOperator internally
+    c_ops_2 = [QobjEvo(a, γ1) + QobjEvo(a, γ1)] # This generates a ComposedOperator internally
+
+    e_ops = [a' * a]
+
+    tlist = range(0, 10, 100)
+    params = [γ0]
+    sol_1 = mesolve(H, ψ0, tlist, c_ops_1, e_ops = e_ops; progress_bar = Val(false), params = params, saveat = tlist)
+    sol_2 = mesolve(H, ψ0, tlist, c_ops_2, e_ops = e_ops; progress_bar = Val(false), params = params, saveat = tlist)
+
+    @test sol_1.expect[1, :] ≈ sol_2.expect[1, :] atol=1e-10
+    @test all(x -> isapprox(x[1].data, x[2].data; atol = 1e-10), zip(sol_1.states, sol_2.states))
+end