using KrylovKit.exponentiate with ITensors

[orialb]

Hi,

As a continuation of the discussion in #98 regarding using KrylovKit, I just wanted to update that I implemented TDVP using ITensors + KrylovKit.exponentiate (here ) and it seems to work nicely.

It works almost smoothly out of the box, the only thing missing is implementation of similar(A::ITensor, T) when T is a different type than the storage type of A (to create a similar ComplexF64 ITensor from a Float64 ITensor ). In the meantime I just make sure the MPS is complex before starting time-evolution.

I did some benchmarks vs. the recently added applyExp function in the C++ version of ITensor, and as far as I can tell exponentiate is actually significantly faster.
Using the Julia code below and the C++ code given here, I get:

• ITensors.jl + KrylovKit - 22.2ms
• ITensor C++ - 32ms

Both ITensor and Julia (v1.2) were compiled with MKL, and I used 1 MKL thread.
I wouldn't expect to see such a big difference between the implementations, so I hope I benchmarked it correctly (otherwise maybe there is some problem in applyExp).

(Note that I also compared my actual TDVP code vs. the code from ITensor/tdvp branch and also there I see a similar performance advantage to Julia)

p.s - Just wanted to add that in general it was super easy to implement TDVP with ITensors.jl, it took me something like 20 mins to write. So, great job with ITensors.jl and also great job by @Jutho with KrylovKit!

using ITensors, BenchmarkTools, KrylovKit

struct LocalOp
    A::ITensor
    B::ITensor
end

(L::LocalOp)(t::ITensor) = noprime(L.A*t*L.B)

m=100
i,j = Index(m,"i"), Index(m,"j")
v = randomITensor(i,j)
A= randomITensor(i',i)
B = randomITensor(j,j')
M =LocalOp(A/norm(A),B/norm(B) )

@btime exponentiate($M,0.1,$v,tol=1e-12)

[orialb]

I wouldn't expect to see such a big difference between the implementations, so I hope I benchmarked it correctly

One possibility for a mistake in the benchmark is that I am not sure that the tolerance parameter tol in exponentiate is equivalent to the ErrGoal parameter in applyExp. So comparing runs where those two parameters are equal might not be the right thing to do.

[mtfishman]

Thanks for the update! It would be interesting to do an in-depth comparison of the C++ implementation of applyExp and the KrylovKit version. @mingruyang may be interested in this discussion (since she implemented the applyExp and TDVP code in C++ ITensor).

I think the best comparison would be the amount of time to reach a certain error from the exact application of the exponential, since as you say the convergence criteria may be different. Jutho's version definitely looked a bit more sophisticated. One main difference I noticed was that he implemented some sort of scheme for timestep adjustment that is not in the C++ version, so that may account for the faster convergence.

It would be good to add some of the functionality you are using to write TDVP into ITensors.jl. For example, we would be happy to have wrappers of some of the KrylovKit methods (like the exponentiate method you show). Also, the similar function you mention sounds useful as well (something like similar(::ITensor,::Type{<:Number}).

[orialb]

we would be happy to have wrappers of some of the KrylovKit methods (like the exponentiate method you show).

I'm not sure what exactly a wrapper adds here, it already just works with exponetiate(PH::ProjMPO, dt, wf::ITensor,...).
Is the motivation for a wrapper that users could use it without using KrylovKit? This could probably be done just by reexporting exponentiate (or is it somehow not a nice thing to do when using someone else's package?).
Or would you like the function to be named applyExp ?

Also, the similar function you mention sounds useful as well (something like similar(::ITensor,::Type{<:Number}).

I could take a look at adding this, but probably I should wait for #109 ?

The third extra feature needed is the ability to change the size of the projected MPO.
In my code I just defined :

singlesite!(PH::ProjMPO) = (PH.nsite = 1)
twosite!(PH::ProjMPO) = (PH.nsite = 2)

then one has to still call position! afterwards to get a single-site ProjMPO if it was previously a two-site ProjMPO.
One could also define singlesite!,twosite! to call position!, but it is not clear what should be the new position of the projected MPO (e.g. if current position of two-site ProjMPO is b should the position of single-site ProjMPO be b or b+1), we can have the user provide it as an argument. Any thoughts regarding this functionality?

[orialb]

I'm not sure what exactly a wrapper adds here

Ok, I guess another motivation for a wrapper would be the option to later change the exponentiation implementation (if it is needed for some reason) in a way which is more transparent to users.

[mtfishman]

Right after I posted I looked at the code and realized you were calling exponentiate directly from KrylovKit. I am actually surprised that works right now, I didn't realize we provided all of the proper overloads of ITensors in order for that to work.

The only reason for a wrapper would be if we wanted a different interface from KrylovKit, but the interface looks reasonable now.

Regarding setting the number of sites of ProjMPO, that is essentially the same interface as the C++ version: https://github.com/ITensor/TDVP/blob/master/tdvp.h#L220. Instead of two separate functions, we could define setNumCenter!(::ProjMPO,::Int). That would be a good thing to add to the interface.

[mtfishman]

Yes that is also a good reason for a wrapper. I think I prefer the name applyExp since it is a bit clearer what the function is actually doing.

[mingruyang]

@mtfishman The ITensor C++ has a similar version of time step adjusting just like what @Jutho did in the KrylovKit.jl. But that doesn't make it faster. The use is to make the algorithm converge to the desired error when the krylov subspace is not large enough for a desired t.

In addition to make the error match, I think there are still several things to check in the c++ version:

1. the way to calculte anorm. In https://github.com/Jutho/KrylovKit.jl/blob/62289670681a3bfdd858c9341108df046124795d/src/matrixfun/expintegrator.jl#L83, it use the norm(Amu_0), in ITensor c++, I calculate (mu_0, Amu_0), which might be slower.

2. Actually in ITensor C++, the applyExp not only calculate the matrix exponential application but also the "energy" after applying the matrix. This can lead to extra cost.

3. I was wondering the julia implementaion of LinearAlgebra.exp!. What is the algorithm they are using. If it is Chebyshev or the pade for symmetric matrix (so that a faster LU decomposition is used), then it can be much faster than the current expPade in ITensor c++ which is used for a general dense matrix.

4. I didn't like the way the current multTElts function in algs_impl.h did. I think there should be a much faster way to do it, like using the lapack function axpy directly? Also I would like to see the detailed implementation of mul! function in https://github.com/Jutho/KrylovKit.jl/blob/62289670681a3bfdd858c9341108df046124795d/src/matrixfun/expintegrator.jl#L156

5. In the ITensor C++, inside the function applyExp I calculate something related to the condition number which is not used but maybe I want to output in the future. This can lead to a tiny extra cost.

[mingruyang]

6. what is the default algorithm to do orthogonalization in KrylovKit.jl? Looks like they check if the AbstractMatrix is hermitian or not? Using Lanczos can be significantly faster than Arnoldi...

[mtfishman]

Thanks for chiming in, Mingru.

Here is the method Julia uses for the exponential: https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/#footnote-H05 which looks like a modification of the Pade approximation with some scaling to decrease the number of matrix multiplications required for convergence.

In general, these Krylov-based methods should be dominated by the number of matrix-vector multiplications, so I would expect a discrepancy as large as Ori is seeing is just a matter of the C++ version doing more matrix-vector multiplications (either from a difference in the algorithm, a difference in the convergence criteria, or both). I think the first thing to do is to make sure we are doing a fair comparison between the two (i.e. compare time to reach a fixed error).

[mtfishman]

I would guess Lanczos is used if the Matrix is Hermitian. We should make sure to compare against the C++ version using Lanczos as well.

Anyway, it is good to have comparisons like this so we can find performance discrepancies on the C++ side!

[mingruyang]

I think for a general dense matrix, the method Julia use is just the same as us (we also use the scaling and squaring technique). But the implementation details like how do they use the LU decomposition exactly might be different. Also, for hermitian matrix, they do eigendecomposition instead of Pade, and that might be faster than Pade?

I agree that the part about the dense matrix might just have tiny cost anyway.

Also I didn't open the access for the user to set the break tolerance of the orthogonalization process. I just set it to be 1E-14. I didn't check the settings of Jutho yet.

[mingruyang]

Although some cost is tiny, when the matrix is as small as 100*100 and the running time is just of order "ms", they can look "significant".

[mtfishman]

That part of the docs says that Julia uses eigendecompositions for Hermitian/Symmetric matrices.

I totally agree about the matrix size, a quick check that the discrepancy is due to things besides matrix-vector multiplications would be to increase the matrix size and see if the difference starts to disappear.

[orialb]

what is the default algorithm to do orthogonalization in KrylovKit.jl? Looks like they check if the AbstractMatrix is hermitian or not? Using Lanczos can be significantly faster than Arnoldi...

By default KrylovKit is using Arnoldi, and in the benchmark code both KrylovKit and ITensor were using Arnoldi (assuming I understood correctly and ITensor also uses Arnoldi by default).
As I wrote above, I am also getting similar performance difference for the TDVP code itself, which is using Lancosz in both cases.

I totally agree about the matrix size, a quick check that the discrepancy is due to things besides matrix-vector multiplications would be to increase the matrix size and see if the difference starts to disappear.

I enlarged the "bond dimension" of the LocalOp (parameter m in the Julia code above) to 1000, with this I get (with the same code I previously used, for which it's not clear if comparison is fair):

• ITensor C++ ~ 11.6 s
• KrylovKit + ITensors.jl ~ 6.13 s

I think the first thing to do is to make sure we are doing a fair comparison between the two (i.e. compare time to reach a fixed error).

It would also be useful if someone else verifies the timings I get on their machine (by just running the code I provided and maybe reading the C++ code) to eliminate the possibility that I made some stupid mistake while benchmarking. If someone wants to try the benchmark it is important to make sure Julia is using the same BLAS as ITensor.

[mtfishman]

Interesting, that makes me think that the C++ version is just doing more matrix-vector multiplies (for whatever reason, possibly just caused by a different convergence criteria).