lasp/src/lasp/pybind11/lasp_dsp_pybind.cpp

#include <carma>
#include "lasp_avpowerspectra.h"
#include "lasp_biquadbank.h"
#include "lasp_fft.h"
#include "lasp_slm.h"
#include "lasp_window.h"
#include <iostream>
#include <pybind11/pybind11.h>

using std::cerr;
namespace py = pybind11;

/**
 * \ingroup pybind
 * @{
 *
 */

/**
 * @brief Initialize DSP code
 *
 * @param m The Python module to add classes and methods to
 */
void init_dsp(py::module &m) {

  py::class_<Fft> fft(m, "Fft");
  fft.def(py::init<us>());
  fft.def("fft", py::overload_cast<const vd &>(&Fft::fft));
  fft.def("fft", py::overload_cast<const dmat &>(&Fft::fft));

  fft.def("ifft", py::overload_cast<const vc &>(&Fft::ifft));
  fft.def("ifft", py::overload_cast<const cmat &>(&Fft::ifft));
  fft.def_static("load_fft_wisdom", &Fft::load_fft_wisdom);
  fft.def_static("store_fft_wisdom", &Fft::store_fft_wisdom);

  /// Window 
  py::class_<Window> w(m, "Window");

  py::enum_<Window::WindowType>(w, "WindowType")
      .value("Hann", Window::WindowType::Hann)
      .value("Hamming", Window::WindowType::Hamming)
      .value("Bartlett", Window::WindowType::Bartlett)
      .value("Blackman", Window::WindowType::Bartlett)
      .value("Rectangular", Window::WindowType::Rectangular);

  w.def_static("toTxt", &Window::toText);


  /// SeriesBiquad 
  py::class_<SeriesBiquad> sbq(m, "SeriesBiquad");
  sbq.def(py::init<const vd &>());
  sbq.def("filter", [](SeriesBiquad &s, const vd &input) {
    vd res = input;
    s.filter(res);
    return res;
  });
  
  /// BiquadBank
  py::class_<BiquadBank, std::shared_ptr<BiquadBank>> bqb(m, "BiquadBank");
  bqb.def(py::init<const dmat&,const vd*>());
  bqb.def("setGains",&BiquadBank::setGains);
  bqb.def("filter",&BiquadBank::filter);

  /// PowerSpectra 
  py::class_<PowerSpectra> ps(m, "PowerSpectra");
  ps.def(py::init<const us, const Window::WindowType>());
  ps.def("compute", &PowerSpectra::compute);

  /// AvPowerSpectra 
  py::class_<AvPowerSpectra> aps(m, "AvPowerSpectra");
  aps.def(py::init<const us, const Window::WindowType, const d, const int>(),
          py::arg("nfft"), py::arg("WindowType"), py::arg("overlap_percentage"),
          py::arg("time_constant"));
  aps.def("compute", [](AvPowerSpectra &aps, const dmat &timedata) {
    std::optional<arma::cx_cube> res = aps.compute(timedata);
      return res.value_or(arma::cx_cube(0,0,0));
    }
  );

  py::class_<SLM> slm(m, "SLM");

  slm.def_static(
      "fromBiquads",
      py::overload_cast<const d, const d, const us, const d, const dmat &>(
          &SLM::fromBiquads));
  slm.def_static(
      "fromBiquads",
      py::overload_cast<const d, const d, const us, const d, const vd &,
                        const dmat &>(&SLM::fromBiquads));
  slm.def("run", &SLM::run);
  slm.def_readonly("Pm", &SLM::Pm);
  slm.def_readonly("Pmax", &SLM::Pmax);
  slm.def_readonly("Ppeak", &SLM::Ppeak);
  slm.def("Lpeak", &SLM::Lpeak);
  slm.def("Leq", &SLM::Leq);
  slm.def("Lmax", &SLM::Lmax);
}
/** @} */