fft.hpp

Go to the documentation of this file.

#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 {
 
template <typename V, require_eigen_vector_vt<is_complex, V>* = nullptr,
          require_not_var_t<base_type_t<value_type_t<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);
}
 
template <typename V, require_eigen_vector_vt<is_complex, V>* = nullptr,
          require_not_var_t<base_type_t<value_type_t<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);
}
 
template <typename M, require_eigen_dense_dynamic_vt<is_complex, M>* = nullptr,
          require_not_var_t<base_type_t<value_type_t<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;
}
 
template <typename M, require_eigen_dense_dynamic_vt<is_complex, M>* = nullptr,
          require_not_var_t<base_type_t<value_type_t<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