How to perform FFT from plane-wave basis function coefficients to real space?

Question

I have a 3D grid in real space of grid spacing $L$ and say 21 grid points in each direction, containing e.g. a charge distribution. This is stored as a numpy array of shape (21, 21, 21). The grid points themselves are of course vectors and so are stored as a numpy array of shape (3, 21, 21, 21).

I also have a set of plane waves that form a basis for reciprocal space. These are of the form $$ \psi_{\alpha}(\vec r) = \exp\left(i\vec G_\alpha\cdot \vec r\right) $$ such that $|\vec G_\alpha|^2 \le \text{cutoff}$, so as to form a finite set. For reference, all the $\vec G_\alpha$ are vertically stacked to form an $\alpha\times3$ matrix, which I'll call $\mathbb G$. Now we can expand our charge density $n(\vec r)$ in terms of the basis: $$ n(\vec r) = c_\alpha\psi_\alpha(\vec r) $$

This means that the computational representation for the charge density is an ordered list of the $\alpha$ numbers $c_\alpha$, i.e. an $\alpha$-length vector $\vec c$. My current objective is to reconstruct the real-space charge density from $\vec c$.

The way I currently do this is through the following code:

I = np.exp(1j * G @ real_grid)
density_reconstruction = c @ I

In other words, I precompute the tensor $\mathbb I$ of shape (α, 21, 21, 21) which is essentially the value of each of the basis functions on each of the grid points. I then matrix multiply/contract along the $\alpha$-dimension with $\vec c$ to get a (21, 21, 21) array - this is the reconstructed charge density. Although this method works, the matrix multiplication in the second step is very expensive.

Question: I have heard that, instead of using a costly matrix multiplication, it is possible to do perform this step using a Fourier transform (in particular, FFT), owing to the form of the plane wave. How can I implement this, e.g. using the scipy.fft package? How does it work for the case where the number of basis functions $\ne$ the number of real-space grid points?

Answer 1

The whole point of the FFT is to avoid the costly matrix-vector multiplication of the DFT, which has an effort of $\mathcal O(N^2)$ (where $N = N_x N_y N_z$), and perform the unitary transformation iteratively to obtain $\mathcal O(N \log N)$. What you do is, in the first step, simply Fast-Fourier-Transform your data to get the frequency components (which are often ordered a bit differently as one would expect, see the documentation of your chosen FFT). Moreover, you can save another factors of two if you have real and/or even/odd grid-data.

For the case, where the number of plane waves $\neq$ grid points, you can apply the normal FFT, and interpolate the frequency points to a new grid using trigonometric interpolation (which should be preferably smaller or equal than the number of grid points, otherwise you fake a resolution which is not there). This scales as long as the number of constrained frequency points is not much smaller than the original number of grid points (if that is the case, one can return to normal DFT again).

No scipy.fft interpolation, sorry.

---

A few more notes about the interpolation approach (the discussion was brought up by @ConvexHull, thanks again for that).

If you do a Fourier transform of a $2\pi$-periodic function $f(x)$, you'll basically calculate:

$$c_k = \frac{1}{2\pi} \int_{0}^{2\pi} f(x)\, e^{-i k x} \, dx$$

to obtain the Fourier transform

$$ f(x)=\sum_{k=-\infty}^\infty c_k\, e^{i k x} = \sum_{k=-\infty}^\infty c_k\, z^k $$

with $z=e^{ix}$. If you'll approximate this for convenience by an odd number of Fourier modes $(2N+1)$, you'll get

$$ f_N(x)=\sum_{k=-N}^N c_k\, z^k $$

The coefficients $c_k$ is what the DFT/FFT yields for a grid of size $(2N+1)$. The way to up- and down-sample this term to a function $f_M(x), M\neq N$ is particularly clear: for $M>N$, set all $c_{k}=0$ where $|k|>N$. And for $M<N$, simply truncate the $c_{k}$ with $|k|>M$. This is what was suggested by @ConvexHull, as far as I understood, and it's correct.

The same effect would be obtained by trigonometric interpolation in the real/nodal space down to $M$ equally spaced grid points, and thereafter performing the FFT on this reduced grid. The trigonometric interpolation in barycentric form is given, e.g., here. However, as this trigonometric interpolation requires an effort of $\mathcal O(MN)$, there is no advantage in using it over the truncation of the Fourier series.