l2_forward_backward.h

Go to the documentation of this file.

#ifndef SOPT_L2_FORWARD_BACKWARD_H
#define SOPT_L2_FORWARD_BACKWARD_H
 
#include "sopt/config.h"
#include <limits> // for std::numeric_limits<>
#include <numeric>
#include <tuple>
#include <utility>
#include "sopt/exception.h"
#include "sopt/forward_backward.h"
#include "sopt/linear_transform.h"
#include "sopt/logging.h"
#include "sopt/proximal.h"
#include "sopt/relative_variation.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 L2ForwardBackward {
  using FB = ForwardBackward<SCALAR>;
 
 public:
  using value_type = typename FB::value_type;
  using Scalar = typename FB::Scalar;
  using Real = typename FB::Real;
  using t_Vector = typename FB::t_Vector;
  using t_LinearTransform = typename FB::t_LinearTransform;
  template <typename T>
  using t_Proximal = std::function<void(t_Vector &, const T &, const t_Vector &)>;
  using t_Gradient = typename FB::t_Gradient;
  using t_IsConverged = typename FB::t_IsConverged;
 
  struct Diagnostic : public FB::Diagnostic {
    Diagnostic(t_uint niters = 0u, bool good = false) : FB::Diagnostic(niters, good) {}
    Diagnostic(t_uint niters, bool good, t_Vector &&residual)
        : FB::Diagnostic(niters, good, std::move(residual)) {}
  };
  struct DiagnosticAndResult : public Diagnostic {
    t_Vector x;
  };
 
  template <typename DERIVED>
  L2ForwardBackward(Eigen::MatrixBase<DERIVED> const &target)
      : l2_proximal_([](t_Vector &output, const t_real &regulariser_strength, const t_Vector &x) -> void {
          proximal::l2_norm(output, regulariser_strength, x);
        }),
        l2_proximal_weighted_(
            [](t_Vector &output, const Vector<Real> &regulariser_strength, const t_Vector &x) -> void {
              proximal::l2_norm(output, regulariser_strength, x);
            }),
        l2_proximal_weights_(Vector<Real>::Ones(1)),
        l2_gradient_([](t_Vector &output, t_Vector const &image, const t_Vector &residual, const t_LinearTransform &Phi) -> void {
          output = Phi.adjoint()*residual;
        }),  // gradient of 1/2 * r^2 = r;
        tight_frame_(false),
        residual_tolerance_(0.),
        relative_variation_(1e-4),
        residual_convergence_(nullptr),
        objective_convergence_(nullptr),
        itermax_(std::numeric_limits<t_uint>::max()),
        regulariser_strength_(1e-8),
        step_size_(1),
        sigma_(1),
        is_converged_(),
        Phi_(linear_transform_identity<Scalar>()),
        target_(target) {}
  virtual ~L2ForwardBackward() {}
 
// Macro helps define properties that can be initialized as in
// auto padmm = L2ForwardBackward<float>().prop0(value).prop1(value);
#define SOPT_MACRO(NAME, TYPE)                             \
  TYPE const &NAME() const { return NAME##_; }             \
  L2ForwardBackward<SCALAR> &NAME(TYPE const &(NAME)) { \
    NAME##_ = NAME;                                        \
    return *this;                                          \
  }                                                        \
                                                           \
 protected:                                                \
  TYPE NAME##_;                                            \
                                                           \
 public:
 
  SOPT_MACRO(l2_proximal, t_Proximal<Real>);
  SOPT_MACRO(l2_proximal_weighted, t_Proximal<Vector<Real>>);
  SOPT_MACRO(l2_proximal_weights, Vector<Real>);
  SOPT_MACRO(l2_gradient, t_Gradient);
  SOPT_MACRO(tight_frame, bool);
  SOPT_MACRO(residual_tolerance, Real);
  SOPT_MACRO(relative_variation, Real);
  SOPT_MACRO(residual_convergence, t_IsConverged);
  SOPT_MACRO(objective_convergence, t_IsConverged);
  SOPT_MACRO(itermax, t_uint);
  SOPT_MACRO(regulariser_strength, Real);
  SOPT_MACRO(step_size, Real);
  SOPT_MACRO(sigma, Real);
  SOPT_MACRO(is_converged, t_IsConverged);
  SOPT_MACRO(Phi, t_LinearTransform);
#ifdef SOPT_MPI
  SOPT_MACRO(obj_comm, mpi::Communicator);
#endif
 
#undef SOPT_MACRO
  t_Vector const &target() const { return target_; }
  Real objmin() const { return objmin_; }
  template <typename DERIVED>
  L2ForwardBackward<Scalar> &target(Eigen::MatrixBase<DERIVED> const &target) {
    target_ = target;
    return *this;
  }
 
  Diagnostic operator()(t_Vector &out) const {
    return operator()(out, ForwardBackward<SCALAR>::initial_guess(target(), Phi()));
  }
  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, ForwardBackward<SCALAR>::initial_guess(target(), Phi()));
    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, L2ForwardBackward &>::type Phi(
      ARGS &&... args) {
    Phi_ = linear_transform(std::forward<ARGS>(args)...);
    return *this;
  }
 
  t_Proximal<Real> &l2_proximal() { return l2_proximal_; }
  t_Proximal<Vector<Real>> &l2_proximal_weighted() { return l2_proximal_weighted_; }
  t_Gradient &l2_gradient() { return l2_gradient_; }
 
  L2ForwardBackward<Scalar> &residual_convergence(Real const &tolerance) {
    return residual_convergence(nullptr).residual_tolerance(tolerance);
  }
  L2ForwardBackward<Scalar> &objective_convergence(Real const &tolerance) {
    return objective_convergence(nullptr).relative_variation(tolerance);
  }
  L2ForwardBackward<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); });
  }
 
 protected:
  t_Vector target_;
  mutable Real objmin_;
 
  Diagnostic operator()(t_Vector &out, t_Vector const &guess, t_Vector const &res) const;
 
  bool residual_convergence(t_Vector const &x, t_Vector const &residual) const;
 
  bool objective_convergence(ScalarRelativeVariation<Scalar> &scalvar, t_Vector const &x,
                             t_Vector const &residual) const;
#ifdef SOPT_MPI
  bool objective_convergence(mpi::Communicator const &obj_comm,
                             ScalarRelativeVariation<Scalar> &scalvar, t_Vector const &x,
                             t_Vector const &residual) const;
#endif
 
  bool is_converged(ScalarRelativeVariation<Scalar> &scalvar, t_Vector const &x,
                    t_Vector const &residual) const;
};
 
template <typename SCALAR>
typename L2ForwardBackward<SCALAR>::Diagnostic L2ForwardBackward<SCALAR>::operator()(
    t_Vector &out, t_Vector const &guess, t_Vector const &res) const {
  SOPT_HIGH_LOG("Performing Forward Backward with L2 and L2 norms");
  // The f proximal is an L2 proximal
  Diagnostic result;
  auto const g_proximal = [this](t_Vector &out, Real regulariser_strength, t_Vector const &x) {
    if (this->l2_proximal_weights().size() > 1)
      this->l2_proximal_weighted()(out, this->l2_proximal_weights() * regulariser_strength, x);
    else
      this->l2_proximal()(out, this->l2_proximal_weights()(0) * regulariser_strength, x);
  };
  const Real sigma_factor = sigma() * sigma();
  const t_Gradient f_gradient = [sigma_factor](t_Vector &out, t_Vector const &image, t_Vector const &res, t_LinearTransform const &Phi) {
    t_Vector temp;
    temp = res / sigma_factor;
    out = Phi.adjoint() * temp;
  };
  ScalarRelativeVariation<Scalar> scalvar(relative_variation(), relative_variation(),
                                          "Objective function");
  auto const convergence = [this, &scalvar](t_Vector const &x, t_Vector const &residual) mutable {
    const bool result = this->is_converged(scalvar, x, residual);
    this->objmin_ = std::real(scalvar.previous());
    return result;
  };
  auto fb = ForwardBackward<SCALAR>(f_gradient, g_proximal, target())
                      .itermax(itermax())
                      .step_size(step_size())
                      .regulariser_strength(regulariser_strength())
                      .Phi(Phi())
                      .is_converged(convergence);
  static_cast<typename ForwardBackward<SCALAR>::Diagnostic &>(result) =
      fb(out, std::tie(guess, res));
  return result;
}
 
template <typename SCALAR>
bool L2ForwardBackward<SCALAR>::residual_convergence(t_Vector const &x,
                                                          t_Vector const &residual) const {
  if (static_cast<bool>(residual_convergence())) return residual_convergence()(x, residual);
  if (residual_tolerance() <= 0e0) return true;
  auto const residual_norm = sopt::l2_norm(residual);
  SOPT_LOW_LOG("    - [FB] Residuals: {} <? {}", residual_norm, residual_tolerance());
  return residual_norm < residual_tolerance();
}
 
template <typename SCALAR>
bool L2ForwardBackward<SCALAR>::objective_convergence(ScalarRelativeVariation<Scalar> &scalvar,
                                                           t_Vector const &x,
                                                           t_Vector const &residual) const {
  if (static_cast<bool>(objective_convergence())) return objective_convergence()(x, residual);
  if (scalvar.relative_tolerance() <= 0e0) return true;
  auto const current = ((regulariser_strength() > 0) ? sopt::l2_norm(x, l2_proximal_weights()) * regulariser_strength() : 0) +
                       std::pow(sopt::l2_norm(residual), 2) / (2 * sigma() * sigma());
  return scalvar(current);
}
 
#ifdef SOPT_MPI
template <typename SCALAR>
bool L2ForwardBackward<SCALAR>::objective_convergence(mpi::Communicator const &obj_comm,
                                                           ScalarRelativeVariation<Scalar> &scalvar,
                                                           t_Vector const &x,
                                                           t_Vector const &residual) const {
  if (static_cast<bool>(objective_convergence())) return objective_convergence()(x, residual);
  if (scalvar.relative_tolerance() <= 0e0) return true;
  auto const current = obj_comm.all_sum_all<t_real>(
      ((regulariser_strength() > 0) ? sopt::l2_norm(x, l2_proximal_weights()) * regulariser_strength() : 0) +
      std::pow(sopt::l2_norm(residual), 2) / (2 * sigma() * sigma()));
  return scalvar(current);
}
#endif
 
template <typename SCALAR>
bool L2ForwardBackward<SCALAR>::is_converged(ScalarRelativeVariation<Scalar> &scalvar,
                                                  t_Vector const &x,
                                                  t_Vector const &residual) const {
  auto const user = static_cast<bool>(is_converged()) == false or is_converged()(x, residual);
  auto const res = residual_convergence(x, residual);
#ifdef SOPT_MPI
  auto const obj = objective_convergence(obj_comm(), scalvar, x, residual);
#else
  auto const obj = objective_convergence(scalvar, x, residual);
#endif
  // beware of short-circuiting!
  // better evaluate each convergence function everytime, especially with mpi
  return user and res and obj;
}
} // namespace sopt::algorithm
#endif