convolution.hh

#ifndef __WT_CONVOLUTION_HH__
#define __WT_CONVOLUTION_HH__

#include "types.hh"
#include "fft.hh"

namespace wt {

class Convolution
{
public:
  Convolution(const CMatrix &kernels, size_t subSample = 1);
  virtual ~Convolution();

  template <class iDerived, class oDerived>
  void apply(const Eigen::DenseBase<iDerived> &signal, Eigen::DenseBase<oDerived> &out)
  {
    // First, clear _lastRes matrix
    this->_lastRes.setConstant(0);

    /*
     * Perform overlap-add FFT
     */
    size_t N = signal.size();
    size_t steps = N/this->_M;
    size_t rem   = N%this->_M;
    size_t out_offset = 0;

    // Compute the first complete steps
    for (size_t i=0; i<steps; i++) {
      // Store piece into forward-trafo buffer
      this->_part.head(this->_M).noalias() =
          signal.block(i*this->_M,0, _M,1).template cast<CScalar>();
      // 0-pad
      this->_part.tail(this->_M).setConstant(0);

      // perform forward FFT
      this->_fwd.exec();

      // Multiply result of forward FFT of the signal piece with every (transformed) kernel
      for (size_t j=0; j<this->_K; j++) {
        this->_work.col(j).noalias() =
            this->_part.cwiseProduct(this->_kernelF.col(j));
      }

      // Peform backward trafo
      this->_rev.exec();

      /*
       * Compute result of convolution and store it into the output buffer
       */
      if (0 == out_offset) { // first block
        out.topRows(this->_M/2).noalias() =
            ( ( this->_work.block(this->_M/2, 0, this->_M/2, this->_K)
                + this->_lastRes.topRows(this->_M/2) )/(2*this->_M) );
        out_offset += ( this->_M/2 );
      } else { // intermediate blocks
        out.block(out_offset, 0, this->_M, this->_K).noalias() =
            ( ( this->_work.topRows(this->_M) + this->_lastRes)/(2*this->_M) );
        out_offset += this->_M;
      }
      // Store remaining part for next step unless we are at the last step
      this->_lastRes.noalias() = this->_work.bottomRows(this->_M);
    }

    // Perform final step (if rem>0)
    if (rem) {
      // Store remaining samples
      this->_part.head(rem).noalias() = signal.block(steps*this->_M,0, rem,1);
      // 0-pad
      this->_part.tail(2*this->_M-rem).setConstant(0);

      // perform forward FFT
      this->_fwd.exec();

      // Multiply result of forward FFT of the signal piece with every (transformed) kernel
      for (size_t j=0; j<this->_K; j++) {
        this->_work.col(j).noalias() =
            this->_part.cwiseProduct(this->_kernelF.col(j));
      }

      // Peform backward trafo
      this->_rev.exec();

      // store result if _M/2+rem < _M
      if (this->_M >= (rem+this->_M/2)) {
        out.block(out_offset, 0, rem+this->_M/2, this->_K).noalias() =
            ( ( this->_work.topRows(rem+this->_M/2) +
                this->_lastRes.topRows(rem+this->_M/2) ) / (2*this->_M) );
      } else {
        out.block(out_offset, 0, this->_M, this->_K).noalias() =
            ( ( this->_work.topRows(this->_M) +
                this->_lastRes.topRows(this->_M) ) / (2*this->_M) );
        out_offset += this->_M;
        size_t n = rem+this->_M/2-this->_M;
        out.block(out_offset, 0, n, this->_K).noalias() =
            (this->_work.block(this->_M, 0, n, this->_K) / (2*this->_M));
      }
    } else {
      // store last _M/2 samples if rem==0
      out.block(out_offset, 0, this->_M/2, this->_K).noalias() =
          ( this->_work.block(this->_M-rem, 0, this->_M/2, this->_K) / (2*this->_M) );
    }
  }

  inline size_t kernelLength() const { return this->_M; }
  inline size_t numKernels() const { return this->_K; }

  inline size_t subSampling() const { return _subSampling; }
  void setSubSampling(size_t subSample) { _subSampling = subSample; }

protected:
  size_t _K;
  size_t _M;
  CMatrix _kernelF;

  CVector _part;
  FFT _fwd;

  CMatrix _lastRes;
  CMatrix _work;
  FFT _rev;

  size_t _subSampling;
};

}

#endif // __WT_CONVOLUTION_HH__