Line data Source code
1 : /** \file fftT.hpp
2 : * \brief The Fast Fourier Transform interface
3 : * \ingroup ft_files
4 : * \author Jared R. Males (jaredmales@gmail.com)
5 : *
6 : */
7 :
8 : //***********************************************************************//
9 : // Copyright 2015-2025 Jared R. Males (jaredmales@gmail.com)
10 : //
11 : // This file is part of mxlib.
12 : //
13 : // mxlib is free software: you can redistribute it and/or modify
14 : // it under the terms of the GNU General Public License as published by
15 : // the Free Software Foundation, either version 3 of the License, or
16 : // (at your option) any later version.
17 : //
18 : // mxlib is distributed in the hope that it will be useful,
19 : // but WITHOUT ANY WARRANTY; without even the implied warranty of
20 : // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
21 : // GNU General Public License for more details.
22 : //
23 : // You should have received a copy of the GNU General Public License
24 : // along with mxlib. If not, see <http://www.gnu.org/licenses/>.
25 : //***********************************************************************//
26 :
27 : #ifndef fftT_hpp
28 : #define fftT_hpp
29 :
30 : #pragma GCC system_header
31 : #include <Eigen/Dense>
32 :
33 : #include <fftw3.h>
34 :
35 : #include "fftwTemplates.hpp"
36 : #include "../../meta/trueFalseT.hpp"
37 :
38 : #include "ftTypes.hpp"
39 :
40 : namespace mx
41 : {
42 : namespace math
43 : {
44 : namespace ft
45 : {
46 :
47 : template <typename _inputT, typename _outputT, size_t _rank, int _cudaGPU = 0>
48 : class fftT;
49 :
50 : template <int rank>
51 : std::vector<int> fftwDimVec( int szX, int szY, int szZ );
52 :
53 : /// Fast Fourier Transforms using the \ref fftw "FFTW library"
54 : /** The \ref fftw_templates "FFTW Templates" type resolution system is used to allow the compiler
55 : * to access the right plan and types for the transforms based on inputT and outputT.
56 : *
57 : * Planning is done with "measure" to achieve optimum performance at the cost of extra time during
58 : * the planning step. This can be mitigated using "wisdom", managed by \ref fftwEnvironment.
59 : *
60 : * Calls to the FFTW plan functions are protected by `\#pragma omp critical` directives
61 : * unless MX_FFTW_NOOMP is defined prior to the first inclusion of this file. This means that
62 : * you do not need to use the `fftw_make_planner_thread_safe` facility in FFTW.
63 : *
64 : * \note I recommend not using fftw_make_planner_thread_safe or specifying number of threads in
65 : * \ref fftwEnvironment.
66 : *
67 : * \todo add execute interface with fftw like signature
68 : * \todo add plan interface where user passes in pointers (to avoid allocations)
69 : *
70 : * \tparam _inputT is the input type of the transform, can be either real or complex
71 : * \tparam _outputT is the output type of the transform, can be either real or complex
72 : * \tparam _rank is the rank of the transform. Limited to 1, 2, or 3.
73 : *
74 : * \ingroup fft
75 : */
76 : template <typename _inputT, typename _outputT, size_t _rank>
77 : class fftT<_inputT, _outputT, _rank, 0>
78 : {
79 :
80 : public:
81 : typedef _inputT inputT;
82 : typedef _outputT outputT;
83 :
84 : static const size_t rank = _rank;
85 :
86 : typedef typename fftwTypeSpec<inputT, outputT>::realT realT;
87 :
88 : typedef typename fftwTypeSpec<inputT, outputT>::complexT complexT;
89 :
90 : typedef Eigen::Array<inputT, -1, -1> eigenArrayInT;
91 :
92 : typedef Eigen::Array<outputT, -1, -1> eigenArrayOutT;
93 :
94 : typedef typename fftwPlanSpec<realT>::planT planT;
95 :
96 : protected:
97 : dir m_direction{ dir::forward }; ///< Direction of this FFT, either dir::forward (default) or dir::backward
98 :
99 : int m_szX{ 0 }; ///< Size of the x dimension
100 : int m_szY{ 0 }; ///< Size of the y dimension
101 : int m_szZ{ 0 }; ///< size of the z dimension
102 :
103 : planT m_plan{ nullptr }; ///< The FFTW plan object. This is a pointer, allocated by FFTW library calls.
104 :
105 : public:
106 : /// Default c'tor
107 : fftT();
108 :
109 : /// Constructor for rank 1 FFT.
110 : template <int crank = _rank>
111 : fftT( int nx, ///< [in] the desired size of the FFT
112 : dir ndir = dir::forward, ///< [in] [optional] direction of this FFT, either dir::forward (default) or
113 : ///< dir::backward
114 : bool inPlace = false, /**< [in] [optional] whether or not this is an in-place transform.
115 : Default is false, out-of-place.*/
116 : typename std::enable_if<crank == 1>::type * = 0 );
117 :
118 : /// Constructor for rank 2 FFT.
119 : template <int crank = _rank>
120 : fftT( int nx, ///< [in] the desired x size of the FFT
121 : int ny, ///< [in] the desired y size of the FFT
122 : dir ndir = dir::forward, /**< [in] [optional] direction of this FFT, either dir::forward
123 : (default) or dir::backward */
124 : bool inPlace = false, /**< [in] [optional] whether or not this is an in-place transform.
125 : Default is false, out-of-place. */
126 : typename std::enable_if<crank == 2>::type * = 0 );
127 :
128 : /// Constructor for rank 3 FFT.
129 : template <int crank = _rank>
130 : fftT( int nx, ///< [in] the desired x size of the FFT
131 : int ny, ///< [in] the desired y size of the FFT
132 : int nz, ///< [in] the desired z size of the FFT
133 : dir ndir = dir::forward, /**< [in] [optional] direction of this FFT, either dir::forward (default)
134 : or dir::backward*/
135 : bool inPlace = false, /**< [in] [optional] whether or not this is an in-place
136 : transform. Default is false, out-of-place. */
137 : typename std::enable_if<crank == 3>::type * = 0 );
138 :
139 : /// Destructor
140 : ~fftT();
141 :
142 : /// Destroy (de-allocate) the plan
143 : void destroyPlan();
144 :
145 : /// Get the direction of this FFT
146 : /** The direction is either dir::forward or dir::backward.
147 : *
148 : * \returns the current value of m_direction.
149 : */
150 : ft::dir direction();
151 :
152 : protected:
153 : /// Call the FFTW planning routine
154 : void doPlan( bool inPlace /**< [in] whether or not the tranform is in-place */ );
155 :
156 : public:
157 :
158 : /// Planning routine for rank 1 transforms.
159 : template <int crank = _rank>
160 : void plan( int nx, ///< [in] the desired size of the FFT
161 : ft::dir ndir = dir::forward, /**< [in] [optional] direction of this FFT, either dir::forward (default)
162 : or dir::backward */
163 : bool inPlace = false, /**< [in] [optional] whether or not this is an in-place transform.
164 : Default is false, out-of-place. */
165 : typename std::enable_if<crank == 1>::type * = 0 );
166 :
167 : /// Planning routine for rank 2 transforms.
168 : template <int crank = _rank>
169 : void plan( int nx, ///< [in] the desired x size of the FFT
170 : int ny, ///< [in] the desired y size of the FFT
171 : ft::dir ndir = dir::forward, /**< [in] [optional] direction of this FFT, either dir::forward (default)
172 : or dir::backward */
173 : bool inPlace = false, /**< [in] [optional] whether or not this is an in-place transform.
174 : Default is false, out-of-place. */
175 : typename std::enable_if<crank == 2>::type * = 0 );
176 :
177 : /// Planning routine for rank 3 transforms.
178 : template <int crank = _rank>
179 : void plan( int nx, ///< [in] the desired x size of the FFT
180 : int ny, ///< [in] the desired y size of the FFT
181 : int nz, ///< [in] the desired z size of the FFT
182 : ft::dir ndir = dir::forward, /**< [in] [optional] direction of this FFT, either dir::forward
183 : (default) or dir::backward */
184 : bool inPlace = false, /**< [in] [optional] whether or not this is an in-place transform.
185 : Default is false, out-of-place. */
186 : typename std::enable_if<crank == 3>::type * = 0 );
187 :
188 : /// Conduct the FFT
189 : void operator()( outputT *out, ///< [out] the output of the FFT, must be pre-allocated
190 : inputT *in ///< [in] the input to the FFT
191 : ) const;
192 :
193 : /// Conduct the FFT
194 : void operator()( eigenArrayOutT &out, /**< [out] the output of the DFT */
195 : eigenArrayInT &in /**< [in] the input to the DFT */ ) const;
196 : };
197 :
198 : template <typename inputT, typename outputT, size_t rank>
199 54 : fftT<inputT, outputT, rank, 0>::fftT()
200 : {
201 54 : }
202 :
203 : template <typename inputT, typename outputT, size_t rank>
204 : template <int crank>
205 : fftT<inputT, outputT, rank, 0>::fftT( int nx, ft::dir ndir, bool inPlace, typename std::enable_if<crank == 1>::type * )
206 : {
207 : m_direction = ndir;
208 :
209 : plan( nx, ndir, inPlace );
210 : }
211 :
212 : template <typename inputT, typename outputT, size_t rank>
213 : template <int crank>
214 : fftT<inputT, outputT, rank, 0>::fftT(
215 : int nx, int ny, ft::dir ndir, bool inPlace, typename std::enable_if<crank == 2>::type * )
216 : {
217 : m_direction = ndir;
218 :
219 : plan( nx, ny, ndir, inPlace );
220 : }
221 :
222 : template <typename inputT, typename outputT, size_t rank>
223 : template <int crank>
224 : fftT<inputT, outputT, rank, 0>::fftT(
225 : int nx, int ny, int nz, ft::dir ndir, bool inPlace, typename std::enable_if<crank == 3>::type * )
226 : {
227 : m_direction = ndir;
228 :
229 : plan( nx, ny, nz, ndir, inPlace );
230 : }
231 :
232 : template <typename inputT, typename outputT, size_t rank>
233 54 : fftT<inputT, outputT, rank, 0>::~fftT()
234 : {
235 54 : destroyPlan();
236 54 : }
237 :
238 : template <typename inputT, typename outputT, size_t rank>
239 108 : void fftT<inputT, outputT, rank, 0>::destroyPlan()
240 : {
241 108 : if( m_plan )
242 : {
243 54 : fftw_destroy_plan<realT>( m_plan );
244 : }
245 :
246 108 : m_plan = 0;
247 :
248 108 : m_szX = 0;
249 108 : m_szY = 0;
250 108 : m_szZ = 0;
251 108 : }
252 :
253 : template <typename inputT, typename outputT, size_t rank>
254 : ft::dir fftT<inputT, outputT, rank, 0>::direction()
255 : {
256 : return m_direction;
257 : }
258 :
259 : template <typename inputT, typename outputT, size_t rank>
260 54 : void fftT<inputT, outputT, rank, 0>::doPlan( bool inPlace )
261 : {
262 :
263 54 : inputT *forplan1 = nullptr;
264 54 : outputT *forplan2 = nullptr;
265 :
266 : int sz;
267 :
268 : if( rank == 1 )
269 : {
270 10 : sz = m_szX;
271 : }
272 : if( rank == 2 )
273 : {
274 34 : sz = m_szX * m_szY;
275 : }
276 : if( rank == 3 )
277 : {
278 10 : sz = m_szX * m_szY * m_szZ;
279 : }
280 :
281 54 : forplan1 = fftw_malloc<inputT>( sz );
282 :
283 54 : if( inPlace )
284 : {
285 30 : forplan2 = reinterpret_cast<outputT *>(forplan1);
286 : }
287 : else
288 : {
289 24 : forplan2 = fftw_malloc<outputT>( sz );
290 : }
291 :
292 54 : int pdir = FFTW_FORWARD;
293 54 : if( m_direction == dir::backward )
294 : {
295 27 : pdir = FFTW_BACKWARD;
296 : }
297 :
298 : // clang-format off
299 : #ifndef MX_FFTW_NOOMP
300 108 : #pragma omp critical
301 : {
302 : #endif // clang-format on
303 :
304 54 : m_plan = fftw_plan_dft<inputT, outputT>( fftwDimVec<rank>( m_szX, m_szY, m_szZ ),
305 : forplan1,
306 : forplan2,
307 : pdir,
308 : FFTW_MEASURE );
309 : // clang-format off
310 : #ifndef MX_FFTW_NOOMP
311 : }
312 : #endif // clang-format on
313 :
314 : //Only free forplan2 if it's distinct. Have to do comparison as same type.
315 54 : if( forplan2 != nullptr && reinterpret_cast<char *>(forplan2) != reinterpret_cast<char *>(forplan1) )
316 : {
317 24 : fftw_free<outputT>( forplan2 );
318 : }
319 :
320 54 : if( forplan1 != nullptr )
321 : {
322 54 : fftw_free<inputT>( forplan1 );
323 : }
324 54 : }
325 :
326 : template <typename inputT, typename outputT, size_t rank>
327 : template <int crank>
328 10 : void fftT<inputT, outputT, rank, 0>::plan( int nx,
329 : ft::dir ndir,
330 : bool inPlace,
331 : typename std::enable_if<crank == 1>::type * )
332 : {
333 10 : if( m_szX == nx && m_direction == ndir && m_plan )
334 : {
335 0 : return;
336 : }
337 :
338 10 : destroyPlan();
339 :
340 10 : m_direction = ndir;
341 :
342 10 : m_szX = nx;
343 10 : m_szY = 0;
344 10 : m_szZ = 0;
345 :
346 10 : doPlan( inPlace );
347 : }
348 :
349 : template <typename inputT, typename outputT, size_t rank>
350 : template <int crank>
351 34 : void fftT<inputT, outputT, rank, 0>::plan(
352 : int nx, int ny, ft::dir ndir, bool inPlace, typename std::enable_if<crank == 2>::type * )
353 : {
354 34 : if( m_szX == nx && m_szY == ny && m_direction == ndir && m_plan )
355 : {
356 0 : return;
357 : }
358 :
359 34 : destroyPlan();
360 :
361 34 : m_direction = ndir;
362 :
363 34 : m_szX = nx;
364 34 : m_szY = ny;
365 34 : m_szZ = 0;
366 :
367 34 : doPlan( inPlace );
368 : }
369 :
370 : template <typename inputT, typename outputT, size_t rank>
371 : template <int crank>
372 10 : void fftT<inputT, outputT, rank, 0>::plan(
373 : int nx, int ny, int nz, ft::dir ndir, bool inPlace, typename std::enable_if<crank == 3>::type * )
374 : {
375 10 : if( m_szX == nx && m_szY == ny && m_szZ == nz && m_direction == ndir && m_plan )
376 : {
377 0 : return;
378 : }
379 :
380 10 : destroyPlan();
381 :
382 10 : m_direction = ndir;
383 :
384 10 : m_szX = nx;
385 10 : m_szY = ny;
386 10 : m_szZ = nz;
387 :
388 10 : doPlan( inPlace );
389 : }
390 :
391 : template <typename inputT, typename outputT, size_t rank>
392 1236 : void fftT<inputT, outputT, rank, 0>::operator()( outputT *out, inputT *in ) const
393 : {
394 1236 : fftw_execute_dft<inputT, outputT>( m_plan, in, out );
395 1236 : }
396 :
397 : template <typename inputT, typename outputT, size_t rank>
398 24 : void fftT<inputT, outputT, rank, 0>::operator()( eigenArrayOutT &out, eigenArrayInT &in ) const
399 : {
400 24 : fftw_execute_dft<inputT, outputT>( m_plan, in.data(), out.data() );
401 24 : }
402 :
403 : template <>
404 : std::vector<int> fftwDimVec<1>( int szX, int szY, int szZ );
405 :
406 : template <>
407 : std::vector<int> fftwDimVec<2>( int szX, int szY, int szZ );
408 :
409 : template <>
410 : std::vector<int> fftwDimVec<3>( int szX, int szY, int szZ );
411 :
412 :
413 : } // namespace ft
414 : } // namespace math
415 : } // namespace mx
416 :
417 : #ifdef MXLIB_CUDA
418 :
419 : #include "fftTcuda.hpp"
420 :
421 : #endif
422 :
423 : #endif // fftT_hpp
|