efft.hpp

Go to the documentation of this file.

 
#ifndef EFFT_HPP
#define EFFT_HPP
 
#include <algorithm>
#include <array>
#include <cmath>
#include <complex>
#include <cstddef>
#include <eigen3/Eigen/Core>
#include <set>
#include <utility>
#include <vector>
 
#ifdef EFFT_USE_FFTW3
#include <fftw3.h>
#endif
 
class Stimulus {
public:
  unsigned int row{0};
  unsigned int col{0};
  bool state{true};
  Stimulus() = default;
  Stimulus(const unsigned int row, const unsigned int col) : row{row}, col{col} {};
  Stimulus(const unsigned int row, const unsigned int col, const bool state) : row{row}, col{col}, state{state} {};
  bool operator==(const Stimulus &p) const { return (row == p.row && col == p.col); }
  bool operator!=(const Stimulus &p) const { return (row != p.row || col != p.col); }
};
 
class Stimuli : public std::vector<Stimulus> {
  using std::vector<Stimulus>::vector;
 
public:
  void filter() {
    std::sort(begin(), end(), [](const Stimulus &left, const Stimulus &right) {
      if(left.row != right.row) {
        return left.row < right.row;
      }
      if(left.col != right.col) {
        return left.col < right.col;
      }
      return left.state && !right.state;
    });
    resize(std::unique(begin(), end(), [](const Stimulus &left, const Stimulus &right) { return (left.row == right.row && left.col == right.col); }) - begin());
  }
 
  void setState(const bool state) {
    std::transform(begin(), end(), begin(), [state](Stimulus &p) { p.state = state; return p; });
  }
};
 
using cfloat = std::complex<float>;
using cfloatmat = Eigen::Matrix<cfloat, Eigen::Dynamic, Eigen::Dynamic>;
 
constexpr unsigned int LOG2(const unsigned int n) { return ((n < 2) ? 0 : 1 + LOG2(n >> 1U)); }
 
#ifdef __has_builtin
#if __has_builtin(__builtin_ffs)
#define EFFT_HAS_BUILTIN_FFS
#endif
#endif
#ifdef EFFT_HAS_BUILTIN_FFS
inline unsigned int log2i(const unsigned int n) { return __builtin_ffs(static_cast<int>(n)) - 1; }
#else
inline unsigned int log2i(const unsigned int n) { return static_cast<unsigned int>(std::log2f(n)); }
#endif
 
template <unsigned int N>
class eFFT {
private:
  static constexpr unsigned int LOG2_N = LOG2(N);
  std::array<std::vector<cfloatmat>, LOG2_N + 1> tree_;
  std::array<cfloat, static_cast<std::size_t>(N) * static_cast<std::size_t>(N + 1)> twiddle_;
#ifdef EFFT_USE_FFTW3
  std::array<fftw_complex, static_cast<std::size_t>(N) * static_cast<std::size_t>(N)> fftwInput_;
  std::array<fftw_complex, static_cast<std::size_t>(N) * static_cast<std::size_t>(N)> fftwOutput_;
  fftw_plan plan_{nullptr};
#endif
 
public:
  eFFT() {
    constexpr float MINUS_TWO_PI = -2 * M_PI;
    for(unsigned int i = 0; i < N; i++) {
      for(unsigned int n = 0; n <= N; n++) {
        twiddle_[i + N * n] = static_cast<cfloat>(std::polar(1.0F, MINUS_TWO_PI * static_cast<float>(i) / static_cast<float>(n)));
      }
    }
  }
 
  ~eFFT() {
#ifdef EFFT_USE_FFTW3
    if(plan_ != nullptr) {
      fftw_destroy_plan(plan_);
    }
#endif
  }
 
  eFFT(const eFFT &) = default;
  eFFT(eFFT &&) noexcept = default;
  eFFT &operator=(const eFFT &) = default;
  eFFT &operator=(eFFT &&) noexcept = default;
 
  void initialize() {
    cfloatmat f0(cfloatmat::Zero(N, N));
    initialize(f0);
  }
 
  void initialize(cfloatmat &x, const int offset = 0) {
    const unsigned int n = x.rows();
    if(n == 1) {
      tree_[0].emplace_back(x);
      return;
    }
    const unsigned int ndiv2 = n >> 1U;
    const unsigned int idx = log2i(ndiv2);
 
    cfloatmat s00(x(Eigen::seq(Eigen::fix<0>, Eigen::last, Eigen::fix<2>), Eigen::seq(Eigen::fix<0>, Eigen::last, Eigen::fix<2>)));
    cfloatmat s01(x(Eigen::seq(Eigen::fix<0>, Eigen::last, Eigen::fix<2>), Eigen::seq(Eigen::fix<1>, Eigen::last, Eigen::fix<2>)));
    cfloatmat s10(x(Eigen::seq(Eigen::fix<1>, Eigen::last, Eigen::fix<2>), Eigen::seq(Eigen::fix<0>, Eigen::last, Eigen::fix<2>)));
    cfloatmat s11(x(Eigen::seq(Eigen::fix<1>, Eigen::last, Eigen::fix<2>), Eigen::seq(Eigen::fix<1>, Eigen::last, Eigen::fix<2>)));
    initialize(s00, 4 * offset);
    initialize(s01, 4 * offset + 4);
    initialize(s10, 4 * offset + 8);
    initialize(s11, 4 * offset + 12);
 
    const cfloatmat &x00 = tree_[idx][offset];
    const cfloatmat &x01 = tree_[idx][offset + 1];
    const cfloatmat &x10 = tree_[idx][offset + 2];
    const cfloatmat &x11 = tree_[idx][offset + 3];
    const unsigned int Nn = N * n;
    for(unsigned int i = 0; i < ndiv2; i++) {
      for(unsigned int j = 0; j < ndiv2; j++) {
        const cfloat tu = twiddle_[j + Nn] * x01(i, j);
        const cfloat td = twiddle_[(i + j) + Nn] * x11(i, j);
        const cfloat ts = twiddle_[i + Nn] * x10(i, j);
 
        const cfloat a = x00(i, j) + tu;
        const cfloat b = x00(i, j) - tu;
        const cfloat c = ts + td;
        const cfloat d = ts - td;
 
        x(i, j) = a + c;
        x(i, j + ndiv2) = b + d;
        x(i + ndiv2, j) = a - c;
        x(i + ndiv2, j + ndiv2) = b - d;
      }
    }
    tree_[log2i(n)].emplace_back(x);
  }
 
  bool update(const Stimulus &p) { return update(tree_[LOG2_N][0], p); }
 
  bool update(cfloatmat &x, const Stimulus &p, const unsigned int offset = 0) {
    const unsigned int n = x.rows();
    if(n == 1) {
      return std::exchange(x(0, 0), p.state).real() != static_cast<float>(p.state);
    }
    const unsigned int ndiv2 = n >> 1U;
    const unsigned int nndiv2 = n * ndiv2;
    const unsigned int idx = log2i(ndiv2);
 
    bool changed = false;
    if(static_cast<bool>(p.row & 1U)) {
      if(static_cast<bool>(p.col & 1U)) {
        changed = update(tree_[idx][offset + 3], {p.row >> 1U, p.col >> 1U, p.state}, 4 * offset + 12); // odd-odd
      } else {
        changed = update(tree_[idx][offset + 2], {p.row >> 1U, p.col >> 1U, p.state}, 4 * offset + 8); // odd-even
      }
    } else {
      if(static_cast<bool>(p.col & 1U)) {
        changed = update(tree_[idx][offset + 1], {p.row >> 1U, p.col >> 1U, p.state}, 4 * offset + 4); // even-odd
      } else {
        changed = update(tree_[idx][offset], {p.row >> 1U, p.col >> 1U, p.state}, 4 * offset); // even-even
      }
    }
 
    if(changed) {
      const cfloat *x00 = tree_[idx][offset].data();
      const cfloat *x01 = tree_[idx][offset + 1].data();
      const cfloat *x10 = tree_[idx][offset + 2].data();
      const cfloat *x11 = tree_[idx][offset + 3].data();
      cfloat *xp = x.data();
      const unsigned int Nn = N * n;
 
      for(unsigned int j = 0; j < ndiv2; j++) {
        const unsigned int ndiv2j = ndiv2 * j, nj = n * j;
        for(unsigned int i = 0; i < ndiv2; i++) {
          const unsigned int k = i + ndiv2j, k1 = i + nj, k2 = k1 + ndiv2;
 
          const cfloat tu = twiddle_[j + Nn] * x01[k];
          const cfloat td = twiddle_[i + j + Nn] * x11[k];
          const cfloat ts = twiddle_[i + Nn] * x10[k];
 
          const cfloat x00_k = x00[k];
          const cfloat a = x00_k + tu;
          const cfloat b = x00_k - tu;
          const cfloat c = ts + td;
          const cfloat d = ts - td;
 
          xp[k1] = a + c;
          xp[k1 + nndiv2] = b + d;
          xp[k2] = a - c;
          xp[k2 + nndiv2] = b - d;
        }
      }
    }
    return changed;
  }
 
  bool update(Stimuli &pv) { return update(tree_[LOG2_N][0], pv.begin(), pv.end()); }
 
  bool update(cfloatmat &x, Stimuli::iterator b0, Stimuli::iterator e0, const unsigned int offset = 0) {
    const unsigned int n = x.rows();
    if(n == 1) {
      if(e0 - b0 == 1) {
        return std::exchange(x(0, 0), b0->state).real() != static_cast<float>(b0->state);
      }
      const bool state = std::any_of(b0, e0, [](const Stimulus &p) { return p.state; });
      return std::exchange(x(0, 0), state).real() != static_cast<float>(state);
    }
    const unsigned int ndiv2 = n >> 1U;
    const unsigned int nndiv2 = n * ndiv2;
    const unsigned int idx = log2i(ndiv2);
 
    Stimuli::iterator e1, e2, e3;
    e2 = std::partition(b0, e0, [](const Stimulus &p) { return p.row & 1U; });
    e1 = std::partition(b0, e2, [](const Stimulus &p) { return p.col & 1U; });
    e3 = std::partition(e2, e0, [](const Stimulus &p) { return p.col & 1U; });
    std::transform(b0, e1, b0, [](const Stimulus &p) { return Stimulus{p.row >> 1U, p.col >> 1U, p.state}; });
    std::transform(e1, e2, e1, [](const Stimulus &p) { return Stimulus{p.row >> 1U, p.col >> 1U, p.state}; });
    std::transform(e2, e3, e2, [](const Stimulus &p) { return Stimulus{p.row >> 1U, p.col >> 1U, p.state}; });
    std::transform(e3, e0, e3, [](const Stimulus &p) { return Stimulus{p.row >> 1U, p.col >> 1U, p.state}; });
 
    bool changed = false;
    if(b0 != e1) {
      changed = update(tree_[idx][offset + 3], b0, e1, 4 * offset + 12) || changed; // odd-odd
    }
    if(e1 != e2) {
      changed = update(tree_[idx][offset + 2], e1, e2, 4 * offset + 8) || changed; // odd-even
    }
    if(e2 != e3) {
      changed = update(tree_[idx][offset + 1], e2, e3, 4 * offset + 4) || changed; // even-odd
    }
    if(e3 != e0) {
      changed = update(tree_[idx][offset], e3, e0, 4 * offset) || changed; // even-even
    }
 
    if(changed) {
      const cfloat *x00 = tree_[idx][offset].data();
      const cfloat *x01 = tree_[idx][offset + 1].data();
      const cfloat *x10 = tree_[idx][offset + 2].data();
      const cfloat *x11 = tree_[idx][offset + 3].data();
      cfloat *xp = x.data();
      const unsigned int Nn = N * n;
 
      for(unsigned int j = 0; j < ndiv2; j++) {
        const unsigned int ndiv2j = ndiv2 * j, nj = n * j;
        for(unsigned int i = 0; i < ndiv2; i++) {
          const unsigned int k = i + ndiv2j, k1 = i + nj, k2 = k1 + ndiv2;
 
          const cfloat tu = twiddle_[j + Nn] * x01[k];
          const cfloat td = twiddle_[i + j + Nn] * x11[k];
          const cfloat ts = twiddle_[i + Nn] * x10[k];
 
          const cfloat x00_k = x00[k];
          const cfloat a = x00_k + tu;
          const cfloat b = x00_k - tu;
          const cfloat c = ts + td;
          const cfloat d = ts - td;
 
          xp[k1] = a + c;
          xp[k1 + nndiv2] = b + d;
          xp[k2] = a - c;
          xp[k2 + nndiv2] = b - d;
        }
      }
    }
    return changed;
  }
 
#ifdef EFFT_USE_FFTW3
  void initializeGroundTruth(const cfloatmat &image = cfloatmat::Zero(N, N)) {
    plan_ = fftw_plan_dft_2d(N, N, fftwInput_.data(), fftwOutput_.data(), FFTW_FORWARD, FFTW_ESTIMATE | FFTW_UNALIGNED);
    for(Eigen::Index i = 0; i < image.rows(); i++) {
      for(Eigen::Index j = 0; j < image.cols(); j++) {
        fftwInput_[N * i + j][0] = image(i, j).real();
        fftwInput_[N * i + j][1] = image(i, j).imag();
      }
    }
    fftw_execute(plan_);
  }
 
  void updateGroundTruth(const Stimulus &p) {
    fftwInput_[N * p.row + p.col][0] = static_cast<double>(p.state);
    fftwInput_[N * p.row + p.col][1] = 0;
    fftw_execute(plan_);
  }
 
  void updateGroundTruth(const Stimuli &pv) {
    std::set<std::pair<unsigned int, unsigned int>> activated;
    for(const Stimulus &p : pv) {
      if(p.state) {
        activated.insert({p.row, p.col});
      } else if(activated.find({p.row, p.col}) != activated.end()) {
        continue;
      }
      fftwInput_[N * p.row + p.col][0] = static_cast<double>(p.state);
      fftwInput_[N * p.row + p.col][1] = 0;
    }
    fftw_execute(plan_);
  }
#endif
 
  [[nodiscard]] inline Eigen::Matrix<cfloat, N, N> getFFT() const {
    return tree_[LOG2_N][0];
  }
 
#ifdef EFFT_USE_FFTW3
  [[nodiscard]] inline Eigen::Matrix<cfloat, N, N> getGroundTruthFFT() const {
    return Eigen::Map<const Eigen::Matrix<std::complex<double>, N, N, Eigen::RowMajor>>(reinterpret_cast<const std::complex<double> *>(fftwOutput_.data())).template cast<cfloat>();
  }
 
  [[nodiscard]] inline double check() const {
    return (getFFT() - getGroundTruthFFT()).norm();
  }
#endif
};
 
#endif // EFFT_HPP