diff --git a/cpp/include/cones.h b/cpp/include/cones.h index ddde2d8..f93230c 100644 --- a/cpp/include/cones.h +++ b/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 &cones, + bool dual); + /* Compute the derivative, at `x`, of a projection onto a cone. * * Args: diff --git a/cpp/include/deriv.h b/cpp/include/deriv.h index 4a512c2..9888e9b 100644 --- a/cpp/include/deriv.h +++ b/cpp/include/deriv.h @@ -4,6 +4,12 @@ #include "eigen_includes.h" #include "linop.h" +LinearOperator dpi(const Vector &u, const Vector &v, double w, + const std::vector &cones); + +Matrix dpi_dense(const Vector &u, const Vector &v, double w, + const std::vector &cones); + LinearOperator M_operator(const SparseMatrix &Q, const std::vector &cones, const Vector &u, const Vector &v, double w); diff --git a/cpp/include/eigen_includes.h b/cpp/include/eigen_includes.h index 9e6d095..d6ca6b2 100644 --- a/cpp/include/eigen_includes.h +++ b/cpp/include/eigen_includes.h @@ -4,6 +4,7 @@ #include "Eigen/Sparse" using Vector = Eigen::VectorXd; +using VectorRef = Eigen::Ref; using Array = Eigen::Array; using Matrix = Eigen::MatrixXd; using MatrixRef = Eigen::Ref; diff --git a/cpp/src/cones.cpp b/cpp/src/cones.cpp index 25fbe9b..3ebed4e 100644 --- a/cpp/src/cones.cpp +++ b/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(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 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(std::sqrt(2 * lower_tri.size())); - Matrix matrix = Matrix::Zero(n, n); + +Vector projection(const Vector &x, const std::vector &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 &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) { diff --git a/cpp/src/wrapper.cpp b/cpp/src/wrapper.cpp index cefc691..86f200c 100644 --- a/cpp/src/wrapper.cpp +++ b/cpp/src/wrapper.cpp @@ -47,8 +47,11 @@ PYBIND11_MODULE(_diffcp, m) { m.def("_solve_derivative_dense", &_solve_derivative_dense, py::call_guard()); m.def("_solve_adjoint_derivative_dense", &_solve_adjoint_derivative_dense, py::call_guard()); - m.def("dprojection", &dprojection); - m.def("dprojection_dense", &dprojection_dense); + m.def("projection", &projection, py::call_guard()); + m.def("dprojection", &dprojection, py::call_guard()); + m.def("dprojection_dense", &dprojection_dense, py::call_guard()); + m.def("dpi", &dpi, py::call_guard()); + m.def("dpi_dense", &dpi_dense, py::call_guard()); m.def("project_exp_cone", &project_exp_cone); m.def("in_exp", &in_exp); m.def("in_exp_dual", &in_exp_dual); diff --git a/diffcp/cones.py b/diffcp/cones.py index 410b241..b04bba4 100644 --- a/diffcp/cones.py +++ b/diffcp/cones.py @@ -3,7 +3,7 @@ import scipy.sparse.linalg as splinalg import warnings -from _diffcp import dprojection, project_exp_cone, Cone, ConeType +from _diffcp import projection, dprojection, project_exp_cone, Cone, ConeType ZERO = "f" POS = "l" @@ -127,18 +127,6 @@ def pi(x, cones, dual=False): Returns: NumPy array that is the projection of `x` onto the (dual) cones """ - projection = np.zeros(x.shape) - offset = 0 - for cone, sz in cones: - sz = sz if isinstance(sz, (tuple, list)) else (sz,) - if sum(sz) == 0: - continue - for dim in sz: - if cone == PSD: - dim = vec_psd_dim(dim) - elif cone == EXP or cone == EXP_DUAL: - dim *= 3 - projection[offset:offset + dim] = _proj( - x[offset:offset + dim], cone, dual=dual) - offset += dim - return projection + + cone_list_cpp = parse_cone_dict_cpp(cones) + return projection(x, cone_list_cpp, dual)