11.3: The Split-Step Fourier Method

As an example of the usefulness of the DFT, let us discuss a DFT-based method for performing numerical integration of a partial differential equation, known as the split-step Fourier method. Here, the method will be presented in the context of the time-dependent Schrödinger equation in 1D space:

\[i \frac{d\psi(x,t)}{dt} = \left[-\frac{1}{2}\frac{\partial^2}{\partial x^2} + V(x,t)\right]\psi(x,t)\]

We have taken \(\hbar = m = 1\) for simplicity. At each time, the wavefunction is a continuous function of \(x\). Let us truncate and discretize this spatial coordinate, by defining a computational domain of length \(L\) containing \(N\) discretization points:

\[\psi_n(t) = \psi(x_n,t), \;\;\;\textrm{where}\;\; x_n = -\frac{L}{2} + n \Delta x, \;\; \Delta x = \frac{L}{N}.\]

Thus, the wavefunction at each time is represented by a complex vector, which we call a "state vector":

\[\vec{\psi}(t) = \begin{pmatrix} \psi_0(t) \\ \psi_1(t) \\ \vdots \\ \psi_{N-1}(t)\end{pmatrix}.\]

Given an initial state vector \(\vec{\psi}(t_a)\), the problem is to compute \(\vec{\psi}(t_b)\) at a later time \(t_{b}\). Note that this differs from previously-studied numerical ODE problems in one important respect: evolving in time involves taking second-order spatial derivatives. We won't go into the details, but it turns out that standard methods for time-stepping and discretizing space don't work very well here, because the errors from time-stepping and spatial discretization interact badly with one another. The split-step Fourier method provides a better way to solve the problem.

11.3.1 Factorizing the Time-Evolution Operator

The split-step Fourier method is based on the concept of the time-evolution operator. Given a wavefunction \(\psi(x,t_j)\), the wavefunction after a small time step \(\tau\) is

\[\psi(x,t_j + \tau) = U(t_j + \tau|t_j)\, \psi(x,t_j), \quad\mathrm{where}\;\;\; U(t_j + \tau|t_j) \approx \exp\left\{-i \tau \left[-\frac{1}{2}\frac{\partial^2}{\partial x^2} + V(x,t_j+\tau/2)\right]\right\}\;.\]

Here, the \(\exp(\cdots)\) refers to the exponential of an operator (one involving spatial derivatives). We call \(U(t_b|t_a)\) the time-evolution operator, which evolves the system from time \(t_{a}\) to \(t_{b}\). The exponential of any operator \(\mathcal{A}\) is defined as the infinite series

\[\exp(\mathcal{A}) = I + \mathcal{A} + \frac{1}{2} \mathcal{A}^2 + \frac{1}{6} \mathcal{A}^3 + \cdots\]

In this case, the exponential contains the Hamiltonian, which consists of a kinetic energy term and a potential energy term that do not generally commute. Due to this non-commutivity, the exponential cannot be simplified by factorization:

\[\exp\left\{-i \tau \left[-\frac{1}{2}\frac{\partial^2}{\partial x^2} + V(x,t_j+\tau/2)\right]\right\} \;\; \ne\;\;\exp\Bigg\{\frac{i \tau}{2} \left[\frac{\partial^2}{\partial x^2}\right]\Bigg\} \;\exp\Bigg\{-i \tau V(x,t_j+\tau/2)\Bigg\}.\]

However, we can obtain an approximate factorization by making use of the series definition of the exponential of an operator. One can show that

\[\exp(\mathcal{A} + \mathcal{B}) = \exp(\mathcal{A}/2) \exp(\mathcal{B}) \exp(\mathcal{A}/2) + O\Big((\mathcal{A},\mathcal{B})^3\Big),\]

which is a variant of an important formula known as the Baker–Campbell–Hausdorff formula. On the right-hand side, note that \(\exp(\mathcal{B})\) is sandwiched "symmetrically" between two copies of \(\exp(\mathcal{A}/2)\). This symmetric arrangement reduces the approximation error to third order, by the cancellation of lower-order errors (in a manner similar to the mid-point formula for the discretized derivative). Applying this factorization to the time-evolution operator gives

\[\boxed{U(t_j + \tau|t_j) \;\; \approx\;\; U_\mathcal{K} \,\cdot\, U_\mathcal{V}(t_j+\tau/2) \,\cdot\, U_\mathcal{K}, \quad \mathrm{where}\;\; \begin{cases} \;U_\mathcal{K} &= \exp\left[\frac{i\tau}{4} \frac{d^2}{dx^2}\right] \\ \; U_\mathcal{V}(t) &= \exp\big[-i\tau\, V(x,\,t)\big]\end{cases}}\]

In other words, the time-evolution operator decomposes into three pieces. That's why we call this a "split-step" algorithm: each time step from \(t_{j}\) to \(t_{j}+\tau\) consists of applying a kinetic step, then applying a potential step, then applying another kinetic step, in sequence. As previously noted, we'll be working with state vectors (complex \(N\)-component vectors), defined through spatial discretization of the wavefunction. So we need to figure out how the above stepping operators act on these state vectors:

\[U_\mathcal{K}\; \vec{\psi} \;= \;\;??, \qquad U_\mathcal{V}\, \vec{\psi} \;=\;\; ??\]

The potential stepping operator is simple to deal with. Since the state vector represents the wavefunction at different points in space, the potential operator is represented by a diagonal matrix, and its exponential is also diagonal:

\[U_\mathcal{V} \begin{pmatrix}\psi_0\\ \vdots \\ \psi_{N-1}\end{pmatrix} \;=\; \begin{pmatrix}\exp\left[-i\tau V(x_0,t+\frac{\tau}{2})\right] & & \\ &\ddots&\\ && \exp\left[-i\tau V(x_{N-1},t+\frac{\tau}{2})\right] \end{pmatrix} \begin{pmatrix}\psi_0\\ \vdots \\ \psi_{N-1}\end{pmatrix} \]

11.3.2 Kinetic Step

The kinetic stepping operator, \(U_\mathcal{K}\), is less obvious. It contains spatial derivatives and is thus not diagonal in the current basis. The key thing to realize, however, is that this operator is diagonal in wavenumber space. Let us return to the continuous wavefunction, and write its Fourier representation:

\[\psi(x) = \int_{-\infty}^\infty \frac{dk}{2\pi} \; e^{ikx}\; \Psi_k.\]

Then

\[U_\mathcal{K} \; \psi(x) = \int_{-\infty}^\infty \frac{dk}{2\pi} \; \exp\left(-\frac{i\tau}{4}k^2\right) \; e^{ikx}\; \Psi_k.\]

Let us discretize space in steps of \(\Delta x\), as discussed earlier, and also discretize the Fourier integrals by steps of \(\Delta k\):

\[\begin{align}x_n &= -\frac{L}{2} + n \Delta x, \\k_n &= -\frac{K}{2} + n\Delta k\end{align}\]

The values of \(K\) and \(\Delta k\) will be chosen shortly. The discretized integrals become

\[\begin{align}\psi_m &\approx \sum_{n=0}^{N-1} \frac{\Delta k}{2\pi} \; e^{ik_n x_m}\; \Psi_{k_n}, \\\left(U_\mathcal{K} \psi\right)_m &\approx \sum_{n=0}^{N-1} \frac{\Delta k}{2\pi} \; e^{-\frac{i\tau}{4}k_n^2} \;e^{ik_n x_m}\; \Psi_{k_n}. \end{align}\]

Let us now choose the \(k\)-space discretization parameters to be

\[\Delta k = \frac{2\pi}{N\Delta x} = \frac{2\pi}{L} , \quad K = N\Delta k = \frac{2\pi}{\Delta x} \;\;\;\Rightarrow \;\;\; k_n = -\frac{\pi}{\Delta x} + \frac{2\pi n}{N\Delta x}.\]

With this choice, we can show with a bit of algebra that the integral for \(\psi_{m}\) reduces to an IDFT:

\[\begin{align}\psi_m &= (-1)^m \frac{1}{N} \sum_{n=0}^{N-1} \left(\frac{1}{\Delta x} e^{iN\pi/2}\;e^{-in\pi}\; \Psi_{k_n}\right) e^{\frac{2\pi i mn}{N}}\\ &= (-1)^m\; \mathrm{IDFT}\left\{\frac{1}{\Delta x} e^{iN\pi/2}\;(-1)^n \;\Psi_{k_n}\right\}_m \\ \Rightarrow \Psi_{k_n} &= \Delta x\, e^{-iN\pi/2}\; (-1)^n \; \;\mathrm{DFT}\Big\{(-1)^m\; \psi_m \Big\}_n \end{align}\]

Likewise,

\[(U_\mathcal{K}\; \psi)_m = (-1)^m \; \mathrm{IDFT}\left\{\frac{1}{\Delta x} e^{iN\pi/2}\;(-1)^n e^{-\frac{i\tau}{4}k_n^2}\; \Psi_{k_n}\right\}_m\]

Putting these results together, we get

\[\begin{align}(U_\mathcal{K}\; \psi)_m = (-1)^m \;\mathrm{IDFT}\Bigg\{ e^{-\frac{i\tau}{4}k_n^2}\; \mathrm{DFT}\Big\{(-1)^{p}\, \psi_{p}\Big\}_n \Bigg\}_m \\ \mathrm{where}\;\; k_n = - \frac{\pi N}{L} + \frac{2\pi n}{L} \end{align}\]

Hence, the \(U_{\mathcal{K}}\) kinetic stepping operator can be implemented by taking a DFT, multiplying the resulting vector elements by \(\exp(-i\tau k_n^2/4)\) phase factors, and taking an IDFT. The runtime of the stepping process is \(O(N\log(N))\). The \(m\), \(n\), and \(p\) indices all run over the range \([0, 1, \cdots, N-1]\), consistent with the standard definition of the DFT and IDFT.