wavelets.h

Go to the documentation of this file.

#ifndef SOPT_WAVELETS_H
#define SOPT_WAVELETS_H
 
#include <iostream>
// Convenience header to include wavelets headers and additional utilities
#include "sopt/config.h"
#include "sopt/linear_transform.h"
#ifdef SOPT_MPI
#include "sopt/mpi/communicator.h"
#endif
#include "sopt/wavelets/sara.h"
#include "sopt/wavelets/wavelets.h"
 
namespace sopt {
namespace details {
namespace {
template <typename T, typename OP>
LinearTransform<Vector<T>> linear_transform(OP const &op) {
  return LinearTransform<Vector<T>>(
      [op](Vector<T> &out, Vector<T> const &x) { op.indirect(x.array(), out.array()); },
      [op](Vector<T> &out, Vector<T> const &x) { op.direct(out.array(), x.array()); });
}
template <typename T, typename OP>
LinearTransform<Vector<T>> linear_transform(OP const &op, t_uint rows, t_uint cols,
                                            t_uint factor = 1) {
  return LinearTransform<Vector<T>>(
      [op, rows, cols, factor](Vector<T> &out, Vector<T> const &x) {
        assert(static_cast<t_uint>(x.size()) == rows * cols * factor);
        out.resize(rows * cols);
        auto signal = Image<T>::Map(out.data(), rows, cols);
        auto const coeffs = Image<T>::Map(x.data(), rows, cols * factor);
        op.indirect(coeffs, signal);
      },
      {{0, 1, static_cast<t_int>(rows * cols)}},
      [op, rows, cols, factor](Vector<T> &out, Vector<T> const &x) {
        assert(static_cast<t_uint>(x.size()) == rows * cols);
        out.resize(rows * cols * factor);
        auto const signal = Image<T>::Map(x.data(), rows, cols);
        auto coeffs = Image<T>::Map(out.data(), rows, cols * factor);
        op.direct(coeffs, signal);
      },
      {{0, 1, static_cast<t_int>(factor * rows * cols)}});
}
}  // namespace
}  // namespace details
 
namespace utilities {
template <typename T>
Vector<T> &get_wavelet_basis_coefficients(Vector<T> &coeffs, const t_uint basis_index,
                                          const t_uint size);
template <typename T>
Vector<T> &get_wavelet_levels_1d(Vector<T> &coeffs, const t_uint level, const t_uint size);
template <typename T>
Vector<T> &get_wavelet_levels(Vector<T> &coeffs, const t_uint level, const t_uint rows,
                              const t_uint cols);
template <typename T>
Vector<T> &get_wavelet_low_high_pass(Vector<T> &coeffs, const t_uint level, const t_uint rows,
                                     const t_uint cols);
template <typename T>
Vector<T> &get_wavelet_high_high_pass(Vector<T> &coeffs, const t_uint level, const t_uint rows,
                                      const t_uint cols);
template <typename T>
Vector<T> &get_wavelet_high_low_pass(Vector<T> &coeffs, const t_uint level, const t_uint rows,
                                     const t_uint cols);
template <typename T>
Vector<T> &get_wavelet_low_low_pass(Vector<T> &coeffs, const t_uint level, const t_uint rows,
                                    const t_uint cols);
template <typename T>
Vector<T> &get_wavelet_high_pass_1d(Vector<T> &coeffs, const t_uint level, const t_uint size);
// macro to add version to work with a wavelet dictionary
#define SOPT_WAVELET_MACRO(NAME)                                                                 \
  template <typename T>                                                                             \
  Vector<T> &NAME(Vector<T> &coeffs, const t_uint basis_index, const t_uint level,               \
                  const t_uint rows, const t_uint cols);                                         \
  template <typename T>                                                                             \
  Vector<T> &NAME(Vector<T> &coeffs, const t_uint basis_index, const t_uint level,               \
                  const t_uint rows, const t_uint cols) {                                        \
    return NAME(get_wavelet_basis_coefficients(coeffs, basis_index, coeffs.size()), level, rows, \
                cols);                                                                           \
  };
SOPT_WAVELET_MACRO(get_wavelet_levels)
SOPT_WAVELET_MACRO(get_wavelet_low_high_pass)
SOPT_WAVELET_MACRO(get_wavelet_high_high_pass)
SOPT_WAVELET_MACRO(get_wavelet_high_low_pass)
SOPT_WAVELET_MACRO(get_wavelet_low_low_pass)
#undef SOPT_WAVELET_MACRO
// implimentations
template <typename T>
Vector<T> &get_wavelet_basis_coefficients(Vector<T> &coeffs, const t_uint basis_index,
                                          const t_uint size) {
  assert(coeffs.size() > basis_index * size);
  return coeffs.segment(basis_index * size, size);
}
template <typename T>
Vector<T> &get_wavelet_levels_1d(Vector<T> &coeffs, const t_uint level, const t_uint size) {
  auto const N = static_cast<t_uint>(coeffs.size()) >> level;  // bitshift to divide by 2^level
  return coeffs.head(N);
}
template <typename T>
Vector<T> &get_wavelet_high_pass_1d(Vector<T> &coeffs, const t_uint level, const t_uint size) {
  auto const N = static_cast<t_uint>(coeffs.size()) >> level;  // bitshift to divide by 2^level
  return get_wavelet_levels(coeffs, level, size).tail(N / 2);
}
template <typename T>
Vector<T> &get_wavelet_levels(Vector<T> &coeffs, const t_uint level, const t_uint rows,
                              const t_uint cols) {
  const Matrix<T> signal = Matrix<T>::Map(coeffs.data(), rows, cols);
  auto const Nx = static_cast<t_uint>(signal.rows()) >> level;  // bitshift to divide by 2^level
  auto const Ny = static_cast<t_uint>(signal.cols()) >> level;
  return Vector<T>::Map(signal.topLeftCorner(Nx, Ny).data(), Nx * Ny);
}
template <typename T>
Vector<T> &get_wavelet_low_high_pass(Vector<T> &coeffs, const t_uint level, const t_uint rows,
                                     const t_uint cols) {
  auto const Nx = rows >> level;  // bitshift to divide by 2^level
  auto const Ny = cols >> level;
  const Matrix<T> signal =
      Matrix<T>::Map(get_wavelet_levels(coeffs, level, rows, cols).data(), Nx, Ny);
  return Vector<T>::Map(signal.topRightCorner(signal.rows() / 2, signal.cols() / 2).data(),
                        signal.size() / 4);
}
template <typename T>
Vector<T> &get_wavelet_high_high_pass(Vector<T> &coeffs, const t_uint level, const t_uint rows,
                                      const t_uint cols) {
  auto const Nx = rows >> level;  // bitshift to divide by 2^level
  auto const Ny = cols >> level;
  const Matrix<T> signal =
      Matrix<T>::Map(get_wavelet_levels(coeffs, level, rows, cols).data(), Nx, Ny);
  return Vector<T>::Map(signal.bottomRightCorner(signal.rows() / 2, signal.cols() / 2).data(),
                        signal.size() / 4);
}
template <typename T>
Vector<T> &get_wavelet_high_low_pass(Vector<T> &coeffs, const t_uint level, const t_uint rows,
                                     const t_uint cols) {
  auto const Nx = rows >> level;  // bitshift to divide by 2^level
  auto const Ny = cols >> level;
  const Matrix<T> signal =
      Matrix<T>::Map(get_wavelet_levels(coeffs, level, rows, cols).data(), Nx, Ny);
  return Vector<T>::Map(signal.bottomLeftCorner(signal.rows() / 2, signal.cols() / 2).data(),
                        signal.size() / 4);
}
template <typename T>
Vector<T> &get_wavelet_low_low_pass(Vector<T> &coeffs, const t_uint level, const t_uint rows,
                                    const t_uint cols) {
  auto const Nx = rows >> level;  // bitshift to divide by 2^level
  auto const Ny = cols >> level;
  const Matrix<T> signal =
      Matrix<T>::Map(get_wavelet_levels(coeffs, level, rows, cols).data(), Nx, Ny);
  return Vector<T>::Map(signal.topLeftCorner(signal.rows() / 2, signal.cols() / 2).data(),
                        signal.size() / 4);
}
 
}  // namespace utilities
 
template <typename T>
LinearTransform<Vector<T>> linear_transform(wavelets::Wavelet const &wavelet) {
  return details::linear_transform<T, wavelets::Wavelet>(wavelet);
}
 
template <typename T>
LinearTransform<Vector<T>> linear_transform(wavelets::SARA const &sara) {
  return details::linear_transform<T, wavelets::SARA>(sara);
}
 
template <typename T>
LinearTransform<Vector<T>> linear_transform(wavelets::Wavelet const &wavelet, t_uint rows,
                                            t_uint cols = 1) {
  return details::linear_transform<T, wavelets::Wavelet>(wavelet, rows, cols);
}
template <typename T>
LinearTransform<Vector<T>> linear_transform(wavelets::SARA const &sara, t_uint rows,
                                            t_uint cols = 1) {
  return details::linear_transform<T, wavelets::SARA>(sara, rows, cols, sara.size());
}
#ifdef SOPT_MPI
template <typename T>
LinearTransform<Vector<T>> linear_transform(wavelets::SARA const &sara, t_uint rows, t_uint cols,
                                            sopt::mpi::Communicator const &comm) {
  auto const factor = sara.size();
  auto const normalization = std::sqrt(sara.size()) / std::sqrt(comm.all_sum_all(sara.size()));
  return LinearTransform<Vector<T>>(
      [sara, rows, cols, factor, comm, normalization](Vector<T> &out, Vector<T> const &x) {
        assert(static_cast<t_uint>(x.size()) == rows * cols * factor);
        out.resize(rows * cols);
        if (sara.empty())
          out.fill(0);
        else {
          auto signal = Image<T>::Map(out.data(), rows, cols);
          auto const coeffs = Image<T>::Map(x.data(), rows, cols * factor);
          sara.indirect(coeffs, signal);
          out *= normalization;
        }
        comm.all_sum_all(out);
      },
      {{0, 1, static_cast<t_int>(rows * cols)}},
      [sara, rows, cols, factor, normalization](Vector<T> &out, Vector<T> const &x) {
        assert(static_cast<t_uint>(x.size()) == rows * cols);
        out.resize(rows * cols * factor);
        auto const signal = Image<T>::Map(x.data(), rows, cols);
        auto coeffs = Image<T>::Map(out.data(), rows, cols * factor);
        sara.direct(coeffs, signal);
        out *= normalization;
      },
      {{0, 1, static_cast<t_int>(factor * rows * cols)}});
}
#endif
 
}  // namespace sopt
 
#endif