# main.cpp

a guest
Apr 2nd, 2017
207
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
1. #include <iostream>
2. #include <vector>
3. #include <cassert>
4. #include <complex>
5. #include <cmath>
6.
7. #include "fftw3.h"
8.
9. #ifndef M_PI
10. #define M_PI           3.14159265358979323846
11. #endif
12.
13. std::complex<double> compute_twiddle(size_t i, size_t len, bool inverse) {
14.     double constant;
15.     if(inverse) {
16.         constant = 2 * M_PI / len;
17.     } else {
18.         constant = -2 * M_PI / len;
19.     }
20.
21.     double real = std::cos(i * constant);
22.     double imag = std::sin(i * constant);
23.
24.     return std::complex<double>(real, imag);
25. }
26.
27. std::vector<double> compute_dct4(std::vector<double> x) {
28.
29.     assert(x.size()%2 == 0);
30.
31.     size_t N = x.size();
32.
33.     //pre-process the input by splitting into into two arrays, "w" and "v"
34.     std::vector<double> w(N/2), v(N/2);
35.
36.     w[0] = x[0];
37.     v[N/2 - 1] = x[N - 1];
38.
39.     for(size_t k = 1; k < N/2; k++) {
40.         w[k] =     x[2 * k - 1] + x[2 * k];
41.         v[k - 1] = x[2 * k - 1] - x[2 * k];
42.
43.         //NOTE: in the paper, the array "v" is 1-indexed rather than 0-indexed, so we have to compensate for this by subtracting 1 from every index
44.         //IE, the loop starts with k = 1, but since v is 1-indexed, this corresponds to v[0] in a 0-indexed array
45.     }
46.
47.     //using FFTW, run a DCT-3 on array "w", and a DST-3 on array "v"
48.     std::vector<double> W(N/2), V(N/2);
49.
50.     auto dct3_plan = fftw_plan_r2r_1d(N/2, w.data(), W.data(), FFTW_REDFT01, FFTW_ESTIMATE);
51.     auto dst3_plan = fftw_plan_r2r_1d(N/2, v.data(), V.data(), FFTW_RODFT01, FFTW_ESTIMATE);
52.
53.     fftw_execute(dct3_plan);
54.     fftw_execute(dst3_plan);
55.
56.     //post-process the data by combining it back into a single array
57.     std::vector<std::complex<double>> Z(N*2);
58.     for(size_t k = 0; k < N/2; k++) {
59.         Z[2 * k + 1] = std::complex<double>(W[k], V[k]);
60.         Z[2 * (x.size() - 1 - k) + 1] = std::complex<double>(W[k], -V[k]);
61.     }
62.
63.     std::vector<double> final_result(N);
64.     for(size_t k = 0; k < N; k++) {
65.         std::complex<double> twiddle = compute_twiddle(2 * k + 1, 8*N, false);
66.
67.         final_result[k] = (Z[2 * k + 1] * twiddle).real();
68.     }
69.
70.     return final_result;
71. }
72.
73. void print_vec(std::vector<double> vec) {
74.     if(vec.size() > 0) {
75.         std::cout << vec[0];
76.     }
77.     for(size_t i = 1; i < vec.size(); i++) {
78.         std::cout << ", " << vec[i];
79.     }
80.     std::cout << std::endl;
81. }
82.
83. int main()
84. {
85.     std::vector<double> input{-1.847365, 4.444461, -18.9616, 20.626879};
86.     std::vector<double> actual = compute_dct4(input);
87.
88.     std::vector<double> expected(input.size());
89.
90.     auto expected_plan = fftw_plan_r2r_1d(input.size(), input.data(), expected.data(), FFTW_REDFT11, FFTW_ESTIMATE);
91.     fftw_execute(expected_plan);
92.
93.     std::cout << "Input:    ";
94.     print_vec(input);
95.
96.     std::cout << "Expected: ";
97.     print_vec(expected);
98.
99.     std::cout << "Actual:   ";
100.     print_vec(actual);
101.
102.     return 0;
103. }