core/fourier_transform.hh

#ifndef VTESTBED_CORE_FOURIER_TRANSFORM_HH
#define VTESTBED_CORE_FOURIER_TRANSFORM_HH

#include <complex>
#include <limits>
#include <stdexcept>

#include <vtestbed/base/blitz.hh>
#include <vtestbed/base/for_loop.hh>
#include <vtestbed/core/math.hh>

namespace vtb {

    namespace core {

        template <class T>
        class Wave_table_base {

        public:
            typedef blitz::Array<int,1> indices_type;
            typedef blitz::Array<T,1> array_type;

        protected:
            array_type _waves;
            indices_type _indices;

        public:

            inline void
            free() {
                this->_waves.free();
                this->_indices.free();
            }

            inline array_type
            waves() const {
                return this->_waves;
            }

            inline indices_type
            indices() const {
                return this->_indices;
            }

        };

        template <class T>
        class Wave_table;

        template <class T>
        class Wave_table: public Wave_table_base<T> {

            static_assert(std::is_floating_point<T>::value, "bad type");

        public:
            typedef blitz::Array<T,1> array_type;

        public:

            void
            shape(int n) {
                if (n < 1 || !blitz::is_power_of_two(n)) {
                    throw std::invalid_argument(__func__);
                }
                auto& indices = this->_indices;
                auto& waves = this->_waves;
                const int nw = n/4;
                waves.resize(2*nw);
                indices.resize(indices_length_real(n));
                this->makewt(nw, indices.data(), waves.data());
                this->makect(nw, indices.data(), waves.data() + nw);
            }

            void
            fourier_transform(array_type& x, int sign) {
                auto& indices = this->_indices;
                auto& waves = this->_waves;
                // TODO the sign???
                this->rdft(
                    x.extent(0),
                    -sign,
                    x.data(),
                    indices.data(),
                    waves.data()
                );
            }

        private:

            inline int
            indices_length_real(int n) {
                return 2 + static_cast<int>(std::sqrt(n/2));
            }

            static void
            makewt(int nw, int *ip, T *w);

            static void
            makect(int nc, int *ip, T *c);

            static void
            rdft(int n, int isgn, T *a, int *ip, const T *w);

        };

        template <class T>
        class Wave_table<std::complex<T>>:
            public Wave_table_base<std::complex<T>> {

            static_assert(std::is_floating_point<T>::value, "bad type");

        public:
            typedef blitz::Array<std::complex<T>,1> array_type;

        public:

            void
            shape(int n) {
                if (n < 1 || !blitz::is_power_of_two(n)) {
                    throw std::invalid_argument(__func__);
                }
                auto& indices = this->_indices;
                auto& waves = this->_waves;
                const int nw = n/2;
                waves.resize(nw);
                indices.resize(indices_length_complex(n));
                this->makewt(nw, indices.data(), from_complex(waves.data()));
            }

            void
            fourier_transform(array_type& x, int sign) {
                auto& indices = this->_indices;
                auto& waves = this->_waves;
                this->cdft(
                    2*x.extent(0),
                    sign,
                    from_complex(x.data()),
                    indices.data(),
                    from_complex(waves.data())
                );
            }

        private:

            inline int
            indices_length_complex(int n) {
                return 2 + static_cast<int>(std::sqrt(n));
            }

            inline T*
            from_complex(std::complex<T>* x) {
                return reinterpret_cast<T*>(x);
            }

            inline const T*
            from_complex(const std::complex<T>* x) {
                return reinterpret_cast<const T*>(x);
            }

            static void
            makewt(int nw, int *ip, T *w);

            static void
            cdft(int n, int isgn, T* a, int* ip, const T* w);

        };

        template <class T, int N>
        inline void
        fourier_transform(
            blitz::Array<T,N>& x,
            int sign,
            Wave_table<T> waves[N]
        ) {
            const int nelems = x.numElements();
            for (int i=0; i<N; ++i) {
                const int stride = x.stride(i);
                const int extent = x.extent(i);
                const int block_size = extent*stride;
                const int nblocks = nelems / block_size;
                #if defined(VTB_DEBUG_FFT)
                std::clog
                    << "n=" << i
                    << ",bs=" << block_size
                    << ",nblocks=" << nblocks
                    << std::endl;
                #endif
                for (int k=0; k<nblocks; ++k) {
                    for (int j=0; j<stride; ++j) {
                        const int offset = block_size*k + j;
                        #if defined(VTB_DEBUG_FFT)
                        std::clog << "FFT: extent=" << extent
                            << ",stride=" << stride
                            << ",offset=" << offset
                            << ",check=" << (offset + stride*(extent-1))
                            << std::endl;
                        #endif
                        blitz::Array<T,1> slice(blitz::shape(extent));
                        for (int idx=0; idx<extent; ++idx) {
                            slice(idx) = *(x.data() + offset + idx*stride);
                        }
                        waves[i].fourier_transform(slice, sign);
                        for (int idx=0; idx<extent; ++idx) {
                            *(x.data() + offset + idx*stride) = slice(idx);
                        }
                    }
                }
            }
        }


        template <class T, int N>
        class Fourier_transform_base {

        public:
            typedef blitz::TinyVector<int,N> shape_type;
            typedef blitz::Array<T,N> array_type;
            typedef Wave_table<T> wavetable_type;

        private:
            shape_type _shape;

        protected:
            wavetable_type _waves[N];

        public:

            inline explicit
            Fourier_transform_base(const shape_type& shp) {
                this->_shape = 0;
                this->shape(shp);
            }

            inline
            Fourier_transform_base() {
                this->_shape = 0;
            }

            inline const shape_type&
            shape() const noexcept {
                return this->_shape;
            }

            inline void
            shape(const shape_type& rhs) {
                for (int i=0; i<N; ++i) {
                    if (this->_shape(i) != rhs(i)) {
                        this->_shape(i) = rhs(i);
                        this->_waves[i].shape(rhs(i));
                    }
                }
            }

            inline void
            clear() {
                this->_shape = 0;
                for (int i=0; i<N; ++i) {
                    this->_waves[i].free();
                }
            }

        protected:

            inline void
            check(const array_type& x) {
                if (!blitz::all(x.shape() == this->shape())) {
                    throw std::invalid_argument("Fourier_transform_base::check: bad shape");
                }
            }

        };

        template <class T, int N>
        class Fourier_transform: public Fourier_transform_base<T,N> {

        private:
            typedef Fourier_transform_base<T,N> base_type;

        public:
            using typename base_type::shape_type;
            using typename base_type::array_type;

        public:

            inline
            Fourier_transform():
            base_type{} {}

            inline explicit
            Fourier_transform(const shape_type& shp):
            base_type{shp} {}

            inline void
            forward(array_type& x) {
                this->transform(x, -1);
            }

            inline void
            backward(array_type& x) {
                this->transform(x, 1);
            }

            inline void
            transform(array_type& x, int dir) {
                #if defined(VTB_DEBUG)
                this->check(x);
                #endif
                fourier_transform<T,N>(x, dir, this->_waves);
            }

        };

        template <class T, int N>
        class Chirp_Z_transform {

        public:
            typedef blitz::TinyVector<int,N> shape_type;
            typedef blitz::Array<T,N> array_type;
            typedef Fourier_transform<T,N> fft_type;

        private:
            typedef blitz::TinyVector<T,N> vec;
            typedef blitz::RectDomain<N> domain;
            typedef typename T::value_type R;

        private:
            fft_type _fft;
            shape_type _shape{0};
            bool _power_of_two = false;
            array_type _chirp;

        public:

            Chirp_Z_transform() = default;

            inline explicit
            Chirp_Z_transform(const shape_type& shp) {
                this->shape(shp);
            }

            inline const shape_type&
            shape() const noexcept {
                return this->_shape;
            }

            inline const shape_type&
            fourier_transform_shape() const noexcept {
                return this->_fft.shape();
            }

            inline void
            shape(const shape_type& rhs) {
                using blitz::all;
                if (all(this->_shape == rhs)) {
                    return;
                }
                this->_shape = rhs;
                shape_type new_shape;
                if (all(is_power_of_two(rhs))) {
                    new_shape = rhs;
                    _power_of_two = true;
                } else {
                    new_shape = next_power_of_two(rhs*2 - 1);
                    _power_of_two = false;
                }
                _fft.shape(new_shape);
                this->make_chirp();
            }

            inline void
            forward(array_type& x) {
                this->transform(x, -1);
            }

            inline void
            backward(array_type& x) {
                this->transform(x, 1);
            }

            inline void
            transform(array_type& x, int dir) {
                using blitz::product;
                if (_power_of_two) {
                    _fft.transform(x, dir);
                } else {
                    #if defined(VTB_DEBUG_CHIRP_Z)
                    std::clog << "chirp=" << _chirp << std::endl;
                    #endif
                    array_type xp{_fft.shape()};
                    xp = 0;
                    shape_type offset{_shape-1};
                    domain rect{offset, _chirp.shape()-1};
                    domain lrect{xp.base(), offset};
                    #if defined(VTB_DEBUG_CHIRP_Z)
                    std::clog << "offset=" << offset << std::endl;
                    #endif
                    if (dir < 0) {
                        xp(lrect) = x*_chirp(rect);
                    } else {
                        // flip the sign of the imaginary part
                        xp(lrect) = x*blitz::conj(_chirp(rect));
                    }
                    #if defined(VTB_DEBUG_CHIRP_Z)
                    std::clog << "xp=" << xp << std::endl;
                    #endif
                    array_type ichirp{_fft.shape()};
                    ichirp = 0;
                    if (dir < 0) {
                        ichirp(_chirp.domain()) = R{1} / _chirp;
                    } else {
                        ichirp(_chirp.domain()) = R{1} / blitz::conj(_chirp);
                    }
                    #if defined(VTB_DEBUG_CHIRP_Z)
                    std::clog << "ichirp=" << ichirp << std::endl;
                    #endif
                    _fft.transform(xp, dir);
                    #if defined(VTB_DEBUG_CHIRP_Z)
                    std::clog << "fft(xp)=" << xp << std::endl;
                    #endif
                    _fft.transform(ichirp, dir);
                    #if defined(VTB_DEBUG_CHIRP_Z)
                    std::clog << "fft(ichirp)=" << ichirp << std::endl;
                    #endif
                    ichirp *= xp;
                    #if defined(VTB_DEBUG_CHIRP_Z)
                    std::clog << "mult2=" << ichirp << std::endl;
                    #endif
                    _fft.transform(ichirp, -dir);
                    #if defined(VTB_DEBUG_CHIRP_Z)
                    std::clog << "ifft=" << ichirp << std::endl;
                    #endif
                    x = ichirp(rect) * _chirp(rect) / R(product(_fft.shape()));
                }
            }

        private:

            void
            make_chirp() {
                using blitz::exp;
                using blitz::pow2;
                using blitz::pow;
                using blitz::product;
                const R pi2 = R{M_PI}*2;
                shape_type chirp_shape{_shape*2-1};
                _chirp.resize(chirp_shape);
                T i{0,1};
                R d{-1};
                vec w = exp(d*i*pi2/vec{_shape});
                #if defined(VTB_DEBUG_CHIRP_Z)
                std::clog << "w=" << w << std::endl;
                #endif
                shape_type offset{_shape-1};
                for_loop<N>(
                    chirp_shape,
                    [&offset,this,&w] (const shape_type& idx) {
                        vec k = idx - offset;
                        _chirp(idx) = product(pow(w, pow2(k)*T{0.5,0}));
                    }
                );
            }

        };

    }

}

#endif // vim:filetype=cpp