ENH: Implement fft.fftshift/ifftshift with np.roll for improved performance (#10073)

See the PR for benchmarking information

1 files changed, 17 insertions, 26 deletions

diff --git a/numpy/fft/helper.py b/numpy/fft/helper.py
index 0856d6759..1a1266e12 100644
--- a/numpy/fft/helper.py
+++ b/numpy/fft/helper.py
@@ -6,11 +6,8 @@ from __future__ import division, absolute_import, print_function
 
 import collections
 import threading
-
 from numpy.compat import integer_types
-from numpy.core import (
-        asarray, concatenate, arange, take, integer, empty
-        )
+from numpy.core import integer, empty, arange, asarray, roll
 
 # Created by Pearu Peterson, September 2002
 
@@ -63,19 +60,16 @@ def fftshift(x, axes=None):
            [-1., -3., -2.]])
 
     """
-    tmp = asarray(x)
-    ndim = tmp.ndim
+    x = asarray(x)
     if axes is None:
-        axes = list(range(ndim))
+        axes = tuple(range(x.ndim))
+        shift = [dim // 2 for dim in x.shape]
     elif isinstance(axes, integer_types):
-        axes = (axes,)
-    y = tmp
-    for k in axes:
-        n = tmp.shape[k]
-        p2 = (n+1)//2
-        mylist = concatenate((arange(p2, n), arange(p2)))
-        y = take(y, mylist, k)
-    return y
+        shift = x.shape[axes] // 2
+    else:
+        shift = [x.shape[ax] // 2 for ax in axes]
+
+    return roll(x, shift, axes)
 
 
 def ifftshift(x, axes=None):
@@ -112,19 +106,16 @@ def ifftshift(x, axes=None):
            [-3., -2., -1.]])
 
     """
-    tmp = asarray(x)
-    ndim = tmp.ndim
+    x = asarray(x)
     if axes is None:
-        axes = list(range(ndim))
+        axes = tuple(range(x.ndim))
+        shift = [-(dim // 2) for dim in x.shape]
     elif isinstance(axes, integer_types):
-        axes = (axes,)
-    y = tmp
-    for k in axes:
-        n = tmp.shape[k]
-        p2 = n-(n+1)//2
-        mylist = concatenate((arange(p2, n), arange(p2)))
-        y = take(y, mylist, k)
-    return y
+        shift = -(x.shape[axes] // 2)
+    else:
+        shift = [-(x.shape[ax] // 2) for ax in axes]
+
+    return roll(x, shift, axes)
 
 
 def fftfreq(n, d=1.0):