slope.h

Go to the documentation of this file.

 
#pragma once
 
#include "clusters.h"
#include "constants.h"
#include "diagnostics.h"
#include "estimate_alpha.h"
#include "logger.h"
#include "losses/loss.h"
#include "losses/setup_loss.h"
#include "math.h"
#include "normalize.h"
#include "regularization_sequence.h"
#include "screening.h"
#include "slope_fit.h"
#include "slope_path.h"
#include "solvers/hybrid_cd.h"
#include "solvers/setup_solver.h"
#include "sorted_l1_norm.h"
#include "timer.h"
#include <Eigen/Core>
#include <Eigen/SparseCore>
#include <cassert>
#include <memory>
#include <numeric>
#include <optional>
 
namespace slope {
 
inline bool
defaultInterruptChecker()
{
  return false;
}
 
class Slope
{
public:
  Slope() = default;
 
  void setSolver(const std::string& solver);
 
  void setIntercept(bool intercept);
 
  void setNormalization(const std::string& type);
 
  void setUpdateClusters(bool update_clusters);
 
  void setReturnClusters(const bool return_clusters);
 
  void setAlphaMinRatio(double alpha_min_ratio);
 
  void setAlphaType(const std::string& alpha_type);
 
  void setLearningRateDecr(double learning_rate_decr);
 
  void setQ(double q);
 
  void setOscarParameters(const double theta1, const double theta2);
 
  void setTol(double tol);
 
  void setRelaxTol(double tol);
 
  void setRelaxMaxOuterIterations(int max_it);
 
  void setRelaxMaxInnerIterations(int max_it);
 
  void setMaxIterations(int max_it);
 
  void setPathLength(int path_length);
 
  void setHybridCdIterations(int cd_iterations);
 
  void setHybridCdType(const std::string& cd_type);
 
  void setLambdaType(const std::string& lambda_type);
 
  void setLoss(const std::string& loss_type);
 
  void setScreening(const std::string& screening_type);
 
  void setModifyX(const bool modify_x);
 
  void setDevChangeTol(const double dev_change_tol);
 
  void setDevRatioTol(const double dev_ratio_tol);
 
  void setMaxClusters(const int max_clusters);
 
  void setCentering(const std::string& type);
 
  void setCentering(const Eigen::VectorXd& x_centers);
 
  void setScaling(const std::string& type);
 
  void setDiagnostics(const bool collect_diagnostics);
 
  void setScaling(const Eigen::VectorXd& x_scales);
 
  void setAlphaEstimationMaxIterations(const int alpha_est_maxit);
 
  void setRandomSeed(const int seed);
 
  void setRandomSeed(std::optional<int> seed);
 
  bool hasRandomSeed() const;
 
  int getRandomSeed() const;
 
  int getAlphaEstimationMaxIterations() const;
 
  bool getFitIntercept() const;
 
  const std::string& getLossType();
 
  template<typename T>
  SlopePath path(
    Eigen::EigenBase<T>& x,
    const Eigen::MatrixXd& y_in,
    Eigen::ArrayXd alpha = Eigen::ArrayXd::Zero(0),
    Eigen::ArrayXd lambda = Eigen::ArrayXd::Zero(0),
    std::function<bool()> check_interrupt = defaultInterruptChecker)
  {
    using Eigen::MatrixXd;
    using Eigen::VectorXd;
 
    const int n = x.rows();
    const int p = x.cols();
 
    const int INTERRUPT_FREQ = 100;
    bool interrupt = false;
 
    if (n != y_in.rows()) {
      throw std::invalid_argument(
        "x and y_in must have the same number of rows");
    }
 
    auto jit_normalization = normalize(x.derived(),
                                       this->x_centers,
                                       this->x_scales,
                                       this->centering_type,
                                       this->scaling_type,
                                       this->modify_x);
 
    std::unique_ptr<Loss> loss = setupLoss(this->loss_type);
 
    MatrixXd y = loss->preprocessResponse(y_in);
 
    const int m = y.cols();
 
    std::vector<int> full_set(p * m);
    std::iota(full_set.begin(), full_set.end(), 0);
 
    VectorXd beta0 = VectorXd::Zero(m);
    VectorXd beta = VectorXd::Zero(p * m);
 
    MatrixXd eta = MatrixXd::Zero(n, m); // linear predictor
 
    if (this->intercept) {
      beta0 = loss->link(y.colwise().mean()).transpose();
      eta.rowwise() = beta0.transpose();
    }
 
    MatrixXd residual = loss->residual(eta, y);
    VectorXd gradient(beta.size());
 
    // Path data
    bool user_alpha = alpha.size() > 0;
    bool user_lambda = lambda.size() > 0;
 
    if (!user_lambda) {
      lambda = lambdaSequence(
        p * m, this->q, this->lambda_type, n, this->theta1, this->theta2);
    } else {
      if (lambda.size() != beta.size()) {
        throw std::invalid_argument(
          "lambda must be the same length as the number of coefficients");
      }
      if (lambda.minCoeff() < 0) {
        throw std::invalid_argument("lambda must be non-negative");
      }
      if (!lambda.isFinite().all()) {
        throw std::invalid_argument("lambda must be finite");
      }
    }
 
    // Setup the regularization sequence and path
    SortedL1Norm sl1_norm;
 
    // TODO: Make this part of the slope class
    auto solver = setupSolver(this->solver_type,
                              this->loss_type,
                              jit_normalization,
                              this->intercept,
                              this->update_clusters,
                              this->cd_iterations,
                              this->cd_type,
                              this->random_seed);
 
    updateGradient(gradient,
                   x.derived(),
                   residual,
                   full_set,
                   this->x_centers,
                   this->x_scales,
                   Eigen::VectorXd::Ones(n),
                   jit_normalization);
 
    int alpha_max_ind = whichMax(gradient.cwiseAbs());
    double alpha_max = sl1_norm.dualNorm(gradient, lambda);
 
    if (alpha_type == "path" ||
        (alpha_type == "estimate" && alpha_estimate != 1)) {
      if (alpha_min_ratio < 0) {
        alpha_min_ratio = n > gradient.size() ? 1e-4 : 1e-2;
      }
 
      alpha =
        regularizationPath(alpha, path_length, alpha_min_ratio, alpha_max);
      path_length = alpha.size();
    } else if (alpha_type == "estimate" && alpha_estimate == -1) {
      if (loss_type != "quadratic") {
        throw std::invalid_argument("Automatic alpha estimation is only "
                                    "available for the quadratic loss");
      }
    }
 
    // Screening setup
    std::unique_ptr<ScreeningRule> screening_rule =
      createScreeningRule(this->screening_type);
    std::vector<int> working_set =
      screening_rule->initialize(full_set, alpha_max_ind);
 
    // Path variables
    double null_deviance = loss->deviance(eta, y);
    double dev_prev = null_deviance;
 
    Timer timer;
 
    double alpha_prev = std::max(alpha_max, alpha(0));
 
    std::vector<SlopeFit> fits;
 
    // Regularization path loop
    for (int path_step = 0; path_step < this->path_length; ++path_step) {
      // Check for interrupt at the start of each path step
      bool local_interrupt = false;
#ifdef _OPENMP
#pragma omp critical(check_interrupt)
#endif
      {
        local_interrupt = check_interrupt();
      }
      if (local_interrupt) {
        interrupt = true;
        break;
      }
 
      double alpha_curr = alpha(path_step);
 
      assert(alpha_curr <= alpha_prev && "Alpha must be decreasing");
 
      Eigen::ArrayXd lambda_curr = alpha_curr * lambda;
      Eigen::ArrayXd lambda_prev = alpha_prev * lambda;
 
      std::vector<double> duals, primals, time;
      timer.start();
 
      // Update gradient for the full set
      // TODO: Only update for non-working set since gradient is updated before
      // the convergence check in the inner loop for the working set
      updateGradient(gradient,
                     x.derived(),
                     residual,
                     full_set,
                     x_centers,
                     x_scales,
                     Eigen::VectorXd::Ones(x.rows()),
                     jit_normalization);
 
      working_set = screening_rule->screen(
        gradient, lambda_curr, lambda_prev, beta, full_set);
 
      int it = 0;
      int total_it = 0;
      for (; it < this->max_it; ++it, ++total_it) {
        // Compute primal, dual, and gap
        residual = loss->residual(eta, y);
        updateGradient(gradient,
                       x.derived(),
                       residual,
                       working_set,
                       this->x_centers,
                       this->x_scales,
                       Eigen::VectorXd::Ones(n),
                       jit_normalization);
 
        double primal = loss->loss(eta, y) +
                        sl1_norm.eval(beta(working_set),
                                      lambda_curr.head(working_set.size()));
 
        MatrixXd theta = residual;
 
        // First compute gradient with potential offset for intercept case
        VectorXd dual_gradient = gradient;
 
        // TODO: Can we avoid this copy? Maybe revert offset afterwards or,
        // alternatively, solve intercept until convergence and then no longer
        // need the offset at all.
        if (this->intercept) {
          VectorXd theta_mean = theta.colwise().mean();
          theta.rowwise() -= theta_mean.transpose();
 
          offsetGradient(dual_gradient,
                         x.derived(),
                         theta_mean,
                         working_set,
                         this->x_centers,
                         this->x_scales,
                         jit_normalization);
        }
 
        // Common scaling operation
        double dual_norm = sl1_norm.dualNorm(
          dual_gradient(working_set), lambda_curr.head(working_set.size()));
        theta.array() /= std::max(1.0, dual_norm);
 
        double dual = loss->dual(theta, y, Eigen::VectorXd::Ones(n));
 
        if (collect_diagnostics) {
          timer.pause();
          double true_dual = computeDual(beta,
                                         residual,
                                         loss,
                                         sl1_norm,
                                         lambda_curr,
                                         x.derived(),
                                         y,
                                         this->x_centers,
                                         this->x_scales,
                                         jit_normalization,
                                         this->intercept);
          timer.resume();
 
          time.emplace_back(timer.elapsed());
          primals.emplace_back(primal);
          duals.emplace_back(true_dual);
        }
 
        double dual_gap = primal - dual;
 
        assert(dual_gap > -1e-6 && "Dual gap should be positive");
 
        double tol_scaled = (std::abs(primal) + constants::EPSILON) * this->tol;
 
        if (dual_gap <= tol_scaled || it == this->max_it) {
          bool no_violations =
            screening_rule->checkKktViolations(gradient,
                                               beta,
                                               lambda_curr,
                                               working_set,
                                               x.derived(),
                                               residual,
                                               this->x_centers,
                                               this->x_scales,
                                               jit_normalization,
                                               full_set);
          if (no_violations) {
            break;
          } else {
            it = 0; // Restart if there are KKT violations
          }
        }
 
        if (it % INTERRUPT_FREQ == 0) {
          bool local_interrupt = false;
#ifdef _OPENMP
#pragma omp critical(check_interrupt)
#endif
          {
            local_interrupt = check_interrupt();
          }
          if (local_interrupt) {
            interrupt = true;
            break;
          }
        }
 
        solver->run(beta0,
                    beta,
                    eta,
                    lambda_curr,
                    loss,
                    sl1_norm,
                    gradient,
                    working_set,
                    x.derived(),
                    this->x_centers,
                    this->x_scales,
                    y);
      }
 
      if (it == this->max_it) {
        WarningLogger::addWarning(
          WarningCode::MAXIT_REACHED,
          "Maximum number of iterations reached at step = " +
            std::to_string(path_step) + ".");
      }
 
      alpha_prev = alpha_curr;
 
      // Compute early stopping criteria
      double dev = loss->deviance(eta, y);
      double dev_ratio = 1 - dev / null_deviance;
      double dev_change = path_step == 0 ? 1.0 : 1 - dev / dev_prev;
      dev_prev = dev;
 
      Clusters clusters;
 
      if (return_clusters) {
        clusters.update(beta);
      }
 
      SlopeFit fit{ beta0,
                    beta.reshaped(p, m).sparseView(),
                    clusters,
                    alpha_curr,
                    lambda,
                    dev,
                    null_deviance,
                    primals,
                    duals,
                    time,
                    total_it,
                    this->centering_type,
                    this->scaling_type,
                    this->intercept,
                    this->x_centers,
                    this->x_scales };
 
      fits.emplace_back(std::move(fit));
 
      if (interrupt) {
        break;
      }
 
      if (!user_alpha) {
        int n_unique = unique(beta.cwiseAbs()).size();
        if (dev_ratio > dev_ratio_tol || dev_change < dev_change_tol ||
            n_unique >= this->max_clusters.value_or(n + 1)) {
          break;
        }
      }
    }
 
    return fits;
  }
 
  template<typename T>
  SlopeFit fit(
    Eigen::EigenBase<T>& x,
    const Eigen::MatrixXd& y_in,
    const double alpha = 1.0,
    Eigen::ArrayXd lambda = Eigen::ArrayXd::Zero(0),
    std::function<bool()> check_interrupt = defaultInterruptChecker)
  {
    Eigen::ArrayXd alpha_arr(1);
    alpha_arr(0) = alpha;
    SlopePath res = path(x, y_in, alpha_arr, lambda, check_interrupt);
 
    return { res(0) };
  };
 
  template<typename T>
  SlopePath estimateAlpha(
    Eigen::EigenBase<T>& x,
    Eigen::MatrixXd& y,
    std::function<bool()> check_interrupt = defaultInterruptChecker)
  {
    int n = x.rows();
    int p = x.cols();
 
    // Create a copy with alpha type set to path to avoid recursion
    Slope model_copy = *this;
    model_copy.setAlphaType("path");
 
    std::vector<int> selected;
    Eigen::ArrayXd alpha(1);
    SlopePath result;
 
    // Estimate the noise level, if possible
    if (n >= p + 30) {
      alpha(0) = estimateNoise(x, y, this->intercept) / n;
      this->alpha_estimate = alpha(0);
      result =
        model_copy.path(x, y, alpha, Eigen::ArrayXd::Zero(0), check_interrupt);
    } else {
      for (int it = 0; it < this->alpha_est_maxit; ++it) {
        T x_selected = subsetCols(x.derived(), selected);
 
        std::vector<int> selected_prev = selected;
        selected.clear();
 
        alpha(0) = estimateNoise(x_selected, y, this->intercept) / n;
        this->alpha_estimate = alpha(0);
 
        result = model_copy.path(
          x, y, alpha, Eigen::ArrayXd::Zero(0), check_interrupt);
        auto coefs = result.getCoefs().back();
 
        for (typename Eigen::SparseMatrix<double>::InnerIterator it(coefs, 0);
             it;
             ++it) {
          selected.emplace_back(it.row());
        }
 
        if (selected == selected_prev) {
          return result;
        }
 
        if (static_cast<int>(selected.size()) >= n + this->intercept) {
          throw std::runtime_error(
            "selected >= n - 1 variables, cannot estimate variance");
        }
      }
 
      slope::WarningLogger::addWarning(
        slope::WarningCode::MAXIT_REACHED,
        "Maximum iterations reached in alpha estimation");
    }
 
    return result;
  }
 
  template<typename T>
  SlopeFit relax(const SlopeFit& fit,
                 T& x,
                 const Eigen::VectorXd& y_in,
                 const double gamma = 0.0,
                 Eigen::VectorXd beta0 = Eigen::VectorXd(0),
                 Eigen::VectorXd beta = Eigen::VectorXd(0))
  {
    using Eigen::MatrixXd;
    using Eigen::VectorXd;
 
    int n = x.rows();
    int p = x.cols();
 
    if (beta0.size() == 0) {
      beta0 = fit.getIntercepts(false);
    }
 
    if (beta.size() == 0) {
      beta = fit.getCoefs(false);
    }
 
    double alpha = fit.getAlpha();
 
    Timer timer;
 
    std::vector<double> primals, duals, time;
    timer.start();
 
    auto jit_normalization =
      normalize(x, x_centers, x_scales, centering_type, scaling_type, modify_x);
 
    bool update_clusters = false;
 
    std::unique_ptr<Loss> loss = setupLoss(this->loss_type);
 
    MatrixXd y = loss->preprocessResponse(y_in);
 
    int m = y.cols();
 
    Eigen::ArrayXd lambda_cumsum_relax = Eigen::ArrayXd::Zero(p * m + 1);
 
    auto working_set = activeSet(beta);
 
    Eigen::MatrixXd eta = linearPredictor(x,
                                          working_set,
                                          beta0,
                                          beta,
                                          x_centers,
                                          x_scales,
                                          jit_normalization,
                                          intercept);
    VectorXd gradient = VectorXd::Zero(p * m);
    MatrixXd residual(n, m);
    MatrixXd working_residual(n, m);
 
    MatrixXd w = MatrixXd::Ones(n, m);
    MatrixXd w_ones = MatrixXd::Ones(n, m);
    MatrixXd z = y;
 
    slope::Clusters clusters(beta);
 
    std::mt19937 rng;
 
    if (random_seed.has_value()) {
      rng.seed(*random_seed);
    } else {
      rng.seed(std::random_device{}());
    }
 
    int passes = 0;
 
    for (int irls_it = 0; irls_it < max_it_outer_relax; irls_it++) {
      residual = loss->residual(eta, y);
 
      if (collect_diagnostics) {
        primals.push_back(loss->loss(eta, y));
        duals.push_back(0.0);
        time.push_back(timer.elapsed());
      }
 
      Eigen::VectorXd cluster_gradient = clusterGradient(beta,
                                                         residual,
                                                         clusters,
                                                         x,
                                                         w_ones,
                                                         x_centers,
                                                         x_scales,
                                                         jit_normalization);
 
      double norm_grad = cluster_gradient.lpNorm<Eigen::Infinity>();
 
      if (norm_grad < tol_relax) {
        break;
      }
 
      loss->updateWeightsAndWorkingResponse(w, z, eta, y);
      working_residual = eta - z;
 
      for (int inner_it = 0; inner_it < max_it_inner_relax; ++inner_it) {
        passes++;
 
        double max_abs_gradient = coordinateDescent(beta0,
                                                    beta,
                                                    working_residual,
                                                    clusters,
                                                    lambda_cumsum_relax,
                                                    x,
                                                    w,
                                                    x_centers,
                                                    x_scales,
                                                    intercept,
                                                    jit_normalization,
                                                    update_clusters,
                                                    rng,
                                                    cd_type);
 
        if (max_abs_gradient < tol_relax) {
          break;
        }
      }
 
      eta = working_residual + z;
 
      if (irls_it == max_it_outer_relax) {
        WarningLogger::addWarning(WarningCode::MAXIT_REACHED,
                                  "Maximum number of IRLS iterations reached.");
      }
    }
 
    double dev = loss->deviance(eta, y);
 
    if (gamma > 0) {
      Eigen::VectorXd old_coefs = fit.getCoefs(false);
      Eigen::VectorXd old_intercept = fit.getIntercepts(false);
      beta = (1 - gamma) * beta + gamma * old_coefs;
    }
 
    SlopeFit fit_out{ beta0,
                      beta.reshaped(p, m).sparseView(),
                      clusters,
                      alpha,
                      fit.getLambda(),
                      dev,
                      fit.getNullDeviance(),
                      primals,
                      duals,
                      time,
                      passes,
                      centering_type,
                      scaling_type,
                      intercept,
                      x_centers,
                      x_scales };
 
    return fit_out;
  }
 
  template<typename T>
  SlopePath relax(const SlopePath& path,
                  T& x,
                  const Eigen::VectorXd& y,
                  const double gamma = 0.0)
  {
    std::vector<SlopeFit> fits;
 
    for (size_t i = 0; i < path.size(); i++) {
      // TODO: Reinstate warm starts. Need to be careful about
      // the warm started values though since they have to
      // agree with the cluster or we will run into trouble.
      // We can probably fix this by using the signs
      // of the cluster object rather than the betas though.
      auto relaxed_fit = relax(path(i), x, y, gamma);
 
      fits.emplace_back(relaxed_fit);
    }
 
    return fits;
  }
 
private:
  // Parameters
  bool collect_diagnostics = false;
  bool intercept = true;
  bool modify_x = false;
  bool return_clusters = true;
  bool update_clusters = true;
  double alpha_min_ratio = -1;
  double dev_change_tol = 1e-5;
  double dev_ratio_tol = 0.999;
  double learning_rate_decr = 0.5;
  double q = 0.1;
  double theta1 = 1.0;
  double theta2 = 0.5;
  double tol = 1e-4;
  double tol_relax = 1e-4;
  double alpha_estimate = -1;
  int alpha_est_maxit = 1000;
  int cd_iterations = 10;
  int max_it = 1e5;
  int max_it_inner_relax = 1e5;
  int max_it_outer_relax = 50;
  int path_length = 100;
  std::optional<int> max_clusters = std::nullopt;
  std::optional<int> random_seed = 0;
  std::string alpha_type = "path";
  std::string cd_type = "permuted";
  std::string centering_type = "mean";
  std::string lambda_type = "bh";
  std::string loss_type = "quadratic";
  std::string scaling_type = "sd";
  std::string screening_type = "strong";
  std::string solver_type = "auto";
 
  // Data
  Eigen::VectorXd x_centers;
  Eigen::VectorXd x_scales;
};
 
} // namespace slope