Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

KrylovKit spectrum solver #367

Merged
merged 8 commits into from
Jul 17, 2023
Merged

KrylovKit spectrum solver #367

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
95a278c
Implement DiagStrategy struct to make fine-tuning possible
aryavorskiy Jun 29, 2023
b328991
Final fixes
aryavorskiy Jun 29, 2023
ef8bc76
Untrack Manifest.toml
aryavorskiy Jun 29, 2023
70c0f9d
Better dispatch and all
aryavorskiy Jul 12, 2023
f13946c
Fix tests erroring on info keyword
aryavorskiy Jul 12, 2023
ef04421
Merge branch 'master' into master
Krastanov Jul 16, 2023
f029487
bump version number, do not export new types, turn info kw to `@debug`
Krastanov Jul 16, 2023
b9dca8c
fix the info keyword
Krastanov Jul 16, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
name = "QuantumOptics"
uuid = "6e0679c1-51ea-5a7c-ac74-d61b76210b0c"
version = "v1.0.10"
version = "1.0.10"

[deps]
Arpack = "7d9fca2a-8960-54d3-9f78-7d1dccf2cb97"
Expand All @@ -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"
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@david-pl , I am in favor of adding KrylovKit as a direct dependency -- it is a very mature and well supported Julia library, and an important part of the iterative linear algebra ecosystem. And we already depend on it because of SciML https://juliahub.com/ui/Packages/KrylovKit/JPeTr/0.6.0?page=2

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Krastanov fine by me! @aryavorskiy Thanks for the work here.

LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
LinearMaps = "7a12625a-238d-50fd-b39a-03d52299707e"
OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
Expand Down
1 change: 1 addition & 0 deletions src/QuantumOptics.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ using SparseArrays, LinearAlgebra
export
ylm,
eigenstates, eigenenergies, simdiag,
LapackDiag, ArpackDiag, KrylovDiag,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Given that the library mostly automatically decides which method to use, I think it would be best to not export these. If nonetheless we want them exported, they should at least have a docstring each.

timeevolution, diagonaljumps, @skiptimechecks,
steadystate,
timecorrelations,
Expand Down
78 changes: 62 additions & 16 deletions src/spectralanalysis.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,27 @@
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."

mutable struct DiagStrategy{DT}
n::Int
v0::Union{Nothing, AbstractVector}
end
DiagStrategy{T}(n::Int) where T = DiagStrategy{T}(n, nothing)
const LapackDiag = DiagStrategy{:lapack}
const ArpackDiag = DiagStrategy{:arpack}
const KrylovDiag = DiagStrategy{:krylov}
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could you consider structuring this a bit differently:

abstract type DiagStrategy end
type LapackDiag{T} <: DiagStrategy
    n::Int
    v0::T
end

My main reasons for this suggestion:

  • It is a bit closer in style to how SciML treats their solver types
  • It is a bit closer in style to Holy Traits https://invenia.github.io/blog/2019/11/06/julialang-features-part-2/
  • It parameterizes v0 instead of leaving it of abstract type, which avoids boxing (but to be fair, you can do that to your version either way)
  • It is not mutable anymore

Please feel free to voice dissent to this suggestion (as always!)

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was wrong on the last one: it still needs to be mutable because you might change n as I see below

Copy link
Contributor Author

@aryavorskiy aryavorskiy Jul 8, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree, this is a better way. Initially I used this structure to make custom subtyping way easier, so that LapackDiag, ArpackDiag and KrylovDiag inherited the same internal structure.
Probably a good way to avoid the DiagStrategy struct being mutable is making the DiagStrategy constructor process the keyword arguments from the eigenstates method. This explanation is probably kinda messy, will clarify in the nearest commit.


arithmetic_unary_error = QuantumOpticsBase.arithmetic_unary_error
DiagStrategy(op::AbstractOperator) = arithmetic_unary_error("DiagStrategy", op)
function DiagStrategy(op::DataOperator)
basis(op) # Checks basis match
DiagStrategy(op.data)
end
DiagStrategy(m::Matrix) = LapackDiag(size(m)[1])
DiagStrategy(::SparseMatrixCSC) = KrylovDiag(6)
DiagStrategy(m::AbstractMatrix) = ArgumentError("Cannot detect DiagStrategy for array type $(typeof(m))")
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you elaborate on this error message. For the last year or more SciML has undertaken a really great process of writing as helpful as possible error messages. For instance, could you add something along the lines of This error probably is raised because you have tried to ... To fix it, consider doing ...


"""
eigenstates(op::AbstractOperator[, n::Int; warning=true])

Expand All @@ -20,10 +40,10 @@ 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`.
"""
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
Expand All @@ -33,26 +53,53 @@ 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,
function eigenstates(op::Operator, ds::ArpackDiag; warning::Bool=true,
info::Bool=true, kwargs...)
b = basis(op)
# TODO: Change to sparese-Hermitian specific algorithm if more efficient
# TODO: Change to sparse-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...)
if ds.v0 === nothing
D, V = eigs(op.data; which=:SR, nev=ds.n, kwargs...)
else
D, V = eigs(op.data; which=:SR, nev=ds.n, v0=ds.v0, kwargs...)
end
states = [Ket(b, V[:, k]) for k=1:length(D)]
D, states
end

function eigenstates(op::Operator, ds::KrylovDiag; warning::Bool=true,
info::Bool=true, kwargs...)
b = basis(op)
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.")
if ds.v0 === nothing
D, Vs = eigsolve(op.data, ds.n, :SR; kwargs...)
else
D, Vs = eigsolve(op.data, ds.v0, ds.n, :SR; kwargs...)
end
states = [Ket(b, Vs[k]) for k=1:ds.n]
D[1:ds.n], states
end

function eigenstates(op::AbstractOperator, n::Union{Int,Nothing}=nothing,
v0::Union{AbstractVector,Nothing}=nothing; warning=true, kw...)
ds = DiagStrategy(op)
n !== nothing && (ds.n = n)
v0 !== nothing && (ds.v0 = v0)
eigenstates(op, ds; warning=warning, kw...)
end

"""
eigenenergies(op::AbstractOperator[, n::Int; warning=true])
Expand All @@ -67,29 +114,28 @@ More details can be found at
If the given operator is non-hermitian a warning is given. This behavior
can be turned off using the keyword `warning=false`.
"""
function eigenenergies(op::DenseOpType, n::Int=length(basis(op)); warning=true)
b = basis(op)
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)
eigenenergies(op::Operator, ds::DiagStrategy; kwargs...) = eigenstates(op, ds; kwargs...)[1]

function eigenenergies(op::AbstractOperator, n::Union{Int,Nothing}=nothing; kw...)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This seems like the main entry to most of the code paths. Could you consider reorganizing the docstrings a bit, because now they seem to be at least partially outdated. A few options:

  • move the information from most docstrings to only here and explain how the dispatch works. Then add very short docstrings to the other methods
  • declare a function eigenenergies end without a method, and add the docstring to it.

Could you also add a brief sentence to the appropriate docstring explaining what is happening and that the library tries to detect the best choice on its own. And that the choice can be overwritten by using the appropriate type.

ds = DiagStrategy(op)
n !== nothing && (ds.n = n)
eigenenergies(op, ds; kw...)
end

"""
simdiag(ops; atol, rtol)
Expand Down