convolution.cc

#include <vtestbed/base/blitz.hh>
#include <vtestbed/base/for_loop.hh>
#include <vtestbed/config/openmp.hh>
#include <vtestbed/config/real_type.hh>
#include <vtestbed/core/convolution.hh>

#include <complex>

#if defined(VTB_WITH_OPENMP)
#include <omp.h>
#endif

template <class T, int N>
void
vtb::core::Convolution<T,N>::shape(
    const shape_type& signal_shape,
    const shape_type& kernel_shape
) {
    using blitz::abs;
    using blitz::min;
    using blitz::next_power_of_two;
    #if defined(VTB_WITH_OPENMP)
    const int parallelism = omp_get_max_threads();
    #else
    const int parallelism = 2;
    #endif
    shape_type guess1 = min(signal_shape, 4*kernel_shape);
    shape_type guess2 = max(kernel_shape, signal_shape / parallelism);
    shape_type block_size = min(guess1, guess2) + abs(guess1 - guess2) / 2;
    block_size = next_power_of_two(block_size + kernel_shape) - kernel_shape;
    this->padded_block_shape(block_size, kernel_shape);
}

template <class T, int N>
auto
vtb::core::Convolution<T,N>::convolve(array_type signal, array_type kernel) -> array_type {
    using blitz::all;
    using blitz::product;
    using blitz::div_ceil;
    using blitz::min;
    if (!all(kernel.shape() <= this->_blockshape)) {
        throw std::length_error("bad kernel shape");
    }
    const shape_type padded_block = this->padded_block_shape();
    domain_type orig_domain(shape_type(0), kernel.shape()-1);
    array_type padded_kernel(padded_block);
    padded_kernel(orig_domain) = kernel;
    const auto nelements = padded_kernel.numElements();
    this->_fft.forward(padded_kernel);
    const shape_type bs = this->_blockshape;
    const shape_type pad = this->_padding;
    const shape_type nparts = div_ceil(signal.shape(), bs);
    if (!all(bs <= signal.shape())) {
        throw std::length_error("bad block size");
    }
    array_type out_signal(signal.shape());
    for_loop<N>(
        nparts,
        [&] (const shape_type& idx) {
            const shape_type offset = idx*bs;
            const shape_type from = offset;
            const shape_type to = min(signal.shape(), offset+bs) - 1;
            const domain_type part_domain(from, to);
            const domain_type dom_to(from-offset, to-offset);
            array_type padded_part(padded_block);
            padded_part(dom_to) = signal(part_domain);
            this->_fft.forward(padded_part);
            padded_part *= padded_kernel;
            this->_fft.backward(padded_part);
            padded_part /= nelements;
            const domain_type padded_from(from, min(to+pad, signal.shape()-1));
            const domain_type padded_to(from-offset, padded_from.ubound()-offset);
            #if defined(VTB_WITH_OPENMP)
            #pragma omp critical
            #endif
            {
                out_signal(padded_from) += padded_part(padded_to);
            }
        }
    );
    return out_signal;
}

template class vtb::core::Convolution<std::complex<VTB_REAL_TYPE>, 1>;
template class vtb::core::Convolution<std::complex<VTB_REAL_TYPE>, 2>;
template class vtb::core::Convolution<std::complex<VTB_REAL_TYPE>, 3>;