add support of QobjEvo for SteadyStateODESolver

Author: cailixun

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 support of `QobjEvo` for `steadystate` (ODE solver only). ([#536])
+
 ## [v0.34.1]
 Release date: 2025-08-23
 
@@ -301,3 +303,4 @@ Release date: 2024-11-13
 [#517]: https://github.com/qutip/QuantumToolbox.jl/issues/517
 [#520]: https://github.com/qutip/QuantumToolbox.jl/issues/520
 [#531]: https://github.com/qutip/QuantumToolbox.jl/issues/531
+[#536]: https://github.com/qutip/QuantumToolbox.jl/issues/536

src/steadystate.jl

@@ -96,7 +96,7 @@ end
 
 @doc raw"""
     steadystate(
-        H::QuantumObject{OpType},
+        H::AbstractQuantumObject{OpType},
         c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
         solver::SteadyStateSolver = SteadyStateDirectSolver(),
         kwargs...,
@@ -111,7 +111,7 @@ Solve the stationary state based on different solvers.
 - `kwargs`: The keyword arguments for the solver.
 """
 function steadystate(
-    H::QuantumObject{OpType},
+    H::AbstractQuantumObject{OpType},
     c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
     solver::SteadyStateSolver = SteadyStateDirectSolver(),
     kwargs...,
@@ -199,7 +199,7 @@ function _steadystate(L::QuantumObject{SuperOperator}, solver::SteadyStateDirect
     return QuantumObject(ρss, Operator(), L.dimensions)
 end
 
-function _steadystate(L::QuantumObject{SuperOperator}, solver::SteadyStateODESolver; kwargs...)
+function _steadystate(L::AbstractQuantumObject{SuperOperator}, solver::SteadyStateODESolver; kwargs...)
     tmax = solver.tmax
 
     ψ0 = isnothing(solver.ψ0) ? rand_ket(L.dimensions) : solver.ψ0
@@ -224,6 +224,9 @@ function _steadystate(L::QuantumObject{SuperOperator}, solver::SteadyStateODESol
     return ρss
 end
 
+_steadystate(L::QuantumObjectEvolution{SuperOperator}, solver::SteadyStateSolver; kwargs...) =
+    throw(ArgumentError("$(get_typename_wrapper(solver)) does not support QobjEvo."))
+
 struct SteadyStateODECondition{CT<:AbstractArray}
     cache::CT
 end

test/core-test/steady_state.jl

@@ -3,6 +3,7 @@
     a = destroy(N)
     a_d = a'
     H = a_d * a + 0.1 * (a + a_d)
+    Ht = QobjEvo(H, (p, t) -> 1) # for test throw
     c_ops = [sqrt(0.1) * a]
     e_ops = [a_d * a]
     psi0 = fock(N, 3)
@@ -17,18 +18,17 @@
     solver = SteadyStateDirectSolver()
     ρ_ss = steadystate(H, c_ops, solver = solver)
     @test tracedist(rho_me, ρ_ss) < 1e-4
+    @test_throws ArgumentError steadystate(Ht, c_ops, solver = solver)
 
     solver = SteadyStateLinearSolver()
     ρ_ss = steadystate(H, c_ops, solver = solver)
     @test tracedist(rho_me, ρ_ss) < 1e-4
-
-    solver = SteadyStateLinearSolver()
-    ρ_ss = steadystate(H, c_ops, solver = solver)
-    @test tracedist(rho_me, ρ_ss) < 1e-4
+    @test_throws ArgumentError steadystate(Ht, c_ops, solver = solver)
 
     solver = SteadyStateEigenSolver()
     ρ_ss = steadystate(H, c_ops, solver = solver)
     @test tracedist(rho_me, ρ_ss) < 1e-4
+    @test_throws ArgumentError steadystate(Ht, c_ops, solver = solver)
 
     @testset "Type Inference (steadystate)" begin
         L = liouvillian(H, c_ops)

test/ext-test/cpu/autodiff/autodiff.jl

@@ -67,6 +67,17 @@ function my_f_mesolve(p)
     return real(expect(a' * a, sol.states[end]))
 end
 
+function my_f_steadystate(p)
+    ρss = steadystate(
+        L,
+        SteadyStateODESolver(ψ0 = ψ0_mesolve, tmax = tlist_mesolve[end]);
+        params = p,
+        sensealg = BacksolveAdjoint(autojacvec = EnzymeVJP()),
+    )
+
+    return real(expect(a' * a, ρss))
+end
+
 # Analytical solution
 n_ss(Δ, F, γ) = abs2(F / (Δ + 1im * γ / 2))
 
@@ -113,8 +124,12 @@ n_ss(Δ, F, γ) = abs2(F / (Δ + 1im * γ / 2))
 
         my_f_mesolve_direct(params)
         my_f_mesolve(params)
+        my_f_steadystate(params)
 
+        # calculate exact solution and check if steadystate works
         grad_exact = Zygote.gradient((p) -> n_ss(p[1], p[2], p[3]), params)[1]
+        grad_ss = Zygote.gradient(my_f_steadystate, params)[1]
+        @test grad_ss ≈ grad_exact atol=1e-6
 
         @testset "ForwardDiff.jl" begin
             grad_qt = ForwardDiff.gradient(my_f_mesolve_direct, params)