xref: /aosp_15_r20/external/eigen/unsupported/test/cxx11_tensor_fft.cpp (revision bf2c37156dfe67e5dfebd6d394bad8b2ab5804d4)
1 // This file is part of Eigen, a lightweight C++ template library
2 // for linear algebra.
3 //
4 // Copyright (C) 2014 Jianwei Cui <[email protected]>
5 //
6 // This Source Code Form is subject to the terms of the Mozilla
7 // Public License v. 2.0. If a copy of the MPL was not distributed
8 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
9 
10 #include "main.h"
11 #include <Eigen/CXX11/Tensor>
12 
13 using Eigen::Tensor;
14 
15 template <int DataLayout>
test_fft_2D_golden()16 static void test_fft_2D_golden() {
17   Tensor<float, 2, DataLayout> input(2, 3);
18   input(0, 0) = 1;
19   input(0, 1) = 2;
20   input(0, 2) = 3;
21   input(1, 0) = 4;
22   input(1, 1) = 5;
23   input(1, 2) = 6;
24 
25   array<ptrdiff_t, 2> fft;
26   fft[0] = 0;
27   fft[1] = 1;
28 
29   Tensor<std::complex<float>, 2, DataLayout> output = input.template fft<Eigen::BothParts, Eigen::FFT_FORWARD>(fft);
30 
31   std::complex<float> output_golden[6]; // in ColMajor order
32   output_golden[0] = std::complex<float>(21, 0);
33   output_golden[1] = std::complex<float>(-9, 0);
34   output_golden[2] = std::complex<float>(-3, 1.73205);
35   output_golden[3] = std::complex<float>( 0, 0);
36   output_golden[4] = std::complex<float>(-3, -1.73205);
37   output_golden[5] = std::complex<float>(0 ,0);
38 
39   std::complex<float> c_offset = std::complex<float>(1.0, 1.0);
40 
41   if (DataLayout == ColMajor) {
42     VERIFY_IS_APPROX(output(0) + c_offset, output_golden[0] + c_offset);
43     VERIFY_IS_APPROX(output(1) + c_offset, output_golden[1] + c_offset);
44     VERIFY_IS_APPROX(output(2) + c_offset, output_golden[2] + c_offset);
45     VERIFY_IS_APPROX(output(3) + c_offset, output_golden[3] + c_offset);
46     VERIFY_IS_APPROX(output(4) + c_offset, output_golden[4] + c_offset);
47     VERIFY_IS_APPROX(output(5) + c_offset, output_golden[5] + c_offset);
48   }
49   else {
50     VERIFY_IS_APPROX(output(0)+ c_offset, output_golden[0]+ c_offset);
51     VERIFY_IS_APPROX(output(1)+ c_offset, output_golden[2]+ c_offset);
52     VERIFY_IS_APPROX(output(2)+ c_offset, output_golden[4]+ c_offset);
53     VERIFY_IS_APPROX(output(3)+ c_offset, output_golden[1]+ c_offset);
54     VERIFY_IS_APPROX(output(4)+ c_offset, output_golden[3]+ c_offset);
55     VERIFY_IS_APPROX(output(5)+ c_offset, output_golden[5]+ c_offset);
56   }
57 }
58 
test_fft_complex_input_golden()59 static void test_fft_complex_input_golden() {
60   Tensor<std::complex<float>, 1, ColMajor> input(5);
61   input(0) = std::complex<float>(1, 1);
62   input(1) = std::complex<float>(2, 2);
63   input(2) = std::complex<float>(3, 3);
64   input(3) = std::complex<float>(4, 4);
65   input(4) = std::complex<float>(5, 5);
66 
67   array<ptrdiff_t, 1> fft;
68   fft[0] = 0;
69 
70   Tensor<std::complex<float>, 1, ColMajor> forward_output_both_parts = input.fft<BothParts, FFT_FORWARD>(fft);
71   Tensor<std::complex<float>, 1, ColMajor> reverse_output_both_parts = input.fft<BothParts, FFT_REVERSE>(fft);
72 
73   Tensor<float, 1, ColMajor> forward_output_real_part = input.fft<RealPart, FFT_FORWARD>(fft);
74   Tensor<float, 1, ColMajor> reverse_output_real_part = input.fft<RealPart, FFT_REVERSE>(fft);
75 
76   Tensor<float, 1, ColMajor> forward_output_imag_part = input.fft<ImagPart, FFT_FORWARD>(fft);
77   Tensor<float, 1, ColMajor> reverse_output_imag_part = input.fft<ImagPart, FFT_REVERSE>(fft);
78 
79   VERIFY_IS_EQUAL(forward_output_both_parts.dimension(0), input.dimension(0));
80   VERIFY_IS_EQUAL(reverse_output_both_parts.dimension(0), input.dimension(0));
81 
82   VERIFY_IS_EQUAL(forward_output_real_part.dimension(0), input.dimension(0));
83   VERIFY_IS_EQUAL(reverse_output_real_part.dimension(0), input.dimension(0));
84 
85   VERIFY_IS_EQUAL(forward_output_imag_part.dimension(0), input.dimension(0));
86   VERIFY_IS_EQUAL(reverse_output_imag_part.dimension(0), input.dimension(0));
87 
88   std::complex<float> forward_golden_result[5];
89   std::complex<float> reverse_golden_result[5];
90 
91   forward_golden_result[0] = std::complex<float>(15.000000000000000,+15.000000000000000);
92   forward_golden_result[1] = std::complex<float>(-5.940954801177935, +0.940954801177934);
93   forward_golden_result[2] = std::complex<float>(-3.312299240582266, -1.687700759417735);
94   forward_golden_result[3] = std::complex<float>(-1.687700759417735, -3.312299240582266);
95   forward_golden_result[4] = std::complex<float>( 0.940954801177934, -5.940954801177935);
96 
97   reverse_golden_result[0] = std::complex<float>( 3.000000000000000, + 3.000000000000000);
98   reverse_golden_result[1] = std::complex<float>( 0.188190960235587, - 1.188190960235587);
99   reverse_golden_result[2] = std::complex<float>(-0.337540151883547, - 0.662459848116453);
100   reverse_golden_result[3] = std::complex<float>(-0.662459848116453, - 0.337540151883547);
101   reverse_golden_result[4] = std::complex<float>(-1.188190960235587, + 0.188190960235587);
102 
103   for(int i = 0; i < 5; ++i) {
104     VERIFY_IS_APPROX(forward_output_both_parts(i), forward_golden_result[i]);
105     VERIFY_IS_APPROX(forward_output_real_part(i), forward_golden_result[i].real());
106     VERIFY_IS_APPROX(forward_output_imag_part(i), forward_golden_result[i].imag());
107   }
108 
109   for(int i = 0; i < 5; ++i) {
110     VERIFY_IS_APPROX(reverse_output_both_parts(i), reverse_golden_result[i]);
111     VERIFY_IS_APPROX(reverse_output_real_part(i), reverse_golden_result[i].real());
112     VERIFY_IS_APPROX(reverse_output_imag_part(i), reverse_golden_result[i].imag());
113   }
114 }
115 
test_fft_real_input_golden()116 static void test_fft_real_input_golden() {
117   Tensor<float, 1, ColMajor> input(5);
118   input(0) = 1.0;
119   input(1) = 2.0;
120   input(2) = 3.0;
121   input(3) = 4.0;
122   input(4) = 5.0;
123 
124   array<ptrdiff_t, 1> fft;
125   fft[0] = 0;
126 
127   Tensor<std::complex<float>, 1, ColMajor> forward_output_both_parts = input.fft<BothParts, FFT_FORWARD>(fft);
128   Tensor<std::complex<float>, 1, ColMajor> reverse_output_both_parts = input.fft<BothParts, FFT_REVERSE>(fft);
129 
130   Tensor<float, 1, ColMajor> forward_output_real_part = input.fft<RealPart, FFT_FORWARD>(fft);
131   Tensor<float, 1, ColMajor> reverse_output_real_part = input.fft<RealPart, FFT_REVERSE>(fft);
132 
133   Tensor<float, 1, ColMajor> forward_output_imag_part = input.fft<ImagPart, FFT_FORWARD>(fft);
134   Tensor<float, 1, ColMajor> reverse_output_imag_part = input.fft<ImagPart, FFT_REVERSE>(fft);
135 
136   VERIFY_IS_EQUAL(forward_output_both_parts.dimension(0), input.dimension(0));
137   VERIFY_IS_EQUAL(reverse_output_both_parts.dimension(0), input.dimension(0));
138 
139   VERIFY_IS_EQUAL(forward_output_real_part.dimension(0), input.dimension(0));
140   VERIFY_IS_EQUAL(reverse_output_real_part.dimension(0), input.dimension(0));
141 
142   VERIFY_IS_EQUAL(forward_output_imag_part.dimension(0), input.dimension(0));
143   VERIFY_IS_EQUAL(reverse_output_imag_part.dimension(0), input.dimension(0));
144 
145   std::complex<float> forward_golden_result[5];
146   std::complex<float> reverse_golden_result[5];
147 
148 
149   forward_golden_result[0] = std::complex<float>(  15, 0);
150   forward_golden_result[1] = std::complex<float>(-2.5, +3.44095480117793);
151   forward_golden_result[2] = std::complex<float>(-2.5, +0.81229924058227);
152   forward_golden_result[3] = std::complex<float>(-2.5, -0.81229924058227);
153   forward_golden_result[4] = std::complex<float>(-2.5, -3.44095480117793);
154 
155   reverse_golden_result[0] = std::complex<float>( 3.0, 0);
156   reverse_golden_result[1] = std::complex<float>(-0.5, -0.688190960235587);
157   reverse_golden_result[2] = std::complex<float>(-0.5, -0.162459848116453);
158   reverse_golden_result[3] = std::complex<float>(-0.5, +0.162459848116453);
159   reverse_golden_result[4] = std::complex<float>(-0.5, +0.688190960235587);
160 
161   std::complex<float> c_offset(1.0, 1.0);
162   float r_offset = 1.0;
163 
164   for(int i = 0; i < 5; ++i) {
165     VERIFY_IS_APPROX(forward_output_both_parts(i) + c_offset, forward_golden_result[i] + c_offset);
166     VERIFY_IS_APPROX(forward_output_real_part(i)  + r_offset, forward_golden_result[i].real() + r_offset);
167     VERIFY_IS_APPROX(forward_output_imag_part(i)  + r_offset, forward_golden_result[i].imag() + r_offset);
168   }
169 
170   for(int i = 0; i < 5; ++i) {
171     VERIFY_IS_APPROX(reverse_output_both_parts(i) + c_offset, reverse_golden_result[i] + c_offset);
172     VERIFY_IS_APPROX(reverse_output_real_part(i)  + r_offset, reverse_golden_result[i].real() + r_offset);
173     VERIFY_IS_APPROX(reverse_output_imag_part(i)  + r_offset, reverse_golden_result[i].imag() + r_offset);
174   }
175 }
176 
177 
178 template <int DataLayout, typename RealScalar, bool isComplexInput, int FFTResultType, int FFTDirection, int TensorRank>
test_fft_real_input_energy()179 static void test_fft_real_input_energy() {
180 
181   Eigen::DSizes<ptrdiff_t, TensorRank> dimensions;
182   ptrdiff_t total_size = 1;
183   for (int i = 0; i < TensorRank; ++i) {
184     dimensions[i] = rand() % 20 + 1;
185     total_size *= dimensions[i];
186   }
187   const DSizes<ptrdiff_t, TensorRank> arr = dimensions;
188 
189   typedef typename internal::conditional<isComplexInput == true, std::complex<RealScalar>, RealScalar>::type InputScalar;
190 
191   Tensor<InputScalar, TensorRank, DataLayout> input;
192   input.resize(arr);
193   input.setRandom();
194 
195   array<ptrdiff_t, TensorRank> fft;
196   for (int i = 0; i < TensorRank; ++i) {
197     fft[i] = i;
198   }
199 
200   typedef typename internal::conditional<FFTResultType == Eigen::BothParts, std::complex<RealScalar>, RealScalar>::type OutputScalar;
201   Tensor<OutputScalar, TensorRank, DataLayout> output;
202   output = input.template fft<FFTResultType, FFTDirection>(fft);
203 
204   for (int i = 0; i < TensorRank; ++i) {
205     VERIFY_IS_EQUAL(output.dimension(i), input.dimension(i));
206   }
207 
208   RealScalar energy_original = 0.0;
209   RealScalar energy_after_fft = 0.0;
210 
211   for (int i = 0; i < total_size; ++i) {
212     energy_original += numext::abs2(input(i));
213   }
214 
215   for (int i = 0; i < total_size; ++i) {
216     energy_after_fft += numext::abs2(output(i));
217   }
218 
219   if(FFTDirection == FFT_FORWARD) {
220     VERIFY_IS_APPROX(energy_original, energy_after_fft / total_size);
221   }
222   else {
223     VERIFY_IS_APPROX(energy_original, energy_after_fft * total_size);
224   }
225 }
226 
227 template <typename RealScalar>
test_fft_non_power_of_2_round_trip(int exponent)228 static void test_fft_non_power_of_2_round_trip(int exponent) {
229   int n = (1 << exponent) + 1;
230 
231   Eigen::DSizes<ptrdiff_t, 1> dimensions;
232   dimensions[0] = n;
233   const DSizes<ptrdiff_t, 1> arr = dimensions;
234   Tensor<RealScalar, 1, ColMajor, ptrdiff_t> input;
235 
236   input.resize(arr);
237   input.setRandom();
238 
239   array<int, 1> fft;
240   fft[0] = 0;
241 
242   Tensor<std::complex<RealScalar>, 1, ColMajor> forward =
243       input.template fft<BothParts, FFT_FORWARD>(fft);
244 
245   Tensor<RealScalar, 1, ColMajor, ptrdiff_t> output =
246       forward.template fft<RealPart, FFT_REVERSE>(fft);
247 
248   for (int i = 0; i < n; ++i) {
249     RealScalar tol = test_precision<RealScalar>() *
250                      (std::abs(input[i]) + std::abs(output[i]) + 1);
251     VERIFY_IS_APPROX_OR_LESS_THAN(std::abs(input[i] - output[i]), tol);
252   }
253 }
254 
EIGEN_DECLARE_TEST(cxx11_tensor_fft)255 EIGEN_DECLARE_TEST(cxx11_tensor_fft) {
256     test_fft_complex_input_golden();
257     test_fft_real_input_golden();
258 
259     test_fft_2D_golden<ColMajor>();
260     test_fft_2D_golden<RowMajor>();
261 
262     test_fft_real_input_energy<ColMajor, float,  true,  Eigen::BothParts, FFT_FORWARD, 1>();
263     test_fft_real_input_energy<ColMajor, double, true,  Eigen::BothParts, FFT_FORWARD, 1>();
264     test_fft_real_input_energy<ColMajor, float,  false,  Eigen::BothParts, FFT_FORWARD, 1>();
265     test_fft_real_input_energy<ColMajor, double, false,  Eigen::BothParts, FFT_FORWARD, 1>();
266 
267     test_fft_real_input_energy<ColMajor, float,  true,  Eigen::BothParts, FFT_FORWARD, 2>();
268     test_fft_real_input_energy<ColMajor, double, true,  Eigen::BothParts, FFT_FORWARD, 2>();
269     test_fft_real_input_energy<ColMajor, float,  false,  Eigen::BothParts, FFT_FORWARD, 2>();
270     test_fft_real_input_energy<ColMajor, double, false,  Eigen::BothParts, FFT_FORWARD, 2>();
271 
272     test_fft_real_input_energy<ColMajor, float,  true,  Eigen::BothParts, FFT_FORWARD, 3>();
273     test_fft_real_input_energy<ColMajor, double, true,  Eigen::BothParts, FFT_FORWARD, 3>();
274     test_fft_real_input_energy<ColMajor, float,  false,  Eigen::BothParts, FFT_FORWARD, 3>();
275     test_fft_real_input_energy<ColMajor, double, false,  Eigen::BothParts, FFT_FORWARD, 3>();
276 
277     test_fft_real_input_energy<ColMajor, float,  true,  Eigen::BothParts, FFT_FORWARD, 4>();
278     test_fft_real_input_energy<ColMajor, double, true,  Eigen::BothParts, FFT_FORWARD, 4>();
279     test_fft_real_input_energy<ColMajor, float,  false,  Eigen::BothParts, FFT_FORWARD, 4>();
280     test_fft_real_input_energy<ColMajor, double, false,  Eigen::BothParts, FFT_FORWARD, 4>();
281 
282     test_fft_real_input_energy<RowMajor, float,  true,  Eigen::BothParts, FFT_FORWARD, 1>();
283     test_fft_real_input_energy<RowMajor, double, true,  Eigen::BothParts, FFT_FORWARD, 1>();
284     test_fft_real_input_energy<RowMajor, float,  false,  Eigen::BothParts, FFT_FORWARD, 1>();
285     test_fft_real_input_energy<RowMajor, double, false,  Eigen::BothParts, FFT_FORWARD, 1>();
286 
287     test_fft_real_input_energy<RowMajor, float,  true,  Eigen::BothParts, FFT_FORWARD, 2>();
288     test_fft_real_input_energy<RowMajor, double, true,  Eigen::BothParts, FFT_FORWARD, 2>();
289     test_fft_real_input_energy<RowMajor, float,  false,  Eigen::BothParts, FFT_FORWARD, 2>();
290     test_fft_real_input_energy<RowMajor, double, false,  Eigen::BothParts, FFT_FORWARD, 2>();
291 
292     test_fft_real_input_energy<RowMajor, float,  true,  Eigen::BothParts, FFT_FORWARD, 3>();
293     test_fft_real_input_energy<RowMajor, double, true,  Eigen::BothParts, FFT_FORWARD, 3>();
294     test_fft_real_input_energy<RowMajor, float,  false,  Eigen::BothParts, FFT_FORWARD, 3>();
295     test_fft_real_input_energy<RowMajor, double, false,  Eigen::BothParts, FFT_FORWARD, 3>();
296 
297     test_fft_real_input_energy<RowMajor, float,  true,  Eigen::BothParts, FFT_FORWARD, 4>();
298     test_fft_real_input_energy<RowMajor, double, true,  Eigen::BothParts, FFT_FORWARD, 4>();
299     test_fft_real_input_energy<RowMajor, float,  false,  Eigen::BothParts, FFT_FORWARD, 4>();
300     test_fft_real_input_energy<RowMajor, double, false,  Eigen::BothParts, FFT_FORWARD, 4>();
301 
302     test_fft_non_power_of_2_round_trip<float>(7);
303     test_fft_non_power_of_2_round_trip<double>(7);
304 }
305