diff --git a/stan/math/prim/fun.hpp b/stan/math/prim/fun.hpp index e8d6fdf4273..8a7bc8d8ec0 100644 --- a/stan/math/prim/fun.hpp +++ b/stan/math/prim/fun.hpp @@ -99,6 +99,7 @@ #include #include #include +#include #include #include #include diff --git a/stan/math/prim/fun/fft.hpp b/stan/math/prim/fun/fft.hpp new file mode 100644 index 00000000000..ffca15a0108 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include +#include + +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 * = nullptr> +inline Eigen::Matrix, -1, 1> fft(const V& x) { + // copy because fft() requires Eigen::Matrix type + Eigen::Matrix, -1, 1> xv = x; + if (xv.size() <= 1) + return xv; + Eigen::FFT> 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 * = nullptr> +inline Eigen::Matrix, -1, 1> inv_fft(const V& y) { + // copy because fft() requires Eigen::Matrix type + Eigen::Matrix, -1, 1> yv = y; + if (y.size() <= 1) + return yv; + Eigen::FFT> 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 * = nullptr> +inline Eigen::Matrix, -1, -1> fft2(const M& x) { + Eigen::Matrix, -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 * = nullptr> +inline Eigen::Matrix, -1, -1> inv_fft2(const M& y) { + Eigen::Matrix, -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 diff --git a/test/unit/math/ad_tolerances.hpp b/test/unit/math/ad_tolerances.hpp index 9d8e84e4b35..a9a2087ab64 100644 --- a/test/unit/math/ad_tolerances.hpp +++ b/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::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 diff --git a/test/unit/math/mix/fun/fft_test.cpp b/test/unit/math/mix/fun/fft_test.cpp new file mode 100644 index 00000000000..78ec560693a --- /dev/null +++ b/test/unit/math/mix/fun/fft_test.cpp @@ -0,0 +1,177 @@ +#include + +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); +} diff --git a/test/unit/math/prim/fun/fft_test.cpp b/test/unit/math/prim/fun/fft_test.cpp new file mode 100644 index 00000000000..78a97bee665 --- /dev/null +++ b/test/unit/math/prim/fun/fft_test.cpp @@ -0,0 +1,191 @@ +#include +#include +#include + +TEST(primFun, fft) { + using c_t = std::complex; + using cv_t = Eigen::Matrix, -1, 1>; + using stan::math::fft; + + // reference answers calculated using Scipy.fft with double precision + + cv_t x0(0); + cv_t y0 = fft(x0); + EXPECT_EQ(0, y0.size()); + + cv_t x1(1); + x1 << c_t(-3.247, 1.98555); + cv_t y1 = fft(x1); + EXPECT_EQ(1, y1.size()); + EXPECT_EQ(real(y1[0]), -3.247); + EXPECT_EQ(imag(y1[0]), 1.98555); + + cv_t y1b = fft(x1 + x1); + EXPECT_NEAR(real(y1b[0]), -3.247 * 2, 1e-6); + EXPECT_NEAR(imag(y1b[0]), 1.98555 * 2, 1e-6); + + cv_t x(3); + x << c_t(1, -2), c_t(-3, 5), c_t(-7, 11); + cv_t y = fft(x); + EXPECT_EQ(3, y.size()); + EXPECT_NEAR(real(y[0]), -9, 1e-6); + EXPECT_NEAR(imag(y[0]), 14, 1e-6); + EXPECT_NEAR(real(y[1]), 0.80384758, 1e-6); + EXPECT_NEAR(imag(y[1]), -13.46410162, 1e-6); + EXPECT_NEAR(real(y[2]), 11.19615242, 1e-6); + EXPECT_NEAR(imag(y[2]), -6.53589838, 1e-6); + + cv_t yb = fft(x + x); + EXPECT_EQ(3, yb.size()); + EXPECT_NEAR(real(yb[0]), 2 * -9, 1e-6); + EXPECT_NEAR(imag(yb[0]), 2 * 14, 1e-6); + EXPECT_NEAR(real(yb[1]), 2 * 0.80384758, 1e-6); + EXPECT_NEAR(imag(yb[1]), 2 * -13.46410162, 1e-6); + EXPECT_NEAR(real(yb[2]), 2 * 11.19615242, 1e-6); + EXPECT_NEAR(imag(yb[2]), 2 * -6.53589838, 1e-6); +} + +template +void expect_complex_mat_eq(const T& x, const U& y, double tol = 1e-8) { + EXPECT_EQ(x.rows(), y.rows()); + EXPECT_EQ(x.cols(), y.cols()); + for (int j = 0; j < x.cols(); ++j) { + for (int i = 0; i < x.rows(); ++i) { + EXPECT_FLOAT_EQ(real(x(i, j)), real(y(i, j))); + EXPECT_FLOAT_EQ(imag(x(i, j)), imag(y(i, j))); + } + } +} + +TEST(primFun, inv_fft) { + using c_t = std::complex; + using cv_t = Eigen::Matrix, -1, 1>; + using stan::math::inv_fft; + + // reference answers calculated using Scipy.fft with double precision + + cv_t y0(0); + cv_t x0 = inv_fft(y0); + cv_t x0_expected(0); + EXPECT_EQ(x0_expected, x0); + + cv_t y1(1); + y1 << c_t(-3.247, 1.98555); + cv_t x1 = inv_fft(y1); + cv_t x1_expected(1); + x1_expected << c_t(-3.247, 1.98555); + expect_complex_mat_eq(x1_expected, x1); + + EXPECT_EQ(1, x1.size()); + EXPECT_EQ(real(x1[0]), -3.247); + EXPECT_EQ(imag(x1[0]), 1.98555); + + cv_t y(3); + y << c_t(-9, 14), c_t(0.80384758, -13.46410162), + c_t(11.19615242, -6.53589838); + cv_t x = inv_fft(y); + EXPECT_EQ(3, y.size()); + Eigen::VectorXcd x_expected(3); + x_expected << c_t(1, -2), c_t(-3, 5), c_t(-7, 11); + expect_complex_mat_eq(x_expected, x); +} + +TEST(primFun, fft2) { + using c_t = std::complex; + using cm_t = Eigen::Matrix, -1, -1>; + using stan::math::fft2; + using stan::math::inv_fft2; + + // reference answers calculated using Scipy.fft with double + // precision, and verified with R + + cm_t x(0, 0); + cm_t y = fft2(x); + EXPECT_EQ(0, y.rows()); + EXPECT_EQ(0, y.cols()); + + cm_t x11(1, 1); + x11 << c_t(1.0, -3.9); + cm_t y11 = fft2(x11); + EXPECT_NEAR(1.0, std::real(y11(0, 0)), 1e-6); + EXPECT_NEAR(-3.9, std::imag(y11(0, 0)), 1e-6); + + cm_t x12(1, 2); + x12 << c_t(1.0, -3.9), c_t(-8.6, 0.2); + cm_t y12 = fft2(x12); + cm_t y12_expected(1, 2); + y12_expected << c_t(-7.6, -3.7), c_t(9.6, -4.1); + expect_complex_mat_eq(y12_expected, y12); + + cm_t x33(3, 3); + x33 << c_t(1, 2), c_t(3, -1.4), c_t(2, 1), c_t(3, -9), c_t(2, -1.3), + c_t(3.9, -1.8), c_t(13, -1.8), c_t(1.3, 1.9), c_t(-2.2, -2.2); + cm_t y33 = fft2(x33); + cm_t y33_expected(3, 3); + y33_expected << c_t(27, -12.6), c_t(13.90525589, -9.15166605), + c_t(10.09474411, -4.64833395), c_t(-13.160254038, 11.471281292), + c_t(-13.29326674, 20.88153533), c_t(-13.25262794, 15.82794549), + c_t(4.160254038, 5.928718708), c_t(-11.34737206, -7.72794549), + c_t(4.89326674, -1.98153533); + expect_complex_mat_eq(y33_expected, y33); +} + +TEST(primFunFFT, invfft2) { + using c_t = std::complex; + using cm_t = Eigen::Matrix, -1, -1>; + using stan::math::fft2; + using stan::math::inv_fft; + using stan::math::inv_fft2; + + // reference answers calculated using R + cm_t x(0, 0); + cm_t y = inv_fft2(x); + EXPECT_EQ(0, y.rows()); + EXPECT_EQ(0, y.cols()); + + cm_t x11(1, 1); + x11 << c_t(1.0, -3.9); + cm_t y11 = inv_fft2(x11); + EXPECT_NEAR(1.0, std::real(y11(0, 0)), 1e-6); + EXPECT_NEAR(-3.9, std::imag(y11(0, 0)), 1e-6); + + cm_t x13(1, 3); + x13 << c_t(-2.3, 1.82), c_t(1.18, 9.32), c_t(1.15, -14.1); + cm_t y13 = inv_fft2(x13); + cm_t y13copy = inv_fft(x13.row(0)); + expect_complex_mat_eq(y13, y13copy.transpose()); + + cm_t x33(3, 3); + x33 << c_t(1, 2), c_t(3, -1.4), c_t(2, 1), c_t(3, -9), c_t(2, -1.3), + c_t(3.9, -1.8), c_t(13, -1.8), c_t(1.3, 1.9), c_t(-2.2, -2.2); + cm_t y33 = fft2(x33); + + // check versus results from R + EXPECT_NEAR(27.000000, real(y33(0, 0)), 1e-5); + EXPECT_NEAR(-12.600000, imag(y33(0, 0)), 1e-5); + EXPECT_NEAR(13.90526, real(y33(0, 1)), 1e-5); + EXPECT_NEAR(-9.15167, imag(y33(0, 1)), 1e-5); + EXPECT_NEAR(10.094744, real(y33(0, 2)), 1e-5); + EXPECT_NEAR(-4.648334, imag(y33(0, 2)), 1e-5); + EXPECT_NEAR(-13.160254, real(y33(1, 0)), 1e-5); + EXPECT_NEAR(11.471281, imag(y33(1, 0)), 1e-5); + EXPECT_NEAR(-13.29327, real(y33(1, 1)), 1e-5); + EXPECT_NEAR(20.88154, imag(y33(1, 1)), 1e-5); + EXPECT_NEAR(-13.252628, real(y33(1, 2)), 1e-5); + EXPECT_NEAR(15.827945, imag(y33(1, 2)), 1e-5); + EXPECT_NEAR(4.160254, real(y33(2, 0)), 1e-5); + EXPECT_NEAR(5.928719, imag(y33(2, 0)), 1e-5); + EXPECT_NEAR(-11.34737, real(y33(2, 1)), 1e-5); + EXPECT_NEAR(-7.72795, imag(y33(2, 1)), 1e-5); + EXPECT_NEAR(4.893267, real(y33(2, 2)), 1e-5); + EXPECT_NEAR(-1.981535, imag(y33(2, 2)), 1e-5); + + // check round trips inv_fft(fft(x)) + cm_t x33copy = inv_fft2(y33); + expect_complex_mat_eq(x33, x33copy); + + // check round trip fft(inv_fft(x)) + cm_t z33 = inv_fft2(x33); + cm_t x33copy2 = fft2(z33); + expect_complex_mat_eq(x33, x33copy2); +} diff --git a/test/unit/math/relative_tolerance.hpp b/test/unit/math/relative_tolerance.hpp index d6ea28b2293..16eaca8cc4d 100644 --- a/test/unit/math/relative_tolerance.hpp +++ b/test/unit/math/relative_tolerance.hpp @@ -3,6 +3,7 @@ #include #include +#include namespace stan { namespace test { @@ -79,6 +80,13 @@ class relative_tolerance { return std::max(tol_ * 0.5 * (fabs(x) + fabs(y)), tol_min_); } + /** + * Returns `true` if the tolerance is infinite. + * + * @return `true` if the tolernace is infinite. + */ + bool is_inf() const { return std::isinf(tol_); } + private: /** * The relative tolerance diff --git a/test/unit/math/test_ad.hpp b/test/unit/math/test_ad.hpp index c5bf2732a6b..5a8fff9ec1d 100644 --- a/test/unit/math/test_ad.hpp +++ b/test/unit/math/test_ad.hpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -333,12 +334,25 @@ template void expect_ad_derivatives(const ad_tolerances& tols, const G& g, const Eigen::VectorXd& x) { double gx = g(x); - test_gradient(tols, g, x, gx); + if (!tols.gradient_val_.is_inf() || !tols.gradient_grad_.is_inf()) { + test_gradient(tols, g, x, gx); + } #ifndef STAN_MATH_TESTS_REV_ONLY - test_gradient_fvar(tols, g, x, gx); - test_hessian(tols, g, x, gx); - test_hessian_fvar(tols, g, x, gx); - test_grad_hessian(tols, g, x, gx); + if (!tols.gradient_fvar_val_.is_inf() || !tols.gradient_fvar_grad_.is_inf()) { + test_gradient_fvar(tols, g, x, gx); + } + if (!tols.hessian_val_.is_inf() || !tols.hessian_grad_.is_inf() + || !tols.hessian_hessian_.is_inf()) { + test_hessian(tols, g, x, gx); + } + if (!tols.hessian_fvar_val_.is_inf() || !tols.hessian_fvar_grad_.is_inf() + || !tols.hessian_hessian_.is_inf()) { + test_hessian_fvar(tols, g, x, gx); + } + if (!tols.grad_hessian_val_.is_inf() || !tols.grad_hessian_hessian_.is_inf() + || !tols.grad_hessian_grad_hessian_.is_inf()) { + test_grad_hessian(tols, g, x, gx); + } #endif }