Guest User

Discrete Cosine Transform Type 4 implementation

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