Implementing Lanczos algorithm as a new solver method

[rupeshknn]

Summary

Lanczos algorithm is an approximate diagonalisation method. It is implemented as an LMDE method and is a considerable speedup compared to scipy.expm

Details and comments

This PR adds a new fixed-step solver method lanczos_diag. This method only works with hermitian generators and works best in sparse evaluation mode. A follow up with a Jax implementation of the same is in progress.

[CLAassistant]

All committers have signed the CLA.

[rupeshknn]

Regarding y0 being a matrix: When Lanczos is initialised using y0 as the seed, we can skip computing the matrix exponential since the equation simplifies to something in terms of the Lanczos basis and tridiagonal matrix. However, when this is not possible, I could seed using a random vector -> compute the matrix exponential, -> find action on expm on the y0 matrix.

This change would affect the speed only a little bit, but this also takes a toll on the accuracy. This is because the rows of Q (the basis vectors) are not perfectly orthogonal.

To mitigate this, one has to perform a QR decomposition which adds to the runtime. However, since we QR decompose, one could make do with a smaller k_dim value. @miamico what do you think?

[miamico]

Regarding y0 being a matrix: When Lanczos is initialised using y0 as the seed, we can skip computing the matrix exponential since the equation simplifies to something in terms of the Lanczos basis and tridiagonal matrix. However, when this is not possible, I could seed using a random vector -> compute the matrix exponential, -> find action on expm on the y0 matrix.

This change would affect the speed only a little bit, but this also takes a toll on the accuracy. This is because the rows of Q (the basis vectors) are not perfectly orthogonal.

To mitigate this, one has to perform a QR decomposition which adds to the runtime. However, since we QR decompose, one could make do with a smaller k_dim value. @miamico what do you think?

I'm still thinking how to deal with an initial matrix instead of a vector. If it's just about batching many initial vectors which are put together into a matrix, couldn't we just do Lanczos on each of the columns (or rows) of the matrix which represent a vector?

[rupeshknn]

@DanPuzzuoli, I've made the requested changes. As for rectangular arrays, I've modified the code so that we run a Lanczos iteration for each row, as @miamico suggested. Let me know what you think.

[DanPuzzuoli]

@DanPuzzuoli, I've made the requested changes. As for rectangular arrays, I've modified the code so that we run a Lanczos iteration for each row, as @miamico suggested. Let me know what you think.

Okay awesome, TY. I'm ending early today, but will take a look next week.

[rupeshknn]

@DanPuzzuoli Thanks for those examples. The if: break condition in the loop was supposed to handle these div by zero errors. However, it comes only after one run of the loop. In those examples you mentioned, It hit zero before the loop starts. I've fixed it now by moving the if condition to the top. The JAX version (not part of this PR) does not have that if: break condition. I might text you on slack when I get to that.

I believe the CI is not automatically running the tests yet. Can you check that?

I've made all the fixes as suggested. Although I'm not sure about the names I've given to the test functions, let me know if they're okay.

[rupeshknn]

I've made the changes and changed lanczos_expm to store the norm and multiply it later. Thanks, @DanPuzzuoli. Now I'll add the Jax code to this branch

[rupeshknn]

Hi @DanPuzzuoli, I've got the changes done. However, I implemented the scan based looping in a slightly different way than you had suggested with some inspiration from fixed_step_solver_template_jax. This method uses only one scan; I guess that's a good thing. The cost is a bunch of iterations that return zeros.

I tested the grad transformation using this tutorial How-to use JAX, and for certain values, thejax_lanczos_diag method gives nan values.

For a modified version of the spin chain tutorial, this doesn't happen.

JAX has an open thread about issues with grad for degenerate eigenvalues https://github.com/google/jax/issues/669#issuecomment-489331458

[rupeshknn]

Hi @DanPuzzuoli I fixed all the inline comments.

Regarding using sparse mode for regular tests, I ran into a bunch of errors for the NumPy backend (aka scipy.sparse.csr_matrix). sparse mode doesn't seem compatible with rotating_frame.

Here's an example (on the main branch)

import numpy as np
from qiskit_dynamics import Solver, Signal
from qiskit_dynamics.array import Array

sparse_solver = Solver(static_hamiltonian=np.array([[0,1],[1,0]]),
                       hamiltonian_operators=[np.array([[0,1],[1,0]])],
                       rotating_frame=np.array([[0,1],[1,0]]),
                       evaluation_mode='sparse'
                      )

res = sparse_solver.solve(
    t_span=[0., 50],
    y0=np.array([0,1]),
    signals = [Signal(Array(1.), carrier_freq=5)],
    method="odeint"
)

# TypeError: operand type(s) all returned NotImplemented 
# from __array_ufunc__(<ufunc 'matmul'>, '__call__', 
#                      [<2x2 sparse matrix of type '<class 'numpy.complex128'>'with 4 stored elements in Compressed Sparse Row format>],
#                      Array([[-0.70710678-0.j,  0.70710678-0.j],
#                             [ 0.70710678-0.j,  0.70710678-0.j]], backend='numpy')
#                     ): 'list', 'Array'

Things work fine with JAX sparse + rotating_frame for the same example.

I suppose it's a bug since the list has a length of 1. But i don't quite understand rotating frames sufficiently to fix it.

[DanPuzzuoli]

Nice bug find. This is a bit of a weird example that I'll need to think about a little to figure out the proper fix.

The reason it hasn't been found before is due to the expected usage that if a user sets evaluation_mode='sparse', then their frame operator is already diagonal (and specified as a 1d array).

For a bit of background: A user can specify the frame operator as a 1d array as a way of explicitly telling Dynamics that it's already diagonal. The main reason this is "expected usage" for sparse evaluation mode is that internally Dynamics always represents all operators in a basis in which the frame operator is diagonal - if it isn't diagonal, it has to change the basis of all other operators. If the operators are sparse, you don't want to do a change of basis (which could ruin sparsity), so specifying the frame operator as a 1d array is something we put in to allow the user a pathway to prevent Dynamics from trying to do any basis changes.

That all being said, everything should still work if the frame operator is not given as a 1d array, but it seems that we've really baked the "expected usage" into our tests, so we've missed the bug so far. I think this will be a quick fix, and will get back to you.

[DanPuzzuoli]

Okay #118 should solve this issue. The example you provided is working (though you need to change the method to something like 'RK45' rather than 'odeint'.

[DanPuzzuoli]

Didn't look at the tests again but just scanned through everything else and looks good. Let me know if you get the tests working, and I'll do a final holistic review.

[rupeshknn]

Thanks for the quick resolution. The tests are working.