Improve accuracy of ODE solvers for general cases (#586)

Author: ytdHuang

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Change default solver detection in `eigensolve` when using `sigma` keyword argument (shift-inverse algorithm). If the operator is a `SparseMatrixCSC`, the default solver is `UMFPACKFactorization`, otherwise it is automatically chosen by LinearSolve.jl, depending on the type of the operator. ([#580]) 
 - Add keyword argument `assume_hermitian` to `liouvillian`. This allows users to disable the assumption that the Hamiltonian is Hermitian. ([#581])
+- Improve accuracy of ODE solvers for general cases. ([#586])
 - Use LinearSolve's internal methods for preconditioners in `SteadyStateLinearSolver`. ([#588])
 - Use `FillArrays.jl` for handling superoperators. This makes the code cleaner and potentially more efficient. ([#589])
 - Make sure state generating functions return dense array by default. ([#591])
@@ -365,6 +366,7 @@ Release date: 2024-11-13
 [#579]: https://github.com/qutip/QuantumToolbox.jl/issues/579
 [#580]: https://github.com/qutip/QuantumToolbox.jl/issues/580
 [#581]: https://github.com/qutip/QuantumToolbox.jl/issues/581
+[#586]: https://github.com/qutip/QuantumToolbox.jl/issues/586
 [#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

Project.toml

@@ -5,7 +5,6 @@ authors = ["Alberto Mercurio", "Yi-Te Huang"]
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 DiffEqNoiseProcess = "77a26b50-5914-5dd7-bc55-306e6241c503"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
@@ -17,7 +16,8 @@ LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
-OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
+OrdinaryDiffEqLowOrderRK = "1344f307-1e59-4825-a18e-ace9aa3fa4c6"
+OrdinaryDiffEqVerner = "79d7bb75-1356-48c1-b8c0-6832512096c2"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -46,7 +46,6 @@ QuantumToolboxMakieExt = "Makie"
 ArrayInterface = "6, 7"
 CUDA = "5.0 - 5.8, 5.9.4 - 5"
 ChainRulesCore = "1"
-DiffEqBase = "6"
 DiffEqCallbacks = "4.2.1 - 4"
 DiffEqNoiseProcess = "5"
 Distributed = "1"
@@ -61,7 +60,8 @@ LinearAlgebra = "1"
 LinearSolve = "2, 3"
 Makie = "0.24"
 OrdinaryDiffEqCore = "1"
-OrdinaryDiffEqTsit5 = "1"
+OrdinaryDiffEqLowOrderRK = "1"
+OrdinaryDiffEqVerner = "1"
 Pkg = "1"
 ProgressMeter = "1.11.0"
 Random = "1"

src/QuantumToolbox.jl

@@ -1,12 +1,16 @@
 module QuantumToolbox
 
-# Standard Julia libraries
+## Standard Julia libraries
 using LinearAlgebra
-import LinearAlgebra: checksquare
 using SparseArrays
+
+import Distributed: RemoteChannel
+import LinearAlgebra: checksquare
+import Pkg
+import Random: AbstractRNG, default_rng, seed!
 import Statistics: mean, std
 
-# SciML packages (for QobjEvo, OrdinaryDiffEq, and LinearSolve)
+## SciML packages (for QobjEvo, OrdinaryDiffEq, and LinearSolve)
 import SciMLBase:
     solve,
     solve!,
@@ -32,8 +36,10 @@ import SciMLBase:
     DiscreteCallback,
     AbstractSciMLProblem,
     AbstractODEIntegrator,
-    AbstractODESolution
-import StochasticDiffEq: StochasticDiffEqAlgorithm, SRA2, SRIW1
+    AbstractODEAlgorithm,
+    AbstractODESolution,
+    AbstractSDEAlgorithm
+import StochasticDiffEq: SRA2, SRIW1
 import SciMLOperators:
     cache_operator,
     iscached,
@@ -49,44 +55,42 @@ import SciMLOperators:
     concretize
 import LinearSolve:
     SciMLLinearSolveAlgorithm, KrylovJL_MINRES, KrylovJL_GMRES, UMFPACKFactorization, OperatorAssumptions
-import DiffEqBase: get_tstops
 import DiffEqCallbacks: PeriodicCallback, FunctionCallingCallback, FunctionCallingAffect, TerminateSteadyState
-import OrdinaryDiffEqCore: OrdinaryDiffEqAlgorithm
-import OrdinaryDiffEqTsit5: Tsit5
+import OrdinaryDiffEqVerner: Vern7
+import OrdinaryDiffEqLowOrderRK: DP5
 import DiffEqNoiseProcess: RealWienerProcess!, RealWienerProcess
 
-# other dependencies (in alphabetical order)
+## other dependencies (in alphabetical order)
 import ArrayInterface: allowed_getindex, allowed_setindex!
-import Distributed: RemoteChannel
 import FFTW: fft, ifft, fftfreq, fftshift
 import FillArrays: Eye
 import Graphs: connected_components, DiGraph
 import IncompleteLU: ilu
 import LaTeXStrings: @L_str
-import Pkg
 import ProgressMeter: Progress, next!
-import Random: AbstractRNG, default_rng, seed!
 import SpecialFunctions: loggamma
 import StaticArraysCore: SVector, MVector
 
 # Export functions from the other modules
 
-# LinearAlgebra
+## LinearAlgebra
 export ishermitian, issymmetric, isposdef, dot, tr, svdvals, norm, normalize, normalize!, diag, Hermitian, Symmetric
 
-# SparseArrays
+## SparseArrays
 export permute
 
-# SciMLOperators
+## SciMLOperators
 export cache_operator, iscached, isconstant
 
-# Utility
+# Source files
+
+## Utility
 include("settings.jl")
 include("utilities.jl")
 include("versioninfo.jl")
 include("linear_maps.jl")
 
-# Quantum Object
+## Quantum Object
 include("qobj/space.jl")
 include("qobj/energy_restricted.jl")
 include("qobj/dimensions.jl")
@@ -103,7 +107,7 @@ include("qobj/superoperators.jl")
 include("qobj/synonyms.jl")
 include("qobj/block_diagonal_form.jl")
 
-# time evolution
+## time evolution
 include("time_evolution/time_evolution.jl")
 include("time_evolution/callback_helpers/callback_helpers.jl")
 include("time_evolution/callback_helpers/sesolve_callback_helpers.jl")
@@ -120,7 +124,7 @@ include("time_evolution/ssesolve.jl")
 include("time_evolution/smesolve.jl")
 include("time_evolution/time_evolution_dynamical.jl")
 
-# Others
+## Others
 include("correlations.jl")
 include("wigner.jl")
 include("spin_lattice.jl")
@@ -132,7 +136,7 @@ include("steadystate.jl")
 include("spectrum.jl")
 include("visualization.jl")
 
-# deprecated functions
+## deprecated functions
 include("deprecated.jl")
 
 end

src/qobj/eigsolve.jl

@@ -422,7 +422,7 @@ end
         H::Union{AbstractQuantumObject{HOpType},Tuple},
         T::Real,
         c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
-        alg::OrdinaryDiffEqAlgorithm = Tsit5(),
+        alg::AbstractODEAlgorithm = DP5(),
         params::NamedTuple = NamedTuple(),
         ρ0::AbstractMatrix = rand_dm(prod(H.dimensions)).data,
         eigvals::Int = 1,
@@ -440,7 +440,7 @@ Solve the eigenvalue problem for a Liouvillian superoperator `L` using the Arnol
 - `H`: The Hamiltonian (or directly the Liouvillian) of the system. It can be a [`QuantumObject`](@ref), a [`QuantumObjectEvolution`](@ref), or a tuple of the form supported by [`mesolve`](@ref).
 - `T`: The time at which to evaluate the time evolution.
 - `c_ops`: A vector of collapse operators. Default is `nothing` meaning the system is closed.
-- `alg`: The differential equation solver algorithm. Default is `Tsit5()`.
+- `alg`: The differential equation solver algorithm. Default is `DP5()`.
 - `params`: A `NamedTuple` containing the parameters of the system.
 - `ρ0`: The initial density matrix. If not specified, a random density matrix is used.
 - `eigvals`: The number of eigenvalues to compute.
@@ -465,7 +465,7 @@ function eigsolve_al(
     H::Union{AbstractQuantumObject{HOpType},Tuple},
     T::Real,
     c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
-    alg::OrdinaryDiffEqAlgorithm = Tsit5(),
+    alg::AbstractODEAlgorithm = DP5(),
     params::NamedTuple = NamedTuple(),
     ρ0::AbstractMatrix = rand_dm(prod(H.dimensions)).data,
     eigvals::Int = 1,

src/steadystate.jl

@@ -41,7 +41,7 @@ end
 
 @doc raw"""
     SteadyStateODESolver(
-        alg = Tsit5(),
+        alg = DP5(),
         ψ0 = nothing,
         tmax = Inf,
         terminate_reltol = 1e-4,
@@ -63,7 +63,7 @@ or
 ```
 
 # Arguments
-- `alg::OrdinaryDiffEqAlgorithm=Tsit5()`: The algorithm to solve the ODE.
+- `alg::AbstractODEAlgorithm=DP5()`: 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.
 - `terminate_reltol` = The relative tolerance for stationary state terminate condition. Default to `1e-4`.
@@ -75,13 +75,13 @@ or
 For more details about the solving `alg`orithms, please refer to [`OrdinaryDiffEq.jl`](https://docs.sciml.ai/OrdinaryDiffEq/stable/).
 """
 Base.@kwdef struct SteadyStateODESolver{
-    MT<:OrdinaryDiffEqAlgorithm,
+    MT<:AbstractODEAlgorithm,
     ST<:Union{Nothing,QuantumObject},
     TT<:Real,
     RT<:Real,
     AT<:Real,
 } <: SteadyStateSolver
-    alg::MT = Tsit5()
+    alg::MT = DP5()
     ψ0::ST = nothing
     tmax::TT = Inf
     terminate_reltol::RT = 1e-4

src/time_evolution/callback_helpers/callback_helpers.jl

@@ -6,6 +6,15 @@ This file contains helper functions for callbacks. The affect! function are defi
 
 abstract type AbstractSaveFunc end
 
+function _merge_tstops(kwargs, prob_is_const::Bool, tlist)
+    if prob_is_const
+        return kwargs
+    else
+        tstops = haskey(kwargs, :tstops) ? unique!(sort!(vcat(tlist, kwargs.tstops))) : tlist
+        return merge(kwargs, (tstops = tstops,))
+    end
+end
+
 # Multiple dispatch depending on the progress_bar and e_ops types
 function _generate_se_me_kwargs(e_ops, progress_bar, tlist, kwargs, method)
     cb = _generate_save_callback(e_ops, tlist, progress_bar, method)
@@ -26,8 +35,7 @@ function _generate_stochastic_kwargs(
 
     # Ensure that the noise is stored in tlist. # TODO: Fix this directly in DiffEqNoiseProcess.jl
     # See https://github.com/SciML/DiffEqNoiseProcess.jl/issues/214 for example
-    tstops = haskey(kwargs, :tstops) ? unique!(sort!(vcat(tlist, kwargs.tstops))) : tlist
-    kwargs2 = merge(kwargs, (tstops = tstops,))
+    kwargs2 = _merge_tstops(kwargs, false, tlist) # set 'prob_is_const = false' to force add 'tstops = tlist'
 
     if SF === SaveFuncSSESolve
         cb_normalize = _ssesolve_generate_normalize_cb()

src/time_evolution/callback_helpers/mcsolve_callback_helpers.jl

@@ -107,20 +107,15 @@ function _generate_mcsolve_kwargs(ψ0, T, e_ops, tlist, c_ops, jump_callback, rn
 
     if e_ops isa Nothing
         # We are implicitly saying that we don't have a `Progress`
-        kwargs2 =
-            haskey(kwargs, :callback) ? merge(kwargs, (callback = CallbackSet(cb1, kwargs.callback),)) :
-            merge(kwargs, (callback = cb1,))
-        return kwargs2
+        kwargs2 = _merge_kwargs_with_callback(kwargs, cb1)
     else
         expvals = Array{ComplexF64}(undef, length(e_ops), length(tlist))
 
         _save_func = SaveFuncMCSolve(get_data.(e_ops), Ref(1), expvals)
         cb2 = FunctionCallingCallback(_save_func, funcat = tlist)
-        kwargs2 =
-            haskey(kwargs, :callback) ? merge(kwargs, (callback = CallbackSet(cb1, cb2, kwargs.callback),)) :
-            merge(kwargs, (callback = CallbackSet(cb1, cb2),))
-        return kwargs2
+        kwargs2 = _merge_kwargs_with_callback(kwargs, CallbackSet(cb1, cb2))
     end
+    return kwargs2
 end
 
 function _lindblad_jump_affect!(

src/time_evolution/lr_mesolve.jl

@@ -25,7 +25,7 @@ struct TimeEvolutionLRSol{
     TS<:AbstractVector,
     TE<:Matrix{ComplexF64},
     RetT<:Enum,
-    AlgT<:OrdinaryDiffEqAlgorithm,
+    AlgT<:AbstractODEAlgorithm,
     TolT<:Real,
     TSZB<:AbstractVector,
     TM<:Vector{<:Integer},
@@ -45,7 +45,7 @@ struct TimeEvolutionLRSol{
 end
 
 lr_mesolve_options_default = (
-    alg = Tsit5(),
+    alg = DP5(),
     progress = true,
     err_max = 0.0,
     p0 = 0.0,

src/time_evolution/mcsolve.jl

@@ -274,7 +274,7 @@ end
         ψ0::QuantumObject{Ket},
         tlist::AbstractVector,
         c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
-        alg::OrdinaryDiffEqAlgorithm = Tsit5(),
+        alg::AbstractODEAlgorithm = DP5(),
         e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
         params = NullParameters(),
         rng::AbstractRNG = default_rng(),
@@ -329,7 +329,7 @@ If the environmental measurements register a quantum jump, the wave function und
 - `ψ0`: Initial state of the system ``|\psi(0)\rangle``.
 - `tlist`: List of time points at which to save either the state or the expectation values of the system.
 - `c_ops`: List of collapse operators ``\{\hat{C}_n\}_n``. It can be either a `Vector` or a `Tuple`.
-- `alg`: The algorithm to use for the ODE solver. Default to `Tsit5()`.
+- `alg`: The algorithm to use for the ODE solver. Default to `DP5()`.
 - `e_ops`: List of operators for which to calculate expectation values. It can be either a `Vector` or a `Tuple`.
 - `params`: Parameters to pass to the solver. This argument is usually expressed as a `NamedTuple` or `AbstractVector` of parameters. For more advanced usage, any custom struct can be used.
 - `rng`: Random number generator for reproducibility.
@@ -361,7 +361,7 @@ function mcsolve(
     ψ0::QuantumObject{Ket},
     tlist::AbstractVector,
     c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
-    alg::OrdinaryDiffEqAlgorithm = Tsit5(),
+    alg::AbstractODEAlgorithm = DP5(),
     e_ops::Union{Nothing,AbstractVector,Tuple} = nothing,
     params = NullParameters(),
     rng::AbstractRNG = default_rng(),
@@ -398,7 +398,7 @@ end
 
 function mcsolve(
     ens_prob_mc::TimeEvolutionProblem,
-    alg::OrdinaryDiffEqAlgorithm = Tsit5(),
+    alg::AbstractODEAlgorithm = DP5(),
     ntraj::Int = 500,
     ensemblealg::EnsembleAlgorithm = EnsembleThreads(),
     keep_runs_results = Val(false),