add math module

Signed-off-by: nbaiot <[email protected]>

CMakeLists.txt

@@ -18,6 +18,8 @@ elseif(UNIX)
     add_definitions(-DOS_UNIX)
 endif()
 
+set(CMAKE_MODULE_PATH "${CMAKE_CURRENT_LIST_DIR}/cmake/external")
+
 ### for boost
 set(Boost_NO_BOOST_CMAKE TRUE)
 set(BOOST_ROOT /opt/third_party/boost/boost_1.71)
@@ -31,11 +33,21 @@ if(Boost_FOUND)
     MESSAGE( STATUS "Boost_LIB_VERSION = ${Boost_LIB_VERSION}")
 endif()
 
+find_package(FFTW3 REQUIRED)
+if(FFTW3_FOUND)
+    include_directories(${FFTW3_INCLUDES})
+    message(STATUS "    libraries: ${FFTW3_INCLUDES}")
+    message(STATUS "    libraries: ${FFTW3_LIBRARIES}")
+endif()
+
+
 ### include module
 add_subdirectory(base)
 
 add_subdirectory(memery)
 
+add_subdirectory(math)
+
 add_subdirectory(test)
 
 

math/CMakeLists.txt

@@ -0,0 +1,11 @@
+
+include_directories(${CMAKE_CURRENT_LIST_DIR}/include)
+include_directories(${FFTW3_INCLUDES})
+
+aux_source_directory(. SRCS)
+
+set(LIBRARY_OUTPUT_PATH ${PROJECT_BINARY_DIR}/lib)
+
+add_library(nb_math SHARED ${SRCS})
+
+target_link_libraries(nb_math ${FFTW3_LIBRARIES} nb_base)

math/fft_factory.cc

@@ -0,0 +1,19 @@
+//
+// Created by [email protected] on 2019/12/7.
+//
+
+#include <fft_factory.h>
+#include "fftw.h"
+#include "ifftw.h"
+
+namespace nbsdk::math {
+
+
+sp<FFT> FTTFactory::CreateFFT() {
+  return std::make_shared<FFTW>();
+}
+
+sp<IFFT> FTTFactory::CreateIFFT() {
+  return std::make_shared<IFFTW>();
+}
+}

math/fftw.cc

@@ -1,5 +1,47 @@
 //
-// Created by shajia on 2019/12/7.
+// Created by [email protected] on 2019/12/7.
 //
 
 #include "fftw.h"
+
+namespace nbsdk::math {
+
+FFTW::~FFTW() {
+  Terminate();
+}
+
+void FFTW::Init(int sampleCount) {
+  sample_count_ = sampleCount;
+  in_ = (double*)fftw_malloc(sizeof(double) * sample_count_);
+  out_size_ = (sample_count_ / 2) + 1;
+  out_ = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * out_size_);
+  plan_ = fftw_plan_dft_r2c_1d(sample_count_, in_, out_, FFTW_ESTIMATE);
+}
+
+void FFTW::Execute(std::complex<double>** out, const int16_t* in) {
+  if (out == nullptr || *out == nullptr || in == nullptr)
+    return;
+  /// TODO: fixme
+  for (int i = 0; i < sample_count_; i++) {
+    in_[i] = (double)in[i];
+  }
+  fftw_execute(plan_);
+  std::complex<double>* oo = *out;
+  for (int i = 0; i < out_size_; i++) {
+    oo[i] = std::complex<double>((double)out_[i][0], (double)out_[i][1]);
+  }
+}
+
+void FFTW::Terminate() {
+  fftw_destroy_plan(plan_);
+  if (in_) {
+    fftw_free(in_);
+    in_ = nullptr;
+  }
+  if (out_) {
+    fftw_free(out_);
+    out_ = nullptr;
+  }
+}
+
+}

math/fftw.h

@@ -1,14 +1,37 @@
 //
-// Created by shajia on 2019/12/7.
+// Created by [email protected] on 2019/12/7.
 //
+#pragma once
 
 #ifndef NBSDK_FFTW_H
 #define NBSDK_FFTW_H
 
+#include <fft.h>
+#include <fftw3.h>
 
-class FFTW {
+namespace nbsdk::math {
 
+class FFTW : public FFT {
+
+public:
+
+  ~FFTW() override ;
+
+  void Init(int sampleCount) override;
+
+  void Execute(std::complex<double>** out, const int16_t* in) override;
+
+  void Terminate() override;
+
+private:
+  int sample_count_{0};
+  int out_size_{0};
+  fftw_plan plan_;
+  double* in_{nullptr};
+  fftw_complex* out_{nullptr};
 };
 
+}
+
 
 #endif //NBSDK_FFTW_H

math/gcc_phat.cc

@@ -0,0 +1,81 @@
+//
+// Created by shajia on 2019/12/7.
+//
+
+#include <base/include/nlog.h>
+#include <cfloat>
+#include <algorithm>
+
+#include "gcc_phat.h"
+#include "fft_factory.h"
+
+namespace nbsdk::math {
+
+GCCPhat::GCCPhat(int sampleCount) : sample_count_(sampleCount) {
+  fft_ = FTTFactory::CreateFFT();
+  fft_->Init(sample_count_);
+
+  ifft_ = FTTFactory::CreateIFFT();
+  ifft_->Init(sample_count_);
+}
+
+GCCPhat::~GCCPhat() {
+  LOG_INFO() << "~~~~~~~~~~~~~~~~~GCCPhat";
+};
+
+int GCCPhat::Execute(const int16_t *s1, const int16_t *s2, int count, int margin) {
+  if (s1 == nullptr || s2 == nullptr)
+    return -1;
+  CHECK(count == sample_count_);
+
+  std::vector<std::complex<double>> s1FFT;
+  s1FFT.resize(count);
+  auto pS1FFT = s1FFT.data();
+  fft_->Execute(&pS1FFT, s1);
+
+  std::vector<std::complex<double>> s2FFT;
+  s2FFT.resize(count);
+  auto pS2FFT = s2FFT.data();
+  fft_->Execute(&pS2FFT, s2);
+
+  std::vector<std::complex<double>> R;
+  R.resize(count);
+  for (int i = 0; i < count; i++) {
+    std::complex<double> v = s2FFT[i] * std::conj(s1FFT[i]);
+    v = v / (std::abs(v) + FLT_MIN);
+    R[i] = v;
+  }
+
+  std::vector<double> crossCorrelation;
+  crossCorrelation.resize(count);
+  auto pCC = crossCorrelation.data();
+  ifft_->Execute(&pCC, R.data());
+
+
+  std::vector<double> shifted;
+  //shifted.resize(crossCorrelation.size());
+  auto half = std::round(crossCorrelation.size() / 2.0);
+  auto half_iter = crossCorrelation.begin() + half;
+  std::copy(half_iter, crossCorrelation.end(), std::back_inserter(shifted));
+  std::copy(crossCorrelation.begin(), half_iter, std::back_inserter(shifted));
+
+  auto newMargin = margin;
+  if ((half - newMargin) < 0) {
+    newMargin = half;
+  }
+  if ((half + newMargin) >= sample_count_) {
+    newMargin = (sample_count_ - 1) - half;
+  }
+
+  auto start = half - newMargin;
+  auto len = 2 * newMargin + 1;
+
+  int maxIndex = std::distance(shifted.begin() + start,
+      std::max_element(shifted.begin() + start, shifted.begin() + start + len));
+
+  int argMax = maxIndex - newMargin;
+
+  return argMax;
+}
+
+}

math/ifftw.cc

@@ -0,0 +1,50 @@
+//
+// Created by [email protected] on 2019/12/7.
+//
+
+#include "ifftw.h"
+
+namespace nbsdk::math {
+
+IFFTW::~IFFTW() {
+  Terminate();
+}
+
+void IFFTW::Init(int sampleCount) {
+  if (out_size_ == sampleCount)
+    return;
+  Terminate();
+  out_size_ = sampleCount;
+  out_ = (double*)fftw_malloc(sizeof(double) * out_size_);
+  in_size_ = (out_size_ / 2) + 1;
+  in_ = (fftw_complex*)fftw_malloc(sizeof(fftw_complex) * in_size_);
+  plan_ = fftw_plan_dft_c2r_1d(out_size_, in_, out_, FFTW_ESTIMATE);
+}
+
+void IFFTW::Execute(double** out, const std::complex<double>* in) {
+  if (out == nullptr || *out == nullptr || in == nullptr)
+    return;
+  for(int i = 0; i < in_size_; i++) {
+    in_[i][0] = (double)in[i].real();
+    in_[i][1] = (double)in[i].imag();
+  }
+  fftw_execute(plan_);
+  double* oo = *out;
+  for (int i = 0; i < out_size_; i++) {
+    oo[i] = out_[i];
+  }
+}
+
+void IFFTW::Terminate() {
+  fftw_destroy_plan(plan_);
+  if (in_) {
+    fftw_free(in_);
+    in_ = nullptr;
+  }
+  if (out_) {
+    fftw_free(out_);
+    out_ = nullptr;
+  }
+}
+
+}

math/ifftw.h

@@ -0,0 +1,37 @@
+//
+// Created by [email protected] on 2019/12/7.
+//
+#pragma once
+
+#ifndef NBSDK_IFFTW_H
+#define NBSDK_IFFTW_H
+
+#include <fft.h>
+#include <fftw3.h>
+
+namespace nbsdk::math {
+
+class IFFTW : public IFFT {
+
+public:
+  ~IFFTW() override;
+
+  void Init(int sampleCount) override;
+
+  void Execute(double** out, const std::complex<double>* in) override;
+
+  void Terminate() override;
+
+private:
+  int in_size_{0};
+  int out_size_{0};
+  double* out_{nullptr};
+  fftw_complex* in_{nullptr};
+  fftw_plan plan_;
+
+};
+
+
+}
+
+#endif //NBSDK_IFFTW_H

math/include/fft.h

@@ -0,0 +1,41 @@
+//
+// Created by [email protected] on 2019/12/7.
+//
+#pragma once
+
+#ifndef NBSDK_FFT_H
+#define NBSDK_FFT_H
+
+#include <vector>
+#include <complex>
+
+namespace nbsdk::math {
+
+class FFT {
+
+public:
+  virtual ~FFT() = default;
+
+  virtual void Init(int sampleCount) = 0;
+
+  ///TODO: use Buffer replace out, in
+  virtual void Execute(std::complex<double>** out, const int16_t* in) = 0;
+
+  virtual void Terminate() = 0;
+};
+
+class IFFT {
+public:
+  virtual ~IFFT() = default;
+
+  virtual void Init(int sampleCount) = 0;
+
+  ///TODO: use Buffer replace out, in
+  virtual void Execute(double** out, const std::complex<double>* in) = 0;
+
+  virtual void Terminate() = 0;
+};
+
+}
+
+#endif //NBSDK_FFT_H

math/include/fft_factory.h

@@ -0,0 +1,27 @@
+//
+// Created by [email protected] on 2019/12/7.
+//
+#pragma once
+
+#ifndef NBSDK_FFT_FACTORY_H
+#define NBSDK_FFT_FACTORY_H
+
+#include <base/include/smart_pointer.h>
+
+#include "fft.h"
+
+
+using namespace nbsdk::base;
+
+namespace nbsdk::math {
+
+class FTTFactory {
+public:
+  static sp<FFT> CreateFFT();
+
+  static sp<IFFT> CreateIFFT();
+};
+
+}
+
+#endif //NBSDK_FFT_FACTORY_H