How to efficient solve $e^{-tA} x =b$, where A is a very sparse matrix

Question

I am going to solve an equation containing an exponential matrix $e^{tA} x =b$, which can be obtained naturally through $x=e^{-tA} b$. A is a 1million $\times$ 1 million matrix with stores 7.15 million elements. However, accurately computing a matrix exponential $e^{-tA}$ requires very high computer memory. Although there are methods to quickly solve exponential matrices, such as the Krylov method, they still require high computational costs.

I tried to solve it using the following approximation method: $$ e^{tA} \approx (I + tA + t^2 \frac{A^2}{2!} + t^3 \frac{A^3}{3!}) $$

However, when I try to solve $$(I + tA + t^2 \frac{A^2}{2!} + t^3 \frac{A^3}{3!})x = b$$, the program quickly occupied me Computer memory (64G) and shut down.

A is a positive semidefinite matrix. It seems that diagonalization technology can be used? For example, let $A = P^{-1} L P$, where $L$ is a diagonal matrix. Then the above formula can be transformed into: $$ (I + tA + t^2\frac{A^2}{2!} + t^3\frac{A^3}{3!}) x = b \to (I + t L + t^2\frac{L^2}{2!} + t^3\frac{L^3}{3!}) y = PbP^{-1}, x = P^{-1}yP. $$ Would such diagonalization be costly? Is there any existing library in python that can quickly perform the above operations?

The code I used when solving was pypardiso.spsolve(A,b), where A is the left-hand matrix of the above equation. I would like to ask if there is a fast algorithm to solve such an equation?

EDIT I'm trying to solve a parabolic problem $$u_t = \Delta u + f,$$ Subject to the homogeneous Neumann boundary conditions, first of all, in terms of spatial discretization, the mass-lumped finite element method is adopted, and the following discrete equation is obtained: $$ ((u_h)_t ,v_h) = (\nabla u_h, \nabla v_h) + (f,u_h), \quad \forall v_h \in V_h $$ where $V_h$ is the finite element subspace of $H^{1}(\Omega)$. Due to the mass-concentrated finite element method used, the above format can be written in matrix-vector form $$ \boldsymbol{u}_t = L \boldsymbol{u} + \boldsymbol{f} $$ Where $L = -M^{-1} A$, $A = \int_{\Omega} \psi_i \psi_j dx$. $\boldsymbol{u}$ is the vector constructed by the nodal value. The numerical solution satisfies $$ \boldsymbol{u}(t_{n+1}) = e^{\tau L} \boldsymbol{u}(t_n) + \tau \int_{0}^{\tau} e^{-(\tau-s) L} \boldsymbol{f}(t_n+s) ds. $$ I want to discretize in the time direction using explicit Runge-Kutta, getting $$ \boldsymbol{u}_{ni} = e^{-c_i \tau L} \boldsymbol{u}^n + \tau \sum\limits_{j=0}^{i-1} a_{i,j} e^{-(c_i-c_j)\tau L} \boldsymbol{f}_{n,j} $$ $$ \boldsymbol{u}^{n+1} = e^{- \tau L} \boldsymbol{u}^n + \tau \sum\limits_{j=0}^{i-1} b_j e^{- (1-c_j)\tau L} \boldsymbol{f}_{n,j} $$ Where $a_{i,j}, b_i, c_i$ are defined in Runge-Kutta's Butcher table. However, usually calculating the exponential matrix in the above process will cost a very high computational cost. Therefore, I want to approximate the matrix exponential, and finally The following equation is obtained $$ \phi_i(-\tau L) \boldsymbol{u}_{ni} = \left( \boldsymbol{u}^n + \tau \sum\limits_{j=0}^{i -1} a_{i,j}\phi_j(-\tau L) \boldsymbol{f}_{n,j} \right), $$ Where $\phi_i(x)$ is a polynomial function used to approximate $e^{c_i x}$.

Taking this step eliminates the need to calculate the exponential matrix, but introduces an additional solving process. Even so, this method remains faster than computing the exponential matrix, as observed in 1D and 2D examples. However, as the problem transitions to three dimensions, the process becomes less satisfactory. When dealing with a $100^3$ grid, $\phi_i(\tau L)$ stores a significantly larger number of elements than the matrix $L$，which introduce additionly trouble.

I have seen some people efficiently implement the matrix exponential using a tensor product and a straightforward finite difference method, but I have not attempted the combination of finite difference and FFT,

Thx to the answers @lightxbulb, I want to start implementing such a method, but my understanding of the ordering and implementation of 3D grid points is very limited. If you can give a 3D example code of the heat equation ( A quick implementation based on finite differences and FFT, I would be very grateful.)

Answer 1

$\newcommand{\bm}[1]{\boldsymbol{#1}}$Continuous Formulation

With the new edits in your post I believe this is a diffusion problem with a time-dependent inhomogeneity on a domain with reflecting boundaries:

\begin{alignat}{2} \partial_t u(t,\bm{x}) &= \Delta u(t,\bm{x}) + f(t,\bm{x}), &\quad& \bm{x}\in\Omega, \,t\in(0,\infty), \\ \partial_{\bm{n}}u(t,\bm{x})&=0, &\quad& \bm{x}\in\partial\Omega, t\in(0,\infty),\\ u(0,\bm{x}) &= b(x),&\quad& \bm{x}\in\Omega, \end{alignat} where $\Omega$ is the solution domain, and $\bm{n}$ is the outward facing normal at the boundary. With this modification things become a bit more complicated, mainly because $f$ is also a function of time.

Spatial Discretisation

As far as I understood you use the finite element method to discretise the problem spatially. That is, given basis vectors $\{\psi_i\}_{i=1}^n$ spanning a subspace of the Sobolev space $H^1(\Omega)$, you can compute the mass matrix $M_{ij} = \int_{\Omega} \psi_i \psi_j$, the stiffness matrix $W_{ij} = \int_{\Omega} \bm{\nabla}\psi_i \cdot \bm{\nabla} \psi_j$, the inhomogeneity vector $f_i = \int_{\Omega}\psi_i f$, and the initial conditions vector $$\bm{b} = \bm{M}^{-1}\begin{bmatrix} \int_{\Omega}\psi_1 b \\ \vdots \\ \int_{\Omega} \psi_n b \end{bmatrix}.$$ The spatially discretised problem then reads: \begin{alignat}{2} d_t \bm{M}\bm{u}(t) &= -\bm{W}\bm{u}(t) + \bm{f}(t), &\quad& t\in(0,\infty), \\ \bm{u}(0) &= \bm{b}. \end{alignat} Here $\bm{u}(t)\in\mathbb{R}^n$ is the vector of coefficients at time $t$ for the solution $$u(t,\bm{x}) = \sum_{j=1}^n u_j(t) \psi_j(\bm{x}).$$ The resulting problem is a matrix ODE system with an initial condition, and its solution is given as $$\bm{u}(t) = \bm{\exp}(-t\bm{M}^{-1}\bm{W})\bm{b} + \int_{0}^t \bm{\exp}(-(t-s)\bm{M}^{-1}\bm{W})\bm{f}(s)\,ds.$$

Spectral Solution

The matrix $\bm{M}^{-1}\bm{W}$ can be non-symmetric even if $\bm{W}$ and $\bm{M}$ are symmetric, so we can instead reparametrize $\bm{w}(t) = \bm{M}^{1/2}\bm{u}(t)$, then the system reads \begin{alignat}{2} d_t \bm{M}^{1/2}\bm{w}(t) &= -\bm{W}\bm{M}^{-1/2}\bm{w}(t) + \bm{f}(t), &\quad& t\in(0,\infty), \\ \bm{w}(0) &= \bm{M}^{1/2}\bm{b}. \end{alignat} The solution is now $$\bm{w}(t) = \bm{\exp}(-t\bm{M}^{-1/2}\bm{W}\bm{M}^{-1/2})\bm{M}^{1/2}\bm{b} + \int_{0}^t\bm{\exp}(-(t-s)\bm{M}^{-1/2}\bm{W}\bm{M}^{-1/2})\bm{M}^{-1/2}\bm{f}(s)\,ds.$$ If $\bm{W}$ and $\bm{M}$ are symmetric then the product is a symmetric matrix and the eigendecomposition is $-\bm{M}^{-1/2}\bm{W}\bm{M}^{-1/2} = \bm{Q}\bm{\Lambda}\bm{Q}^T$. You can use this to evaluate the application of the exponential as $$\bm{\exp}(-t\bm{M}^{-1/2}\bm{W}\bm{M}^{-1/2})\bm{x} = \bm{Q}(\bm{\exp}(t\bm{\Lambda})(\bm{Q}^T\bm{x})).$$

Provided you know how to efficiently compute products $\bm{Q}\bm{x}$, $\bm{Q}^T\bm{x}$ you can efficiently evaluate the application of the exponential (e.g. $\bm{Q}\bm{x}$ may be implemented as a fast discrete sine/cosine/Fourier transform in specific cases). If your mass matrix is diagonal/lumped then also the powers and inverse of $\bm{M}$ are easy to compute. The main issue is computing the integral - I do not believe you have an analytic solution for it in the general case, so you could use standard numerical integration methods (each point would cost you an evaluation of the exponential, but if you can do this quickly it shouldn't be too bad). In the special case when $\bm{f}$ is not a function of time there is a nice expression for the integral. If I didn't mess up something it should be:

$$\int_0^t\bm{\exp}((t-s)\bm{A})\,ds = \bm{A}^{-}\bm{\exp}(t\bm{A})+t(\bm{I}-\bm{A}^{-}\bm{A})-\bm{A}^{-}.$$

Here $\bm{A}^{-}$ is the group inverse (in the case $\bm{A}$ is symmetric the Moore-Penrose inverse would do) such that $\bm{A}^{-}\bm{A}^{k+1} = \bm{A}^k$ for $k\geq 1$. Knowing the eigendecomposition of $\bm{A}$ this should also not be too hard to compute. But as far as I understood in your problem $\bm{f}$ is a function of $t$, so you cannot benefit from this, and I added it just for completeness.

Time Discretisation

You may instead opt for discretising time using explicit Euler: $$d_t\bm{M}\bm{u}(t) = \bm{W}\bm{u}(t) + \bm{f}(t) \,\,\approx\,\, \frac{\bm{M}\bm{u}^{k+1}-\bm{M}\bm{u}^k}{\tau} = \bm{W}\bm{u}^k + \bm{f}^k \\ \implies \\ \bm{u}^{k+1} = \bm{M}^{-1}(\bm{M}+\tau\bm{W})\bm{u}^k + \tau\bm{M}^{-1}\bm{f}^k,$$ or implicit Euler: $$d_t\bm{M}\bm{u}(t) = \bm{W}\bm{u}(t) + \bm{f}(t) \approx \frac{\bm{M}\bm{u}^{k+1}-\bm{M}\bm{u}^k}{\tau} = \bm{W}\bm{u}^{k+1} + \bm{f}^{k+1}\\ \implies \\ (\bm{M}-\tau\bm{W})\bm{u}^{k+1}=\bm{M}\bm{u}^k+\tau\bm{f}^{k+1}.$$

The explicit Euler does not require solving linear systems (if $\bm{M}$ is diagonal) but it is stable only if $\|\bm{I}+\tau\bm{M}^{-1}\bm{W}\|<1$. On the other hand implicit Euler requires solving a linear system at each step, but it is unconditionally stable (i.e. you can pick $\tau$ as large as you like, but the solution will deviate from the true one/the spectral one if you have too few iterations). In either case you set $\bm{u}^0 = \bm{b}$ and start the iterations from there. There are other integrators such as Runge-Kutta but I won't discuss those here.

Rectangular Domain and Regular Grid FEM or FDM

For a regular grid discretisation with linear triangular Lagrange elements the lumped mass matrix is the identity times a constant $\bm{M}=\alpha\bm{I}$, and the stiffness matrix is a Toeplitz plus Hankel matrix. For such matrices the eigenvectors and eigenvalues are known, see "Diagonalizing Properties of the Discrete Cosine Transforms" by Sanchez et al., "The Discrete Cosine Transform" by Strang, and "Functions of Difference Matrices are Toeplitz Plus Hankel" by Strang and MacNamara. For reflecting boundary conditions in 1D with FEM or FDM on a regular grid you typically get a matrix $$\Delta = d_{xx} \approx -\bm{W} = \frac{1}{h^2}\begin{bmatrix} -1 & 1 & & \\ 1 & -2 & 1 & \\ & \ddots & \ddots & \ddots \\ & & 1 & -2 & 1 \\ &&&1&-1\end{bmatrix}.$$ This matrix is diagonalized by the orthogonal DCT-II and DCT-III matrices. The eigenvalues and eigenvectors are:

$$\lambda_j = -\frac{4}{h^2}\sin^2\left(\frac{\pi(j-1)}{2n}\right), \quad v_{ij} = \begin{cases} \frac{1}{\sqrt{n}}, &j=1,\\ \sqrt{\frac{2}{n}}\cos\left(\frac{\pi(j-1)(i-0.5)}{n}\right), &j\ne 1. \end{cases}$$

This means that $-\bm{W} = \bm{V}\bm{\Lambda}\bm{V}^T$. If you go to higher dimensions you can write $\Delta = \sum_{k=1}^d \partial^2_{x_k}$. In 2D you can discretise this using the Kronecker product as: $\Delta \approx \bm{D}_{xx}\otimes\bm{I} + \bm{I}\otimes \bm{D}_{yy}$. The eigenvalue is now the sum of eigenvalues along the two dimensions $\lambda_{jl} = \lambda_j + \lambda_l$, and the eigenvectors are the products $v_{ijkl} = v_{ij}v_{kl}$. That is, you can apply the DCT-II and DCT-III first on each row, and then on each column (or vice versa) in order to compute $\bm{V}\bm{x}$ and $\bm{V}^T\bm{x}$ where $\bm{x}$ represents your values on a 2d grid. In either case having the eigendecomposition $-\bm{W} = \bm{V}\bm{\Lambda}\bm{V}^T$ and assuming that $\bm{M} = \alpha \bm{I}$ you can compute the exponential as: $$\bm{\exp}(-t\bm{M}^{-1}\bm{W}) = \bm{\exp}\bm{V}(t\alpha^{-1}\bm{\Lambda})\bm{V}^T.$$

For three dimensions it is similar $$\Delta \approx \bm{D}_{xx} \otimes \bm{I} \otimes \bm{I} + \bm{I} \otimes \bm{D}_{yy} \otimes \bm{I} + \bm{I} \otimes \bm{I} \otimes \bm{D}_{zz},$$ and the eigenvalues are $\lambda_{j_1 j_2 j_3} = \lambda_{j_1} + \lambda_{j_2} + \lambda_{j_3}$ with eigenvectors $\lambda_{i_1j_1i_2j_2i_3j_3} = v_{i_1j_1}v_{i_2j_2}v_{i_3j_3}$. This means you can apply the 1D DCT-II/DCT-III to each dimension separately to get the 3d DCT-II/DCT-III. This generalizes similarly to higher dimensions. Note that I wrote this in terms of matrices just for the sake of the theory. In practice you can work matrix-free in the case of a regular grid.