From a0db3777d2f2a8c9db9e95c6530f68a2760efbe1 Mon Sep 17 00:00:00 2001
From: kradchen <chenjuijun@realcan.cn>
Date: Thu, 8 Jun 2023 15:24:49 +0800
Subject: [PATCH] Add ifftshift

---
 src/Function2D.cpp       | 25 +++++++++++++++++++++----
 src/Function2D.h         |  6 ++++++
 test/Function2D_Test.cpp | 11 +++++++++++
 3 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/src/Function2D.cpp b/src/Function2D.cpp
index a34d1d4..06ca233 100644
--- a/src/Function2D.cpp
+++ b/src/Function2D.cpp
@@ -837,13 +837,30 @@ Matrix Aurora::ifft_symmetric(const Matrix &aMatrix,long length)
 
 void Aurora::fftshift(Matrix &aMatrix){
     int backwardLength = aMatrix.getDimSize(0)/2;
+    if (aMatrix.getDimSize(0)%2!=0)++backwardLength;
     int forwardLength = aMatrix.getDimSize(0) - backwardLength;
     double* buffer = malloc(forwardLength,true);
+    int valueStep = aMatrix.isComplex()?2:1;
     for (int i = 0; i<aMatrix.getDimSize(1); ++i) {
-        double* dataPtr = aMatrix.getData()+aMatrix.getDimSize(0)*i*2;
-        cblas_dcopy(forwardLength*2, dataPtr+backwardLength*2, 1, buffer, 1);
-        cblas_dcopy(backwardLength*2, dataPtr, 1, dataPtr+forwardLength*2, 1);
-        cblas_dcopy(forwardLength*2, buffer, 1, dataPtr, 1);
+        double* dataPtr = aMatrix.getData()+aMatrix.getDimSize(0)*i*valueStep;
+        cblas_dcopy(forwardLength*valueStep, dataPtr+backwardLength*valueStep, 1, buffer, 1);
+        cblas_dcopy(backwardLength*valueStep, dataPtr, 1, dataPtr+forwardLength*valueStep, 1);
+        cblas_dcopy(forwardLength*valueStep, buffer, 1, dataPtr, 1);
+    }
+    Aurora::free(buffer);
+}
+
+void Aurora::ifftshift(Matrix &aMatrix){
+    int forwardLength= aMatrix.getDimSize(0)/2;
+    if (aMatrix.getDimSize(0)%2!=0)++forwardLength;
+    int backwardLength = aMatrix.getDimSize(0) - forwardLength;
+    double* buffer = malloc(forwardLength,true);
+    int valueStep = aMatrix.isComplex()?2:1;
+    for (int i = 0; i<aMatrix.getDimSize(1); ++i) {
+        double* dataPtr = aMatrix.getData()+aMatrix.getDimSize(0)*i*valueStep;
+        cblas_dcopy(forwardLength*valueStep, dataPtr+backwardLength*valueStep, 1, buffer, 1);
+        cblas_dcopy(backwardLength*valueStep, dataPtr, 1, dataPtr+forwardLength*valueStep, 1);
+        cblas_dcopy(forwardLength*valueStep, buffer, 1, dataPtr, 1);
     }
     Aurora::free(buffer);
 }
diff --git a/src/Function2D.h b/src/Function2D.h
index a7877fa..40313ab 100644
--- a/src/Function2D.h
+++ b/src/Function2D.h
@@ -135,6 +135,12 @@ namespace Aurora
     */
     void fftshift(Matrix &aMatrix);
 
+    /**
+    * fftshift,在原有数据上进行修改，将fft的数据的前半部分和后半部分交换，支持2D数据
+    * @param aMatrix 
+    */
+    void ifftshift(Matrix &aMatrix);
+
     /**
      * 逆fft，支持到2维，输入必须是复数，输出必是复数
      * @attention 如有需要可使用real去除虚部
diff --git a/test/Function2D_Test.cpp b/test/Function2D_Test.cpp
index 87ba200..4b0c563 100644
--- a/test/Function2D_Test.cpp
+++ b/test/Function2D_Test.cpp
@@ -369,6 +369,17 @@ TEST_F(Function2D_Test, fftAndComplexAndIfft){
     EXPECT_DOUBLE_EQ(-3.4641, fourDecimalRound(result3[4].imag()));
     EXPECT_DOUBLE_EQ(-1.0, fourDecimalRound(result3[13].real()));
     EXPECT_DOUBLE_EQ(-3.4641, fourDecimalRound(result3[16].imag()));
+    {
+        auto fm = Aurora::Matrix::fromRawData(new double[7]{1,2,3, 4, 5, 6, 7},7,1);
+        Aurora::fftshift(fm);
+        EXPECT_DOUBLE_EQ(5, fm[0]);
+        EXPECT_DOUBLE_EQ(1, fm[3]);
+        EXPECT_DOUBLE_EQ(4, fm[6]);
+        Aurora::ifftshift(fm);
+        EXPECT_DOUBLE_EQ(1, fm[0]);
+        EXPECT_DOUBLE_EQ(4, fm[3]);
+        EXPECT_DOUBLE_EQ(7, fm[6]);
+    }
 }
 
 TEST_F(Function2D_Test, hilbert) {