1 #ifndef STAN_MATH_PRIM_FUN_AUTOCORRELATION_HPP
2 #define STAN_MATH_PRIM_FUN_AUTOCORRELATION_HPP
3
4 #include <stan/math/prim/fun/Eigen.hpp>
5 #include <stan/math/prim/fun/mean.hpp>
6 #include <unsupported/Eigen/FFT>
7 #include <complex>
8 #include <vector>
9
10 namespace stan {
11 namespace math {
12 namespace internal {
13
14 /**
15 * Find the optimal next size for the FFT so that
16 * a minimum number of zeros are padded.
17 */
fft_next_good_size(size_t N)18 inline size_t fft_next_good_size(size_t N) {
19 if (N <= 2) {
20 return 2;
21 }
22 while (true) {
23 size_t m = N;
24 while ((m % 2) == 0) {
25 m /= 2;
26 }
27 while ((m % 3) == 0) {
28 m /= 3;
29 }
30 while ((m % 5) == 0) {
31 m /= 5;
32 }
33 if (m <= 1) {
34 return N;
35 }
36 N++;
37 }
38 }
39 } // namespace internal
40
41 /**
42 * Write autocorrelation estimates for every lag for the specified
43 * input sequence into the specified result using the specified
44 * FFT engine. The return vector be resized to the same length as
45 * the input sequence with lags given by array index.
46 *
47 * <p>The implementation involves a fast Fourier transform,
48 * followed by a normalization, followed by an inverse transform.
49 *
50 * <p>An FFT engine can be created for reuse for type double with:
51 *
52 * <pre>
53 * Eigen::FFT<double> fft;
54 * </pre>
55 *
56 * @tparam T Scalar type.
57 * @param y Input sequence.
58 * @param ac Autocorrelations.
59 * @param fft FFT engine instance.
60 */
61 template <typename T>
autocorrelation(const std::vector<T> & y,std::vector<T> & ac,Eigen::FFT<T> & fft)62 void autocorrelation(const std::vector<T>& y, std::vector<T>& ac,
63 Eigen::FFT<T>& fft) {
64 using std::complex;
65 using std::vector;
66
67 size_t N = y.size();
68 size_t M = internal::fft_next_good_size(N);
69 size_t Mt2 = 2 * M;
70
71 vector<complex<T> > freqvec;
72
73 // centered_signal = y-mean(y) followed by N zeroes
74 vector<T> centered_signal(y);
75 centered_signal.insert(centered_signal.end(), Mt2 - N, 0.0);
76 T mean = stan::math::mean(y);
77 for (size_t i = 0; i < N; i++) {
78 centered_signal[i] -= mean;
79 }
80
81 fft.fwd(freqvec, centered_signal);
82 for (size_t i = 0; i < Mt2; ++i) {
83 freqvec[i] = complex<T>(norm(freqvec[i]), 0.0);
84 }
85
86 fft.inv(ac, freqvec);
87 ac.resize(N);
88
89 for (size_t i = 0; i < N; ++i) {
90 ac[i] /= (N - i);
91 }
92 T var = ac[0];
93 for (size_t i = 0; i < N; ++i) {
94 ac[i] /= var;
95 }
96 }
97
98 /**
99 * Write autocorrelation estimates for every lag for the specified
100 * input sequence into the specified result using the specified
101 * FFT engine. The return vector be resized to the same length as
102 * the input sequence with lags given by array index.
103 *
104 * <p>The implementation involves a fast Fourier transform,
105 * followed by a normalization, followed by an inverse transform.
106 *
107 * <p>An FFT engine can be created for reuse for type double with:
108 *
109 * <pre>
110 * Eigen::FFT<double> fft;
111 * </pre>
112 *
113 * @tparam T Scalar type.
114 * @param y Input sequence.
115 * @param ac Autocorrelations.
116 * @param fft FFT engine instance.
117 */
118 template <typename T, typename DerivedA, typename DerivedB>
autocorrelation(const Eigen::MatrixBase<DerivedA> & y,Eigen::MatrixBase<DerivedB> & ac,Eigen::FFT<T> & fft)119 void autocorrelation(const Eigen::MatrixBase<DerivedA>& y,
120 Eigen::MatrixBase<DerivedB>& ac, Eigen::FFT<T>& fft) {
121 size_t N = y.size();
122 size_t M = internal::fft_next_good_size(N);
123 size_t Mt2 = 2 * M;
124
125 // centered_signal = y-mean(y) followed by N zeros
126 Eigen::Matrix<T, Eigen::Dynamic, 1> centered_signal(Mt2);
127 centered_signal.setZero();
128 centered_signal.head(N) = y.array() - y.mean();
129
130 Eigen::Matrix<std::complex<T>, Eigen::Dynamic, 1> freqvec(Mt2);
131 fft.SetFlag(fft.HalfSpectrum);
132 fft.fwd(freqvec, centered_signal);
133 // cwiseAbs2 == norm
134 freqvec = freqvec.cwiseAbs2();
135
136 Eigen::Matrix<T, Eigen::Dynamic, 1> ac_tmp(Mt2);
137 fft.inv(ac_tmp, freqvec);
138 fft.ClearFlag(fft.HalfSpectrum);
139
140 for (size_t i = 0; i < N; ++i) {
141 ac_tmp(i) /= (N - i);
142 }
143
144 ac = ac_tmp.head(N).array() / ac_tmp(0);
145 }
146
147 /**
148 * Write autocorrelation estimates for every lag for the specified
149 * input sequence into the specified result. The return vector be
150 * resized to the same length as the input sequence with lags
151 * given by array index.
152 *
153 * <p>The implementation involves a fast Fourier transform,
154 * followed by a normalization, followed by an inverse transform.
155 *
156 * <p>This method is just a light wrapper around the three-argument
157 * autocorrelation function.
158 *
159 * @tparam T Scalar type.
160 * @param y Input sequence.
161 * @param ac Autocorrelations.
162 */
163 template <typename T>
autocorrelation(const std::vector<T> & y,std::vector<T> & ac)164 void autocorrelation(const std::vector<T>& y, std::vector<T>& ac) {
165 Eigen::FFT<T> fft;
166 size_t N = y.size();
167 ac.resize(N);
168
169 const Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, 1> > y_map(&y[0], N);
170 Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, 1> > ac_map(&ac[0], N);
171 autocorrelation<T>(y_map, ac_map, fft);
172 }
173
174 /**
175 * Write autocorrelation estimates for every lag for the specified
176 * input sequence into the specified result. The return vector be
177 * resized to the same length as the input sequence with lags
178 * given by array index.
179 *
180 * <p>The implementation involves a fast Fourier transform,
181 * followed by a normalization, followed by an inverse transform.
182 *
183 * <p>This method is just a light wrapper around the three-argument
184 * autocorrelation function
185 *
186 * @tparam T Scalar type.
187 * @param y Input sequence.
188 * @param ac Autocorrelations.
189 */
190 template <typename T, typename DerivedA, typename DerivedB>
autocorrelation(const Eigen::MatrixBase<DerivedA> & y,Eigen::MatrixBase<DerivedB> & ac)191 void autocorrelation(const Eigen::MatrixBase<DerivedA>& y,
192 Eigen::MatrixBase<DerivedB>& ac) {
193 Eigen::FFT<T> fft;
194 autocorrelation(y, ac, fft);
195 }
196
197 } // namespace math
198 } // namespace stan
199
200 #endif
201