Interleaving FFT real & complex parts in a tensor

1

I have a use-case where I have to do FFT for a given tensor as. Here, FFT is applied to each of the 10 rows, in a column-wise manner which gives the dimension (10, 11) post FFT.

# Random data-
x = torch.rand((10, 20))

# Compute RFFT of 'x'-
x_fft = torch.fft.rfft(x)

# Sanity check-
x.shape, x_fft.shape
# (torch.Size([10, 20]), torch.Size([10, 11]))

# FFT for the first 2 rows are-
x_fft[:2, :]
'''
tensor([[12.2561+0.0000j,  0.7551-1.2075j,  1.1119-0.0458j, -0.2814-1.5266j,
          1.4083-0.7302j,  0.6648+0.3311j,  0.3969+0.0632j, -0.8031-0.1904j,
         -0.4206+0.9066j, -0.2149+0.9160j,  0.4800+0.0000j],
        [ 9.8967+0.0000j, -0.5100-0.2377j, -0.6344+2.2406j,  0.4584-1.0705j,
          0.2235+0.4788j, -0.3923+0.8205j, -1.0372-0.0292j, -1.6368+0.5517j,
          1.5093+0.0419j,  0.5755-1.2133j,  2.9269+0.0000j]])
'''

# The goal is to have for each row, 1-D vector (of size 11) as follows:
# So, for first row, the desired 1-D vector (size = 11) is-
[12.2561, 0.0000, 0.7551, -1.2075, 1.1119, -0.0458, -0.2814, -1.5266,
1.4083, -0.7302, 0.6648, 0.3311, 0.3969, 0.0632, -0.8031, -0.1904,
-0.4206, 0.9066, -0.2149, 0.9160, 0.4800, 0.0000]
'''

Here, you are taking the real and imaginary components and placing them adjacent to each other.

Adjacent means:

[a_1_real, a_1_imag, a_2_real, a_2_imag, a_3_real, a_3_imag, ....., a_n_real, a_n_imag]

Since for each row, you get 11 FFT complex numbers, a_n = a_11.

How to go about it?

2

Could you clarify your desired output, considering that x_fft[:-2,:].real and .imag are 2d Tensors with shape (8, 11). What exactly is an from your message? That suggests a 1d Tensor and is ambiguous in the 2d case. Perhaps I misunderstood!

Also in your example above, the first 4 numbers are a bit confusing:

[1.1547e+01, 0.0000, -0.8197, 0.1383, ...
...
[a1_real, a1_imag, a2_real, a2_imag, ...

The 3rd number -0.8197 should be a2_real, but -0.8197 looks like it’s in the imaginary part when looking at the tensors above.

3

I have edited the question. Hope it helps?

4

Got it! Try this (stacking along the 3rd dimension):

torch.dstack([x_fft.real, x_fft.imag]).view(10, 11 * 2)
1 Like