Lanczos-based spectrum method (#476)

Author: matteosecli

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)
 
+- Introduce `Lanczos` solver for `spectrum`. ([#476])
+
 ## [v0.31.1]
 Release date: 2025-05-16
 
@@ -228,3 +230,4 @@ Release date: 2024-11-13
 [#455]: https://github.com/qutip/QuantumToolbox.jl/issues/455
 [#456]: https://github.com/qutip/QuantumToolbox.jl/issues/456
 [#460]: https://github.com/qutip/QuantumToolbox.jl/issues/460
+[#476]: https://github.com/qutip/QuantumToolbox.jl/issues/476

benchmarks/correlations_and_spectrum.jl

@@ -19,6 +19,8 @@ function benchmark_correlations_and_spectrum!(SUITE)
 
     PI_solver = PseudoInverse()
 
+    L_solver = Lanczos()
+
     SUITE["Correlations and Spectrum"]["FFT Correlation"] =
         @benchmarkable _calculate_fft_spectrum($H, $t_l, $c_ops, $(a'), $a)
 
@@ -28,5 +30,8 @@ function benchmark_correlations_and_spectrum!(SUITE)
     SUITE["Correlations and Spectrum"]["Spectrum"]["Pseudo Inverse"] =
         @benchmarkable spectrum($H, $ω_l, $c_ops, $(a'), $a, solver = $PI_solver)
 
+    SUITE["Correlations and Spectrum"]["Spectrum"]["Lanczos"] =
+        @benchmarkable spectrum($H, $ω_l, $c_ops, $(a'), $a, solver = $L_solver)
+
     return nothing
 end

docs/src/resources/api.md

@@ -248,6 +248,7 @@ spectrum_correlation_fft
 spectrum
 ExponentialSeries
 PseudoInverse
+Lanczos
 ```
 
 ## [Entropy and Metrics](@id doc-API:Entropy-and-Metrics)

src/QuantumToolbox.jl

@@ -115,14 +115,14 @@ include("time_evolution/time_evolution_dynamical.jl")
 
 # Others
 include("correlations.jl")
-include("spectrum.jl")
 include("wigner.jl")
 include("spin_lattice.jl")
 include("arnoldi.jl")
 include("entropy.jl")
 include("metrics.jl")
 include("negativity.jl")
 include("steadystate.jl")
+include("spectrum.jl")
 include("visualization.jl")
 
 # deprecated functions

src/spectrum.jl

@@ -1,5 +1,5 @@
 export spectrum, spectrum_correlation_fft
-export SpectrumSolver, ExponentialSeries, PseudoInverse
+export SpectrumSolver, ExponentialSeries, PseudoInverse, Lanczos
 
 abstract type SpectrumSolver end
 
@@ -26,6 +26,20 @@ end
 
 PseudoInverse(; alg::SciMLLinearSolveAlgorithm = KrylovJL_GMRES()) = PseudoInverse(alg)
 
+@doc raw"""
+    Lanczos(; tol = 1e-8, maxiter = 5000, verbose = 0)
+
+A solver which solves [`spectrum`](@ref) by using a non-symmetric Lanczos variant of the algorithm in [Koch2011](https://www.cond-mat.de/events/correl11/manuscripts/koch.pdf).
+The nonsymmetric Lanczos algorithm is adapted from Algorithm 6.6 in [Saad2011](https://www-users.cse.umn.edu/~saad/eig_book_2ndEd.pdf).
+The running estimate is updated via a [Wallis-Euler recursion](https://en.wikipedia.org/wiki/Continued_fraction).
+"""
+Base.@kwdef struct Lanczos{T<:Real,SS<:Union{Nothing,<:SteadyStateSolver}} <: SpectrumSolver
+    tol::T = 1e-8
+    maxiter::Int = 5000
+    verbose::Int = 0
+    steadystate_solver::SS = nothing
+end
+
 @doc raw"""
     spectrum(H::QuantumObject,
         ωlist::AbstractVector,
@@ -44,6 +58,7 @@ S(\omega) = \int_{-\infty}^\infty \lim_{t \rightarrow \infty} \left\langle \hat{
 See also the following list for `SpectrumSolver` docstrings:
 - [`ExponentialSeries`](@ref)
 - [`PseudoInverse`](@ref)
+- [`Lanczos`](@ref)
 """
 function spectrum(
     H::QuantumObject{HOpType},
@@ -144,6 +159,142 @@ function _spectrum(
     return spec
 end
 
+function _spectrum(
+    L::QuantumObject{SuperOperator},
+    ωlist::AbstractVector,
+    A::QuantumObject{Operator},
+    B::QuantumObject{Operator},
+    solver::Lanczos,
+)
+    check_dimensions(L, A, B)
+
+    # Define type shortcuts
+    fT = _FType(L)
+    cT = _CType(L)
+
+    # Calculate |v₁> = B|ρss>
+    ρss =
+        isnothing(solver.steadystate_solver) ? mat2vec(steadystate(L)) :
+        mat2vec(steadystate(L; solver = solver.steadystate_solver))
+    vₖ = (spre(B) * ρss).data
+
+    # Define (possibly GPU) vector type
+    vT = typeof(vₖ)
+
+    # Calculate <w₁| = <I|A
+    D = prod(L.dimensions)
+    Ivec = SparseVector(D^2, [1 + n * (D + 1) for n in 0:(D-1)], ones(cT, D)) # same as vec(system_identity_matrix)
+    wₖ = spre(A).data' * vT(Ivec)
+
+    # Store the norm of the Green's function before renormalizing |v₁> and <w₁|
+    gfNorm = abs(dot(wₖ, vₖ))
+    vₖ ./= sqrt(gfNorm)
+    wₖ ./= sqrt(gfNorm)
+
+    # Handle input frequency range
+    ωList = vT(convert(Vector{fT}, ωlist))  # Make sure they're real frequencies and potentially on GPU
+    Length = length(ωList)
+
+    # Current and previous estimates of the spectrum
+    lanczosFactor = vT(zeros(cT, Length))
+    lanczosFactor₋₁ = vT(zeros(cT, Length))
+
+    # Tridiagonal matrix elements
+    αₖ = cT(0)
+    βₖ = cT(-1)
+    δₖ = fT(+1)
+
+    # Current and up to second-to-last A and B Euler sequences
+    A₋₂ = vT(ones(cT, Length))
+    A₋₁ = vT(zeros(cT, Length))
+    Aₖ = vT(zeros(cT, Length))
+    B₋₂ = vT(zeros(cT, Length))
+    B₋₁ = vT(ones(cT, Length))
+    Bₖ = vT(zeros(cT, Length))
+
+    # Maximum norm and residue
+    maxNorm = vT(zeros(cT, length(ωList)))
+    maxResidue = fT(0.0)
+
+    # Previous and next left/right Krylov vectors
+    v₋₁ = vT(zeros(cT, D^2))
+    v₊₁ = vT(zeros(cT, D^2))
+    w₋₁ = vT(zeros(cT, D^2))
+    w₊₁ = vT(zeros(cT, D^2))
+
+    # Frequency of renormalization
+    renormFrequency = 1
+
+    # Loop over the Krylov subspace(s)
+    for k in 1:solver.maxiter
+        # k-th diagonal element
+        mul!(w₊₁, L.data', wₖ)
+        αₖ = dot(w₊₁, vₖ)
+
+        # Update A(k), B(k) and continuous fraction; normalization avoids overflow
+        Aₖ .= (-1im .* ωList .+ αₖ) .* A₋₁ .- (βₖ * δₖ) .* A₋₂
+        Bₖ .= (-1im .* ωList .+ αₖ) .* B₋₁ .- (βₖ * δₖ) .* B₋₂
+        lanczosFactor₋₁ .= lanczosFactor
+        lanczosFactor .= Aₖ ./ Bₖ
+
+        # Renormalize Euler sequences to avoid overflow
+        if k % renormFrequency == 0
+            map!((x, y) -> max(abs(x), abs(y)), maxNorm, Aₖ, Bₖ)
+            Aₖ ./= maxNorm
+            Bₖ ./= maxNorm
+            A₋₁ ./= maxNorm
+            B₋₁ ./= maxNorm
+        end
+
+        # Check for convergence
+
+        residueNorm = max(maximum(abs, lanczosFactor), maximum(abs, lanczosFactor₋₁))
+        lanczosFactor₋₁ .-= lanczosFactor
+        maxResidue = maximum(abs, lanczosFactor₋₁) / residueNorm
+        if maxResidue <= solver.tol
+            if solver.verbose > 1
+                println("spectrum(): solver::Lanczos converged after $(k) iterations")
+            end
+            break
+        end
+
+        # (k+1)-th left/right vectors, orthogonal to previous ones
+        mul!(v₊₁, L.data, vₖ)
+        v₊₁ .= v₊₁ .- αₖ .* vₖ .- βₖ .* v₋₁
+        w₊₁ .= w₊₁ .- conj(αₖ) .* wₖ .- conj(δₖ) .* w₋₁
+        v₋₁, vₖ = vₖ, v₋₁
+        vₖ, v₊₁ = v₊₁, vₖ
+        w₋₁, wₖ = wₖ, w₋₁
+        wₖ, w₊₁ = w₊₁, wₖ
+
+        # k-th off-diagonal elements
+        buf = dot(wₖ, vₖ)
+        δₖ = sqrt(abs(buf))
+        βₖ = buf / δₖ
+
+        # Normalize (k+1)-th left/right vectors
+        vₖ ./= δₖ
+        wₖ ./= conj(βₖ)
+
+        # Update everything for the next cycle
+        A₋₂, A₋₁ = A₋₁, A₋₂
+        A₋₁, Aₖ = Aₖ, A₋₁
+        B₋₂, B₋₁ = B₋₁, B₋₂
+        B₋₁, Bₖ = Bₖ, B₋₁
+    end
+
+    if solver.verbose > 0 && maxResidue > solver.tol
+        println("spectrum(): maxiter = $(solver.maxiter) reached before convergence!")
+        println("spectrum(): Max residue = $maxResidue")
+        println("spectrum(): Consider increasing maxiter and/or tol")
+    end
+
+    # Restore the norm
+    lanczosFactor .= gfNorm .* lanczosFactor
+
+    return -2 .* real(lanczosFactor)
+end
+
 @doc raw"""
     spectrum_correlation_fft(tlist, corr; inverse=false)
 

test/core-test/correlations_and_spectrum.jl

@@ -13,10 +13,12 @@
     ω_l2 = range(0, 3, length = 1000)
     spec2 = spectrum(H, ω_l2, c_ops, a', a)
     spec3 = spectrum(H, ω_l2, c_ops, a', a; solver = PseudoInverse())
+    spec4 = spectrum(H, ω_l2, c_ops, a', a; solver = Lanczos())
 
     spec1 = spec1 ./ maximum(spec1)
     spec2 = spec2 ./ maximum(spec2)
     spec3 = spec3 ./ maximum(spec3)
+    spec4 = spec4 ./ maximum(spec4)
 
     test_func1 = maximum(real.(spec1)) * (0.1 / 2)^2 ./ ((ω_l1 .- 1) .^ 2 .+ (0.1 / 2)^2)
     test_func2 = maximum(real.(spec2)) * (0.1 / 2)^2 ./ ((ω_l2 .- 1) .^ 2 .+ (0.1 / 2)^2)
@@ -26,12 +28,26 @@
     @test sum(abs2.(spec2[idxs2] .- test_func2[idxs2])) / sum(abs2.(test_func2[idxs2])) < 0.01
     @test all(corr1 .≈ corr2)
     @test all(spec2 .≈ spec3)
+    @test all(spec2 .≈ spec4)
 
     @testset "Type Inference spectrum" begin
         @inferred correlation_2op_1t(H, nothing, t_l, c_ops, a', a; progress_bar = Val(false))
         @inferred spectrum_correlation_fft(t_l, corr1)
         @inferred spectrum(H, ω_l2, c_ops, a', a)
         @inferred spectrum(H, ω_l2, c_ops, a', a; solver = PseudoInverse())
+        @inferred spectrum(H, ω_l2, c_ops, a', a; solver = Lanczos())
+    end
+
+    @testset "Verbose mode Lanczos" begin
+        cout = stdout
+        r, w = redirect_stdout()
+        nout = @async read(r, String)
+        spectrum(H, ω_l2, c_ops, a', a; solver = Lanczos(verbose = 2, maxiter = 2, tol = 1e-16));
+        redirect_stdout(cout)
+        close(w)
+        out = fetch(nout)
+        outlines = split(out, '\n', keepempty = false)
+        @test last(outlines) == "spectrum(): Consider increasing maxiter and/or tol"
     end
 
     # tlist and τlist checks

test/ext-test/gpu/cuda_ext.jl

@@ -154,6 +154,29 @@ end
     @test ρ_ss_cpu.data ≈ Array(ρ_ss_gpu_csr.data) atol = 1e-8 * length(ρ_ss_cpu)
 end
 
+@testset "CUDA spectrum" begin
+    N = 10
+    a = cu(destroy(N))
+    H = a' * a
+    c_ops = [sqrt(0.1 * (0.01 + 1)) * a, sqrt(0.1 * (0.01)) * a']
+    solver = Lanczos(steadystate_solver = SteadyStateLinearSolver())
+
+    ω_l = range(0, 3, length = 1000)
+    spec = spectrum(H, ω_l, c_ops, a', a; solver = solver)
+
+    spec = collect(spec)
+    spec = spec ./ maximum(spec)
+
+    test_func = maximum(real.(spec)) * (0.1 / 2)^2 ./ ((ω_l .- 1) .^ 2 .+ (0.1 / 2)^2)
+    idxs = test_func .> 0.05
+    @test sum(abs2.(spec[idxs] .- test_func[idxs])) / sum(abs2.(test_func[idxs])) < 0.01
+
+    # TODO: Fix this
+    # @testset "Type Inference spectrum" begin
+    #     @inferred spectrum(H, ω_l, c_ops, a', a; solver = solver)
+    # end
+end
+
 @testset "CUDA ptrace" begin
     g = fock(2, 1)
     e = fock(2, 0)