Extend MatrixBlock::LeftMultiplyByBlockDiagonal() (RobotLocomotion#20312

)

When the MatrixBlock is Block3x3SparseMatrix, we now allow the blocks in
the block diagonal to be of size 3n-by-3n instead of restricted to be
3x3.

multibody/contact_solvers/matrix_block.cc

@@ -108,47 +108,98 @@ MatrixBlock<T> MatrixBlock<T>::LeftMultiplyByBlockDiagonal(
   DRAKE_DEMAND(end >= start);
   DRAKE_DEMAND(static_cast<int>(Gs.size()) > end);
   /* Verify that the sizes of G and M is compatible. */
-  int G_rows = 0;
+  std::vector<int>
+      row_starts;  // starting row index for each diagonal block in G
+  row_starts.reserve(end - start + 1);
+  int row = 0;
   for (int i = start; i <= end; ++i) {
     DRAKE_DEMAND(Gs[i].rows() == Gs[i].cols());
-    G_rows += Gs[i].rows();
+    row_starts.emplace_back(row);
+    row += Gs[i].rows();
     if (!is_dense_) {
-      DRAKE_DEMAND(Gs[i].rows() == 3);
+      DRAKE_DEMAND(Gs[i].rows() % 3 == 0);
     }
   }
-  DRAKE_DEMAND(G_rows == rows());
+  DRAKE_DEMAND(row == rows());
 
   if (is_dense_) {
     const MatrixX<T>& M_dense = std::get<MatrixX<T>>(data_);
     MatrixX<T> GM(rows(), cols());
-    int row_offset = 0;
     for (int i = start; i <= end; ++i) {
       const int rows = Gs[i].rows();
-      GM.middleRows(row_offset, rows).noalias() =
-          Gs[i] * M_dense.middleRows(row_offset, rows);
-      row_offset += rows;
+      const int row_start = row_starts[i - start];
+      GM.middleRows(row_start, rows).noalias() =
+          Gs[i] * M_dense.middleRows(row_start, rows);
     }
     return MatrixBlock<T>(std::move(GM));
   }
 
   const Block3x3SparseMatrix<T>& M_sparse =
       std::get<Block3x3SparseMatrix<T>>(data_);
-  Block3x3SparseMatrix<T> GM(M_sparse.block_rows(), M_sparse.block_cols());
+  /* At this point, we know that M is block 3x3 sparse and G is block diagonal,
+   G is block diagonal, with block h being 3nₕ x 3nₕ for positive integers nₕ.
+   Therefore, all quantities in the computation of A = G * M can be represented
+   as 3x3 blocks. We loop over all 3x3 non-zero blocks of M, and for each Mₖⱼ
+   (the k,j-th 3x3 block of M), we find all non-zero 3x3 blocks Gᵢₖ that
+   multiplies with Mₖⱼ, and add their product to the resulting 3x3 block, Aᵢⱼ.
+
+   Because G is block diagonal, to achieve the above, we only need:
+   (1) The index h into `Gs` to obtain the correct diagonal block of G that
+       multiplies with Mₖⱼ.
+   (2) The block row indices (i.e. the i's) of these non-zero 3x3 G subblocks.
+
+   Note that, for (2), the block row indices are contiguous, so we only need the
+   starting block row index and the size of that G block to determine all the
+   block row indices. */
+  std::vector<int> G_indices(
+      M_sparse.block_rows());  // `G_indices[k]` gives the relevant index h into
+                               // `Gs` for block column k.
+  std::vector<int> i_start(
+      M_sparse.block_rows());  // `i_start[k]` gives the starting block row
+                               // index from (2) above.
+  {
+    int i = 0;
+    int h = start;
+    for (int k = 0; k < M_sparse.block_rows(); ++k) {
+      if (3 * k >= row_starts[h - start] + Gs[h].rows()) {
+        i += Gs[h].cols() / 3;
+        ++h;
+      }
+      i_start[k] = i;
+      G_indices[k] = h;
+    }
+  }
+  /* We also record, for each k, the local block column index, l, such that the
+   `Gs[G_indices[k]].col(l)` gives the (non-zero entries of) k-th column of G.
+   */
+  std::vector<int> local_block_cols(M_sparse.block_rows());
+  for (int k = 0; k < M_sparse.block_rows(); ++k) {
+    local_block_cols[k] = k - i_start[k];
+  }
+
+  Block3x3SparseMatrix<T> A(M_sparse.block_rows(), M_sparse.block_cols());
   using Triplet = typename Block3x3SparseMatrix<T>::Triplet;
   const std::vector<std::vector<Triplet>>& M_triplets = M_sparse.get_triplets();
-  std::vector<Triplet> GM_triplets;
-  GM_triplets.reserve(M_sparse.num_blocks());
+  std::vector<Triplet> A_triplets;
+  A_triplets.reserve(M_sparse.num_blocks());
   for (const auto& row_data : M_triplets) {
     for (const Triplet& t : row_data) {
-      const int block_row = std::get<0>(t);
-      const int block_col = std::get<1>(t);
+      const int k = std::get<0>(t);
+      const int j = std::get<1>(t);
       const Matrix3<T>& M_block = std::get<2>(t);
-      GM_triplets.emplace_back(block_row, block_col,
-                               Gs[start + block_row] * M_block);
+      const int h = G_indices[k];
+      /* Given k and j, perform Aᵢⱼ+= Gᵢₖ * Mₖⱼ for all relevant i's. */
+      for (int l = 0; 3 * l < Gs[h].cols(); ++l) {
+        const int i = i_start[k] + l;
+        A_triplets.emplace_back(
+            i, j,
+            Gs[h].template block<3, 3>(3 * l, 3 * local_block_cols[k]) *
+                M_block);
+      }
     }
   }
-  GM.SetFromTriplets(GM_triplets);
-  return MatrixBlock<T>(std::move(GM));
+  A.SetFromTriplets(A_triplets);
+  return MatrixBlock<T>(std::move(A));
 }
 
 template <class T>

multibody/contact_solvers/matrix_block.h

@@ -87,7 +87,8 @@ class MatrixBlock {
    @pre 0 <= start <= end < Gs.size().
    @pre All matrices in Gs are square.
    @pre Gs[start].rows() + ... + Gs[end].rows() == M.rows().
-   @pre If M is sparse, all matrices in Gs are 3x3. */
+   @pre If M is sparse, all matrices in Gs are 3n-by-3n for some positive
+   integer n. */
   MatrixBlock<T> LeftMultiplyByBlockDiagonal(const std::vector<MatrixX<T>>& Gs,
                                              int start, int end) const;
 

multibody/contact_solvers/test/matrix_block_test.cc

@@ -132,6 +132,44 @@ GTEST_TEST(MatrixBlockTest, LeftMultiplyByBlockDiagonal) {
   EXPECT_TRUE(CompareMatrices(dense_result.MakeDenseMatrix(), expected));
 }
 
+GTEST_TEST(MatrixBlockTest, LeftMultiplyByBlockDiagonalWithNon3x3GBlocks) {
+  Block3x3SparseMatrix<double> sparse_matrix = MakeBlockSparseMatrix();
+  const MatrixBlock<double> sparse_block(std::move(sparse_matrix));
+
+  /* Here we construct a block diagonal G with 7 square blocks of sizes
+   3 3 3 [3 6 3] 3. We'll use only the three blocks marked with brackets to
+   multiply with the sparse matrix which has 12 rows. */
+  MatrixXd matrix6(6, 6);
+  for (int i = 0; i < 6; ++i) {
+    for (int j = 0; j < 6; ++j) {
+      matrix6(i, j) = 0.123 * i + 0.456 * j;
+    }
+  }
+  const int num_Gs = 7;
+  std::vector<MatrixXd> Gs;
+  for (int i = 0; i < num_Gs; ++i) {
+    if (i != 4) {
+      Gs.emplace_back(Matrix3d::Constant(3.14 * i));
+    } else {
+      Gs.emplace_back(matrix6);
+    }
+  }
+  const int start = 3;
+  const int end = 5;
+  const MatrixBlock<double> result =
+      sparse_block.LeftMultiplyByBlockDiagonal(Gs, start, end);
+
+  /* Compute the expected results using dense matrices. */
+  MatrixXd dense_G = MatrixXd::Zero(12, 12);
+  dense_G.topLeftCorner(3, 3) = Gs[3];
+  dense_G.block<6, 6>(3, 3) = Gs[4];
+  dense_G.bottomRightCorner(3, 3) = Gs[5];
+
+  MatrixXd expected = dense_G * sparse_block.MakeDenseMatrix();
+  EXPECT_TRUE(CompareMatrices(result.MakeDenseMatrix(), expected,
+                              16 * std::numeric_limits<double>::epsilon()));
+}
+
 GTEST_TEST(MatrixBlockTest, MultiplyWithScaledTransposeAndAddTo) {
   Block3x3SparseMatrix<double> sparse_matrix = MakeBlockSparseMatrix();
   const MatrixXd dense_matrix = sparse_matrix.MakeDenseMatrix();

multibody/plant/deformable_driver.cc

@@ -501,10 +501,7 @@ void DeformableDriver<T>::AppendDeformableRigidFixedConstraintKinematics(
         p_WQs.template segment<3>(3 * v) = X_WB * spec.p_BQs[v].cast<T>();
       }
       negative_Jv_v_WAp.SetFromTriplets(jacobian_triplets);
-      // TODO(xuchenhan-tri): Use the sparse Jacobian when we support
-      // multiplication of sparse jacobians with weight matrices that are not
-      // 3x3 in size. See MatrixBlock::LeftMultiplyByBlockDiagonal().
-      MatrixBlock<T> jacobian_block_A(negative_Jv_v_WAp.MakeDenseMatrix());
+      MatrixBlock<T> jacobian_block_A(std::move(negative_Jv_v_WAp));
 
       /* Positions of fixed vertices of the deformable body in the deformable
        body's frame which is always assumed to be the world frame. */