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.

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

KrylovKit spectrum solver #367

Merged
merged 8 commits into from
Jul 17, 2023
Merged

KrylovKit spectrum solver #367

merged 8 commits into from
Jul 17, 2023

Conversation

aryavorskiy
Copy link
Contributor

@aryavorskiy aryavorskiy commented Jun 29, 2023

#364
The Arnoldi algorithm implementation from the KrylovKit.jl package seems to be much more performant than the ARPACK one.
Also a DiagStrategy trait was implemented to make selecting the correct method on the fly (or defining a new one) more flexible

@codecov
Copy link

codecov bot commented Jul 5, 2023

Codecov Report

Merging #367 (b9dca8c) into master (eff22f1) will decrease coverage by 0.30%.
The diff coverage is 88.37%.

@@            Coverage Diff             @@
##           master     #367      +/-   ##
==========================================
- Coverage   98.11%   97.82%   -0.30%     
==========================================
  Files          18       18              
  Lines        1488     1516      +28     
==========================================
+ Hits         1460     1483      +23     
- Misses         28       33       +5     
Impacted Files Coverage Δ
src/spectralanalysis.jl 93.42% <88.37%> (-6.58%) ⬇️

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

Copy link
Collaborator

@Krastanov Krastanov left a comment

Choose a reason for hiding this comment

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

Thank you for taking care of all of this, I greatly appreciate the contribution! Attached here are a few comments and suggestions. Let me know what you think and how you would like to proceed.

@@ -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.

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

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.

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 ...


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.

@@ -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.

@aryavorskiy
Copy link
Contributor Author

aryavorskiy commented Jul 8, 2023

@Krastanov thanks for the review! I will implement these requests by Tuesday, 11th, I think. I also have some new ideas:

  • Maybe remove the Arpack diagonalization routine? KrylovKit is much faster (also it is universal and can be applied even to CUDA arrays). Benchmarks attached:
julia> using Arpack, KrylovKit, SparseArrays, BenchmarkTools

julia> A = sprand(1000, 1000, 0.001); v0 = rand(1000);

julia> @benchmark eigs(A, v0=v0, nev=5, which=:SR)  # Arpack
BenchmarkTools.Trial: 589 samples with 1 evaluation.
 Range (min … max):  7.974 ms …  12.672 ms  ┊ GC (min … max): 0.00% … 33.56%
 Time  (median):     8.431 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   8.492 ms ± 514.293 μs  ┊ GC (mean ± σ):  0.64% ±  3.72%

     ▄▆▆█▅▃
  ▃▄▇██████▆▄▄▃▃▂▁▂▂▂▁▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▂▁▁▁▁▁▁▁▁▂▁▁▂ ▃
  7.97 ms         Histogram: frequency by time        12.3 ms <

 Memory estimate: 702.83 KiB, allocs estimate: 949.

julia> @benchmark eigsolve(A, v0, 5, :SR)  # KrylovKit
BenchmarkTools.Trial: 2167 samples with 1 evaluation.
 Range (min … max):  1.931 ms …   6.259 ms  ┊ GC (min … max): 0.00% … 50.37%
 Time  (median):     2.162 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   2.302 ms ± 457.034 μs  ┊ GC (mean ± σ):  2.63% ±  8.01%

    ▄█
  ▃▇███▇▇▆▆▅▄▄▃▃▄▄▃▃▃▃▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▁▁▁▁▂▁▁▂▁▂▂▂▂▂▂▂▂▂▂▁▂▂▂▂ ▃
  1.93 ms         Histogram: frequency by time        4.52 ms <
  • Maybe create an EigenSystem struct that will store computed eigenvalues and eigevectors? In this format they can be easily dumped to a file or used to construct a density matrix or Green's function.
  • Maybe move this functionality to QuantumOpticsBase? Other quantum physics packages depending on it might want to use this as well. I am myself now in progress of creating such one :-) Or is this a generally bad idea?

@Krastanov
Copy link
Collaborator

  • It seems relatively conclusive that KrylovKit is better. I am comfortable with the idea of removing arpack. Here are other comments from the community https://discourse.julialang.org/t/suggestions-needed-diagonalizing-large-hermitian-sparse-matrix/96580
  • I personally do not think an EigenSystem type would be better than simply an array or tuple of results. My reasoning is that an extra type makes sense only if we need some API on top of it, i.e. a bunch of useful methods that are more ergonomic to use with EigenSystem than with a vector or tuple of results. For the serialization question: I think a Tuple is as easy or easier to serialize than a custom type. Do not hesitate to elaborate on an API for EigenSystem that would be useful for you -- that would be a good reason to reverse my opinion.
  • Could you share a bit more about the package you are constructing? Separation like that might make sense (and has happened before):
    • QuantumOpticsBase was split out from QuantumOptics, so that there is a base library for data structures and a separate library for dynamics over these datastructures (and the eigensolver is kinda in between, but in my mental model a bit closer to QuantumOptics than to QuantumOpticsBase)
    • QuantumInterface was split away from QuantumOpticsBase, so that the abstract types and bases are in an interface library that can be reused by simulators that do not care about state vectors and density matrices and master equations. E.g. thanks to that split, the QuantumClifford library for simulating Clifford circuits and the QuantumSymbolics library for symbolic expressions can share an API with QuantumOptics without depending on it
    • As you are contributing more to this ecosystem, we should consider how to make it easier for you, so some rearrangement might make sense

@aryavorskiy
Copy link
Contributor Author

Done.
One final idea: I don't like how two separate keywords are used for nonhermitian warning and info message about the diagonalization method. My proposal is using a verbose keyword for both. This change is breaking, but only a little, so I don't know if it is a good idea.
@Krastanov about my package: it is called LatticeModels.jl, it is designed to perform various computations on tight-binding lattices. I decided to use the bases & operators infrastructure from QuantumOpticsBase when I discovered that it is a separate package, and am currently working on its integration.
My package is still far from finished, I will definitely rename it sometime, the docs are a little outdated, but most important functionality is already implemented.

@Krastanov
Copy link
Collaborator

@aryavorskiy , there was a minor merge conflict that I resolved, but that ended up pushing to your master branch. Be sure to do a pull before you perform other edits. Also, it is a good idea to do your development not to the master branch of your repo, but rather on a separate development branch -- it makes other modifications down the line easier.

I will go through the current changes and review. Thanks for improving this part of the library!

@Krastanov
Copy link
Collaborator

I made two small changes to your PR:

  • The strategies are not exported anymore. I try to be very minimal with what is exported, because the convention is that exported symbols will be supported in the future, which makes refactoring more difficult.
  • I turned your info keyword argument in a @debug statement as this is a more universal julia-wide way to track this type of user messages. That way any function that provides such stdout information can be controlled in the same way (instead of the user needing to learn the conventions of every single library)

I will wait on the tests, give it a final read through, and I should be able to merge it and release it today.

@Krastanov Krastanov mentioned this pull request Jul 16, 2023
@Krastanov
Copy link
Collaborator

My bad, it seems info was something already done to this library so you were just following the established convention. Give me a minute to fix this.

@Krastanov
Copy link
Collaborator

It seems the info string was meant to specifically warn about defaulting to the sparse eigsolver when the dense one might be better. I moved it to that location and cleaned it up a bit.

@Krastanov Krastanov merged commit 35aa334 into qojulia:master Jul 17, 2023
5 of 7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants