Add sortby and rev arguments to eigensolver functions (#546)

Author: albertomercurio

CHANGELOG.md

@@ -21,6 +21,7 @@ Release date: 2025-08-23
 
 - Improve Bloch sphere rendering for animation. ([#520])
 - Add support to `Enzyme.jl` for `sesolve` and `mesolve`. ([#531])
+- Add `sortby` and `rev` keyword arguments to eigensolvers. ([#546])
 
 ## [v0.34.0]
 Release date: 2025-07-29
@@ -314,3 +315,4 @@ Release date: 2024-11-13
 [#536]: https://github.com/qutip/QuantumToolbox.jl/issues/536
 [#537]: https://github.com/qutip/QuantumToolbox.jl/issues/537
 [#539]: https://github.com/qutip/QuantumToolbox.jl/issues/539
+[#546]: https://github.com/qutip/QuantumToolbox.jl/issues/546

src/qobj/eigsolve.jl

@@ -149,12 +149,12 @@ function _permuteschur!(
     return T, Q
 end
 
-function _update_schur_eigs!(Hₘ, Uₘ, Uₘᵥ, f, m, β, sorted_vals)
+function _update_schur_eigs!(Hₘ, Uₘ, Uₘᵥ, f, m, β, sorted_vals, sortby, rev)
     F = hessenberg!(Hₘ)
     copyto!(Uₘ, Hₘ)
     LAPACK.orghr!(1, m, Uₘ, F.τ)
     Tₘ, Uₘ, values = hseqr!(Hₘ, Uₘ)
-    sortperm!(sorted_vals, values, by = abs, rev = true)
+    sortperm!(sorted_vals, values, by = sortby, rev = rev)
     _permuteschur!(Tₘ, Uₘ, sorted_vals)
     mul!(f, Uₘᵥ, β)
 
@@ -170,6 +170,8 @@ function _eigsolve(
     m::Int = max(20, 2 * k + 1);
     tol::Real = 1e-8,
     maxiter::Int = 200,
+    sortby::Function = abs2,
+    rev = true,
 ) where {T<:BlasFloat,ObjType<:Union{Nothing,Operator,SuperOperator}}
     n = size(A, 2)
     V = similar(b, n, m + 1)
@@ -209,7 +211,7 @@ function _eigsolve(
 
     M = typeof(cache0)
 
-    Tₘ, Uₘ = _update_schur_eigs!(Hₘ, Uₘ, Uₘᵥ, f, m, β, sorted_vals)
+    Tₘ, Uₘ = _update_schur_eigs!(Hₘ, Uₘ, Uₘᵥ, f, m, β, sorted_vals, sortby, rev)
 
     numops = m
     iter = 0
@@ -235,7 +237,7 @@ function _eigsolve(
 
         # println( A * Vₘ ≈ Vₘ * M(Hₘ) + qₘ * M(transpose(βeₘ)) )     # SHOULD BE TRUE
 
-        Tₘ, Uₘ = _update_schur_eigs!(Hₘ, Uₘ, Uₘᵥ, f, m, β, sorted_vals)
+        Tₘ, Uₘ = _update_schur_eigs!(Hₘ, Uₘ, Uₘᵥ, f, m, β, sorted_vals, sortby, rev)
 
         numops += m - k - 1
         iter += 1
@@ -263,6 +265,8 @@ end
         tol::Real = 1e-8,
         maxiter::Int = 200,
         solver::Union{Nothing, SciMLLinearSolveAlgorithm} = nothing,
+        sortby::Function = abs2,
+        rev::Bool = true,
         kwargs...)
 
 Solve for the eigenvalues and eigenvectors of a matrix `A` using the Arnoldi method.
@@ -276,6 +280,8 @@ Solve for the eigenvalues and eigenvectors of a matrix `A` using the Arnoldi met
 - `tol::Real`: the tolerance for the Arnoldi method. Default is `1e-8`.
 - `maxiter::Int`: the maximum number of iterations for the Arnoldi method. Default is `200`.
 - `solver::Union{Nothing, SciMLLinearSolveAlgorithm}`: the linear solver algorithm. Default is `nothing`.
+- `sortby::Function`: the function to sort eigenvalues. Default is `abs2`.
+- `rev::Bool`: whether to sort in descending order. Default is `true`.
 - `kwargs`: Additional keyword arguments passed to the solver.
 
 # Notes
@@ -293,6 +299,8 @@ function eigsolve(
     tol::Real = 1e-8,
     maxiter::Int = 200,
     solver::Union{Nothing,SciMLLinearSolveAlgorithm} = nothing,
+    sortby::Function = abs2,
+    rev::Bool = true,
     kwargs...,
 )
     return eigsolve(
@@ -306,6 +314,8 @@ function eigsolve(
         tol = tol,
         maxiter = maxiter,
         solver = solver,
+        sortby = sortby,
+        rev = rev,
         kwargs...,
     )
 end
@@ -321,14 +331,27 @@ function eigsolve(
     tol::Real = 1e-8,
     maxiter::Int = 200,
     solver::Union{Nothing,SciMLLinearSolveAlgorithm} = nothing,
+    sortby::Function = abs2,
+    rev::Bool = true,
     kwargs...,
 )
     T = eltype(A)
     isH = ishermitian(A)
     v0 === nothing && (v0 = normalize!(rand(T, size(A, 1))))
 
     if sigma === nothing
-        res = _eigsolve(A, v0, type, dimensions, eigvals, krylovdim, tol = tol, maxiter = maxiter)
+        res = _eigsolve(
+            A,
+            v0,
+            type,
+            dimensions,
+            eigvals,
+            krylovdim,
+            tol = tol,
+            maxiter = maxiter,
+            sortby = sortby,
+            rev = rev,
+        )
         vals = res.values
     else
         Aₛ = A - sigma * I
@@ -346,7 +369,18 @@ function eigsolve(
 
         Amap = EigsolveInverseMap(T, size(A), linsolve)
 
-        res = _eigsolve(Amap, v0, type, dimensions, eigvals, krylovdim, tol = tol, maxiter = maxiter)
+        res = _eigsolve(
+            Amap,
+            v0,
+            type,
+            dimensions,
+            eigvals,
+            krylovdim,
+            tol = tol,
+            maxiter = maxiter,
+            sortby = sortby,
+            rev = rev,
+        )
         vals = @. (1 + sigma * res.values) / res.values
     end
 
@@ -368,6 +402,8 @@ end
         krylovdim::Int = min(10, size(H, 1)),
         maxiter::Int = 200,
         eigstol::Real = 1e-6,
+        sortby::Function = abs2,
+        rev::Bool = true,
         kwargs...,
     )
 
@@ -384,6 +420,8 @@ Solve the eigenvalue problem for a Liouvillian superoperator `L` using the Arnol
 - `krylovdim`: The dimension of the Krylov subspace.
 - `maxiter`: The maximum number of iterations for the eigsolver.
 - `eigstol`: The tolerance for the eigsolver.
+- `sortby::Function`: the function to sort eigenvalues. Default is `abs2`.
+- `rev::Bool`: whether to sort in descending order. Default is `true`.
 - `kwargs`: Additional keyword arguments passed to the differential equation solver.
 
 # Notes
@@ -407,6 +445,8 @@ function eigsolve_al(
     krylovdim::Int = min(10, size(H, 1)),
     maxiter::Int = 200,
     eigstol::Real = 1e-6,
+    sortby::Function = abs2,
+    rev::Bool = true,
     kwargs...,
 ) where {HOpType<:Union{Operator,SuperOperator}}
     L_evo = _mesolve_make_L_QobjEvo(H, c_ops)
@@ -422,8 +462,18 @@ function eigsolve_al(
 
     Lmap = ArnoldiLindbladIntegratorMap(eltype(H), size(L_evo), integrator)
 
-    res =
-        _eigsolve(Lmap, mat2vec(ρ0), L_evo.type, L_evo.dimensions, eigvals, krylovdim, maxiter = maxiter, tol = eigstol)
+    res = _eigsolve(
+        Lmap,
+        mat2vec(ρ0),
+        L_evo.type,
+        L_evo.dimensions,
+        eigvals,
+        krylovdim,
+        maxiter = maxiter,
+        tol = eigstol,
+        sortby = sortby,
+        rev = rev,
+    )
 
     vals = similar(res.values)
     vecs = similar(res.vectors)

test/core-test/eigenvalues_and_operators.jl

@@ -1,4 +1,4 @@
-@testitem "Eigenvalues and Operators" begin
+@testitem "Eigenvalues" begin
     σx = sigmax()
     result = eigenstates(σx, sparse = false)
     λd, ψd, Td = result
@@ -33,12 +33,10 @@
 
     vals_d, vecs_d, mat_d = eigenstates(H_d)
     vals_c, vecs_c, mat_c = eigenstates(H_c)
-    vals2, vecs2, mat2 = eigenstates(H_d, sparse = true, sigma = -0.9, eigvals = 10, krylovdim = 30)
-    sort!(vals_c, by = real)
-    sort!(vals2, by = real)
+    vals2, vecs2, mat2 = eigenstates(H_d, sparse = true, sigma = -0.9, eigvals = 10, krylovdim = 30, by = real)
 
-    @test sum(real.(vals_d[1:20]) .- real.(vals_c[1:20])) / 20 < 1e-3
-    @test sum(real.(vals_d[1:10]) .- real.(vals2[1:10])) / 20 < 1e-3
+    @test real.(vals_d[1:20]) ≈ real.(vals_c[1:20])
+    @test real.(vals_d[1:10]) ≈ real.(vals2[1:10])
 
     N = 5
     a = kron(destroy(N), qeye(N))
@@ -58,16 +56,15 @@
 
     # eigen solve for general matrices
     vals, _, vecs = eigsolve(L.data, sigma = 0.01, eigvals = 10, krylovdim = 50)
-    vals2, _, vecs2 = eigenstates(L)
+    vals2, _, vecs2 = eigenstates(L; sortby = abs)
     vals3, state3, vecs3 = eigsolve_al(L, 1 \ (40 * κ), eigvals = 10, krylovdim = 50)
-    idxs = sortperm(vals2, by = abs)
-    vals2 = vals2[idxs][1:10]
-    vecs2 = vecs2[:, idxs][:, 1:10]
+    vals2 = vals2[1:10]
+    vecs2 = vecs2[:, 1:10]
 
-    @test isapprox(sum(abs2, vals), sum(abs2, vals2), atol = 1e-7)
-    @test isapprox(abs2(vals2[1]), abs2(vals3[1]), atol = 1e-7)
-    @test isapprox(vec2mat(vecs[:, 1]) * exp(-1im * angle(vecs[1, 1])), vec2mat(vecs2[:, 1]), atol = 1e-7)
-    @test isapprox(vec2mat(vecs[:, 1]) * exp(-1im * angle(vecs[1, 1])), vec2mat(vecs3[:, 1]), atol = 1e-5)
+    @test sum(abs2, vals) ≈ sum(abs2, vals2)
+    @test abs2(vals2[1]) ≈ abs2(vals3[1]) atol=1e-7
+    @test vec2mat(vecs[:, 1]) * exp(-1im * angle(vecs[1, 1])) ≈ vec2mat(vecs2[:, 1]) atol=1e-7
+    @test vec2mat(vecs[:, 1]) * exp(-1im * angle(vecs[1, 1])) ≈ vec2mat(vecs3[:, 1]) atol=1e-5
 
     # eigen solve for QuantumObject
     result = eigenstates(L, sparse = true, sigma = 0.01, eigvals = 10, krylovdim = 50)
@@ -78,19 +75,18 @@
     @test resstring ==
           "EigsolveResult:   type=$(SuperOperator())   dims=$(result.dims)\nvalues:\n$(valstring)\nvectors:\n$vecsstring"
 
-    vals2, vecs2 = eigenstates(L, sparse = false)
-    idxs = sortperm(vals2, by = abs)
-    vals2 = vals2[idxs][1:10]
-    vecs2 = vecs2[idxs][1:10]
+    vals2, vecs2 = eigenstates(L, sortby = abs)
+    vals2 = vals2[1:10]
+    vecs2 = vecs2[1:10]
 
     @test result.type isa SuperOperator
     @test result.dims == L.dims
     @test all([v.type isa OperatorKet for v in vecs])
     @test typeof(result.vectors) <: AbstractMatrix
-    @test isapprox(sum(abs2, vals), sum(abs2, vals2), atol = 1e-7)
-    @test isapprox(abs2(vals2[1]), abs2(vals3[1]), atol = 1e-7)
-    @test isapprox(vec2mat(vecs[1]).data * exp(-1im * angle(vecs[1][1])), vec2mat(vecs2[1]).data, atol = 1e-7)
-    @test isapprox(vec2mat(vecs[1]).data * exp(-1im * angle(vecs[1][1])), vec2mat(state3[1]).data, atol = 1e-5)
+    @test sum(abs2, vals) ≈ sum(abs2, vals2)
+    @test abs2(vals2[1]) ≈ abs2(vals3[1]) atol=1e-7
+    @test vec2mat(vecs[1]).data * exp(-1im * angle(vecs[1][1])) ≈ vec2mat(vecs2[1]).data
+    @test vec2mat(vecs[1]).data * exp(-1im * angle(vecs[1][1])) ≈ vec2mat(state3[1]).data atol=1e-5
 
     @testset "Type Inference (eigen)" begin
         N = 5