Align steadystate ODE solver and improve GPU support

CHANGELOG.md

@@ -22,6 +22,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - `hellinger_dist`
   - `bures_dist`
   - `bures_angle`
+- Align `steadystate` ODE solver to other methods and improve GPU support. ([#421])
 
 ## [v0.27.0]
 Release date: 2025-02-14
@@ -165,3 +166,4 @@ Release date: 2024-11-13
 [#418]: https://github.com/qutip/QuantumToolbox.jl/issues/418
 [#419]: https://github.com/qutip/QuantumToolbox.jl/issues/419
 [#420]: https://github.com/qutip/QuantumToolbox.jl/issues/420
+[#421]: https://github.com/qutip/QuantumToolbox.jl/issues/421

docs/src/users_guide/steadystate.md

@@ -37,10 +37,6 @@ To change a solver, use the keyword argument `solver`, for example:
 ρ_ss = steadystate(H, c_ops; solver = SteadyStateLinearSolver())
 ```
 
-!!! note "Initial state for `SteadyStateODESolver()`"
-    It is necessary to provide the initial state `ψ0` for ODE solving method, namely
-    `steadystate(H, ψ0, tspan, c_ops, solver = SteadyStateODESolver())`, where `tspan::Real` represents the final time step, defaults to `Inf` (infinity).
-
 Although it is not obvious, the [`SteadyStateDirectSolver()`](@ref SteadyStateDirectSolver) and [`SteadyStateEigenSolver()`](@ref SteadyStateEigenSolver) methods all use an LU decomposition internally and thus can have a large memory overhead. In contrast, for [`SteadyStateLinearSolver()`](@ref SteadyStateLinearSolver), iterative algorithms provided by [`LinearSolve.jl`](https://docs.sciml.ai/LinearSolve/stable/solvers/solvers/), such as `KrylovJL_GMRES()` and `KrylovJL_BICGSTAB()`, do not factor the matrix and thus take less memory than the LU methods and allow, in principle, for extremely large system sizes. The downside is that these methods can take much longer than the direct method as the condition number of the Liouvillian matrix is large, indicating that these iterative methods require a large number of iterations for convergence. 
 
 To overcome this, one can provide preconditioner that solves for an approximate inverse for the (modified) Liouvillian, thus better conditioning the problem, leading to faster convergence. The left and right preconditioner can be specified as the keyword argument `Pl` and `Pr`, respectively:

ext/QuantumToolboxCUDAExt.jl

@@ -2,6 +2,7 @@ module QuantumToolboxCUDAExt
 
 using QuantumToolbox
 using QuantumToolbox: makeVal, getVal
+import QuantumToolbox: _sparse_similar
 import CUDA: cu, CuArray, allowscalar
 import CUDA.CUSPARSE: CuSparseVector, CuSparseMatrixCSC, CuSparseMatrixCSR, AbstractCuSparseArray
 import SparseArrays: SparseVector, SparseMatrixCSC
@@ -106,4 +107,7 @@ QuantumToolbox.to_dense(A::MT) where {MT<:AbstractCuSparseArray} = CuArray(A)
 QuantumToolbox.to_dense(::Type{T1}, A::CuArray{T2}) where {T1<:Number,T2<:Number} = CuArray{T1}(A)
 QuantumToolbox.to_dense(::Type{T}, A::AbstractCuSparseArray) where {T<:Number} = CuArray{T}(A)
 
+QuantumToolbox._sparse_similar(A::CuSparseMatrixCSC, args...) = sparse(args..., fmt = :csc)
+QuantumToolbox._sparse_similar(A::CuSparseMatrixCSR, args...) = sparse(args..., fmt = :csr)
+
 end

src/steadystate.jl

@@ -27,7 +27,7 @@ struct SteadyStateEigenSolver <: SteadyStateSolver end
 
 A solver which solves [`steadystate`](@ref) by finding the inverse of Liouvillian [`SuperOperator`](@ref) using the `alg`orithms given in [`LinearSolve.jl`](https://docs.sciml.ai/LinearSolve/stable/).
 
-# Parameters
+# Arguments
 - `alg::SciMLLinearSolveAlgorithm=KrylovJL_GMRES()`: algorithms given in [`LinearSolve.jl`](https://docs.sciml.ai/LinearSolve/stable/)
 - `Pl::Union{Function,Nothing}=nothing`: left preconditioner, see documentation [Solving for Steady-State Solutions](@ref doc:Solving-for-Steady-State-Solutions) for more details.
 - `Pr::Union{Function,Nothing}=nothing`: right preconditioner, see documentation [Solving for Steady-State Solutions](@ref doc:Solving-for-Steady-State-Solutions) for more details.
@@ -43,16 +43,40 @@ Base.@kwdef struct SteadyStateLinearSolver{
 end
 
 @doc raw"""
-    SteadyStateODESolver(alg = Tsit5())
+    SteadyStateODESolver(
+        alg = Tsit5(),
+        ψ0 = nothing,
+        tmax = Inf,
+        )
 
 An ordinary differential equation (ODE) solver for solving [`steadystate`](@ref).
 
-It includes a field (attribute) `SteadyStateODESolver.alg` that specifies the solving algorithm. Default to `Tsit5()`.
+Solve the stationary state based on time evolution (ordinary differential equations; `OrdinaryDiffEq.jl`) with a given initial state.
+
+The termination condition of the stationary state ``|\rho\rangle\rangle`` is that either the following condition is `true`:
+
+```math
+\lVert\frac{\partial |\hat{\rho}\rangle\rangle}{\partial t}\rVert \leq \textrm{reltol} \times\lVert\frac{\partial |\hat{\rho}\rangle\rangle}{\partial t}+|\hat{\rho}\rangle\rangle\rVert
+```
+
+or
 
-For more details about the solvers, please refer to [`OrdinaryDiffEq.jl`](https://docs.sciml.ai/OrdinaryDiffEq/stable/)
+```math
+\lVert\frac{\partial |\hat{\rho}\rangle\rangle}{\partial t}\rVert \leq \textrm{abstol}
+```
+
+# Arguments
+- `alg::OrdinaryDiffEqAlgorithm=Tsit5()`: The algorithm to solve the ODE.
+- `ψ0::Union{Nothing,QuantumObject}=nothing`: The initial state of the system. If not specified, a random pure state will be generated.
+- `tmax::Real=Inf`: The final time step for the steady state problem.
+
+For more details about the solvers, please refer to [`OrdinaryDiffEq.jl`](https://docs.sciml.ai/OrdinaryDiffEq/stable/).
 """
-Base.@kwdef struct SteadyStateODESolver{MT<:OrdinaryDiffEqAlgorithm} <: SteadyStateSolver
+Base.@kwdef struct SteadyStateODESolver{MT<:OrdinaryDiffEqAlgorithm,ST<:Union{Nothing,QuantumObject},T<:Real} <:
+                   SteadyStateSolver
     alg::MT = Tsit5()
+    ψ0::ST = nothing
+    tmax::T = Inf
 end
 
 @doc raw"""
@@ -108,18 +132,18 @@ function _steadystate(L::QuantumObject{SuperOperatorQuantumObject}, solver::Stea
     N = prod(L.dimensions)
     weight = norm(L_tmp, 1) / length(L_tmp)
 
-    v0 = _get_dense_similar(L_tmp, N^2)
+    v0 = _dense_similar(L_tmp, N^2)
     fill!(v0, 0)
     allowed_setindex!(v0, weight, 1) # Because scalar indexing is not allowed on GPU arrays
 
     idx_range = collect(1:N)
-    rows = _get_dense_similar(L_tmp, N)
-    cols = _get_dense_similar(L_tmp, N)
-    vals = _get_dense_similar(L_tmp, N)
+    rows = _dense_similar(L_tmp, N)
+    cols = _dense_similar(L_tmp, N)
+    vals = _dense_similar(L_tmp, N)
     fill!(rows, 1)
     copyto!(cols, N .* (idx_range .- 1) .+ idx_range)
     fill!(vals, weight)
-    Tn = sparse(rows, cols, vals, N^2, N^2)
+    Tn = _sparse_similar(L_tmp, rows, cols, vals, N^2, N^2)
     L_tmp = L_tmp + Tn
 
     (haskey(kwargs, :Pl) || haskey(kwargs, :Pr)) && error("The use of preconditioners must be defined in the solver.")
@@ -155,14 +179,14 @@ function _steadystate(L::QuantumObject{SuperOperatorQuantumObject}, solver::Stea
     N = prod(L.dimensions)
     weight = norm(L_tmp, 1) / length(L_tmp)
 
-    v0 = _get_dense_similar(L_tmp, N^2)
+    v0 = _dense_similar(L_tmp, N^2)
     fill!(v0, 0)
     allowed_setindex!(v0, weight, 1) # Because scalar indexing is not allowed on GPU arrays
 
     idx_range = collect(1:N)
-    rows = _get_dense_similar(L_tmp, N)
-    cols = _get_dense_similar(L_tmp, N)
-    vals = _get_dense_similar(L_tmp, N)
+    rows = _dense_similar(L_tmp, N)
+    cols = _dense_similar(L_tmp, N)
+    vals = _dense_similar(L_tmp, N)
     fill!(rows, 1)
     copyto!(cols, N .* (idx_range .- 1) .+ idx_range)
     fill!(vals, weight)
@@ -175,68 +199,20 @@ function _steadystate(L::QuantumObject{SuperOperatorQuantumObject}, solver::Stea
     return QuantumObject(ρss, Operator, L.dimensions)
 end
 
-_steadystate(L::QuantumObject{SuperOperatorQuantumObject}, solver::SteadyStateODESolver; kwargs...) = throw(
-    ArgumentError(
-        "The initial state ψ0 is required for SteadyStateODESolver, use the following call instead: `steadystate(H, ψ0, tmax, c_ops)`.",
-    ),
-)
+function _steadystate(L::QuantumObject{SuperOperatorQuantumObject}, solver::SteadyStateODESolver; kwargs...)
+    tmax = solver.tmax
 
-@doc raw"""
-    steadystate(
-        H::QuantumObject{HOpType},
-        ψ0::QuantumObject{StateOpType},
-        tmax::Real = Inf,
-        c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
-        solver::SteadyStateODESolver = SteadyStateODESolver(),
-        reltol::Real = 1.0e-8,
-        abstol::Real = 1.0e-10,
-        kwargs...,
-    )
-
-Solve the stationary state based on time evolution (ordinary differential equations; `OrdinaryDiffEq.jl`) with a given initial state.
+    ψ0 = isnothing(solver.ψ0) ? rand_ket(L.dimensions) : solver.ψ0
+    abstol = haskey(kwargs, :abstol) ? kwargs[:abstol] : DEFAULT_ODE_SOLVER_OPTIONS.abstol
+    reltol = haskey(kwargs, :reltol) ? kwargs[:reltol] : DEFAULT_ODE_SOLVER_OPTIONS.reltol
 
-The termination condition of the stationary state ``|\rho\rangle\rangle`` is that either the following condition is `true`:
-
-```math
-\lVert\frac{\partial |\hat{\rho}\rangle\rangle}{\partial t}\rVert \leq \textrm{reltol} \times\lVert\frac{\partial |\hat{\rho}\rangle\rangle}{\partial t}+|\hat{\rho}\rangle\rangle\rVert
-```
-
-or
-
-```math
-\lVert\frac{\partial |\hat{\rho}\rangle\rangle}{\partial t}\rVert \leq \textrm{abstol}
-```
-
-# Parameters
-- `H`: The Hamiltonian or the Liouvillian of the system.
-- `ψ0`: The initial state of the system.
-- `tmax=Inf`: The final time step for the steady state problem.
-- `c_ops=nothing`: The list of the collapse operators.
-- `solver`: see [`SteadyStateODESolver`](@ref) for more details.
-- `reltol=1.0e-8`: Relative tolerance in steady state terminate condition and solver adaptive timestepping.
-- `abstol=1.0e-10`: Absolute tolerance in steady state terminate condition and solver adaptive timestepping.
-- `kwargs`: The keyword arguments for the ODEProblem.
-"""
-function steadystate(
-    H::QuantumObject{HOpType},
-    ψ0::QuantumObject{StateOpType},
-    tmax::Real = Inf,
-    c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
-    solver::SteadyStateODESolver = SteadyStateODESolver(),
-    reltol::Real = 1.0e-8,
-    abstol::Real = 1.0e-10,
-    kwargs...,
-) where {
-    HOpType<:Union{OperatorQuantumObject,SuperOperatorQuantumObject},
-    StateOpType<:Union{KetQuantumObject,OperatorQuantumObject},
-}
     ftype = _FType(ψ0)
-    cb = TerminateSteadyState(abstol, reltol, _steadystate_ode_condition)
+    _terminate_func = SteadyStateODECondition(similar(mat2vec(ket2dm(ψ0)).data))
+    cb = TerminateSteadyState(abstol, reltol, _terminate_func)
     sol = mesolve(
-        H,
+        L,
         ψ0,
         [ftype(0), ftype(tmax)],
-        c_ops,
         progress_bar = Val(false),
         save_everystep = false,
         saveat = ftype[],
@@ -247,12 +223,17 @@ function steadystate(
     return ρss
 end
 
-function _steadystate_ode_condition(integrator, abstol, reltol, min_t)
+struct SteadyStateODECondition{CT<:AbstractArray}
+    cache::CT
+end
+
+function (f::SteadyStateODECondition)(integrator, abstol, reltol, min_t)
     # this condition is same as DiffEqBase.NormTerminationMode
 
-    du_dt = (integrator.u - integrator.uprev) / integrator.dt
-    norm_du_dt = norm(du_dt)
-    if (norm_du_dt <= reltol * norm(du_dt + integrator.u)) || (norm_du_dt <= abstol)
+    f.cache .= (integrator.u .- integrator.uprev) ./ integrator.dt
+    norm_du_dt = norm(f.cache)
+    f.cache .+= integrator.u
+    if norm_du_dt <= reltol * norm(f.cache) || norm_du_dt <= abstol
         return true
     else
         return false

src/utilities.jl

@@ -130,8 +130,10 @@
 
 get_typename_wrapper(A) = Base.typename(typeof(A)).wrapper
 
-_get_dense_similar(A::AbstractArray, args...) = similar(A, args...)
-_get_dense_similar(A::AbstractSparseMatrix, args...) = similar(nonzeros(A), args...)
+_dense_similar(A::AbstractArray, args...) = similar(A, args...)
+_dense_similar(A::AbstractSparseMatrix, args...) = similar(nonzeros(A), args...)
+
+_sparse_similar(A::AbstractArray, args...) = sparse(args...)
 
 _Ginibre_ensemble(n::Int, rank::Int = n) = randn(ComplexF64, n, rank) / sqrt(n)
 

test/core-test/steady_state.jl

@@ -11,9 +11,8 @@
     rho_me = sol_me.states[end]
 
     solver = SteadyStateODESolver()
-    ρ_ss = steadystate(H, psi0, t_l[end], c_ops, solver = solver)
+    ρ_ss = steadystate(H, c_ops, solver = solver)
     @test tracedist(rho_me, ρ_ss) < 1e-4
-    @test_throws ArgumentError steadystate(H, c_ops, solver = solver)
 
     solver = SteadyStateDirectSolver()
     ρ_ss = steadystate(H, c_ops, solver = solver)
@@ -34,9 +33,9 @@
     @testset "Type Inference (steadystate)" begin
         L = liouvillian(H, c_ops)
 
-        solver = SteadyStateODESolver()
-        @inferred steadystate(H, psi0, t_l[end], c_ops, solver = solver)
-        @inferred steadystate(L, psi0, t_l[end], solver = solver)
+        solver = SteadyStateODESolver(tmax = t_l[end])
+        @inferred steadystate(H, c_ops, solver = solver)
+        @inferred steadystate(L, solver = solver)
 
         solver = SteadyStateDirectSolver()
         @inferred steadystate(H, c_ops, solver = solver)

test/ext-test/gpu/cuda_ext.jl

@@ -125,6 +125,31 @@
     @test all([isapprox(sol_cpu.expect[i], sol_gpu32.expect[i]; atol = 1e-6) for i in 1:length(tlist)])
 end
 
+@testset "CUDA steadystate" begin
+    N = 50
+    Δ = 0.01
+    F = 0.1
+    γ = 0.1
+    nth = 2
+
+    a = destroy(N)
+    H = Δ * a' * a + F * (a + a')
+    c_ops = [sqrt(γ * (nth + 1)) * a, sqrt(γ * nth) * a']
+
+    ρ_ss_cpu = steadystate(H, c_ops)
+
+    H_gpu_csc = cu(H)
+    c_ops_gpu_csc = [cu(c_op) for c_op in c_ops]
+    ρ_ss_gpu_csc = steadystate(H_gpu_csc, c_ops_gpu_csc, solver = SteadyStateLinearSolver())
+
+    H_gpu_csr = CuSparseMatrixCSR(H_gpu_csc)
+    c_ops_gpu_csr = [CuSparseMatrixCSR(c_op) for c_op in c_ops_gpu_csc]
+    ρ_ss_gpu_csr = steadystate(H_gpu_csr, c_ops_gpu_csr, solver = SteadyStateLinearSolver())
+
+    @test ρ_ss_cpu.data ≈ Array(ρ_ss_gpu_csc.data) atol = 1e-8 * length(ρ_ss_cpu)
+    @test ρ_ss_cpu.data ≈ Array(ρ_ss_gpu_csr.data) atol = 1e-8 * length(ρ_ss_cpu)
+end
+
 @testset "CUDA ptrace" begin
     g = fock(2, 1)
     e = fock(2, 0)