l1_proximal.h

Go to the documentation of this file.

#ifndef SOPT_L1_PROXIMAL_H
#define SOPT_L1_PROXIMAL_H
 
#include "sopt/config.h"
#include <array>
#include <type_traits>
#include <utility> // for std::forward<>
#include <Eigen/Core>
#include "sopt/linear_transform.h"
#include "sopt/maths.h"
#include "sopt/proximal_expression.h"
#ifdef SOPT_MPI
#include "sopt/mpi/communicator.h"
#include "sopt/mpi/utilities.h"
#endif
 
namespace sopt::proximal {
 
template <typename SCALAR>
class L1TightFrame {
 public:
  using Scalar = SCALAR;
  using Real = typename real_type<Scalar>::type;
 
#ifdef SOPT_MPI
  L1TightFrame(mpi::Communicator const &direct_comm = mpi::Communicator(),
               mpi::Communicator const &adjoint_comm = mpi::Communicator())
      : Psi_(linear_transform_identity<Scalar>()),
        nu_(1e0),
        direct_space_comm_(direct_comm),
        adjoint_space_comm_(adjoint_comm),
        weights_(Vector<Real>::Ones(1)) {}
#else
  L1TightFrame()
      : Psi_(linear_transform_identity<Scalar>()), nu_(1e0), weights_(Vector<Real>::Ones(1)) {}
#endif
 
#define SOPT_MACRO(NAME, TYPE)                   \
  TYPE const &NAME() const { return NAME##_; }   \
  L1TightFrame<Scalar> &NAME(TYPE const &(NAME)) { \
    NAME##_ = NAME;                              \
    return *this;                                \
  }                                              \
                                                 \
 protected:                                      \
  TYPE NAME##_;                                  \
                                                 \
 public:
  SOPT_MACRO(Psi, LinearTransform<Vector<Scalar>>);
  SOPT_MACRO(nu, Real);
#ifdef SOPT_MPI
  SOPT_MACRO(direct_space_comm, mpi::Communicator);
  SOPT_MACRO(adjoint_space_comm, mpi::Communicator);
#endif
#undef SOPT_MACRO
  Vector<Real> const &weights() const { return weights_; }
  template <typename T>
  L1TightFrame<Scalar> &weights(Eigen::MatrixBase<T> const &w) {
    if ((w.array() < 0e0).any()) SOPT_THROW("Weights cannot be negative");
    if (w.stableNorm() < 1e-12) SOPT_THROW("Weights cannot be null");
    weights_ = w;
    return *this;
  }
 
  L1TightFrame<Scalar> &weights(Real const &value) {
    if (value <= 0e0) SOPT_THROW("Weight cannot be negative or null");
    weights_ = Vector<Real>::Ones(1) * value;
    return *this;
  }
 
  template <typename... ARGS>
  typename std::enable_if<sizeof...(ARGS) >= 1, L1TightFrame &>::type Psi(ARGS &&... args) {
    Psi_ = linear_transform(std::forward<ARGS>(args)...);
    return *this;
  }
 
  template <typename T0, typename T1>
  typename std::enable_if<is_complex<Scalar>::value == is_complex<typename T0::Scalar>::value and
                          is_complex<Scalar>::value == is_complex<typename T1::Scalar>::value>::type
  operator()(Eigen::MatrixBase<T0> &out, Real gamma, Eigen::MatrixBase<T1> const &x) const;
 
  template <typename T0>
  ProximalExpression<L1TightFrame<Scalar> const &, T0> operator()(
      Real const &gamma, Eigen::MatrixBase<T0> const &x) const {
    return {*this, gamma, x};
  }
 
  template <typename T0, typename T1>
  typename std::enable_if<is_complex<Scalar>::value == is_complex<typename T0::Scalar>::value and
                              is_complex<Scalar>::value == is_complex<typename T1::Scalar>::value,
                          Real>::type
  objective(Eigen::MatrixBase<T0> const &x, Eigen::MatrixBase<T1> const &z,
            Real const &gamma) const;
 
 protected:
  Vector<Real> weights_;
};
 
template <typename SCALAR>
template <typename T0, typename T1>
typename std::enable_if<is_complex<SCALAR>::value == is_complex<typename T0::Scalar>::value and
                        is_complex<SCALAR>::value == is_complex<typename T1::Scalar>::value>::type
L1TightFrame<SCALAR>::operator()(Eigen::MatrixBase<T0> &out, Real gamma,
                                 Eigen::MatrixBase<T1> const &x) const {
  Vector<Scalar> const psit_x = Psi().adjoint() * x;
  if (weights().size() == 1)
    out = static_cast<Vector<Scalar>>(
              Psi() * (soft_threshhold(psit_x, nu() * gamma * weights()(0)) - psit_x)) /
              nu() +
          x;
  else
    out = static_cast<Vector<Scalar>>(
              Psi() * (soft_threshhold(psit_x, nu() * gamma * weights()) - psit_x)) /
              nu() +
          x;
  SOPT_LOW_LOG("Prox L1: objective = {}", objective(x, out, gamma));
}
 
template <typename SCALAR>
template <typename T0, typename T1>
typename std::enable_if<is_complex<SCALAR>::value == is_complex<typename T0::Scalar>::value and
                            is_complex<SCALAR>::value == is_complex<typename T1::Scalar>::value,
                        typename real_type<SCALAR>::type>::type
L1TightFrame<SCALAR>::objective(Eigen::MatrixBase<T0> const &x, Eigen::MatrixBase<T1> const &z,
                                Real const &gamma) const {
#ifdef SOPT_MPI
  auto const adj = gamma * sopt::mpi::l1_norm(static_cast<T1>(Psi().adjoint() * z), weights(),
                                              adjoint_space_comm());
  auto const dir = direct_space_comm().all_sum_all(0.5 * (x - z).squaredNorm());
  return adj + dir;
#else
  return 0.5 * (x - z).squaredNorm() +
         gamma * sopt::l1_norm(static_cast<T1>(Psi().adjoint() * z), weights());
#endif
}
 
template <typename SCALAR>
class L1 : protected L1TightFrame<SCALAR> {
 public:
  class FistaMixing;
  class NoMixing;
  class Breaker;
 
  using L1TightFrame<SCALAR>::objective;
#ifdef SOPT_MPI
  using L1TightFrame<SCALAR>::direct_space_comm;
  using L1TightFrame<SCALAR>::adjoint_space_comm;
#endif
 
  using Scalar = typename L1TightFrame<SCALAR>::Scalar;
  using Real = typename L1TightFrame<SCALAR>::Real;
 
  struct Diagnostic {
    t_uint niters;
    Real relative_variation;
    Real objective;
    bool good;
    Diagnostic(t_uint niters = 0, Real relative_variation = 0, Real objective = 0,
               bool good = false)
        : niters(niters),
          relative_variation(relative_variation),
          objective(objective),
          good(good) {}
  };
 
  struct DiagnosticAndResult : public Diagnostic {
    Vector<SCALAR> proximal;
  };
 
  template <typename T0>
  Diagnostic operator()(Eigen::MatrixBase<T0> &out, Real gamma, Vector<Scalar> const &x) const {
    // Note that we *must* call eval on x, in case it is an expression involving out
    if (gamma <= 0) {
      apply_constraints(out, x);
      return Diagnostic(0, 0, 0.5 * (out - x).squaredNorm(), true);
    }
 
    if (fista_mixing())
      return operator()(out, gamma, x, FistaMixing());
    else
      return operator()(out, gamma, x, NoMixing());
  }
 
  template <typename T0>
  DiagnosticAndResult operator()(Real const &gamma, Eigen::MatrixBase<T0> const &x) const {
    DiagnosticAndResult result;
    static_cast<Diagnostic &>(result) = operator()(result.proximal, gamma, x);
    return result;
  }
 
#ifdef SOPT_MPI
  L1(mpi::Communicator const &direct_comm = mpi::Communicator(),
     mpi::Communicator const &adjoint_comm = mpi::Communicator())
      : L1TightFrame<SCALAR>(direct_comm, adjoint_comm),
        itermax_(0),
        tolerance_(1e-8),
        positivity_constraint_(false),
        real_constraint_(false),
        fista_mixing_(true) {}
#else
  L1()
      : L1TightFrame<SCALAR>(),
        itermax_(0),
        tolerance_(1e-8),
        positivity_constraint_(false),
        real_constraint_(false),
        fista_mixing_(true) {}
#endif
 
#define SOPT_MACRO(NAME, TYPE)                 \
  TYPE const &NAME() const { return NAME##_; } \
  L1<Scalar> &NAME(TYPE const &(NAME)) {         \
    NAME##_ = NAME;                            \
    return *this;                              \
  }                                            \
                                               \
 protected:                                    \
  TYPE NAME##_;                                \
                                               \
 public:
  SOPT_MACRO(itermax, t_uint);
  SOPT_MACRO(tolerance, Real);
  SOPT_MACRO(positivity_constraint, bool);
  SOPT_MACRO(real_constraint, bool);
  SOPT_MACRO(fista_mixing, bool);
#undef SOPT_MACRO
 
  Vector<Real> const &weights() const { return L1TightFrame<Scalar>::weights(); }
  template <typename T>
  L1<Scalar> &weights(Eigen::MatrixBase<T> const &w) {
    L1TightFrame<Scalar>::weights(w);
    return *this;
  }
  L1<Scalar> &weights(Real const &w) {
    L1TightFrame<Scalar>::weights(w);
    return this;
  }
 
  Real nu() const { return L1TightFrame<Scalar>::nu(); }
  L1<Scalar> &nu(Real const &nu) {
    L1TightFrame<SCALAR>::nu(nu);
    return *this;
  }
 
  LinearTransform<Vector<Scalar>> const &Psi() const { return L1TightFrame<Scalar>::Psi(); }
  template <typename... ARGS>
  typename std::enable_if<sizeof...(ARGS) >= 1, L1<Scalar> &>::type Psi(ARGS &&... args) {
    L1TightFrame<Scalar>::Psi(std::forward<ARGS>(args)...);
    return *this;
  }
 
  template <typename... T>
  auto tight_frame(T &&... args) const
      -> decltype(this->L1TightFrame<Scalar>::operator()(std::forward<T>(args)...)) {
    return this->L1TightFrame<Scalar>::operator()(std::forward<T>(args)...);
  }
 
 protected:
  template <typename T1>
  Vector<SCALAR> apply_soft_threshhold(Real gamma, Eigen::MatrixBase<T1> const &x) const;
  template <typename T0, typename T1>
  void apply_constraints(Eigen::MatrixBase<T0> &out, Eigen::MatrixBase<T1> const &x) const;
 
  template <typename T0, typename MIXING>
  Diagnostic operator()(Eigen::MatrixBase<T0> &out, Real gamma, Vector<Scalar> const &x,
                        MIXING mixing) const;
};
 
template <typename SCALAR>
template <typename T0, typename MIXING>
typename L1<SCALAR>::Diagnostic L1<SCALAR>::operator()(Eigen::MatrixBase<T0> &out, Real gamma,
                                                       Vector<Scalar> const &x,
                                                       MIXING mixing) const {
  SOPT_MEDIUM_LOG("Starting Proximal L1 operator:");
  t_uint niters = 0;
  out = x;
 
  Breaker breaker(objective(out, x, gamma), tolerance(), false);  // not fista_mixing());
  SOPT_LOW_LOG("    - [ProxL1] iter {}, prox_fval = {}", niters, breaker.current());
  Vector<Scalar> const res = Psi().adjoint() * out;
  Vector<Scalar> u_l1 = 1e0 / nu() * (res - apply_soft_threshhold(gamma, res));
  apply_constraints(out, x - Psi() * u_l1);
 
  // Move on to other iterations
  for (++niters; niters < itermax() or itermax() == 0; ++niters) {
    auto const do_break = breaker(objective(x, out, gamma));
    SOPT_LOW_LOG("    - [ProxL1] iter {}, prox_fval = {}, rel_fval = {}", niters, breaker.current(),
                 breaker.relative_variation());
    if (do_break) break;
 
    Vector<Scalar> const res = u_l1 * nu() + Psi().adjoint() * out;
    mixing(u_l1, 1e0 / nu() * (res - apply_soft_threshhold(gamma, res)), niters);
    apply_constraints(out, x - Psi() * u_l1);
  }
 
  if (breaker.two_cycle()) SOPT_WARN("Two-cycle detected when computing L1");
 
  if (breaker.converged()) {
    SOPT_LOW_LOG("Proximal L1 operator converged at {} in {} iterations", breaker.current(),
                 niters);
  } else
    SOPT_ERROR("Proximal L1 operator did not converge after {} iterations", niters);
  return {niters, breaker.relative_variation(), breaker.current(), breaker.converged()};
}
 
template <typename SCALAR>
template <typename T1>
Vector<SCALAR> L1<SCALAR>::apply_soft_threshhold(Real gamma, Eigen::MatrixBase<T1> const &x) const {
  if (weights().size() == 1)
    return soft_threshhold(x, gamma * weights()(0));
  else
    return soft_threshhold(x, gamma * weights());
}
 
template <typename SCALAR>
template <typename T0, typename T1>
void L1<SCALAR>::apply_constraints(Eigen::MatrixBase<T0> &out,
                                   Eigen::MatrixBase<T1> const &x) const {
  if (positivity_constraint())
    out = sopt::positive_quadrant(x);
  else if (real_constraint())
    out = x.real().template cast<SCALAR>();
  else
    out = x;
}
 
template <typename SCALAR>
class L1<SCALAR>::FistaMixing {
 public:
  using Real = typename real_type<SCALAR>::type;
  FistaMixing() : t(1){};
  template <typename T1>
  void operator()(Vector<SCALAR> &previous, Eigen::MatrixBase<T1> const &unmixed, t_uint iter) {
    // reset
    if (iter == 0) {
      previous = unmixed;
      return;
    }
    if (iter <= 1) t = next(1);
    auto const prior_t = t;
    t = next(t);
    auto const alpha = (prior_t - 1) / t;
    previous = (1e0 + alpha) * unmixed.derived() - alpha * previous;
  }
  static Real next(Real t) { return 0.5 + 0.5 * std::sqrt(1e0 + 4e0 * t * t); }
 
 private:
  Real t;
};
 
template <typename SCALAR>
class L1<SCALAR>::NoMixing {
 public:
  template <typename T1>
  void operator()(Vector<SCALAR> &previous, Eigen::MatrixBase<T1> const &unmixed, t_uint) {
    previous = unmixed;
  }
};
 
template <typename SCALAR>
class L1<SCALAR>::Breaker {
 public:
  using Real = typename real_type<SCALAR>::type;
  Breaker(Real objective, Real tolerance = 1e-8, bool do_two_cycle = true)
      : tolerance_(tolerance),
        iter(0),
        objectives({{objective, 0, 0, 0}}),
        do_two_cycle(do_two_cycle) {}
  bool operator()(Real objective) {
    ++iter;
    objectives = {{objective, objectives[0], objectives[1], objectives[2]}};
    return converged() or two_cycle();
  }
  Real current() const { return objectives[0]; }
  Real previous() const { return objectives[1]; }
  Real relative_variation() const { return std::abs((current() - previous()) / current()); }
  bool two_cycle() const {
    return do_two_cycle and iter > 3 and std::abs(objectives[0] - objectives[2]) < tolerance() and
           std::abs(objectives[1] - objectives[3]) < tolerance();
  }
 
  bool converged() const {
    // If current ~ 0, then defaults to absolute convergence
    // This is mainly to avoid a division by zero
    if (std::abs(current() * 1000) < tolerance()) return std::abs(previous() * 1000) < tolerance();
    return relative_variation() < tolerance();
  }
  Real tolerance() const { return tolerance_; }
  L1<SCALAR>::Breaker &tolerance(Real tol) const {
    tolerance_ = tol;
    return *this;
  }
 
 protected:
  Real tolerance_;
  t_uint iter;
  std::array<Real, 4> objectives;
  bool do_two_cycle;
};
} // namespace sopt::proximal
 
#endif