forward_backward.h

Go to the documentation of this file.

#ifndef SOPT_FORWARD_BACKWARD_H
#define SOPT_FORWARD_BACKWARD_H
 
#include "sopt/config.h"
#include <functional>
#include <limits>
#include <tuple> // for tuple<>
#include <utility> // for std::move<>
#include "sopt/exception.h"
#include "sopt/linear_transform.h"
#include "sopt/logging.h"
#include "sopt/types.h"
 
#include "sopt/gradient_utils.h"
 
namespace sopt::algorithm {
 
template <typename SCALAR>
class ForwardBackward {
 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_Proximal = ProximalFunction<Scalar>;
  // The first argument is the output vector, the others are inputs
  using t_Gradient = std::function<void(t_Vector &gradient, const t_Vector &image, const t_Vector &residual, const t_LinearTransform& Phi)>;
  using t_randomUpdater = std::function<std::shared_ptr<IterationState<t_Vector>>()>;
 
  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;
  };
 
  ForwardBackward(t_Gradient const &f_gradient, t_Proximal const &g_proximal,
                  t_Vector const &target)
      : itermax_(std::numeric_limits<t_uint>::max()),
        regulariser_strength_(1e-8),
        step_size_(1),
        is_converged_(),
          fista_(true),
        f_gradient_(f_gradient),
        g_proximal_(g_proximal)
        {
          std::shared_ptr<t_LinearTransform> Id = std::make_shared<t_LinearTransform>(linear_transform_identity<Scalar>());
          problem_state = std::make_shared<IterationState<t_Vector>>(target, Id);
        }
  virtual ~ForwardBackward() {}
 
// Macro helps define properties that can be initialized as in
// auto sdmm  = ForwardBackward<float>().prop0(value).prop1(value);
#define SOPT_MACRO(NAME, TYPE)                      \
  TYPE const &NAME() const { return NAME##_; }      \
  ForwardBackward<SCALAR> &NAME(TYPE const &(NAME)) { \
    NAME##_ = NAME;                                 \
    return *this;                                   \
  }                                                 \
                                                    \
 protected:                                         \
  TYPE NAME##_;                                     \
                                                    \
 public:
 
  SOPT_MACRO(itermax, t_uint);
  SOPT_MACRO(regulariser_strength, Real);
  SOPT_MACRO(step_size, Real);
  SOPT_MACRO(fista, bool);
  SOPT_MACRO(is_converged, t_IsConverged);
  SOPT_MACRO(f_gradient, t_Gradient);
  SOPT_MACRO(g_proximal, t_Proximal);
 
  t_LinearTransform const &Phi() const { return problem_state->Phi(); }
  ForwardBackward<SCALAR> &Phi(t_LinearTransform const &new_phi) {
    problem_state->Phi(new_phi);
    return *this;
  }
 
  ForwardBackward<SCALAR> &random_updater(t_randomUpdater &rU)
  {
    random_updater_ = rU;
    return *this;
  }
 
  ForwardBackward<SCALAR> &set_problem_state(std::shared_ptr<IterationState<t_Vector>> pS)
  {
    problem_state = pS;
    return *this;
  }
 
#undef SOPT_MACRO
  void f_gradient(t_Vector &out, t_Vector const &x, t_Vector const &res, t_LinearTransform const &Phi) const { f_gradient()(out, x, res, Phi); }
  void g_proximal(t_Vector &out, Real regulariser_strength, t_Vector const &x) const {
    g_proximal()(out, regulariser_strength, x);
  }
 
  ForwardBackward<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 problem_state->target(); }
  ForwardBackward<Scalar> &target(t_Vector const &target) {
    problem_state->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) { return operator()(out, initial_guess()); }
  Diagnostic operator()(t_Vector &out, std::tuple<t_Vector, t_Vector> const &guess) {
    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) {
    return operator()(out, std::get<0>(guess), std::get<1>(guess));
  }
  DiagnosticAndResult operator()(std::tuple<t_Vector, t_Vector> const &guess) {
    return operator()(std::tie(std::get<0>(guess), std::get<1>(guess)));
  }
  DiagnosticAndResult operator()(
      std::tuple<t_Vector const &, t_Vector const &> const &guess) {
    DiagnosticAndResult result;
    static_cast<Diagnostic &>(result) = operator()(result.x, guess);
    return result;
  }
  DiagnosticAndResult operator()() {
    DiagnosticAndResult result;
    static_cast<Diagnostic &>(result) = operator()(result.x, initial_guess());
    return result;
  }
  DiagnosticAndResult operator()(DiagnosticAndResult const &warmstart) {
    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, ForwardBackward &>::type Phi(ARGS &&... args) {
    problem_state->Phi(linear_transform(std::forward<ARGS>(args)...));
    return *this;
  }
 
  std::tuple<t_Vector, t_Vector> initial_guess() const {
    return ForwardBackward<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) = phi * std::get<0>(guess) - target;
    return guess;
  }
 
 protected:
  void iteration_step(t_Vector &out, t_Vector &residual, t_Vector &p, t_Vector &z,
                      const t_real lambda);
 
  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);
 
  std::shared_ptr<IterationState<t_Vector>> problem_state;
  t_randomUpdater random_updater_;
};
 
template <typename SCALAR>
void ForwardBackward<SCALAR>::iteration_step(t_Vector &image, t_Vector &residual, t_Vector &auxilliary_image,
                                             t_Vector &gradient_current, const t_real FISTA_step) {
  t_Vector prev_image = image;
  f_gradient(gradient_current, auxilliary_image, residual, Phi());  // assigns gradient_current (non normalised)
  t_Vector auxilliary_with_step = auxilliary_image - step_size() / Phi().sq_norm() * gradient_current;  // step to new image using gradient
  const Real weight = regulariser_strength() * step_size();
  g_proximal(image, weight, auxilliary_with_step);  // apply proximal operator to new image
  auxilliary_image = image + FISTA_step * (image - prev_image);  // update auxilliary vector with FISTA acceleration step  
  
  // set up next iteration
  if(random_updater_)
  {
    problem_state = random_updater_();
  }
  residual = (Phi() * auxilliary_image) - target();  // updates the residual for the NEXT iteration (new image).
}
 
template <typename SCALAR>
typename ForwardBackward<SCALAR>::Diagnostic ForwardBackward<SCALAR>::operator()(
    t_Vector &out, t_Vector const &x_guess, t_Vector const &res_guess) {
  SOPT_HIGH_LOG("Performing Forward Backward Splitting");
  if (fista()) {
    SOPT_HIGH_LOG("Using FISTA algorithm");
  } else {
    SOPT_HIGH_LOG("Using standard FB algorithm");
  }
  sanity_check(x_guess, res_guess);
 
  const size_t image_size = x_guess.size();
 
  t_Vector auxilliary_image = x_guess;
  t_Vector residual = res_guess;
  t_Vector gradient_current = t_Vector::Zero(image_size);
  out = x_guess;
 
  t_uint niters(0);
  bool converged = false;
  Real theta = 1.0;
  Real theta_new = 1.0;
  Real FISTA_step = 0.0;
  for (; (not converged) && (niters < itermax()); ++niters) {
    SOPT_MEDIUM_LOG("    - [FB] Iteration {}/{}", niters, itermax());
    if (fista()) {
      theta_new = (1 + std::sqrt(1 + 4 * theta * theta)) / 2.;
      FISTA_step = (theta - 1) / (theta_new);
      theta = theta_new;
    }
    SOPT_LOW_LOG("      - Call iteration step");
    iteration_step(out, residual, auxilliary_image, gradient_current, FISTA_step);
    SOPT_LOW_LOG("      - [FB] Sum of residuals: {}", residual.array().abs().sum());
    converged = is_converged(out, residual);
  }
 
  if (converged) {
    SOPT_MEDIUM_LOG("    - [FB] converged in {} of {} iterations", niters, itermax());
  } else if (static_cast<bool>(is_converged())) {
    // not meaningful if not convergence function
    SOPT_ERROR("    - [FB] did not converge within {} iterations", itermax());
  }
  return {niters, converged, std::move(residual)};
}
} // namespace sopt::algorithm
#endif