primal_dual.h

Go to the documentation of this file.

#ifndef SOPT_PRIMAL_DUAL_H
#define SOPT_PRIMAL_DUAL_H
 
#include "sopt/config.h"
#include <functional>
#include <limits>
#include <tuple> // for std::tuple<>
#include <utility> // for std::move<>
#include "sopt/exception.h"
#include "sopt/linear_transform.h"
#include "sopt/logging.h"
#include "sopt/types.h"
 
#ifdef SOPT_MPI
#include "sopt/mpi/communicator.h"
#include "sopt/mpi/utilities.h"
#endif
 
namespace sopt::algorithm {
 
template <typename SCALAR>
class PrimalDual {
 public:
  using value_type = SCALAR;
  using Scalar = value_type;
  using Real = typename real_type<Scalar>::type;
  using t_Vector = Vector<Scalar>;
  using t_LinearTransform = LinearTransform<t_Vector>;
  using t_IsConverged = std::function<bool (const t_Vector &, const t_Vector &)>;
  using t_Constraint = std::function<void (t_Vector &, const t_Vector &)>;
  using t_Random_Updater = std::function<bool ()>;
  using t_Proximal = ProximalFunction<Scalar>;
 
  struct Diagnostic {
    t_uint niters;
    bool good;
    t_Vector residual;
 
    Diagnostic(t_uint niters = 0u, bool good = false)
        : niters(niters), good(good), residual(t_Vector::Zero(0)) {}
    Diagnostic(t_uint niters, bool good, t_Vector &&residual)
        : niters(niters), good(good), residual(std::move(residual)) {}
  };
  struct DiagnosticAndResult : public Diagnostic {
    t_Vector x;
  };
 
  template <typename DERIVED>
  PrimalDual(t_Proximal const &f_proximal, t_Proximal const &g_proximal,
             Eigen::MatrixBase<DERIVED> const &target)
      : itermax_(std::numeric_limits<t_uint>::max()),
        sigma_(1),
        tau_(0.5),
        regulariser_strength_(0.5),
        update_scale_(1),
        xi_(1),
        rho_(1),
        is_converged_(),
        constraint_([](t_Vector &out, t_Vector const &x) { out = x; }),
        Phi_(linear_transform_identity<Scalar>()),
        Psi_(linear_transform_identity<Scalar>()),
        f_proximal_(f_proximal),
        g_proximal_(g_proximal),
        random_measurement_updater_([]() { return true; }),
        random_wavelet_updater_([]() { return true; }),
#ifdef SOPT_MPI
        v_all_sum_all_comm_(mpi::Communicator()),
        u_all_sum_all_comm_(mpi::Communicator()),
#endif
        target_(target) {
  }
  virtual ~PrimalDual() {}
 
// Macro helps define properties that can be initialized as in
// auto sdmm  = PrimalDual<float>().prop0(value).prop1(value);
#define SOPT_MACRO(NAME, TYPE)                 \
  TYPE const &NAME() const { return NAME##_; } \
  PrimalDual<SCALAR> &NAME(TYPE const &(NAME)) { \
    NAME##_ = NAME;                            \
    return *this;                              \
  }                                            \
                                               \
 protected:                                    \
  TYPE NAME##_;                                \
                                               \
 public:
 
  SOPT_MACRO(itermax, t_uint);
  SOPT_MACRO(update_scale, Real);
  SOPT_MACRO(regulariser_strength, Real);
  SOPT_MACRO(sigma, Real);
  SOPT_MACRO(xi, Real);
  SOPT_MACRO(rho, Real);
  SOPT_MACRO(tau, Real);
  SOPT_MACRO(is_converged, t_IsConverged);
  SOPT_MACRO(constraint, t_Constraint);
  SOPT_MACRO(Phi, t_LinearTransform);
  SOPT_MACRO(Psi, t_LinearTransform);
  SOPT_MACRO(f_proximal, t_Proximal);
  SOPT_MACRO(g_proximal, t_Proximal);
  SOPT_MACRO(random_measurement_updater, t_Random_Updater);
  SOPT_MACRO(random_wavelet_updater, t_Random_Updater);
#ifdef SOPT_MPI
  SOPT_MACRO(v_all_sum_all_comm, mpi::Communicator);
  SOPT_MACRO(u_all_sum_all_comm, mpi::Communicator);
#endif
#undef SOPT_MACRO
  void f_proximal(t_Vector &out, Real regulariser_strength, t_Vector const &x) const {
    f_proximal()(out, regulariser_strength, x);
  }
  void g_proximal(t_Vector &out, Real regulariser_strength, t_Vector const &x) const {
    g_proximal()(out, regulariser_strength, x);
  }
 
  PrimalDual<Scalar> &is_converged(std::function<bool(t_Vector const &x)> const &func) {
    return is_converged([func](t_Vector const &x, t_Vector const &) { return func(x); });
  }
 
  t_Vector const &target() const { return target_; }
  template <typename DERIVED>
  PrimalDual<Scalar> &target(Eigen::MatrixBase<DERIVED> const &target) {
    target_ = target;
    return *this;
  }
 
  bool is_converged(t_Vector const &x, t_Vector const &residual) const {
    return static_cast<bool>(is_converged()) and is_converged()(x, residual);
  }
 
  Diagnostic operator()(t_Vector &out) const { return operator()(out, initial_guess()); }
  Diagnostic operator()(t_Vector &out, std::tuple<t_Vector, t_Vector> const &guess) const {
    return operator()(out, std::get<0>(guess), std::get<1>(guess));
  }
  Diagnostic operator()(t_Vector &out,
                        std::tuple<t_Vector const &, t_Vector const &> const &guess) const {
    return operator()(out, std::get<0>(guess), std::get<1>(guess));
  }
  DiagnosticAndResult operator()(std::tuple<t_Vector, t_Vector> const &guess) const {
    return operator()(std::tie(std::get<0>(guess), std::get<1>(guess)));
  }
  DiagnosticAndResult operator()(
      std::tuple<t_Vector const &, t_Vector const &> const &guess) const {
    DiagnosticAndResult result;
    static_cast<Diagnostic &>(result) = operator()(result.x, guess);
    return result;
  }
  DiagnosticAndResult operator()() const {
    DiagnosticAndResult result;
    static_cast<Diagnostic &>(result) = operator()(result.x, initial_guess());
    return result;
  }
  DiagnosticAndResult operator()(DiagnosticAndResult const &warmstart) const {
    DiagnosticAndResult result = warmstart;
    static_cast<Diagnostic &>(result) = operator()(result.x, warmstart.x, warmstart.residual);
    return result;
  }
  template <typename... ARGS>
  typename std::enable_if<sizeof...(ARGS) >= 1, PrimalDual &>::type Phi(ARGS &&... args) {
    Phi_ = linear_transform(std::forward<ARGS>(args)...);
    return *this;
  }
  template <typename... ARGS>
  typename std::enable_if<sizeof...(ARGS) >= 1, PrimalDual &>::type Psi(ARGS &&... args) {
    Psi_ = linear_transform(std::forward<ARGS>(args)...);
    return *this;
  }
 
  std::tuple<t_Vector, t_Vector> initial_guess() const {
    return PrimalDual<SCALAR>::initial_guess(target(), Phi());
  }
 
  static std::tuple<t_Vector, t_Vector> initial_guess(t_Vector const &target,
                                                      t_LinearTransform const &phi) {
    std::tuple<t_Vector, t_Vector> guess;
    std::get<0>(guess) = static_cast<t_Vector>(phi.adjoint() * target) / phi.sq_norm();
    std::get<1>(guess) = target;
    return guess;
  }
 
 protected:
  void iteration_step(t_Vector &out, t_Vector &out_hold, t_Vector &u, t_Vector &u_hold, t_Vector &v,
                      t_Vector &v_hold, t_Vector &residual, t_Vector &q, t_Vector &r,
                      bool &random_measurement_update, bool &random_wavelet_update,
                      t_Vector &u_update, t_Vector &v_update) const;
 
  void sanity_check(t_Vector const &x_guess, t_Vector const &res_guess) const {
    if ((Phi().adjoint() * target()).size() != x_guess.size())
      SOPT_THROW("target, adjoint measurement operator and input vector have inconsistent sizes");
    if (target().size() != res_guess.size())
      SOPT_THROW("target and residual vector have inconsistent sizes");
    if ((Phi() * x_guess).size() != target().size())
      SOPT_THROW("target, measurement operator and input vector have inconsistent sizes");
    if (not static_cast<bool>(is_converged()))
      SOPT_WARN("No convergence function was provided: algorithm will run for {} steps", itermax());
  }
 
  Diagnostic operator()(t_Vector &out, t_Vector const &guess, t_Vector const &res) const;
 
  t_Vector target_;
};
 
template <typename SCALAR>
void PrimalDual<SCALAR>::iteration_step(t_Vector &out, t_Vector &out_hold, t_Vector &u,
                                        t_Vector &u_hold, t_Vector &v, t_Vector &v_hold,
                                        t_Vector &residual, t_Vector &q, t_Vector &r,
                                        bool &random_measurement_update,
                                        bool &random_wavelet_update, t_Vector &u_update,
                                        t_Vector &v_update) const {
  // dual calculations for measurements
  if (random_measurement_update) {
    g_proximal(v_hold, rho(), v + residual);
    v_hold = v + residual - v_hold;
    v = v + update_scale() * (v_hold - v);
    v_update = static_cast<t_Vector>(Phi().adjoint() * v);
  }
  // dual calculations for wavelet
  if (random_wavelet_update) {
    q = static_cast<t_Vector>(Psi().adjoint() * out_hold) * sigma();
    f_proximal(u_hold, regulariser_strength(), (u + q));
    u_hold = u + q - u_hold;
    u = u + update_scale() * (u_hold - u);
    u_update = static_cast<t_Vector>(Psi() * u);
  }
  // primal calculations
  r = out;
#ifdef SOPT_MPI
  if (v_all_sum_all_comm().size() > 0 and u_all_sum_all_comm().size() > 0)
    constraint()(
        out_hold,
        r - tau() * (u_all_sum_all_comm().all_sum_all(static_cast<const t_Vector>(u_update)) +
                     v_all_sum_all_comm().all_sum_all(static_cast<const t_Vector>(v_update))));
  else
#endif
    constraint()(out_hold, r - tau() * (u_update + v_update));
  out = r + update_scale() * (out_hold - r);
  out_hold = 2 * out_hold - r;
  random_measurement_update = random_measurement_updater_();
  random_wavelet_update = random_wavelet_updater_();
  // update residual
  if (random_measurement_update)
    residual = static_cast<t_Vector>(Phi() * out_hold) * xi() - target();
}
 
template <typename SCALAR>
typename PrimalDual<SCALAR>::Diagnostic PrimalDual<SCALAR>::operator()(
    t_Vector &out, t_Vector const &x_guess, t_Vector const &res_guess) const {
  SOPT_HIGH_LOG("Performing Primal Dual");
  sanity_check(x_guess, res_guess);
  bool random_measurement_update = random_measurement_updater_();
  bool random_wavelet_update = random_wavelet_updater_();
  t_Vector residual = res_guess;
  out = x_guess;
  t_Vector out_hold = x_guess;
  t_Vector r = x_guess;
  t_Vector v = residual;
  t_Vector v_hold = residual;
  t_Vector v_update = x_guess;
  t_Vector u = Psi().adjoint() * out;
  t_Vector u_hold = u;
  t_Vector u_update = out;
  t_Vector q = u;
 
  t_uint niters(0);
  bool converged = false;
  for (; (not converged) && (niters < itermax()); ++niters) {
    SOPT_LOW_LOG("    - [Primal Dual] Iteration {}/{}", niters, itermax());
    iteration_step(out, out_hold, u, u_hold, v, v_hold, residual, q, r, random_measurement_update,
                   random_wavelet_update, u_update, v_update);
    SOPT_LOW_LOG("      - [Primal Dual] Sum of residuals: {}",
                 static_cast<t_Vector>(residual).array().abs().sum());
    converged = is_converged(out, residual);
  }
 
  if (converged) {
    SOPT_MEDIUM_LOG("    - [Primal Dual] converged in {} of {} iterations", niters, itermax());
  } else if (static_cast<bool>(is_converged())) {
    // not meaningful if not convergence function
    SOPT_ERROR("    - [Primal Dual] did not converge within {} iterations", itermax());
  }
  return {niters, converged, std::move(residual)};
}
} // namespace sopt::algorithm
#endif