Align steadystate ODE solver

Author: albertomercurio

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:

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,46 @@ Base.@kwdef struct SteadyStateLinearSolver{
 end
 
 @doc raw"""
-    SteadyStateODESolver(alg = Tsit5())
+    SteadyStateODESolver(
+        alg = Tsit5(),
+        ψ0 = nothing,
+        tmax = Inf,
+        reltol = 1.0e-6,
+        abstol = 1.0e-8,
+        )
 
 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
+
+```math
+\lVert\frac{\partial |\hat{\rho}\rangle\rangle}{\partial t}\rVert \leq \textrm{abstol}
+```
 
-For more details about the solvers, please refer to [`OrdinaryDiffEq.jl`](https://docs.sciml.ai/OrdinaryDiffEq/stable/)
+# 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.
+- `reltol::Real=1.0e-6`: Relative tolerance in steady state terminate condition. It can be different from the integrator's tolerance.
+- `abstol::Real=1.0e-8`: Absolute tolerance in steady state terminate condition. It can be different from the integrator's tolerance.
+
+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}} <:
+                   SteadyStateSolver
     alg::MT = Tsit5()
+    ψ0::ST = nothing
+    tmax::Float64 = Inf
+    reltol::Float64 = 1.0e-4
+    abstol::Float64 = 1.0e-6
 end
 
 @doc raw"""
@@ -175,68 +205,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)`.",
-    ),
-)
-
-@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.
-
-The termination condition of the stationary state ``|\rho\rangle\rangle`` is that either the following condition is `true`:
+function _steadystate(L::QuantumObject{SuperOperatorQuantumObject}, solver::SteadyStateODESolver; kwargs...)
+    tmax = solver.tmax
+    abstol = solver.abstol
+    reltol = solver.reltol
 
-```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}
-```
+    ψ0 = isnothing(solver.ψ0) ? rand_ket(L.dimensions) : solver.ψ0
 
-# 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 +229,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

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

@@ -146,8 +146,8 @@ end
     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)
+    @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