Add ExtractPrincipalSubmatrix and ToLowerTriangularColumnsFromMatrix …

…utility functions (RobotLocomotion#20646)

bindings/pydrake/math/math_py_monolith.cc

@@ -503,7 +503,21 @@ void DoScalarIndependentDefinitions(py::module m) {
                 lower_triangular_columns);
           },
           py::arg("lower_triangular_columns"),
-          doc.ToSymmetricMatrixFromLowerTriangularColumns.doc_dynamic_size);
+          doc.ToSymmetricMatrixFromLowerTriangularColumns.doc_dynamic_size)
+      .def(
+          "ToLowerTriangularColumnsFromMatrix",
+          [](const Eigen::Ref<const MatrixX<T>>& matrix) {
+            return ToLowerTriangularColumnsFromMatrix(matrix);
+          },
+          py::arg("matrix"), doc.ToLowerTriangularColumnsFromMatrix.doc)
+      .def(
+          "ExtractPrincipalSubmatrix",
+          [](const Eigen::Ref<const MatrixX<T>>& matrix,
+              const std::set<int>& indices) {
+            return ExtractPrincipalSubmatrix(matrix, indices);
+          },
+          py::arg("matrix"), py::arg("indices"),
+          doc.ExtractPrincipalSubmatrix.doc);
 
   // Quadratic Form.
   m  // BR

bindings/pydrake/math/test/math_test.py

@@ -532,6 +532,17 @@ def test_matrix_util(self):
         np.testing.assert_array_equal(
             symmetric_mat, np.array([[1, 2, 3], [2, 4, 5], [3, 5, 6]]))
 
+        lower_triangular2 = mut.ToLowerTriangularColumnsFromMatrix(
+            matrix=symmetric_mat)
+        np.testing.assert_array_equal(lower_triangular, lower_triangular2)
+
+        minor_indices = {0, 2}
+        minor = mut.ExtractPrincipalSubmatrix(matrix=symmetric_mat,
+                                              indices=minor_indices)
+        np.testing.assert_array_equal(
+            minor,
+            symmetric_mat[np.ix_(list(minor_indices), list(minor_indices))])
+
     def test_quadratic_form(self):
         Q = np.diag([1., 2., 3.])
         X = mut.DecomposePSDmatrixIntoXtransposeTimesX(Q, 1e-8)

math/matrix_util.h

@@ -3,6 +3,8 @@
 #include <algorithm>
 #include <cmath>
 #include <functional>
+#include <set>
+#include <utility>
 #include <vector>
 
 #include <Eigen/Dense>
@@ -122,6 +124,25 @@ ToSymmetricMatrixFromLowerTriangularColumns(
   return symmetric_matrix;
 }
 
+/// Given a square matrix, extract the lower triangular part as a stacked column
+/// vector. This is a particularly useful operation when vectorizing symmetric
+/// matrices.
+template <typename Derived>
+drake::VectorX<typename Derived::Scalar> ToLowerTriangularColumnsFromMatrix(
+    const Eigen::MatrixBase<Derived>& matrix) {
+  DRAKE_ASSERT(matrix.rows() == matrix.cols());
+  const int num_rows = (matrix.rows() * (matrix.rows() + 1)) / 2;
+  drake::VectorX<typename Derived::Scalar> stacked_lower_triangular(num_rows);
+  int row_count = 0;
+  for (int c = 0; c < matrix.cols(); ++c) {
+    const int num_elts = matrix.cols() - c;
+    stacked_lower_triangular.segment(row_count, num_elts) =
+        matrix.col(c).tail(num_elts);
+    row_count += num_elts;
+  }
+  return stacked_lower_triangular;
+}
+
 /// Checks if a matrix is symmetric (with tolerance @p symmetry_tolerance --
 /// @see IsSymmetric) and has all eigenvalues greater than @p
 /// eigenvalue_tolerance.  @p eigenvalue_tolerance must be >= 0 -- where 0
@@ -182,5 +203,48 @@ MatrixX<T> StdVectorToEigen(const std::vector<MatrixX<T>>& vec) {
   return mat;
 }
 
+/// Extracts the principal submatrix from the ordered set of indices. The
+/// indices must be in monotonically increasing order. This method makes
+/// no assumptions about the symmetry of the matrix, nor that the matrix is
+/// square. However, all indices must be valid for both rows and columns.
+template <typename Derived>
+MatrixX<typename Derived::Scalar> ExtractPrincipalSubmatrix(
+    const Eigen::MatrixBase<Derived>& mat, const std::set<int>& indices) {
+  // Stores the contiguous intervals of the index set of the minor. These
+  // intervals include the first index but exclude the last, i.e.
+  // [intervals[i][0], intervals[i][1]).
+  std::vector<std::pair<int, int>> intervals;
+  int interval_start{*indices.begin()};
+  int last_idx{*indices.begin()};
+  DRAKE_ASSERT(last_idx >= 0);
+  for (const int& i : indices) {
+    DRAKE_ASSERT(i < mat.rows() && i < mat.cols());
+    if (i - last_idx > 1) {
+      intervals.emplace_back(interval_start, last_idx + 1);
+      interval_start = i;
+    }
+    last_idx = i;
+  }
+  intervals.emplace_back(last_idx, *indices.rbegin() + 1);
+
+  MatrixX<typename Derived::Scalar> minor(indices.size(), indices.size());
+  int minor_row_count = 0;
+  for (auto row_it = intervals.begin(); row_it != intervals.cend(); ++row_it) {
+    int minor_col_count = 0;
+    const int cur_block_row_size = row_it->second - row_it->first;
+    for (auto col_it = intervals.begin(); col_it != intervals.cend();
+         ++col_it) {
+      const int cur_block_col_size = col_it->second - col_it->first;
+      minor.block(minor_row_count, minor_col_count, cur_block_row_size,
+                  cur_block_col_size) =
+          mat.block(row_it->first, col_it->first, cur_block_row_size,
+                    cur_block_col_size);
+      minor_col_count += cur_block_col_size;
+    }
+    minor_row_count += cur_block_row_size;
+  }
+  return minor;
+}
+
 }  // namespace math
 }  // namespace drake

math/test/matrix_util_test.cc

@@ -68,6 +68,30 @@ GTEST_TEST(TestMatrixUtil, TestToSymmetricMatrixFromLowerTriangularColumns) {
       CompareMatrices(ToSymmetricMatrixFromLowerTriangularColumns(x2), X2));
 }
 
+GTEST_TEST(TestMatrixUtil, TestToLowerTriangularColumnFromMatrix) {
+  // Tests a static size vector.
+  Eigen::Vector3d x1(1, 2, 3);
+  Eigen::Matrix2d X1;
+  // clang-format off
+  X1 << 1, 2,
+        2, 3;
+  // clang-format on
+  auto x1_result = ToLowerTriangularColumnsFromMatrix(X1);
+  EIGEN_STATIC_ASSERT_SAME_MATRIX_SIZE(decltype(x1_result), Eigen::Vector3d);
+  EXPECT_TRUE(CompareMatrices(x1, x1_result));
+
+  // Tests a dynamic size vector.
+  Eigen::VectorXd x2(6);
+  x2 << 1, 2, 3, 4, 5, 6;
+  Eigen::Matrix3d X2;
+  // clang-format off
+  X2 << 1, 7, 8,
+        2, 4, 9,
+        3, 5, 6;
+  // clang-format on
+  EXPECT_TRUE(CompareMatrices(ToLowerTriangularColumnsFromMatrix(X2), x2));
+}
+
 GTEST_TEST(TestMatrixUtil, IsPositiveDefiniteTest) {
   Eigen::Matrix3d A;
   A << 1, 2, 4, 2, 3, 5, 4, 5, 6;
@@ -104,6 +128,62 @@ GTEST_TEST(TestMatrixUtil, StdVector) {
       3);
 }
 
+GTEST_TEST(TestMatrixUtil, ExtractPrincipalSubmatrixSquareMatrix) {
+  int n = 8;
+  drake::MatrixX<symbolic::Variable> X(n, n);
+  for (int i = 0; i < n; ++i) {
+    for (int j = 0; j < n; ++j) {
+      X(i, j) = symbolic::Variable(fmt::format("X({},{})", i, j));
+    }
+  }
+  // A non-contiguous set of indices
+  const std::set<int> submatrix_indices{0, 1, 3, 4, 5, 7};
+  drake::MatrixX<symbolic::Variable> submatrix_manual(submatrix_indices.size(),
+                                                      submatrix_indices.size());
+  // clang-format off
+  submatrix_manual << X(0, 0), X(0, 1), X(0, 3), X(0, 4), X(0, 5), X(0, 7),
+                      X(1, 0), X(1, 1), X(1, 3), X(1, 4), X(1, 5), X(1, 7),
+                      X(3, 0), X(3, 1), X(3, 3), X(3, 4), X(3, 5), X(3, 7),
+                      X(4, 0), X(4, 1), X(4, 3), X(4, 4), X(4, 5), X(4, 7),
+                      X(5, 0), X(5, 1), X(5, 3), X(5, 4), X(5, 5), X(5, 7),
+                      X(7, 0), X(7, 1), X(7, 3), X(7, 4), X(7, 5), X(7, 7);
+  // clang-format on
+
+  auto submatrix = ExtractPrincipalSubmatrix(X, submatrix_indices);
+
+  EIGEN_STATIC_ASSERT_SAME_MATRIX_SIZE(decltype(submatrix),
+                                       drake::MatrixX<symbolic::Variable>);
+  EXPECT_EQ(submatrix.rows(), submatrix.cols());
+  EXPECT_EQ(static_cast<int>(submatrix.rows()),
+            static_cast<int>(submatrix_indices.size()));
+  for (int r = 0; r < submatrix.rows(); ++r) {
+    for (int c = 0; c < submatrix.cols(); ++c) {
+      EXPECT_TRUE(submatrix(r, c).equal_to(submatrix_manual(r, c)));
+    }
+  }
+
+  const std::set<int> submatrix_indices2{2, 3, 5, 7};
+  drake::MatrixX<symbolic::Variable> submatrix_manual2(
+      submatrix_indices2.size(), submatrix_indices2.size());
+  auto submatrix2 = ExtractPrincipalSubmatrix(X, submatrix_indices2);
+  // clang-format off
+  submatrix_manual2 << X(2, 2), X(2, 3), X(2, 5), X(2, 7),
+                       X(3, 2), X(3, 3), X(3, 5), X(3, 7),
+                       X(5, 2), X(5, 3), X(5, 5), X(5, 7),
+                       X(7, 2), X(7, 3), X(7, 5), X(7, 7);
+  // clang-format on
+  EIGEN_STATIC_ASSERT_SAME_MATRIX_SIZE(decltype(submatrix2),
+                                       drake::MatrixX<symbolic::Variable>);
+  EXPECT_EQ(submatrix2.rows(), submatrix2.cols());
+  EXPECT_EQ(static_cast<int>(submatrix2.rows()),
+            static_cast<int>(submatrix_indices2.size()));
+  for (int r = 0; r < submatrix2.rows(); ++r) {
+    for (int c = 0; c < submatrix2.cols(); ++c) {
+      EXPECT_TRUE(submatrix2(r, c).equal_to(submatrix_manual2(r, c)));
+    }
+  }
+}
+
 }  // namespace test
 }  // namespace math
 }  // namespace drake

solvers/mathematical_program.cc

@@ -1146,50 +1146,14 @@ MathematicalProgram::AddPositiveSemidefiniteConstraint(
   return AddConstraint(constraint, symmetric_matrix_var);
 }
 
-namespace {
-
-template <typename T>
-// Extract the principal submatrix from the ordered set of minor_indices. The
-// minor_indices must be in monotonically increasing order. This method makes no
-// assumptions about the symmetry of the matrix, nor that the matrix is square.
-// TODO(hongkai.dai)  move this to matrix_utils.h
-MatrixX<T> MakePrincipalSubmatrix(const Eigen::Ref<const MatrixX<T>>& mat,
-                                  const std::set<int>& minor_indices) {
-  // In Debug builds, check if the minor_indices are valid.
-  if (kDrakeAssertIsArmed) {
-    auto elt_is_in_bounds = [&mat](int elt) {
-      return elt >= 0 && elt < mat.rows() && elt < mat.cols();
-    };
-    DRAKE_ASSERT(std::all_of(minor_indices.begin(), minor_indices.end(),
-                             elt_is_in_bounds));
-  }
-
-  MatrixX<T> minor(minor_indices.size(), minor_indices.size());
-  int row_count = 0;
-  for (auto row_it = minor_indices.begin(); row_it != minor_indices.cend();
-       ++row_it) {
-    minor(row_count, row_count) = mat(*row_it, *row_it);
-    int col_count = row_count + 1;
-    for (auto col_it = next(row_it); col_it != minor_indices.cend(); ++col_it) {
-      minor(row_count, col_count) = mat(*row_it, *col_it);
-      minor(col_count, row_count) = mat(*col_it, *row_it);
-      ++col_count;
-    }
-    ++row_count;
-  }
-  return minor;
-}
-}  // namespace
-
 Binding<PositiveSemidefiniteConstraint>
 MathematicalProgram::AddPrincipalSubmatrixIsPsdConstraint(
     const Eigen::Ref<const MatrixXDecisionVariable>& symmetric_matrix_var,
     const std::set<int>& minor_indices) {
   // This function relies on AddPositiveSemidefiniteConstraint to validate the
   // documented symmetry prerequisite.
   return AddPositiveSemidefiniteConstraint(
-      MakePrincipalSubmatrix<symbolic::Variable>(symmetric_matrix_var,
-                                                 minor_indices));
+      math::ExtractPrincipalSubmatrix(symmetric_matrix_var, minor_indices));
 }
 
 Binding<PositiveSemidefiniteConstraint>
@@ -1199,7 +1163,7 @@ MathematicalProgram::AddPrincipalSubmatrixIsPsdConstraint(
   // This function relies on AddPositiveSemidefiniteConstraint to validate the
   // documented symmetry prerequisite.
   return AddPositiveSemidefiniteConstraint(
-      MakePrincipalSubmatrix<symbolic::Expression>(e, minor_indices));
+      math::ExtractPrincipalSubmatrix(e, minor_indices));
 }
 
 Binding<LinearMatrixInequalityConstraint> MathematicalProgram::AddConstraint(