Added helper functions (fftfreq and fftshift) + tests

Author: egpbos

CMakeLists.txt

@@ -109,6 +109,7 @@ endif(FIX_RPATH)
 
 set(XTENSOR_FFTW_HEADERS
         ${XTENSOR_FFTW_INCLUDE_DIR}/xtensor-fftw/basic.hpp
+        ${XTENSOR_FFTW_INCLUDE_DIR}/xtensor-fftw/helper.hpp
         ${XTENSOR_FFTW_INCLUDE_DIR}/xtensor-fftw/xtensor-fftw_config.hpp
 )
 

include/xtensor-fftw/helper.hpp

@@ -0,0 +1,78 @@
+/*
+ * xtensor-fftw
+ * Copyright (c) 2017, Patrick Bos
+ * Distributed under the terms of the BSD 3-Clause License.
+ *
+ * The full license is in the file LICENSE, distributed with this software.
+ */
+
+#ifndef XTENSOR_FFTW_HELPER_HPP
+#define XTENSOR_FFTW_HELPER_HPP
+
+#include <xtensor/xarray.hpp>
+
+namespace xt {
+  namespace fftw {
+
+    template <typename T>
+    xt::xarray<T> fftshift(xt::xarray<T> in) {
+      // partly mimic np.fftshift (only 1D arrays)
+      xt::xarray<T> shifted(in);
+      auto it_in = in.begin();
+      for (std::size_t ix_shifted = in.size()/2; it_in != in.end(); ++it_in, ++ix_shifted) {
+        shifted[ix_shifted % in.size()] = *it_in;
+      }
+      return shifted;
+    }
+
+    template <typename T>
+    xt::xarray<T> ifftshift(xt::xarray<T> shifted) {
+      // partly mimic np.ifftshift (only 1D arrays)
+      xt::xarray<T> out(shifted);
+      auto it_out = out.begin();
+      for (std::size_t ix_shifted = out.size()/2; it_out != out.end(); ++ix_shifted, ++it_out) {
+        *it_out = shifted[ix_shifted % out.size()];
+      }
+      return out;
+    }
+
+
+    template <typename T>
+    xt::xarray<T> fftfreq(unsigned long n, T d=1.0) {
+      // mimic np.fftfreq
+      T df = 1 / (d * static_cast<T>(n));
+      xt::xarray<T> frequencies;
+      if (n % 2 == 0) {
+        frequencies = xt::arange<T>(-static_cast<long>(n/2), n/2) * df;
+      } else {
+        frequencies = xt::arange<T>(-static_cast<long>(n/2), n/2 + 1) * df;
+      }
+      frequencies = ifftshift(frequencies);
+      return frequencies;
+    }
+
+    template <typename T>
+    xt::xarray<T> fftscale(unsigned long n, T d=1.0) {
+      // mimic np.fftfreq, but in scale space instead of frequency space (dk = 2\pi df)
+      return 2 * M_PI * fftfreq(n, d);
+    }
+
+
+    template <typename T>
+    xt::xarray<T> rfftfreq(unsigned long n, T d=1.0) {
+      // mimic np.rfftfreq
+      T df = 1 / (d * static_cast<T>(n));
+      xt::xarray<T> frequencies = xt::arange<T>(0., n/2 + 1) * df;
+      return frequencies;
+    }
+
+    template <typename T>
+    xt::xarray<T> rfftscale(unsigned long n, T d=1.0) {
+      // mimic np.rfftfreq, but in scale space instead of frequency space (dk = 2\pi df)
+      return 2 * M_PI * rfftfreq(n, d);
+    }
+
+  }
+}
+
+#endif //XTENSOR_FFTW_HELPER_HPP

test/CMakeLists.txt

@@ -96,6 +96,7 @@ include_directories(${GTEST_INCLUDE_DIRS})
 
 set(XTENSOR_FFTW_TESTS
         basic_interface.cpp
+        helper.cpp
 )
 
 set(XTENSOR_FFTW_TARGET test_xtensor-fftw)

test/helper.cpp

@@ -0,0 +1,57 @@
+/*
+ * xtensor-fftw
+ * Copyright (c) 2017, Patrick Bos
+ * Distributed under the terms of the BSD 3-Clause License.
+ *
+ * The full license is in the file LICENSE, distributed with this software.
+ */
+
+#include <xtensor/xarray.hpp>
+#include <xtensor-fftw/helper.hpp>
+
+#include "gtest/gtest.h"
+
+
+TEST(helper, fftshift) {
+  xt::xarray<double> odd = {3, 4, 0, 1, 2};
+  xt::xarray<double> even = {3, 4, 5, 0, 1, 2};
+  xt::xarray<double> odd_range = xt::arange<double>(5);
+  xt::xarray<double> even_range = xt::arange<double>(6);
+  EXPECT_EQ(xt::fftw::fftshift(odd_range), odd);
+  EXPECT_EQ(xt::fftw::fftshift(even_range), even);
+}
+
+TEST(helper, ifftshift) {
+  xt::xarray<double> odd = {3, 4, 0, 1, 2};
+  xt::xarray<double> even = {3, 4, 5, 0, 1, 2};
+  EXPECT_EQ(xt::fftw::ifftshift(odd), xt::arange<double>(5));
+  EXPECT_EQ(xt::fftw::ifftshift(even), xt::arange<double>(6));
+}
+
+TEST(helper, fftfreq) {
+  xt::xarray<double> reference9 = {0.,  1.,  2.,  3.,  4., -4., -3., -2., -1.};
+  xt::xarray<double> reference10 = {0.  ,  0.04,  0.08,  0.12,  0.16, -0.2 , -0.16, -0.12, -0.08, -0.04};
+  EXPECT_EQ(xt::fftw::fftfreq(9, 1./9), reference9);
+  EXPECT_EQ(xt::fftw::fftfreq(10, 2.5), reference10);
+}
+
+TEST(helper, fftscale) {
+  xt::xarray<double> reference9 = {0.,  1.,  2.,  3.,  4., -4., -3., -2., -1.};
+  xt::xarray<double> reference10 = {0.  ,  0.04,  0.08,  0.12,  0.16, -0.2 , -0.16, -0.12, -0.08, -0.04};
+  EXPECT_EQ(xt::fftw::fftscale(9, 1./9), 2 * M_PI * reference9);
+  EXPECT_EQ(xt::fftw::fftscale(10, 2.5), 2 * M_PI * reference10);
+}
+
+TEST(helper, rfftfreq) {
+  xt::xarray<double> reference9 = {0.,  1.,  2.,  3.,  4.};
+  xt::xarray<double> reference10 = {0.  ,  0.04,  0.08,  0.12,  0.16, 0.2};
+  EXPECT_EQ(xt::fftw::rfftfreq(9, 1./9), reference9);
+  EXPECT_EQ(xt::fftw::rfftfreq(10, 2.5), reference10);
+}
+
+TEST(helper, rfftscale) {
+  xt::xarray<double> reference9 = {0.,  1.,  2.,  3.,  4.};
+  xt::xarray<double> reference10 = {0.  ,  0.04,  0.08,  0.12,  0.16, 0.2};
+  EXPECT_EQ(xt::fftw::rfftscale(9, 1./9), 2 * M_PI * reference9);
+  EXPECT_EQ(xt::fftw::rfftscale(10, 2.5), 2 * M_PI * reference10);
+}