How should I interpret the output of numpy.fft.rfft2?

Question

Obviously the rfft2 function simply computes the discrete fft of the input matrix. However how do I interpret a given index of the output? Given an index of the output, which Fourier coefficient am I looking at?
I am especially confused by the sizes of the output. For an n by n matrix, the output seems to be an n by (n/2)+1 matrix (for even n). Why does a square matrix ends up with a non-square fourier transform?

Answer 1

The output of numpy.fft.rfft2 is simply the left half (plus one column) of a standard two-dimensional FFT, as computed by numpy.fft.fft2. There's no need for rfft2 to supply the right half of the result, because the FFT of a real array has a natural and simple symmetry, and the right half of the full FFT can therefore be derived from the left half using that symmetry.

Here's an example, to illustrate. First, to make it easy to reproduce and easy to see, I'll set NumPy's random state and printing options:

In [1]: import numpy as np

In [2]: np.set_printoptions(precision=3, suppress=True, linewidth=128)

In [3]: random = np.random.RandomState(seed=15206)

Let's create a real input array, with 6 rows and 6 columns:

In [4]: x = random.randn(6, 6)

In [5]: x
Out[5]: 
array([[ 1.577,  0.426,  0.322, -0.891, -0.793,  0.017],
       [ 0.238,  0.603, -0.094, -0.087, -0.936, -1.139],
       [-0.583,  0.394,  0.323, -1.384,  1.255,  0.457],
       [-0.186,  0.687, -0.815, -0.54 ,  0.762, -0.674],
       [-1.604, -0.557,  1.933, -1.122, -0.516, -1.51 ],
       [-1.683, -0.006, -1.648, -0.016,  1.145,  0.809]])

Now take a look at the full FFT (using fft2, not rfft2):

In [6]: fft2_result = np.fft.fft2(x)

In [7]: fft2_result
Out[7]: 
array([[ -5.834+0.j   ,   1.084-2.33j ,  -6.504-3.884j,   3.228-0.j   ,  -6.504+3.884j,   1.084+2.33j ],
       [  1.475-3.311j,   1.865-3.699j,   2.777-0.095j,  -2.570-1.152j,   4.705-3.373j,   4.555-3.657j],
       [  2.758+3.339j,  -3.512+0.398j,   5.824-4.045j,   1.149-3.705j,   0.661-2.127j,  12.368+1.464j],
       [  1.326-0.j   ,   1.191-4.479j,  -3.263+6.19j ,   8.939-0.j   ,  -3.263-6.19j ,   1.191+4.479j],
       [  2.758-3.339j,  12.368-1.464j,   0.661+2.127j,   1.149+3.705j,   5.824+4.045j,  -3.512-0.398j],
       [  1.475+3.311j,   4.555+3.657j,   4.705+3.373j,  -2.570+1.152j,   2.777+0.095j,   1.865+3.699j]])

Notice that there's a symmetry here: for any indices i and j with 0 <= i < 6 and 0 <= j < 6, fft2_result[i, j] is the complex conjugate of fft_result[-i, -j]. For example:

In [8]: fft2_result[2, 4]
Out[8]: (0.66075993512998199-2.127249005984857j)

In [9]: fft2_result[-2, -4].conj()
Out[9]: (0.66075993512998199-2.127249005984857j)

That means that we don't need to include the right-hand half of the output, since it can be derived from the left half. We can save memory and perhaps also a tiny bit of time by only computing the left half of the full FFT. And that's exactly what rfft2 does:

In [10]: rfft2_result = np.fft.rfft2(x)

In [11]: rfft2_result
Out[11]: 
array([[ -5.834+0.j   ,   1.084-2.33j ,  -6.504-3.884j,   3.228+0.j   ],
       [  1.475-3.311j,   1.865-3.699j,   2.777-0.095j,  -2.570-1.152j],
       [  2.758+3.339j,  -3.512+0.398j,   5.824-4.045j,   1.149-3.705j],
       [  1.326-0.j   ,   1.191-4.479j,  -3.263+6.19j ,   8.939-0.j   ],
       [  2.758-3.339j,  12.368-1.464j,   0.661+2.127j,   1.149+3.705j],
       [  1.475+3.311j,   4.555+3.657j,   4.705+3.373j,  -2.570+1.152j]])

Notice that rfft2_result matches fft2_result[:, :4], at least up to numerical error:

In [12]: np.allclose(rfft2_result, fft2_result[:, :4])
Out[12]: True

We could also choose to keep the top half of the output rather than the left half, by using the axes argument to np.fft.rfft2:

In [13]: np.fft.rfft2(x, axes=[1, 0])
Out[13]: 
array([[ -5.834+0.j   ,   1.084-2.33j ,  -6.504-3.884j,   3.228-0.j   ,  -6.504+3.884j,   1.084+2.33j ],
       [  1.475-3.311j,   1.865-3.699j,   2.777-0.095j,  -2.570-1.152j,   4.705-3.373j,   4.555-3.657j],
       [  2.758+3.339j,  -3.512+0.398j,   5.824-4.045j,   1.149-3.705j,   0.661-2.127j,  12.368+1.464j],
       [  1.326+0.j   ,   1.191-4.479j,  -3.263+6.19j ,   8.939-0.j   ,  -3.263-6.19j ,   1.191+4.479j]])

As the documentation for np.fft.rfftn says, NumPy performs a real FFT over the last axis specified, and complex FFTs over the other axes.

Of course, there's still some redundancy in rfft2_result: we could discard the bottom half of the first column, and the bottom half of the last column, and still be able to reconstruct them using the same symmetry as before. And the entries at positions [0, 0], [0, 3], [3, 0] and [3, 3] are all going to be real, so we could discard their imaginary parts. But that would leave us with a less convenient array representation.

Answer 2

Also note the ordering of the coefficients in the fft output:

According to the doc: by default the 1st element is the coefficient for 0 frequency component (effectively the sum or mean of the array), and starting from the 2nd we have coeffcients for the postive frequencies in increasing order, and starts from n/2+1 they are for negative frequencies in decreasing order. To have a view of the frequencies for a length-10 array:

np.fft.fftfreq(10) the output is: array([ 0. , 0.1, 0.2, 0.3, 0.4, -0.5, -0.4, -0.3, -0.2, -0.1])

Use np.fft.fftshift(cf), where cf=np.fft.fft(array), the output is shifted so that it corresponds to this frequency ordering: array([-0.5, -0.4, -0.3, -0.2, -0.1, 0. , 0.1, 0.2, 0.3, 0.4]) which makes for sense for plotting.

In the 2D case it is the same. And the fft2 and rfft2 difference is as explained by others.