Feature/20 fft

stan/math/prim/fun.hpp

@@ -99,6 +99,7 @@
 #include <stan/math/prim/fun/factor_cov_matrix.hpp>
 #include <stan/math/prim/fun/falling_factorial.hpp>
 #include <stan/math/prim/fun/fdim.hpp>
+#include <stan/math/prim/fun/fft.hpp>
 #include <stan/math/prim/fun/fill.hpp>
 #include <stan/math/prim/fun/finite_diff_stepsize.hpp>
 #include <stan/math/prim/fun/floor.hpp>

stan/math/prim/fun/fft.hpp

@@ -0,0 +1,117 @@
+#ifndef STAN_MATH_PRIM_FUN_FFT_HPP
+#define STAN_MATH_PRIM_FUN_FFT_HPP
+
+#include <stan/math/prim/meta.hpp>
+#include <stan/math/prim/fun/Eigen.hpp>
+#include <unsupported/Eigen/FFT>
+#include <Eigen/Dense>
+#include <complex>
+#include <type_traits>
+#include <vector>
+
+namespace stan {
+namespace math {
+
+/**
+ * Return the discrete Fourier transform of the specified complex
+ * vector.
+ *
+ * Given an input complex vector `x[0:N-1]` of size `N`, the discrete
+ * Fourier transform computes entries of the resulting complex
+ * vector `y[0:N-1]` by
+ *
+ * ```
+ * y[n] = SUM_{i < N} x[i] * exp(-n * i * 2 * pi * sqrt(-1) / N)
+ * ```
+ *
+ * If the input is of size zero, the result is a size zero vector.
+ *
+ * @tparam V type of complex vector argument
+ * @param[in] x vector to transform
+ * @return discrete Fourier transform of `x`
+ */
+template <typename V, require_eigen_vector_vt<is_complex, V>* = nullptr>
+inline Eigen::Matrix<scalar_type_t<V>, -1, 1> fft(const V& x) {
+  // copy because fft() requires Eigen::Matrix type
+  Eigen::Matrix<scalar_type_t<V>, -1, 1> xv = x;
+  if (xv.size() <= 1)
+    return xv;
+  Eigen::FFT<base_type_t<V>> fft;
+  return fft.fwd(xv);
+}
+
+/**
+ * Return the inverse discrete Fourier transform of the specified
+ * complex vector.
+ *
+ * Given an input complex vector `y[0:N-1]` of size `N`, the inverse
+ * discrete Fourier transform computes entries of the resulting
+ * complex vector `x[0:N-1]` by
+ *
+ * ```
+ * x[n] = SUM_{i < N} y[i] * exp(n * i * 2 * pi * sqrt(-1) / N)
+ * ```
+ *
+ * If the input is of size zero, the result is a size zero vector.
+ * The only difference between the discrete DFT and its inverse is
+ * the sign of the exponent.
+ *
+ * @tparam V type of complex vector argument
+ * @param[in] y vector to inverse transform
+ * @return inverse discrete Fourier transform of `y`
+ */
+template <typename V, require_eigen_vector_vt<is_complex, V>* = nullptr>
+inline Eigen::Matrix<scalar_type_t<V>, -1, 1> inv_fft(const V& y) {
+  // copy because fft() requires Eigen::Matrix type
+  Eigen::Matrix<scalar_type_t<V>, -1, 1> yv = y;
+  if (y.size() <= 1)
+    return yv;
+  Eigen::FFT<base_type_t<V>> fft;
+  return fft.inv(yv);
+}
+
+/**
+ * Return the two-dimensional discrete Fourier transform of the
+ * specified complex matrix.  The 2D discrete Fourier transform first
+ * runs the discrete Fourier transform on the each row, then on each
+ * column of the result.
+ *
+ * @tparam M type of complex matrix argument
+ * @param[in] x matrix to transform
+ * @return discrete 2D Fourier transform of `x`
+ */
+template <typename M, require_eigen_dense_dynamic_vt<is_complex, M>* = nullptr>
+inline Eigen::Matrix<scalar_type_t<M>, -1, -1> fft2(const M& x) {
+  Eigen::Matrix<scalar_type_t<M>, -1, -1> y(x.rows(), x.cols());
+  for (int i = 0; i < y.rows(); ++i)
+    y.row(i) = fft(x.row(i));
+  for (int j = 0; j < y.cols(); ++j)
+    y.col(j) = fft(y.col(j));
+  return y;
+}
+
+/**
+ * Return the two-dimensional inverse discrete Fourier transform of
+ * the specified complex matrix.  The 2D inverse discrete Fourier
+ * transform first runs the 1D inverse Fourier transform on the
+ * columns, and then on the resulting rows.  The composition of the
+ * FFT and inverse FFT (or vice-versa) is the identity.
+ *
+ * @tparam M type of complex matrix argument
+ * @param[in] y matrix to inverse trnasform
+ * @return inverse discrete 2D Fourier transform of `y`
+ */
+template <typename M, require_eigen_dense_dynamic_vt<is_complex, M>* = nullptr>
+inline Eigen::Matrix<scalar_type_t<M>, -1, -1> inv_fft2(const M& y) {
+  Eigen::Matrix<scalar_type_t<M>, -1, -1> x(y.rows(), y.cols());
+  for (int j = 0; j < x.cols(); ++j)
+    x.col(j) = inv_fft(y.col(j));
+  for (int i = 0; i < x.rows(); ++i)
+    x.row(i) = inv_fft(x.row(i));
+  return x;
+}
+
+}  // namespace math
+}  // namespace stan
+
+#endif

test/unit/math/ad_tolerances.hpp

@@ -69,6 +69,34 @@ struct ad_tolerances {
         grad_hessian_grad_hessian_(1e-2) {}
 };
 
+/**
+ * Return tolerances that are infinite other than for the value
+ * tolerance and gradient tolerance for reverse mode.
+ *
+ * @param val_tol value relative tolerance (default `1e-8`)
+ * @param grad_tol gradient relative tolerance (default `1e-4`)
+ */
+ad_tolerances reverse_only_ad_tolerances(double val_tol = 1e-8,
+                                         double grad_tol = 1e-4) {
+  ad_tolerances tols;
+  tols.gradient_val_ = val_tol;
+  tols.gradient_grad_ = grad_tol;
+  constexpr double inf = std::numeric_limits<double>::infinity();
+  tols.gradient_grad_varmat_matvar_ = inf;
+  tols.gradient_fvar_val_ = inf;
+  tols.gradient_fvar_grad_ = inf;
+  tols.hessian_val_ = inf;
+  tols.hessian_grad_ = inf;
+  tols.hessian_hessian_ = inf;
+  tols.hessian_fvar_val_ = inf;
+  tols.hessian_fvar_grad_ = inf;
+  tols.hessian_fvar_hessian_ = inf;
+  tols.grad_hessian_val_ = inf;
+  tols.grad_hessian_hessian_ = inf;
+  tols.grad_hessian_grad_hessian_ = inf;
+  return tols;
+}
+
 }  // namespace test
 }  // namespace stan
 #endif

test/unit/math/mix/fun/fft_test.cpp

@@ -0,0 +1,177 @@
+#include <test/unit/math/test_ad.hpp>
+
+void expect_fft(const Eigen::VectorXcd& x) {
+  stan::test::ad_tolerances tols = stan::test::reverse_only_ad_tolerances();
+  for (int m = 0; m < x.rows(); ++m) {
+    auto g = [m](const auto& x) {
+      using stan::math::fft;
+      return fft(x)(m);
+    };
+    stan::test::expect_ad(tols, g, x);
+  }
+}
+
+TEST(mathMixFun, fft) {
+  using cvec_t = Eigen::VectorXcd;
+
+  cvec_t x0(0);
+  expect_fft(x0);
+
+  cvec_t x1(1);
+  x1[0] = {1, 2};
+  expect_fft(x1);
+
+  cvec_t x2(2);
+  x2[0] = {1, 2};
+  x2[1] = {-1.3, 2.9};
+  expect_fft(x2);
+
+  Eigen::VectorXcd x3(3);
+  x3[0] = {1, -1.3};
+  x3[1] = {2.9, 14.7};
+  x3[2] = {-12.9, -4.8};
+  expect_fft(x3);
+
+  Eigen::VectorXcd x4(4);
+  x4[0] = {1, -1.3};
+  x4[1] = {-2.9, 14.7};
+  x4[2] = {-12.9, -4.8};
+  x4[3] = {8.398, 9.387};
+  expect_fft(x4);
+}
+
+void expect_inv_fft(const Eigen::VectorXcd& x) {
+  stan::test::ad_tolerances tols = stan::test::reverse_only_ad_tolerances();
+  for (int m = 0; m < x.rows(); ++m) {
+    auto g = [m](const auto& x) {
+      using stan::math::inv_fft;
+      return inv_fft(x)(m);
+    };
+    stan::test::expect_ad(tols, g, x);
+  }
+}
+
+TEST(mathMixFun, invFft) {
+  using cvec_t = Eigen::VectorXcd;
+
+  cvec_t x0(0);
+  expect_inv_fft(x0);
+
+  cvec_t x1(1);
+  x1[0] = {1, 2};
+  expect_inv_fft(x1);
+
+  cvec_t x2(2);
+  x2[0] = {1, 2};
+  x2[1] = {-1.3, 2.9};
+  expect_inv_fft(x2);
+
+  Eigen::VectorXcd x3(3);
+  x3[0] = {1, -1.3};
+  x3[1] = {2.9, 14.7};
+  x3[2] = {-12.9, -4.8};
+  expect_inv_fft(x3);
+
+  Eigen::VectorXcd x4(4);
+  x4[0] = {1, -1.3};
+  x4[1] = {-2.9, 14.7};
+  x4[2] = {-12.9, -4.8};
+  x4[3] = {8.398, 9.387};
+  expect_inv_fft(x4);
+}
+
+void expect_fft2(const Eigen::MatrixXcd& x) {
+  stan::test::ad_tolerances tols = stan::test::reverse_only_ad_tolerances();
+  for (int n = 0; n < x.cols(); ++n) {
+    for (int m = 0; m < x.rows(); ++m) {
+      auto g = [m, n](const auto& x) {
+        using stan::math::fft2;
+        return fft2(x)(m, n);
+      };
+      stan::test::expect_ad(tols, g, x);
+    }
+  }
+}
+
+TEST(mathMixFun, fft2) {
+  using cmat_t = Eigen::MatrixXcd;
+
+  cmat_t x00(0, 0);
+  expect_fft2(x00);
+
+  cmat_t x11(1, 1);
+  x11(0, 0) = {1, 2};
+  expect_fft2(x11);
+
+  cmat_t x21(2, 1);
+  x21(0, 0) = {1, 2};
+  x21(1, 0) = {-1.3, 2.9};
+  expect_fft2(x21);
+
+  cmat_t x22(2, 2);
+  x22(0, 0) = {3, 9};
+  x22(0, 1) = {-13.2, 8.345};
+  x22(1, 0) = {-4.23, 7.566};
+  x22(1, 1) = {1, -12.9};
+  expect_fft2(x22);
+
+  cmat_t x33(3, 3);
+  x33(0, 0) = {3, 9};
+  x33(0, 1) = {-13.2, 8.345};
+  x33(0, 2) = {4.333, -1.9};
+  x33(1, 0) = {-4.23, 7.566};
+  x33(1, 1) = {1, -12.9};
+  x33(1, 2) = {-1.01, -4.01};
+  x33(2, 0) = {3.87, 5.89};
+  x33(2, 1) = {-2.875, -2.999};
+  x33(2, 2) = {12.98, 14.5555};
+  expect_fft2(x33);
+}
+
+void expect_inv_fft2(const Eigen::MatrixXcd& x) {
+  stan::test::ad_tolerances tols = stan::test::reverse_only_ad_tolerances();
+  for (int n = 0; n < x.cols(); ++n) {
+    for (int m = 0; m < x.rows(); ++m) {
+      auto g = [m, n](const auto& x) {
+        using stan::math::inv_fft2;
+        return inv_fft2(x)(m, n);
+      };
+      stan::test::expect_ad(tols, g, x);
+    }
+  }
+}
+
+TEST(mathMixFun, inv_fft2) {
+  using cmat_t = Eigen::MatrixXcd;
+
+  cmat_t x00(0, 0);
+  expect_inv_fft2(x00);
+
+  cmat_t x11(1, 1);
+  x11(0, 0) = {1, 2};
+  expect_inv_fft2(x11);
+
+  cmat_t x21(2, 1);
+  x21(0, 0) = {1, 2};
+  x21(1, 0) = {-1.3, 2.9};
+  expect_inv_fft2(x21);
+
+  cmat_t x22(2, 2);
+  x22(0, 0) = {3, 9};
+  x22(0, 1) = {-13.2, 8.345};
+  x22(1, 0) = {-4.23, 7.566};
+  x22(1, 1) = {1, -12.9};
+  expect_inv_fft2(x22);
+
+  cmat_t x33(3, 3);
+  x33(0, 0) = {3, 9};
+  x33(0, 1) = {-13.2, 8.345};
+  x33(0, 2) = {4.333, -1.9};
+  x33(1, 0) = {-4.23, 7.566};
+  x33(1, 1) = {1, -12.9};
+  x33(1, 2) = {-1.01, -4.01};
+  x33(2, 0) = {3.87, 5.89};
+  x33(2, 1) = {-2.875, -2.999};
+  x33(2, 2) = {12.98, 14.5555};
+  expect_inv_fft2(x33);
+}