sdmm.h

Go to the documentation of this file.

#ifndef SOPT_SDMM_H
#define SOPT_SDMM_H
 
#include "sopt/config.h"
#include <limits>
#include <numeric>
#include <utility> // for std::forward<>
#include <vector>
#include "sopt/conjugate_gradient.h"
#include "sopt/exception.h"
#include "sopt/linear_transform.h"
#include "sopt/logging.h"
#include "sopt/proximal.h"
#include "sopt/proximal_expression.h"
#include "sopt/types.h"
#include "sopt/wrapper.h"
 
namespace sopt::algorithm {
 
template <typename SCALAR>
class SDMM {
 public:
  struct Diagnostic {
    t_uint niters;
    bool good;
    ConjugateGradient::Diagnostic cg_diagnostic;
  };
  struct DiagnosticAndResult : public Diagnostic {
    Vector<SCALAR> x;
  };
  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_Proximal = ProximalFunction<SCALAR>;
  using t_IsConverged = ConvergenceFunction<SCALAR>;
 
  SDMM()
      : itermax_(std::numeric_limits<t_uint>::max()),
        gamma_(1e-8),
        conjugate_gradient_(std::numeric_limits<t_uint>::max(), 1e-6),
        is_converged_([](t_Vector const &) { return false; }) {}
  virtual ~SDMM() {}
 
// Macro helps define properties that can be initialized as in
// auto sdmm  = SDMM<float>().prop0(value).prop1(value);
#define SOPT_MACRO(NAME, TYPE)                 \
  TYPE const &NAME() const { return NAME##_; } \
  SDMM<SCALAR> &NAME(TYPE const &(NAME)) {       \
    NAME##_ = NAME;                            \
    return *this;                              \
  }                                            \
                                               \
 protected:                                    \
  TYPE NAME##_;                                \
                                               \
 public:
  SOPT_MACRO(itermax, t_uint);
  SOPT_MACRO(gamma, Real);
  SOPT_MACRO(conjugate_gradient, ConjugateGradient);
  SOPT_MACRO(is_converged, t_IsConverged);
#undef SOPT_MACRO
  SDMM<SCALAR> &conjugate_gradient(t_uint itermax, t_real tolerance) {
    conjugate_gradient_.itermax(itermax);
    conjugate_gradient_.tolerance(tolerance);
    return *this;
  }
  template <typename PROXIMAL, typename T>
  SDMM<SCALAR> &append(PROXIMAL proximal, T args) {
    proximals().emplace_back(proximal);
    transforms().emplace_back(linear_transform(args));
    return *this;
  }
  template <typename PROXIMAL>
  SDMM<SCALAR> &append(PROXIMAL proximal) {
    return append(proximal, linear_transform_identity<Scalar>());
  }
  template <typename PROXIMAL, typename L, typename LADJOINT>
  SDMM<SCALAR> &append(PROXIMAL proximal, L l, LADJOINT ladjoint) {
    return append(proximal, linear_transform<t_Vector>(l, ladjoint));
  }
  template <typename PROXIMAL, typename L, typename LADJOINT>
  SDMM<SCALAR> &append(PROXIMAL proximal, L l, LADJOINT ladjoint, std::array<t_int, 3> sizes) {
    return append(proximal, linear_transform<t_Vector>(l, ladjoint, sizes));
  }
  template <typename PROXIMAL, typename L, typename LADJOINT>
  SDMM<SCALAR> &append(PROXIMAL proximal, L l, std::array<t_int, 3> dsizes, LADJOINT ladjoint,
                       std::array<t_int, 3> isizes) {
    return append(proximal, linear_transform<t_Vector>(l, dsizes, ladjoint, isizes));
  }
 
  Diagnostic operator()(t_Vector &out, t_Vector const &input) const;
  DiagnosticAndResult operator()(t_Vector const &input) const {
    DiagnosticAndResult result;
    static_cast<Diagnostic &>(result) = operator()(result.x, input);
    return result;
  }
  DiagnosticAndResult operator()(DiagnosticAndResult const &warmstart) const {
    DiagnosticAndResult result;
    static_cast<Diagnostic &>(result) = operator()(result.x, warmstart.x);
    return result;
  }
 
  std::vector<t_LinearTransform> const &transforms() const { return transforms_; }
  std::vector<t_LinearTransform> &transforms() { return transforms_; }
  t_LinearTransform const &transforms(t_uint i) const { return transforms_[i]; }
  t_LinearTransform &transforms(t_uint i) { return transforms_[i]; }
 
  std::vector<t_Proximal> const &proximals() const { return proximals_; }
  std::vector<t_Proximal> &proximals() { return proximals_; }
  t_Proximal const &proximals(t_uint i) const { return proximals_[i]; }
  t_Proximal &proximals(t_uint i) { return proximals_[i]; }
  template <typename T0>
  proximal::ProximalExpression<t_Proximal const &, T0> proximals(
      t_uint i, Eigen::MatrixBase<T0> const &x) const {
    return {proximals()[i], gamma(), x};
  }
 
  t_uint size() const { return proximals().size(); }
 
  // We must declare the first argument explicitly so that the function never
  // match the getter with the same name.
  template <typename T0, typename... T>
  auto conjugate_gradient(T0 &&t0, T &&... args) const
      -> decltype(this->conjugate_gradient()(std::forward<T0>(t0), std::forward<T>(args)...)) {
    return conjugate_gradient()(std::forward<T0>(t0), std::forward<T>(args)...);
  }
 
  bool is_converged(t_Vector const &x) const { return is_converged()(x); }
 
 protected:
  std::vector<t_LinearTransform> transforms_;
  std::vector<t_Proximal> proximals_;
 
  using t_Vectors = std::vector<t_Vector>;
  virtual ConjugateGradient::Diagnostic solve_for_xn(t_Vector &out, t_Vectors const &y,
                                                     t_Vectors const &z) const;
  virtual void update_directions(t_Vectors &y, t_Vectors &z, t_Vector const &x) const;
 
  virtual void initialization(t_Vectors &y, t_Vectors &z, t_Vector const &x) const;
 
  virtual void sanity_check(t_Vector const &input) const;
};
 
template <typename SCALAR>
typename SDMM<SCALAR>::Diagnostic SDMM<SCALAR>::operator()(t_Vector &out,
                                                           t_Vector const &input) const {
  sanity_check(input);
  bool convergence = false;
  t_uint niters(0);
  // Figures out where itermax or convergence reached
  auto const has_finished = [&convergence, &niters, this](t_Vector const &out) {
    convergence = is_converged(out);
    return niters >= itermax() or convergence;
  };
 
  SOPT_HIGH_LOG("Performing SDMM ");
  out = input;
  t_Vectors y(transforms().size());
  t_Vectors z(transforms().size());
 
  // Initial step replaces iteration update with initialization
  initialization(y, z, input);
  auto cg_diagnostic = solve_for_xn(out, y, z);
 
  while (not has_finished(out)) {
    SOPT_LOW_LOG("Iteration {}/{}. ", niters, itermax());
    // computes y and z from out and transforms
    update_directions(y, z, out);
    SOPT_LOW_LOG("   - sum z_ij = {}",
                 std::accumulate(z.begin(), z.end(), Scalar(0e0),
                                 [](Scalar const &a, t_Vector const &z) { return a + z.sum(); }));
    // computes x = L^-1 y
    cg_diagnostic = solve_for_xn(out, y, z);
    SOPT_LOW_LOG("   - CG Residual = {} in {}/{} iterations", cg_diagnostic.residual,
                 cg_diagnostic.niters, conjugate_gradient().itermax());
 
    ++niters;
  }
  return {niters, convergence, cg_diagnostic};
}
 
template <typename SCALAR>
ConjugateGradient::Diagnostic SDMM<SCALAR>::solve_for_xn(t_Vector &out, t_Vectors const &y,
                                                         t_Vectors const &z) const {
  assert(z.size() == transforms().size());
  assert(y.size() == transforms().size());
  SOPT_TRACE("Solving for x_n");
 
  // Initialize b of A x = b = sum_i L_i^H(z_i - y_i)
  t_Vector b = out.Zero(out.size());
  for (t_uint i(0); i < transforms().size(); ++i) b += transforms(i).adjoint() * (y[i] - z[i]);
  if (b.stableNorm() < 1e-12) {
    out.fill(0e0);
    return {0, 0, true};
  }
 
  // Then create operator A
  auto A = [this](t_Vector &out, t_Vector const &input) {
    out = out.Zero(input.size());
    for (auto const &transform : this->transforms())
      out += transform.adjoint() * static_cast<t_Vector>(transform * input);
  };
 
  // Call conjugate gradient
  auto const diagnostic = this->conjugate_gradient(out, A, b);
  if (not diagnostic.good) {
    SOPT_ERROR("CG error - iterations: {}/{} - residuals {}\n", diagnostic.niters,
               conjugate_gradient().itermax(), diagnostic.residual);
    SOPT_THROW("Conjugate gradient failed to converge");
  }
 
  return diagnostic;
}
 
template <typename SCALAR>
void SDMM<SCALAR>::update_directions(t_Vectors &y, t_Vectors &z, t_Vector const &x) const {
  SOPT_TRACE("Updating directions");
  for (t_uint i(0); i < transforms().size(); ++i) {
    z[i] += transforms(i) * x;
    y[i] = proximals(i, z[i]);
    z[i] -= y[i];
  }
}
 
template <typename SCALAR>
void SDMM<SCALAR>::initialization(t_Vectors &y, t_Vectors &z, t_Vector const &x) const {
  SOPT_TRACE("Initializing SDMM");
  for (t_uint i(0); i < transforms().size(); i++) {
    y[i] = transforms(i) * x;
    z[i].resize(y[i].size());
    z[i].fill(0);
    assert(z[i].size() == y[i].size());
    SOPT_TRACE("    - transform {}: {}", i, y[i].transpose());
  }
}
 
template <typename SCALAR>
void SDMM<SCALAR>::sanity_check(t_Vector const &x) const {
  bool doexit = false;
  if (proximals().size() != transforms().size()) {
    SOPT_ERROR("Internal error: number of proximals and transforms do not match");
    doexit = true;
  }
  if (x.size() == 0) SOPT_WARN("Input vector has zero size");
  if (size() == 0) SOPT_WARN("No operators - SDMM is empty");
  for (t_uint i(0); i < size(); ++i) {
    auto const xdual = t_Vector::Zero((transforms(i) * x).size());
    auto const r = (transforms(i).adjoint() * xdual).size();
    if (r != x.size()) {
      SOPT_ERROR("Output size of transform {} and input do not match: {} vs {}", i, r, x.size());
      doexit = true;
    }
  }
  if (doexit) SOPT_THROW("Input to SDMM is inconsistent");
}
} // namespace sopt::algorithm
#endif