Iterative Matrix-Vector Product Eigensolvers in Python

Question

I have a very large Hermitian matrix of size (21^6, 21^6) that I need to diagonalize. The matrix is highly sparse with roughly 1E^10 non-zero elements (out of 7.35E^15), and is built up from a direct product of many smaller matrices. For this, I make use of scipy's sparse data structures, building the final matrix up from a dedicated series of sparse matrix operations. However, even in sparse format the total matrix is too large to store in memory (csr format) and diagonalize in the standard fashion using scipy.sparse.linalg.eigsh as,

from scipy.sparse.linalg import eigsh

n = 20 # number of eigenvals/vecs

vals, vecs = eigsh(A, k=n, which='SA')

The scipy.sparse documentation (https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.eigsh.html#scipy.sparse.linalg.eigsh) mentions that one can pass a linear operator representation of the operation (in numpy notation) A@x to solve A@x = w@x. I would like to use this to iteratively solve the eigenvalue problem using a matrix-vector product method so I do not have to store the whole matrix A at any one time. The documentation on this is pretty poor, with the only example given (https://docs.scipy.org/doc/scipy/tutorial/arpack.html#use-of-linearoperator),

from scipy.sparse.linalg import LinearOperator

class Diagonal(LinearOperator):

    def __init__(self, diag, dtype='float32'):
        self.diag = diag
        self.shape = (len(self.diag), len(self.diag)
        self.dtype = np.dtype(dtype)

    def _matvec(self, x):
        return self.diag*x

    def _rmatvec(self, x):
        return self.diag*x

N = 100
rng = np.random.default_rng()
d = rng.normal(0, 1, N).astype(np.float64)
D = np.diag(d)
Dop = Diagonal(d, dtype=np.float64)
evals_all, evecs_all = eigh(D)
evals_large, evecs_large = eigsh(Dop, 3, which='LA', maxiter=1e3)

which gives the same result,

>>> evals_all[-3:]
    array([1.53092498, 1.77243671, 2.00582508])

>>> evals_large
    array([1.53092498, 1.77243671, 2.00582508])

But this is not really useful as the matrix is already diagonal and this is equivalent to constructing a sparse csr representation of D from the diagonal values d and passing that to eigsh directly...

If anyone has any idea on how to use such a LinearOperator to construct a representation of A@x such that the full sparse matrix of A is never allocated in memory at once - I would greatly appreciate some advice.

EDIT - Construction of matrix A:

The sparse matrix A is built from a series of d smaller matrices K. To avoid getting too complicated, let us consider the case when d=3, and A is built up from three matrices Kx, Ky, and Kz. Say each matrix K has dimensions (n, n) with n=21. The matrix elements of A can be written as,

which I compute using scipy's sparse implementation of the Kronecker product,

,

where each matrix I is the identity of shape (21,21). Given that each matrix K is (21,21) the shape of A is then (n^d, n^d) - i.e. (21^3, 21^3). Note in the second expression I have excluded the elements of V, which is just a vector subsequently added to the diagonal of A.

While we only consider d=3 and n=21 here, the implementation in my code is iterative and allows one to construct the sparse matrix A in csr format for any value of d and n (until I run out of memory). The actual problem above I am trying to solve correspond to d=6 and n=21, thus there are not three lots of two Kronecker product operations, but six lots of five. While for d=3 and n=21, the matrix A can easily be allocated in sparse format, for any larger values of d this quickly becomes impossible.

I am not too clued up on how these iterative eigensolvers work, from my understanding I can define a LinearOperator object that returns the dot product of a row of A with a vector - and scipy/ARPACK will iteratively diagonalize A without having to store the whole thing in memory.