From 35aa334390319a78dcf647a13ee6ffc53a5db039 Mon Sep 17 00:00:00 2001 From: Alexander Yavorsky <56161455+aryavorskiy@users.noreply.github.com> Date: Mon, 17 Jul 2023 03:18:45 +0300 Subject: [PATCH] KrylovKit spectrum solver (#367) --------- Co-authored-by: Stefan Krastanov --- Project.toml | 3 +- src/spectralanalysis.jl | 185 ++++++++++++++++++++++++++++++---------- 2 files changed, 143 insertions(+), 45 deletions(-) diff --git a/Project.toml b/Project.toml index 6fdbbe4b..b2a968fa 100644 --- a/Project.toml +++ b/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" diff --git a/src/spectralanalysis.jl b/src/spectralanalysis.jl index be708a8f..64b24194 100644 --- a/src/spectralanalysis.jl +++ b/src/spectralanalysis.jl @@ -1,17 +1,106 @@ 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 @@ -19,11 +108,30 @@ 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)