Accelerating the computation of scipy.sparse.linalg.expm_multiply

Question

I have a tridiagonal antiHermitian matrix ($-i*Hami*t$) with nonzero elements only along the upper diagonal and lower diagonal, and the goal is to know the action of exponential of such matrix on a given vector. Currently, I am using scipy.sparse.linalg.expm_multiply function in Python to check the action of a million by million dimensional matrix on a million dimensional vector:scipy.sparse.linalg.expm_multiply(Hami, psi, start=0, stop=t_end-dt, num=int(t_end/dt), endpoint=True, traceA=0). This gives me the required state at various times $(e^{-i∗Hami∗t}|psi⟩)$ for $t=\{0,dt,...,t_{end}-dt\}$.

This takes multiple days to get the answer. The end goal is to do this computation for 1000s of cases. I was wondering if there is any way to accelerate such computation?

Answer 1

Sorry for an incomplete answer (but it might work in the future if CuPy implements the missing expm_multply)

Side note: LinearOperator fun

If your Hami has some form of pattern that can be used, a LinearOperator could take shortcuts by e.g. never even constructing it in the first place

import scipy.sparse
import scipy.sparse.linalg

n = 1000000

# Simple predicable matrix:
a, b = 1.2, 2.3
Hami = scipy.sparse.eye(n, k=0)*a + scipy.sparse.eye(n, k=1)*b - scipy.sparse.eye(n, k=-1)*b 

# This linear operator does the same, without constructing the matrix:
def matvec(x):
    y = a*x
    y[1:] -= b*x[:-1]
    y[:-1] += b*x[1:]
    return y

HamiL = scipy.sparse.linalg.LinearOperator((n,n), matvec=matvec)

import timeit
import numpy as np

psi = np.random.rand(n)
timeit.timeit('Hami*psi', number=1000, globals=globals())
timeit.timeit('HamiL*psi', number=1000, globals=globals())

Though in the scenario given here, there were basically no speedup (I would expect approaches like this to be well worth it if the size of the matrix is so large that even constructing it at all is a problem).

There are still some things you could possible do here, using e.g. numpexpr to speed up the matvec, though it depends on your scenario.

CuPy

When thinking accelerating, I think accelerators: GPUs. I've found the CuPy library is very easy to use, and, ideally, you would just switch your imports and be pretty much done;

import cupy as cp
import cupyx.scipy
import cupyx.scipy.sparse.linalg

n = 1000000
a, b = 1.2, 2.3
Hami = cupyx.scipy.sparse.eye(n, k=0, dtype=cp.float32)*a + cupyx.scipy.sparse.eye(n, k=1, dtype=cp.float32)*b - cupyx.scipy.sparse.eye(n, k=-1, dtype=cp.float32)*b

psi = cp.random.rand(n, dtype=cp.float32)

import timeit
timeit.timeit('Hami*psi; cp.cuda.Device().synchronize()', number=1000, globals=globals())

Granted that the precision was dropped to float32 (for the benefit of the GPU), this still showed a ~20x speedup (including the synchronizing) for these matrix sizes.

Here is where i would say to just use

psi = cp.array(psi)
r = cupyx.scipy.sparse.linalg.expm_multiply(Hami, psi, start=0, stop=t_end-dt, num=int(t_end/dt), endpoint=True, traceA=0)

but unfortunately at the time of writing CuPy has not yet implemented support for expm_multiply. https://docs.cupy.dev/en/stable/reference/comparison.html

The implementation in SciPy unfortunately tries to convert everything to numpy or assumes there are only scipy sparse arrays in many places, so one really needs a CuPy implementation for this to work.

Wrapping the operation in a LinearOperator might be tempting

from scipy.sparse.linalg import expm_multiply, LinearOperator

def matvec(x):
    return (Hami*cp.array(x)).get()  # costly memory transfers

HamiL = LinearOperator((n,n), matvec=matvec)

but the constant memory transfers eats up all performance gains for these small arrays.