KrylovKit spectrum solver (#367)

---------

Co-authored-by: Stefan Krastanov <[email protected]>

Project.toml

@@ -1,6 +1,6 @@
 name = "QuantumOptics"
 uuid = "6e0679c1-51ea-5a7c-ac74-d61b76210b0c"
-version = "v1.0.11"
+version = "1.0.12"
 
 [deps]
 Arpack = "7d9fca2a-8960-54d3-9f78-7d1dccf2cb97"
@@ -9,6 +9,7 @@ DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
+KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearMaps = "7a12625a-238d-50fd-b39a-03d52299707e"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"

src/spectralanalysis.jl

@@ -1,29 +1,137 @@
 using Arpack
+import KrylovKit: eigsolve
 
 const nonhermitian_warning = "The given operator is not hermitian. If this is due to a numerical error make the operator hermitian first by calculating (x+dagger(x))/2 first."
 
 """
-    eigenstates(op::AbstractOperator[, n::Int; warning=true])
+    abstract type DiagStrategy
 
-Calculate the lowest n eigenvalues and their corresponding eigenstates.
+Represents an algorithm used to find eigenvalues and eigenvectors of some operator.
+Subtypes of this abstract type correspond to concrete routines. See `LapackDiag`,
+`KrylovDiag` for more info.
+"""
+abstract type DiagStrategy end
+
+"""
+    LapackDiag <: DiagStrategy
+
+Represents the `LinearArgebra.eigen` diagonalization routine.
+The only parameter `n` represents the number of (lowest) eigenvectors.
+"""
+struct LapackDiag <: DiagStrategy
+    n::Int
+end
+
+"""
+    KrylovDiag <: DiagStrategy
+
+Represents the `KrylovKit.eigsolve` routine. Implements the Lanczos & Arnoldi algorithms.
+"""
+struct KrylovDiag{VT} <: DiagStrategy
+    n::Int
+    v0::VT
+    krylovdim::Int
+end
+"""
+    KrylovDiag(n::Int, [v0=nothing, krylovdim::Int=n + 30])
+
+Parameters:
+- `n`: The number of eigenvectors to find
+- `v0`: The starting vector. By default it is `nothing`, which means it will be a random dense
+`Vector`. This will not work for non-trivial array types like from `CUDA.jl`, so you might want
+to define a new method for the `QuantumOptics.get_starting_vector` function.
+- `krylovdim`: The upper bound for dimenstion count of the emerging Krylov space.
+"""
+KrylovDiag(n::Int, v0=nothing) = KrylovDiag(n, v0, n + 30)
+Base.print(io::IO, kds::KrylovDiag) =
+    print(io, "KrylovDiag($(kds.n))")
+
+arithmetic_unary_error = QuantumOpticsBase.arithmetic_unary_error
+"""
+    detect_diagstrategy(op::Operator; kw...)
+
+Find a `DiagStrategy` for the given operator; processes the `kw` keyword arguments
+and automatically sets parameters of the resulting `DiagStrategy`object.
+Returns a tuple of the `DiagStrategy` and unprocessed keyword arguments from `kw`.
+"""
+function detect_diagstrategy(op::DataOperator; kw...)
+    QuantumOpticsBase.check_samebases(op)
+    detect_diagstrategy(op.data; kw...)
+end
+detect_diagstrategy(op::AbstractOperator; kw...) = arithmetic_unary_error("detect_diagstrategy", op)
+
+"""
+    get_starting_vector(m::AbstractMatrix)
+
+Generate a default starting vector for Arnoldi-like iterative methods for matrix `m`.
+"""
+get_starting_vector(::SparseMatrixCSC) = nothing
+function detect_diagstrategy(m::AbstractSparseMatrix; kw...)
+    if get(kw, :info, true)
+        @info "Defaulting to sparse diagonalization for sparse operator. If storing the full operator is possible, it might be faster to do `eigenstates(dense(op))`. Set `info=false` to turn off this message."
+    end
+    nev = get(kw, :n, 6)
+    v0 = get(kw, :v0, get_starting_vector(m))
+    krylovdim = get(kw, :krylovdim, nev + 30)
+    new_kw = Base.structdiff(values(kw), NamedTuple{(:n, :v0, :krylovdim, :info)})
+    return KrylovDiag(nev, v0, krylovdim), new_kw
+end
+function detect_diagstrategy(m::Matrix; kw...)
+    nev = get(kw, :n, size(m)[1])
+    new_kw = Base.structdiff(values(kw), NamedTuple{(:n, :info)})
+    return LapackDiag(nev), new_kw
+end
+"""
+    detect_diagstrategy(m::AbstractMatrix; kw...)
 
-This is just a thin wrapper around julia's `eigen` and `eigs` functions. Which
-of them is used depends on the type of the given operator. If more control
-about the way the calculation is done is needed, use the functions directly.
-More details can be found at
-[http://docs.julialang.org/en/stable/stdlib/linalg/].
+Same as above, but dispatches on different internal array types.
+"""
+detect_diagstrategy(m::T; _...) where T<:AbstractMatrix = throw(ArgumentError(
+    """Cannot detect DiagStrategy for array type $(typeof(m)).
+    Consider defining `QuantumOptics.detect_diagstrategy(::$T; kw...)` method.
+    Refer to `QuantumOptics.detect_diagstrategy` docstring for more info."""))
+
+"""
+    eigenstates(op::Operator[, n::Int; warning=true, kw...])
+
+Calculate the lowest n eigenvalues and their corresponding eigenstates. By default `n` is
+equal to the matrix size for dense matrices; for sparse matrices the default value is 6.
+
+This is just a thin wrapper around julia's `LinearArgebra.eigen` and `KrylovKit.eigsolve`
+functions. Which of them is used depends on the type of the given operator. If more control
+about the way the calculation is done is needed, use the method instance with `DiagStrategy`
+(see below).
 
 NOTE: Especially for small systems full diagonalization with Julia's `eigen`
 function is often more desirable. You can convert a sparse operator `A` to a
 dense one using `dense(A)`.
 
 If the given operator is non-hermitian a warning is given. This behavior
 can be turned off using the keyword `warning=false`.
+
+## Optional arguments
+- `n`: It can be a keyword argument too!
+- `v0`: The starting vector for Arnoldi-like iterative methods.
+- `krylovdim`: The upper bound for dimenstion count of the emerging Krylov space.
+"""
+function eigenstates(op::AbstractOperator; kw...)
+    ds, kwargs_rem = detect_diagstrategy(op; kw...)
+    eigenstates(op, ds; kwargs_rem...)
+end
+eigenstates(op::AbstractOperator, n::Int; warning=true, kw...) =
+    eigenstates(op; warning=warning, kw..., n=n)
+
+"""
+    eigenstates(op::Operator, ds::DiagStrategy[; warning=true, kw...])
+
+Calculate the lowest eigenvalues and their corresponding eigenstates of the `op` operator
+using the `ds` diagonalization strategy. The `kw...` arguments can be passed to the exact
+function that does the diagonalization (like `KrylovKit.eigsolve`).
 """
-function eigenstates(op::DenseOpType, n::Int=length(basis(op)); warning=true)
+function eigenstates(op::Operator, ds::LapackDiag; warning=true)
     b = basis(op)
     if ishermitian(op)
-        D, V = eigen(Hermitian(op.data), 1:n)
+        D, V = eigen(Hermitian(op.data), 1:ds.n)
         states = [Ket(b, V[:, k]) for k=1:length(D)]
         return D, states
     else
@@ -33,63 +141,52 @@ function eigenstates(op::DenseOpType, n::Int=length(basis(op)); warning=true)
         perm = sortperm(D, by=real)
         permute!(D, perm)
         permute!(states, perm)
-        return D[1:n], states[1:n]
+        return D[1:ds.n], states[1:ds.n]
     end
 end
 
-"""
-For sparse operators by default it only returns the 6 lowest eigenvalues.
-"""
-function eigenstates(op::SparseOpType, n::Int=6; warning::Bool=true,
-        info::Bool=true, kwargs...)
+function eigenstates(op::Operator, ds::KrylovDiag; warning::Bool=true, kwargs...)
     b = basis(op)
-    # TODO: Change to sparese-Hermitian specific algorithm if more efficient
     ishermitian(op) || (warning && @warn(nonhermitian_warning))
-    info && println("INFO: Defaulting to sparse diagonalization.
-        If storing the full operator is possible, it might be faster to do
-        eigenstates(dense(op)). Set info=false to turn off this message.")
-    D, V = eigs(op.data; which=:SR, nev=n, kwargs...)
-    states = [Ket(b, V[:, k]) for k=1:length(D)]
-    D, states
+    if ds.v0 === nothing
+        D, Vs = eigsolve(op.data, ds.n, :SR; krylovdim = ds.krylovdim, kwargs...)
+    else
+        D, Vs = eigsolve(op.data, ds.v0, ds.n, :SR; krylovdim = ds.krylovdim, kwargs...)
+    end
+    states = [Ket(b, Vs[k]) for k=1:ds.n]
+    D[1:ds.n], states
 end
 
-
 """
-    eigenenergies(op::AbstractOperator[, n::Int; warning=true])
+    eigenenergies(op::AbstractOperator[, n::Int; warning=true, kwargs...])
 
-Calculate the lowest n eigenvalues.
-
-This is just a thin wrapper around julia's `eigvals`. If more control
-about the way the calculation is done is needed, use the function directly.
-More details can be found at
-[http://docs.julialang.org/en/stable/stdlib/linalg/].
+Calculate the lowest n eigenvalues of given operator.
 
 If the given operator is non-hermitian a warning is given. This behavior
 can be turned off using the keyword `warning=false`.
+
+See `eigenstates` for more info.
 """
-function eigenenergies(op::DenseOpType, n::Int=length(basis(op)); warning=true)
-    b = basis(op)
+function eigenenergies(op::AbstractOperator; kw...)
+    ds, kw_rem = detect_diagstrategy(op; kw...)
+    eigenenergies(op, ds; kw_rem...)
+end
+eigenenergies(op::AbstractOperator, n::Int; kw...) = eigenenergies(op; kw..., n=n)
+
+function eigenenergies(op::Operator, ds::LapackDiag; warning=true)
     if ishermitian(op)
-        D = eigvals(Hermitian(op.data), 1:n)
+        D = eigvals(Hermitian(op.data), 1:ds.n)
         return D
     else
         warning && @warn(nonhermitian_warning)
         D = eigvals(op.data)
         sort!(D, by=real)
-        return D[1:n]
+        return D[1:ds.n]
     end
 end
 
-"""
-For sparse operators by default it only returns the 6 lowest eigenvalues.
-"""
-eigenenergies(op::SparseOpType, n::Int=6; kwargs...) = eigenstates(op, n; kwargs...)[1]
-
-
-arithmetic_unary_error = QuantumOpticsBase.arithmetic_unary_error
-eigenstates(op::AbstractOperator, n::Int=0) = arithmetic_unary_error("eigenstates", op)
-eigenenergies(op::AbstractOperator, n::Int=0) = arithmetic_unary_error("eigenenergies", op)
-
+# Call eigenstates
+eigenenergies(op::Operator, ds::DiagStrategy; kwargs...) = eigenstates(op, ds; kwargs...)[1]
 
 """
     simdiag(ops; atol, rtol)