added fftshift and ifftshift

Author: dpilger26

include/NumCpp/FFT/fftshift.hpp

@@ -29,6 +29,7 @@
 
 #include "NumCpp/Core/Internal/StaticAsserts.hpp"
 #include "NumCpp/Core/Types.hpp"
+#include "NumCpp/Functions/roll.hpp"
 #include "NumCpp/NdArray.hpp"
 
 namespace nc::fft
@@ -51,6 +52,25 @@ namespace nc::fft
     {
         STATIC_ASSERT_ARITHMETIC_OR_COMPLEX(dtype);
 
-        return {};
+        switch (inAxis)
+        {
+            case Axis::NONE:
+            {
+                return roll(inX, inX.size() / 2, inAxis);
+            }
+            case Axis::COL:
+            {
+                return roll(inX, inX.numCols() / 2, inAxis);
+            }
+            case Axis::ROW:
+            {
+                return roll(inX, inX.numRows() / 2, inAxis);
+            }
+            default:
+            {
+                THROW_INVALID_ARGUMENT_ERROR("Unimplemented axis type.");
+                return {};
+            }
+        }
     }
 } // namespace nc::fft

include/NumCpp/FFT/ifftshift.hpp

@@ -50,6 +50,31 @@ namespace nc::fft
     {
         STATIC_ASSERT_ARITHMETIC_OR_COMPLEX(dtype);
 
-        return {};
+        switch (inAxis)
+        {
+            case Axis::NONE:
+            {
+                auto shift = inX.size() / 2;
+                shift += inX.size() % 2 == 1 ? 1 : 0;
+                return roll(inX, shift, inAxis);
+            }
+            case Axis::COL:
+            {
+                auto shift = inX.numCols() / 2;
+                shift += inX.numCols() % 2 == 1 ? 1 : 0;
+                return roll(inX, shift, inAxis);
+            }
+            case Axis::ROW:
+            {
+                auto shift = inX.numRows() / 2;
+                shift += inX.numRows() % 2 == 1 ? 1 : 0;
+                return roll(inX, shift, inAxis);
+            }
+            default:
+            {
+                THROW_INVALID_ARGUMENT_ERROR("Unimplemented axis type.");
+                return {};
+            }
+        }
     }
 } // namespace nc::fft

test/pytest/test_fft.py

@@ -611,12 +611,44 @@ def test_fftfreq():
 
 ####################################################################################
 def test_fftshift():
-    assert False
+    for _ in range(50):
+        n = np.random.randint(10, 1000)
+        d = np.random.rand()
+        freqs = np.fft.fftfreq(n, d)
+        cFreqs = NumCpp.NdArray(1, freqs.size)
+        cFreqs.setArray(freqs)
+        assert np.array_equal(np.round(NumCpp.fftshift(cFreqs, NumCpp.Axis.NONE).flatten(), 8), np.round(np.fft.fftshift(freqs), 8))
+
+        dim0 = np.random.randint(10, 100)
+        dim1 = np.random.randint(10, 100)
+        n = dim0 * dim1
+        d = np.random.rand()
+        freqs = np.fft.fftfreq(n, d).reshape(dim0, dim1)
+        cFreqs = NumCpp.NdArray(freqs.shape[0], freqs.shape[1])
+        cFreqs.setArray(freqs)
+        assert np.array_equal(np.round(NumCpp.fftshift(cFreqs, NumCpp.Axis.ROW), 8), np.round(np.fft.fftshift(freqs, axes=0), 8))
+        assert np.array_equal(np.round(NumCpp.fftshift(cFreqs, NumCpp.Axis.COL), 8), np.round(np.fft.fftshift(freqs, axes=1), 8))
 
 
 ####################################################################################
 def test_ifftshift():
-    assert False
+    for _ in range(50):
+        n = np.random.randint(10, 1000)
+        d = np.random.rand()
+        freqs = np.fft.fftfreq(n, d)
+        cFreqs = NumCpp.NdArray(1, freqs.size)
+        cFreqs.setArray(freqs)
+        assert np.array_equal(np.round(NumCpp.ifftshift(cFreqs, NumCpp.Axis.NONE).flatten(), 8), np.round(np.fft.ifftshift(freqs), 8))
+
+        dim0 = np.random.randint(10, 100)
+        dim1 = np.random.randint(10, 100)
+        n = dim0 * dim1
+        d = np.random.rand()
+        freqs = np.fft.fftfreq(n, d).reshape(dim0, dim1)
+        cFreqs = NumCpp.NdArray(freqs.shape[0], freqs.shape[1])
+        cFreqs.setArray(freqs)
+        assert np.array_equal(np.round(NumCpp.ifftshift(cFreqs, NumCpp.Axis.ROW), 8), np.round(np.fft.ifftshift(freqs, axes=0), 8))
+        assert np.array_equal(np.round(NumCpp.ifftshift(cFreqs, NumCpp.Axis.COL), 8), np.round(np.fft.ifftshift(freqs, axes=1), 8))
 
 
 ####################################################################################