# Accelerating the computation of scipy.sparse.linalg.expm_multiply

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?

• One of the references in the documentation says the computation is dominated by computation of matvecs with $A$, so any improvements there will likely scale directly to your problem. Since your system is tridiagonal, parallelization should lead to some big improvements in efficiency Jun 24, 2023 at 17:46
• Note that expm_multiply takes a LinearOperator; depending on exactly how your matrix is constructed, you could potentially take a huge shortcut here and just compute A*v or A^H*v directly: A = LinearOperator((n,n), matvec=your_smart_matvec_func) Jun 24, 2023 at 19:57
• @MikaelÖhman: can you please elaborate by what you mean? Right now, I do scipy.sparse.linalg.expm_multiply(Hami, psi, start=0, stop=t_end-dt, num=int(t_end/dt), endpoint=True, traceA=0) and this gives me the required state at various times $\bigl(e^{-i*Hami*t}|psi\rangle \bigr)$ for $t=\{0,dt,...,t_{end}\}$ Jun 26, 2023 at 19:33
• This function doesn't actually require access to the individual elements of $A$, but instead it just needs a function that computes it's action on a vector, i.e., a map $v \mapsto Av$. Depending on $A$, sometimes you can write a function that executes faster than scipy's sparse matrix multiplication routines, although I don't know the specifics of those Jun 27, 2023 at 2:00
• "matrix-free" and "Krylov" would be the primary keywords here. They are most often used to solve $Ax=b$ without forming $A$ explicitly for giant PDE discretizations. docs.scipy.org/doc/scipy/reference/generated/… contains links to a couple references on the specific applications of applying matrix functions, but Iterative Methods for Linear and Nonlinear Equations by Tim Kelley is an approachable introduction to Krylov methods for solving linear equations and the general idea. Iterative Methods for Sparse Lienar Systems by Y. Saad is also good Jun 27, 2023 at 22:07

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.