Lanczos-based spectrum method

[matteosecli]

Checklist

Thank you for contributing to QuantumToolbox.jl! Please make sure you have finished the following tasks before opening the PR.

• Please read Contributing to Quantum Toolbox in Julia.
• Any code changes were done in a way that does not break public API.
• Appropriate tests were added and tested locally by running: make test.
• Any code changes should be julia formatted by running: make format.
• All documents (in docs/ folder) related to code changes were updated and able to build locally by running: make docs.
• (If necessary) the CHANGELOG.md should be updated (regarding to the code changes) and built by running: make changelog.

Request for a review after you have completed all the tasks. If you have not finished them all, you can also open a Draft Pull Request to let the others know this on-going work.

Description

Lanczos-based spectrum method.

Related issues or PRs

Resolves #463.

Additional context

This PR implements the non-symmetric variant of the Hermitian Lanczos method described in Koch2011.

The pseudo-code for the non-symmetric Lanczos algorithm can be found in Algorithm 6.6 of Saad2011.

At each step, the running estimate for the result is updated via a Wallis-Euler recursion to avoid recomputing the whole continued fraction each time, and the convergence is evaluated accordingly.

[matteosecli]

Note: I still have to add tests, evaluate performance, etc.

[matteosecli]

I've been running some benchmarks (on CPU) in the meantime.

I'm just using the example for the vacuum Rabi splitting, with the following script that writes out the runtime (minus compile time) out in the legend:

using QuantumToolbox
using CairoMakie

# %%
N  = 4           # number of cavity fock states
ωc = 1.0 * 2 * π  # cavity frequency
ωa = 1.0 * 2 * π  # atom frequency
g  = 0.1 * 2 * π  # coupling strength
κ  = 0.75         # cavity dissipation rate
γ  = 0.25         # atom dissipation rate

# Jaynes-Cummings Hamiltonian
a  = tensor(destroy(N), qeye(2))
sm = tensor(qeye(N), destroy(2))
H = ωc * a' * a + ωa * sm' * sm + g * (a' * sm + a * sm')

# collapse operators
n_th = 0.25
c_ops = [
    sqrt(κ * (1 + n_th)) * a,
    sqrt(κ *      n_th)  * a',
    sqrt(γ)              * sm,
];

# calculate the correlation function using mesolve, and then FFT to obtain the spectrum.
# Here we need to make sure to evaluate the correlation function for a sufficient long time and
# sufficiently high sampling rate so that FFT captures all the features in the resulting spectrum.
tlist = LinRange(0, 100, 5000)
corr = correlation_2op_1t(H, nothing, tlist, c_ops, a', a; progress_bar = Val(false))
ωlist1, spec1 = spectrum_correlation_fft(tlist, corr)

# calculate the power spectrum using spectrum
# using Exponential Series (default) method
ωlist2 = LinRange(0.25, 1.75, 200) * 2 * π
(spec2, rt2, _, _, _, _, ct2, _) = @timed spectrum(H, ωlist2, c_ops, a', a; solver = ExponentialSeries())

# calculate the power spectrum using spectrum
# using Pseudo-Inverse method
(spec3, rt3, _, _, _, _, ct3, _) = @timed spectrum(H, ωlist2, c_ops, a', a; solver = PseudoInverse())

# calculate the power spectrum using spectrum
# using Lanczos method
(spec4, rt4, _, _, _, _, ct4, _) = @timed spectrum(H, ωlist2, c_ops, a', a; solver = Lanczos())

# plot by CairoMakie.jl
fig = Figure(size = (600, 420))
ax = Axis(fig[1, 1], title = "Vacuum Rabi splitting", xlabel = "Frequency", ylabel = "Emission power spectrum")
lines!(ax, ωlist1 / (2 * π), spec1, label = "mesolve + FFT", linestyle = :solid)
lines!(ax, ωlist2 / (2 * π), spec2, label = string(rt2-ct2)*": "*"Exponential Series", linestyle = :dash)
lines!(ax, ωlist2 / (2 * π), spec3, label = string(rt3-ct3)*": "*"Pseudo-Inverse", linestyle = :dashdot)
lines!(ax, ωlist2 / (2 * π), spec4, label = string(rt4-ct4)*": "*"Lanczos", linestyle = :dot)

xlims!(ax, ωlist2[1] / (2 * π), ωlist2[end] / (2 * π))
axislegend(ax, position = :rb)

fig

# save the figure
#save("vacuumRabiSplitting.png", fig)
#save("vacuumRabiSplitting.pdf", fig)

I've started by increasing the cutoff to (a useless) N=16, just to be able to appreciate the difference:

PseudoInverse is 10 times faster than ExponentialSeries, and in turn Lanczos is 42 (no pun intended) times faster than PseudoInverse.

But it shines even more if you try to increase your frequency resolution, for example by asking for 5000 points instead of 200 (again, useless in this case, but it's just for demonstration purposes):

Now PseudoInverse actually takes twice as much as ExponentialSeries, becoming 22 times slower, while Lanczos slows down only by a factor 5, thus being almost 200 times faster than PseudoInverse and more than 90 times faster than ExponentialSeries.

If you then try to increase the cutoff too, to e.g. an even more useless N=32, Lanczos becomes more than 3500 times faster than ExponentialSeries and 730 times faster than PseudoInverse:

In summary, at least according to these initial tests, while ExponentialSeries struggles with higher cutoffs and PseudoInverse struggles with higher frequency resolution (because it has to compute a pseudo-inverse for each frequency point), Lanczos seems to offer a competitive advantage in both scenarios.

[albertomercurio]

Nice! I will give a look later.

[ytdHuang]

Thank you @matteosecli !

The speed up is very impressing !

Although you still mark this PR as draft, I just have a few minor comments after I take the first look.

[matteosecli]

I should have implemented GPU support for the Lanczos solver, but testing it as-is (and as a matter of fact, any solver) doesn't work, and you'd have to change this line:

    ρss = mat2vec(steadystate(L))

to

    ρss = mat2vec(steadystate(L; solver = SteadyStateLinearSolver()))

to make it work (and other solvers like PseudoInverse too).

A simple example taken from the README:

using QuantumToolbox
using CUDA
CUDA.allowscalar(false) # Avoid unexpected scalar indexing

N = 20 # cutoff of the Hilbert space dimension
ω = 1.0 # frequency of the harmonic oscillator
γ = 0.1 # damping rate

a_gpu = cu(destroy(N)) # The only difference in the code is the cu() function

H_gpu = ω * a_gpu' * a_gpu

ψ0_gpu = cu(fock(N, 3))

c_ops = [sqrt(γ) * a_gpu]
e_ops = [a_gpu' * a_gpu]

L = liouvillian(H_gpu, c_ops)

ρss = steadystate(L) # doesn't work
# ρss = steadystate(L, solver = SteadyStateLinearSolver()) # this works

I seem to remember that steadystate used to work, maybe something broke in the meantime?

[matteosecli]

Btw maybe it's worth keeping kwargs to let the user specify the steadystate solver, in this case? I'm not sure about the other solvers, but this is basically effortlessly GPU-compatible (it's just spmat/vec products), as long as the steady-state calculation is also GPU-compatible.

[codecov]

Codecov Report

Attention: Patch coverage is 98.71795% with 1 line in your changes missing coverage. Please review.

Project coverage is 93.00%. Comparing base (a1d1627) to head (a6abdbd).
Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
src/spectrum.jl	98.71%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #476      +/-   ##
==========================================
- Coverage   94.08%   93.00%   -1.09%     
==========================================
  Files          49       49              
  Lines        3042     3201     +159     
==========================================
+ Hits         2862     2977     +115     
- Misses        180      224      +44     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[matteosecli]

Btw maybe it's worth keeping kwargs to let the user specify the steadystate solver, in this case? I'm not sure about the other solvers, but this is basically effortlessly GPU-compatible (it's just spmat/vec products), as long as the steady-state calculation is also GPU-compatible.

Or, the steady-state solver could be an additional field in the Lanczos struct.

[albertomercurio]

Or, the steady-state solver could be an additional field in the Lanczos struct.

I think this option is better. Better to keep kwargs for other options, e.g. for LinearSolve.

[matteosecli]

Or, the steady-state solver could be an additional field in the Lanczos struct.

I think this option is better. Better to keep kwargs for other options, e.g. for LinearSolve.

I've done that, although I had to modify the include order in QuantumToolbox.jl so that the SteadyStateSolver abstract type defined in steadystate.jl is available to the Lanczos solver in spectrum.jl.

With this change, something like

solver = Lanczos(steadystate_solver = SteadyStateLinearSolver())
spec = spectrum(H, ωlist, c_ops, A, B; solver = solver)

now works if H, c_ops, A, and B are defined on GPU.

[albertomercurio]

Could you add the changelogs? Remember to do make changelog then

[matteosecli]

Could you add the changelogs? Remember to do make changelog then

Done!

[ytdHuang]

Somehow I think the previous two comments (merging from the main branch) are not done in a correct way...

Those changes are now overwritten again in this PR...

[matteosecli]

Somehow I think the previous two comments (merging from the main branch) are not done in a correct way...

Those changes are now overwritten again in this PR...

Hi @ytdHuang, do you get a merge conflict or something similar in GitHub’s interface? I’ve just rebased my branch against the upstream changes, rather than merging the changes in a separate commit. This was to ensure that my changes were compatible with the updated upstream, especially regarding tests. Is there something missing?

In any case, I can always roll back the rebase and do a merge commit or nothing at all.

[ytdHuang]

@matteosecli

After you force-pushed it, now it looks fine !

The code doesn't seem to have any big problems. I'm gonna start the CI pipeline.