Drop projections to C++ and add cone program refinement

cpp/include/cones.h

@@ -16,6 +16,20 @@ class Cone {
       : type(type), sizes(sizes){};
 };
 
+/* Project `x` onto the primal or dual cone.
+ *
+ *  Args:
+ *    x:     The point at which to evaluate the derivative
+ *    cones: A list of cones; the cone on which to project is the cartesian
+ *           product of these cones
+ *    dual:  whether to project onto the dual cone
+ *
+ *  Returns:
+ *    A Vector with the projection.
+ */
+Vector projection(const Vector &x, const std::vector<Cone> &cones,
+                  bool dual);
+
 /* Compute the derivative, at `x`, of a projection onto a cone.
  *
  *  Args:

cpp/include/deriv.h

@@ -3,13 +3,25 @@
 #include "cones.h"
 #include "eigen_includes.h"
 #include "linop.h"
+#include "lsqr.h"
+
+LinearOperator dpi(const Vector &u, const Vector &v, double w,
+                   const std::vector<Cone> &cones);
+
+Matrix dpi_dense(const Vector &u, const Vector &v, double w,
+                 const std::vector<Cone> &cones);
 
 LinearOperator M_operator(const SparseMatrix &Q, const std::vector<Cone> &cones,
                           const Vector &u, const Vector &v, double w);
 
 Matrix M_dense(const Matrix &Q, const std::vector<Cone> &cones, const Vector &u,
                const Vector &v, double w);
 
+LsqrResult _solve_adjoint_derivative_lsqr(
+    const SparseMatrix &Q, const std::vector<Cone> &cones,
+    const Vector &u, const Vector &v, double w, const Vector &dz);
+
+
 // this function releases the GIL.
 Vector _solve_derivative_dense(const Matrix &M, const Matrix &MT,
                                const Vector &rhs);

cpp/include/eigen_includes.h

@@ -4,6 +4,7 @@
 #include "Eigen/Sparse"
 
 using Vector = Eigen::VectorXd;
+using VectorRef = Eigen::Ref<Eigen::VectorXd>;
 using Array = Eigen::Array<double, Eigen::Dynamic, 1>;
 using Matrix = Eigen::MatrixXd;
 using MatrixRef = Eigen::Ref<Eigen::MatrixXd>;

cpp/include/residual.h

@@ -0,0 +1,81 @@
+#pragma once
+
+#include "cones.h"
+#include "eigen_includes.h"
+#include "linop.h"
+#include "lsqr.h"
+#include "deriv.h"
+
+std::tuple<Vector, Vector> proj_pq(const std::vector<Cone> &cones,
+                                   const Vector &u, const Vector &v, double w);
+
+/**
+ * Refines the solution to a cone program by using LSQR to improve
+ * the normalized residual map:
+ * https://stanford.edu/~boyd/papers/pdf/cone_prog_refine.pdf
+ */
+Vector refine(const SparseMatrix &Q, const std::vector<Cone> &cones,
+              const Vector &u, const Vector &v, double w,
+              const int n_iter,
+              const double lsqr_lambda=1e-8,
+              const double lsqr_iter_lim=30,
+              const int alpha_K=10);
+
+
+/**
+ * Compute the value of the residual map at a point z = (u, v, w)
+ * for a cone program defined by a sparse Q.
+ */
+Vector residual(const SparseMatrix &Q, const std::vector<Cone> &cones,
+                const Vector &u, const Vector &v, double w);
+
+/**
+ * Compute the value of the normalized residual map at a point z = (u, v, w)
+ * for a cone program defined by a sparse Q.
+ */
+Vector N_residual(const SparseMatrix &Q, const std::vector<Cone> &cones,
+                  const Vector &u, const Vector &v, double w);
+
+/**
+ * Compute the value of the residual map at a point z = (u, v, w)
+ * for a cone program defined by a dense Q.
+ */
+Vector residual_dense(const Matrix &Q, const std::vector<Cone> &cones,
+                      const Vector &u, const Vector &v, double w);
+
+/**
+ * Compute the value of the normalized residual map at a point z = (u, v, w)
+ * for a cone program defined by a sparse Q.
+ */
+Vector N_residual_dense(const Matrix &Q, const std::vector<Cone> &cones,
+                        const Vector &u, const Vector &v, double w);
+
+/**
+ * The derivative of the normalized residual map at a point z = (u, v, w)
+ * for a cone program defined by a sparse Q.
+ */
+LinearOperator DN(const SparseMatrix &Q, const std::vector<Cone> &cones,
+                  const Vector &u, const Vector &v, double w);
+
+/**
+ * The derivative of the normalized residual map at a point z = (u, v, w)
+ * for a cone program defined by a dense Q.
+ */
+Matrix DN_dense(const Matrix &Q, const std::vector<Cone> &cones, const Vector &u,
+                const Vector &v, double w);
+
+/**
+ * Left-multiply a point dz = (du, dv, dw) by the derivative of the
+ * normalized residual map.
+ */
+Vector apply_DN_matvec(const SparseMatrix &Q, const std::vector<Cone> &cones,
+                       const Vector &u, const Vector &v, double w,
+                       const Vector &du, const Vector &dv, double dw);
+
+/**
+ * Right-multiply a point dz = (du, dv, dw) by the derivative of the
+ * normalized residual map.
+ */
+Vector apply_DN_rmatvec(const SparseMatrix &Q, const std::vector<Cone> &cones,
+                        const Vector &u, const Vector &v, double w,
+                        const Vector &du, const Vector &dv, double dw);

cpp/src/cones.cpp

@@ -14,6 +14,55 @@
 const double EulerConstant = std::exp(1.0);
 const double sqrt_two = std::sqrt(2.0);
 
+int vectorized_psd_size(int n) { return n * (n + 1) / 2; }
+
+
+bool in_exp(const Eigen::Vector3d &x) {
+  return (x[0] <= 0 && std::abs(x[1]) <= CONE_THRESH && x[2] >= 0) ||
+         (x[1] > 0 && x[1] * exp(x[0] / x[1]) - x[2] <= CONE_THRESH);
+}
+
+bool in_exp_dual(const Eigen::Vector3d &x) {
+  return (std::abs(x[0]) <= CONE_THRESH && x[1] >= 0 && x[2] >= 0) ||
+         (x[0] < 0 &&
+          -x[0] * exp(x[1] / x[0]) - EulerConstant * x[2] <= CONE_THRESH);
+}
+
+Vector lower_triangular_from_matrix(const Matrix &matrix) {
+  int n = matrix.rows();
+  Vector lower_tri = Vector::Zero(vectorized_psd_size(n));
+  int offset = 0;
+  for (int col = 0; col < n; ++col) {
+    for (int row = col; row < n; ++row) {
+      if (row != col) {
+        lower_tri[offset] = matrix(row, col) * sqrt_two;
+      } else {
+        lower_tri[offset] = matrix(row, col);
+      }
+      ++offset;
+    }
+  }
+  return lower_tri;
+}
+
+Matrix matrix_from_lower_triangular(const Vector &lower_tri) {
+  int n = static_cast<int>(std::sqrt(2 * lower_tri.size()));
+  Matrix matrix = Matrix::Zero(n, n);
+  int offset = 0;
+  for (int col = 0; col < n; ++col) {
+    for (int row = col; row < n; ++row) {
+      if (row != col) {
+        matrix(row, col) = lower_tri[offset] / sqrt_two;
+        matrix(col, row) = lower_tri[offset] / sqrt_two;
+      } else {
+        matrix(row, col) = lower_tri[offset];
+      }
+      ++offset;
+    }
+  }
+  return matrix;
+}
+
 double exp_newton_one_d(double rho, double y_hat, double z_hat) {
   double t = std::max(-z_hat, 1e-6);
   double f, fp;
@@ -115,52 +164,91 @@ Eigen::Vector3d project_exp_cone(const Eigen::Vector3d &x) {
   return projection;
 }
 
-bool in_exp(const Eigen::Vector3d &x) {
-  return (x[0] <= 0 && std::abs(x[1]) <= CONE_THRESH && x[2] >= 0) ||
-         (x[1] > 0 && x[1] * exp(x[0] / x[1]) - x[2] <= CONE_THRESH);
-}
 
-bool in_exp_dual(const Eigen::Vector3d &x) {
-  return (std::abs(x[0]) <= CONE_THRESH && x[1] >= 0 && x[2] >= 0) ||
-         (x[0] < 0 &&
-          -x[0] * exp(x[1] / x[0]) - EulerConstant * x[2] <= CONE_THRESH);
-}
-
-int vectorized_psd_size(int n) { return n * (n + 1) / 2; }
-
-Vector lower_triangular_from_matrix(const Matrix &matrix) {
-  int n = matrix.rows();
-  Vector lower_tri = Vector::Zero(vectorized_psd_size(n));
-  int offset = 0;
-  for (int col = 0; col < n; ++col) {
-    for (int row = col; row < n; ++row) {
-      if (row != col) {
-        lower_tri[offset] = matrix(row, col) * sqrt_two;
-      } else {
-        lower_tri[offset] = matrix(row, col);
-      }
-      ++offset;
+void _projection(VectorRef &y, const Vector &x,
+                 ConeType type, bool dual) {
+  if (type == ZERO) {
+    if (dual) {
+      y << x;
+    } else {
+      /* segment is already zero */
+    }
+  } else if (type == POS) {
+    y << x.cwiseMax(0.);
+  } else if (type == SOC) {
+    double t = x[0];
+    const Vector z = x.segment(1, x.rows()-1);
+    double norm_z = z.norm();
+    if (norm_z <= t) {
+      y<< x;
+    } else if (norm_z <= -t) {
+      /* segment is already zero */
+    } else {
+      float c = 0.5 * (1. + t/norm_z);
+      VectorRef y_z = y.segment(1, y.rows()-1);
+      y[0] = c * norm_z;
+      y_z << c * z;
     }
+  } else if (type == PSD) {
+    const Matrix &X = matrix_from_lower_triangular(x);
+    Eigen::SelfAdjointEigenSolver<Matrix> eigen_solver(X.rows());
+    eigen_solver.compute(X);
+    const Vector &eigenvalues = eigen_solver.eigenvalues();
+    const Matrix &Q = eigen_solver.eigenvectors();
+    const Vector &clamped_eigenvalues = eigenvalues.cwiseMax(0.);
+    Matrix result = Q * clamped_eigenvalues.asDiagonal() * Q.transpose();
+    y << lower_triangular_from_matrix(result);
+  } else if (type == EXP) {
+    int num_cones = x.size() / 3;
+    int offset = 0;
+    for (int i = 0; i < num_cones; ++i) {
+        Eigen::Vector3d x_i;
+        if (dual) {
+          x_i = -1. * x.segment(offset, 3);
+        } else {
+          x_i = x.segment(offset, 3);
+        }
+        y.segment(offset, 3) << project_exp_cone(x_i);
+        offset += 3;
+    }
+    // via Moreau: Pi_K*(x) = x + Pi_K(-x)
+    if (dual) {
+        y += x;
+    }
+  } else if (type == EXP_DUAL) {
+    _projection(y, x, EXP, !dual);
+  } else {
+    throw std::invalid_argument("Cone not supported.");
   }
-  return lower_tri;
 }
 
-Matrix matrix_from_lower_triangular(const Vector &lower_tri) {
-  int n = static_cast<int>(std::sqrt(2 * lower_tri.size()));
-  Matrix matrix = Matrix::Zero(n, n);
+
+Vector projection(const Vector &x, const std::vector<Cone> &cones,
+                  bool dual) {
+  Vector y = Vector::Zero(x.rows());
+
   int offset = 0;
-  for (int col = 0; col < n; ++col) {
-    for (int row = col; row < n; ++row) {
-      if (row != col) {
-        matrix(row, col) = lower_tri[offset] / sqrt_two;
-        matrix(col, row) = lower_tri[offset] / sqrt_two;
-      } else {
-        matrix(row, col) = lower_tri[offset];
+  for (const Cone &cone : cones) {
+    const ConeType &type = cone.type;
+    const std::vector<int> &sizes = cone.sizes;
+    if (std::accumulate(sizes.begin(), sizes.end(), 0) == 0) {
+      continue;
+    }
+    for (int size : sizes) {
+      if (type == PSD) {
+        size = vectorized_psd_size(size);
+      } else if (type == EXP) {
+        size *= 3;
+      } else if (type == EXP_DUAL) {
+        size *= 3;
       }
-      ++offset;
+      VectorRef y_segment = y.segment(offset, size);
+      _projection(y_segment, x.segment(offset, size), type, dual);
+      offset += size;
     }
   }
-  return matrix;
+
+  return y;
 }
 
 LinearOperator _dprojection_exp(const Vector &x, bool dual) {

cpp/src/deriv.cpp

@@ -50,12 +50,22 @@ Matrix M_dense(const Matrix &Q, const std::vector<Cone> &cones, const Vector &u,
   return (Q - eye) * dpi_dense(u, v, w, cones) + eye;
 }
 
+
 Vector _solve_derivative_dense(const Matrix &M, const Matrix &MT,
                                const Vector &rhs) {
   // TODO: Factorization could be cached to optimize multiple calls
   return (MT * M).ldlt().solve(MT * rhs);
 }
 
+LsqrResult _solve_adjoint_derivative_lsqr(
+    const SparseMatrix &Q, const std::vector<Cone> &cones,
+    const Vector &u, const Vector &v, double w, const Vector &dz) {
+  LinearOperator M = M_operator(Q, cones, u, v, w);
+  LinearOperator MT = M.transpose();
+  LsqrResult result = lsqr(MT, dz);
+  return result;
+}
+
 Vector _solve_adjoint_derivative_dense(const Matrix &M, const Matrix &MT,
                                        const Vector &dz) {
   // TODO: Factorization could be cached to optimize multiple calls