How to write fftshift and ifftshift in R? [closed]

In numpy, we have the following functions :

import numpy
from numpy.fft import fft2, ifft2, fftshift, ifftshift

I would like to rewrite these functions in R. fft in R works just as fft or fft2 in python. Also for ifft2, we have to do fft(,inverse=T)

Now I would like to know how to rewrite the fftshift and ifftshift functions (for matrices), efficiently in R.

The concept behind fftshift and ifftshift is quite straight forward. Here's a figure that I pulled from MathWorks (creators of MATLAB):

^{Source: MathWorks doc page on fftshift}

Imagine that your input 2D matrix is split up into quadrants. Quadrant #1 is at the top left, quadrant #2 is at the top right, quadrant #3 is at the bottom right and quadrant #4 is at the bottom left. For 2D matrices, fftshift swaps the first and third quadrants and second and fourth quadrants by default. You can override this behaviour where you can simply perform fftshift on one dimension separately. If you do this, you are swapping what are known as half spaces. If you specified to swap along the rows (i.e. dimension 1), then the top half of the matrix gets swapped with the bottom half. If you specified to swap along the columns (i.e. dimension 2), then the right half gets swapped with the left half:

^{Source: MathWorks doc page on fftshift}

Using fftshift by default chains the swapping of dimensions 1 and dimensions 2 in sequence. If you have an even sized matrix where the rows and columns are even, then it's very unambiguous to cut the matrix into the four portions and do the swapping. However, if the matrix is odd sized, it depends on which language you're looking at. In MATLAB and Python numpy for instance, where to perform the switching is defined as (r,c) = ceil(rows/2), ceil(cols/2) where rows and cols are the rows and columns of the matrix. r and c are the row and column of where the swapping takes place.

For ifftshift you simply reverse the actions done on fftshift. Therefore, the default action is to swap dimensions 2 then dimensions 1. However, you have to redefine where the centre of switching is for matrices of odd dimensions. Instead of ceil, you must use floor because this precisely determines where the half spaces were after fftshift was performed and you are now undoing what was done on the original matrix. Therefore, the new centre of switching is (r,c) = floor(rows/2), floor(cols/2). Other than that, the logic to swap among a single dimension is the same - just that the centre of switching has now changed.

As such, here's what I have come up with. These functions take in a matrix defined in R as well as a second optional argument to determine what dimension you want to swap. Omitting this performs the default quadrant swap I just talked about:

fftshift <- function(input_matrix, dim = -1) {

    rows <- dim(input_matrix)[1]    
    cols <- dim(input_matrix)[2]    

    swap_up_down <- function(input_matrix) {
        rows_half <- ceiling(rows/2)
        return(rbind(input_matrix[((rows_half+1):rows), (1:cols)], input_matrix[(1:rows_half), (1:cols)]))
    }

    swap_left_right <- function(input_matrix) {
        cols_half <- ceiling(cols/2)
        return(cbind(input_matrix[1:rows, ((cols_half+1):cols)], input_matrix[1:rows, 1:cols_half]))
    }

    if (dim == -1) {
        input_matrix <- swap_up_down(input_matrix)
        return(swap_left_right(input_matrix))
    }
    else if (dim == 1) {
        return(swap_up_down(input_matrix))
    }
    else if (dim == 2) {
        return(swap_left_right(input_matrix))
    }
    else {
        stop("Invalid dimension parameter")
    }
}

ifftshift <- function(input_matrix, dim = -1) {

    rows <- dim(input_matrix)[1]    
    cols <- dim(input_matrix)[2]    

    swap_up_down <- function(input_matrix) {
        rows_half <- floor(rows/2)
        return(rbind(input_matrix[((rows_half+1):rows), (1:cols)], input_matrix[(1:rows_half), (1:cols)]))
    }

    swap_left_right <- function(input_matrix) {
        cols_half <- floor(cols/2)
        return(cbind(input_matrix[1:rows, ((cols_half+1):cols)], input_matrix[1:rows, 1:cols_half]))
    }

    if (dim == -1) {
        input_matrix <- swap_left_right(input_matrix)
        return(swap_up_down(input_matrix))
    }
    else if (dim == 1) {
        return(swap_up_down(input_matrix))
    }
    else if (dim == 2) {
        return(swap_left_right(input_matrix))
    }
    else {
        stop("Invalid dimension parameter")
    }
}

In each function, I define functions that swap the left and right halves as well as the top and bottom halves. To swap the left and right halves, simply determine what column you're using to perform the swapping and use indexing to create a new matrix by concatenating these two halves together where the first half is the right half and the second half is the left half. You would do the same when swapping the top and bottom halves but find the right row to perform the swap.

The second parameter dim can optionally be set to either 1 or 2 to swap the corresponding dimension. Note that the only difference between fftshift and ifftshift is the centre of swapping that is defined as well as the order of operations when you default to swapping both dimensions. The rest of the code is the same.

As a means of demonstration, suppose I have declared a 5 x 5 numeric matrix like so:

> O <- matrix(1:25, 5, 5)
> O
     [,1] [,2] [,3] [,4] [,5]
[1,]    1    6   11   16   21
[2,]    2    7   12   17   22
[3,]    3    8   13   18   23
[4,]    4    9   14   19   24
[5,]    5   10   15   20   25

Note that the size of the matrix is odd in both dimensions. Performing a fftshift gives us:

> P <- fftshift(O)
> P
     [,1] [,2] [,3] [,4] [,5]
[1,]   19   24    4    9   14
[2,]   20   25    5   10   15
[3,]   16   21    1    6   11
[4,]   17   22    2    7   12
[5,]   18   23    3    8   13

You can see that the corresponding dimensions have been swapped. Reversing this with ifftshift should give us the original matrix back, and it does:

> Q <- ifftshift(P)
> Q
     [,1] [,2] [,3] [,4] [,5]
[1,]    1    6   11   16   21
[2,]    2    7   12   17   22
[3,]    3    8   13   18   23
[4,]    4    9   14   19   24
[5,]    5   10   15   20   25

Of course you can override this and specify which dimension you want to swap, so let's say you wanted to swap just the rows:

> fftshift(O, 1)
     [,1] [,2] [,3] [,4] [,5]
[1,]    4    9   14   19   24
[2,]    5   10   15   20   25
[3,]    1    6   11   16   21
[4,]    2    7   12   17   22
[5,]    3    8   13   18   23

> ifftshift(fftshift(O, 1), 1)
     [,1] [,2] [,3] [,4] [,5]
[1,]    1    6   11   16   21
[2,]    2    7   12   17   22
[3,]    3    8   13   18   23
[4,]    4    9   14   19   24
[5,]    5   10   15   20   25

Take care that when you perform ifftshift on a matrix that was swapped in one dimension, you must use the same dimension that was used for swapping when you used fftshift to ensure you get the original matrix back.

We can also do the same for the second dimension:

> ifftshift(O, 2)
     [,1] [,2] [,3] [,4] [,5]
[1,]   11   16   21    1    6
[2,]   12   17   22    2    7
[3,]   13   18   23    3    8
[4,]   14   19   24    4    9
[5,]   15   20   25    5   10

> ifftshift(fftshift(O, 2), 2)
     [,1] [,2] [,3] [,4] [,5]
[1,]    1    6   11   16   21
[2,]    2    7   12   17   22
[3,]    3    8   13   18   23
[4,]    4    9   14   19   24
[5,]    5   10   15   20   25

As an interesting note, we can verify that numpy does the same thing as what we discussed above. Here's an IPython session:

In [16]: import numpy as np

In [17]: from numpy.fft import fftshift, ifftshift

In [18]: O = np.reshape(np.arange(1,26), (5,5)).T

In [19]: O
Out[19]:
array([[ 1,  6, 11, 16, 21],
       [ 2,  7, 12, 17, 22],
       [ 3,  8, 13, 18, 23],
       [ 4,  9, 14, 19, 24],
       [ 5, 10, 15, 20, 25]])

In [20]: fftshift(O)
Out[20]:
array([[19, 24,  4,  9, 14],
       [20, 25,  5, 10, 15],
       [16, 21,  1,  6, 11],
       [17, 22,  2,  7, 12],
       [18, 23,  3,  8, 13]])

In [21]: ifftshift(fftshift(O))
Out[21]:
array([[ 1,  6, 11, 16, 21],
       [ 2,  7, 12, 17, 22],
       [ 3,  8, 13, 18, 23],
       [ 4,  9, 14, 19, 24],
       [ 5, 10, 15, 20, 25]])

numpy as well as fftshift and ifftshift from the numpy.fft package are imported in and a 2D matrix is created that is the same as what you saw in the example defined in R. We then call fftshift, then fftshift and ifftshift in sequence and we can see that we get the same results as seen in the R code.