sara.h

Go to the documentation of this file.

#ifndef SOPT_WAVELETS_SARA_H
#define SOPT_WAVELETS_SARA_H
 
#include "sopt/config.h"
#include <algorithm> // for std::min<>
#include <cmath>
#include <initializer_list>
#include <string> // for std::string
#include <tuple>
#include <vector>
#include "sopt/logging.h"
#include "sopt/wavelets/wavelets.h"
#ifdef SOPT_MPI
#include "sopt/mpi/communicator.h"
#endif
 
namespace sopt::wavelets {
 
class SARA : public std::vector<Wavelet> {
 public:
#ifndef SOPT_HAS_NOT_USING
  // Constructors
  using std::vector<Wavelet>::vector;
#else
  SARA() : std::vector<Wavelet>(){};
#endif
  SARA(std::initializer_list<std::tuple<std::string, t_uint>> const &init)
      : SARA(init.begin(), init.end()) {}
  template <typename ITERATOR,
            class T = typename std::enable_if<
                std::is_convertible<decltype(std::get<0>(*std::declval<ITERATOR>())),
                                    std::string>::value and
                std::is_convertible<decltype(std::get<1>(*std::declval<ITERATOR>())),
                                    t_uint>::value>::type>
  SARA(ITERATOR first, ITERATOR last) {
    for (; first != last; ++first) emplace_back(std::get<0>(*first), std::get<1>(*first));
  }
 
  SARA(const_iterator first, const_iterator last) : std::vector<Wavelet>(first, last) {}
  virtual ~SARA() {}
 
  template <typename T0>
  typename T0::PlainObject direct(Eigen::ArrayBase<T0> const &signal) const;
  template <typename T0, typename T1>
  void direct(Eigen::ArrayBase<T1> &coefficients, Eigen::ArrayBase<T0> const &signal) const;
  template <typename T0, typename T1>
  void direct(Eigen::ArrayBase<T1> &&coefficients, Eigen::ArrayBase<T0> const &signal) const {
    direct(coefficients, signal);
  }
  template <typename T0>
  typename T0::PlainObject indirect(Eigen::ArrayBase<T0> const &coeffs) const;
  template <typename T0, typename T1>
  void indirect(Eigen::ArrayBase<T1> const &coefficients, Eigen::ArrayBase<T0> &signal) const;
  template <typename T0, typename T1>
  void indirect(Eigen::ArrayBase<T1> const &coeffs, Eigen::ArrayBase<T0> &&signal) const {
    indirect(coeffs, signal);
  }
 
  t_uint max_levels() const {
    if (size() == 0) return 0;
    auto cmp = [](Wavelet const &a, Wavelet const &b) { return a.levels() < b.levels(); };
    return std::max_element(begin(), end(), cmp)->levels();
  }
 
  void emplace_back(std::string const &name, t_uint nlevels) {
    std::vector<Wavelet>::emplace_back(factory(name, nlevels));
  }
};
 
#define SOPT_WAVELET_ERROR_MACRO(INPUT)                                           \
  if ((INPUT).rows() % (1u << max_levels()) != 0)                                   \
    throw std::length_error("Inconsistent number of columns and wavelet levels"); \
  else if ((INPUT).cols() != 1 and (INPUT).cols() % (1u << max_levels()))             \
    throw std::length_error("Inconsistent number of rows and wavelet levels");
 
template <typename T0, typename T1>
void SARA::direct(Eigen::ArrayBase<T1> &coeffs, Eigen::ArrayBase<T0> const &signal) const {
  SOPT_WAVELET_ERROR_MACRO(signal);
  if (coeffs.rows() != signal.rows() or coeffs.cols() != signal.cols() * static_cast<t_int>(size()))
    coeffs.derived().resize(signal.rows(), signal.cols() * size());
  if (coeffs.rows() != signal.rows() or coeffs.cols() != signal.cols() * static_cast<t_int>(size()))
    throw std::length_error("Incorrect size for output matrix(or could not resize)");
  if (size() == 0) return;
  auto const Ncols = signal.cols();
  {
#ifndef SOPT_OPENMP
    SOPT_TRACE("Calling direct sara without threads");
#endif
    if (Ncols == 1)
      for (size_type i = 0; i < size(); ++i) at(i).direct(coeffs.col(i), signal.col(0));
    else
      for (size_type i = 0; i < size(); ++i)
        at(i).direct(coeffs.leftCols((i + 1) * Ncols).rightCols(Ncols), signal);
  }
 
  coeffs /= std::sqrt(size());
}
 
template <typename T0, typename T1>
void SARA::indirect(Eigen::ArrayBase<T1> const &coeffs, Eigen::ArrayBase<T0> &signal) const {
  if (size() == 0) throw std::runtime_error("Empty wavelets: adjoint operation undefined");
  if (signal.cols() == 1) {
    if (coeffs.rows() % (1u << max_levels()) != 0)
      throw std::length_error("Inconsistent number of columns and wavelet levels");
  } else
    SOPT_WAVELET_ERROR_MACRO(coeffs);
  if (coeffs.cols() % size() != 0)
    throw std::length_error(
        "Columns of coefficient matrix and number of wavelets are inconsistent");
  if (coeffs.rows() != signal.rows() or coeffs.cols() != signal.cols() * static_cast<t_int>(size()))
    signal.derived().resize(coeffs.rows(), coeffs.cols() / size());
  if (coeffs.rows() != signal.rows() or coeffs.cols() != signal.cols() * static_cast<t_int>(size()))
    throw std::length_error("Incorrect size for output matrix(or could not resize)");
  auto const Ncols = signal.cols();
  {
#ifndef SOPT_OPENMP
    SOPT_TRACE("Calling indirect sara without threads");
#endif
    signal = T0::Zero(signal.rows(), signal.cols());
    if (Ncols == 1)
      for (size_type i = 0; i < size(); ++i) signal.col(0) += at(i).indirect(coeffs.col(i));
    else
      for (size_type i = 0; i < size(); ++i)
        signal += at(i).indirect(coeffs.leftCols((i + 1) * Ncols).rightCols(Ncols));
  }
  signal /= std::sqrt(size());
}
 
#undef SOPT_WAVELET_ERROR_MACRO
 
template <typename T0>
typename T0::PlainObject SARA::indirect(Eigen::ArrayBase<T0> const &coeffs) const {
  using t_Output = decltype(this->front().indirect(coeffs));
  t_Output signal = t_Output::Zero(coeffs.rows(), coeffs.cols() / size());
  (*this).indirect(coeffs, signal);
  return signal;
}
 
template <typename T0>
typename T0::PlainObject SARA::direct(Eigen::ArrayBase<T0> const &signal) const {
  using t_Output = decltype(this->front().direct(signal));
  t_Output result = t_Output::Zero(signal.rows(), signal.cols() * size());
  (*this).direct(result, signal);
  return result;
}
 
#ifdef SOPT_MPI
template <typename T>
T distribute_sara(T const &sara, t_uint size, t_uint rank) {
  auto const start = [](t_uint size, t_uint ncomms, t_uint rank) {
    return std::min(size, rank * (size / ncomms) + std::min(rank, size % ncomms));
  };
  auto const startw = start(sara.size(), size, rank);
  auto const endw = start(sara.size(), size, rank + 1);
  return T(sara.begin() + startw, sara.begin() + endw);
}
template <typename T>
T distribute_sara(T const &all_wavelets, mpi::Communicator const &comm) {
  return distribute_sara<T>(all_wavelets, comm.size(), comm.rank());
}
#endif
} // namespace sopt::wavelets
#endif