Refactor BlockJacobiPreconditioner

1. Rename BlockJacobiPreconditionet to BlockSparseJacobiPreconditioner.
2. Add CompressedRowSparseJacobiPreconditioner which is a
   block Jacobi preconditioner for CompressedRowSparseMatrix objects.
3. Re-write the tests to be more comprehensive.

Change-Id: Icbc91f9ad2cefaad593c11397f8cdcf805d7e118

internal/ceres/block_jacobi_preconditioner.cc

@@ -30,6 +30,7 @@
 
 #include "ceres/block_jacobi_preconditioner.h"
 
+#include "Eigen/Dense"
 #include "ceres/block_random_access_diagonal_matrix.h"
 #include "ceres/block_sparse_matrix.h"
 #include "ceres/block_structure.h"
@@ -38,7 +39,7 @@
 
 namespace ceres::internal {
 
-BlockJacobiPreconditioner::BlockJacobiPreconditioner(
+BlockSparseJacobiPreconditioner::BlockSparseJacobiPreconditioner(
     const BlockSparseMatrix& A) {
   const CompressedRowBlockStructure* bs = A.block_structure();
   std::vector<int> blocks(bs->cols.size());
@@ -49,10 +50,10 @@ BlockJacobiPreconditioner::BlockJacobiPreconditioner(
   m_ = std::make_unique<BlockRandomAccessDiagonalMatrix>(blocks);
 }
 
-BlockJacobiPreconditioner::~BlockJacobiPreconditioner() = default;
+BlockSparseJacobiPreconditioner::~BlockSparseJacobiPreconditioner() = default;
 
-bool BlockJacobiPreconditioner::UpdateImpl(const BlockSparseMatrix& A,
-                                           const double* D) {
+bool BlockSparseJacobiPreconditioner::UpdateImpl(const BlockSparseMatrix& A,
+                                                 const double* D) {
   const CompressedRowBlockStructure* bs = A.block_structure();
   const double* values = A.values();
   m_->SetZero();
@@ -90,9 +91,91 @@ bool BlockJacobiPreconditioner::UpdateImpl(const BlockSparseMatrix& A,
   return true;
 }
 
-void BlockJacobiPreconditioner::RightMultiplyAndAccumulate(const double* x,
-                                                           double* y) const {
-  m_->RightMultiplyAndAccumulate(x, y);
+BlockCRSJacobiPreconditioner::BlockCRSJacobiPreconditioner(
+    const CompressedRowSparseMatrix& A) {
+  auto& col_blocks = A.col_blocks();
+
+  // Compute the number of non-zeros in the preconditioner. This is needed so
+  // that we can construct the CompressedRowSparseMatrix.
+  int m_nnz = 0;
+  for (int col_block_size : col_blocks) {
+    m_nnz += col_block_size * col_block_size;
+  }
+
+  m_ = std::make_unique<CompressedRowSparseMatrix>(
+      A.num_cols(), A.num_cols(), m_nnz);
+
+  const int num_col_blocks = col_blocks.size();
+
+  // Populate the sparsity structure of the preconditioner matrix.
+  int* m_cols = m_->mutable_cols();
+  int* m_rows = m_->mutable_rows();
+  m_rows[0] = 0;
+  for (int i = 0, col = 0, idx = 0; i < num_col_blocks; ++i) {
+    // For each column block populate a diagonal block in the preconditioner.
+    // Not that the because of the way the CompressedRowSparseMatrix format
+    // works, the entire diagonal block is laid out contiguously in memory as a
+    // row-major matrix. We will use this when updating the block.
+    const int col_block_size = col_blocks[i];
+    for (int j = 0; j < col_block_size; ++j) {
+      for (int k = 0; k < col_block_size; ++k, ++idx) {
+        m_cols[idx] = col + k;
+      }
+      m_rows[col + j + 1] = idx;
+    }
+    col += col_block_size;
+  }
+
+  CHECK_EQ(m_rows[A.num_cols()], m_nnz);
+}
+
+BlockCRSJacobiPreconditioner::~BlockCRSJacobiPreconditioner() = default;
+
+bool BlockCRSJacobiPreconditioner::UpdateImpl(
+    const CompressedRowSparseMatrix& A, const double* D) {
+  const auto& col_blocks = A.col_blocks();
+  const int num_col_blocks = col_blocks.size();
+
+  const int* a_rows = A.rows();
+  const int* a_cols = A.cols();
+  const double* a_values = A.values();
+
+  m_->SetZero();
+  double* m_values = m_->mutable_values();
+  const int* m_rows = m_->rows();
+
+  const int num_rows = A.num_rows();
+  // The following loop can likely be optimized by exploiting the fact that each
+  // row block has exactly the same sparsity structure.
+  for (int r = 0; r < num_rows; ++r) {
+    int idx = a_rows[r];
+    while (idx < a_rows[r + 1]) {
+      const int col = a_cols[idx];
+      const int col_block_size = m_rows[col + 1] - m_rows[col];
+      // We make use of the fact that the entire diagonal block is stored
+      // contiguously in memory as a row-major matrix.
+      MatrixRef m(m_values + m_rows[col], col_block_size, col_block_size);
+      ConstVectorRef b(a_values + idx, col_block_size);
+      m.selfadjointView<Eigen::Upper>().rankUpdate(b);
+      idx += col_block_size;
+    }
+  }
+
+  for (int i = 0, col = 0; i < num_col_blocks; ++i) {
+    const int col_block_size = m_rows[col + 1] - m_rows[col];
+    MatrixRef m(m_values + m_rows[col], col_block_size, col_block_size);
+
+    if (D != nullptr) {
+      m.diagonal() +=
+          ConstVectorRef(D + col, col_block_size).array().square().matrix();
+    }
+
+    m = m.selfadjointView<Eigen::Upper>().llt().solve(
+        Matrix::Identity(col_block_size, col_block_size));
+    col += col_block_size;
+  }
+
+  return true;
 }
 
 }  // namespace ceres::internal

internal/ceres/block_jacobi_preconditioner.h

@@ -41,30 +41,22 @@
 namespace ceres::internal {
 
 class BlockSparseMatrix;
-struct CompressedRowBlockStructure;
+class CompressedRowSparseMatrix;
 
 // A block Jacobi preconditioner. This is intended for use with
-// conjugate gradients, or other iterative symmetric solvers. To use
-// the preconditioner, create one by passing a BlockSparseMatrix "A"
-// to the constructor. This fixes the sparsity pattern to the pattern
-// of the matrix A^TA.
-//
-// Before each use of the preconditioner in a solve with conjugate gradients,
-// update the matrix by running Update(A, D). The values of the matrix A are
-// inspected to construct the preconditioner. The vector D is applied as the
-// D^TD diagonal term.
-class CERES_NO_EXPORT BlockJacobiPreconditioner
+// conjugate gradients, or other iterative symmetric solvers.
+
+// This version of the preconditioner is for use with BlockSparseMatrix
+// Jacobians.
+class CERES_NO_EXPORT BlockSparseJacobiPreconditioner
     : public BlockSparseMatrixPreconditioner {
  public:
   // A must remain valid while the BlockJacobiPreconditioner is.
-  explicit BlockJacobiPreconditioner(const BlockSparseMatrix& A);
-  BlockJacobiPreconditioner(const BlockJacobiPreconditioner&) = delete;
-  void operator=(const BlockJacobiPreconditioner&) = delete;
-
-  ~BlockJacobiPreconditioner() override;
-
-  // Preconditioner interface
-  void RightMultiplyAndAccumulate(const double* x, double* y) const final;
+  explicit BlockSparseJacobiPreconditioner(const BlockSparseMatrix& A);
+  ~BlockSparseJacobiPreconditioner() override;
+  void RightMultiplyAndAccumulate(const double* x, double* y) const final {
+    return m_->RightMultiplyAndAccumulate(x, y);
+  }
   int num_rows() const final { return m_->num_rows(); }
   int num_cols() const final { return m_->num_rows(); }
   const BlockRandomAccessDiagonalMatrix& matrix() const { return *m_; }
@@ -75,6 +67,27 @@ class CERES_NO_EXPORT BlockJacobiPreconditioner
   std::unique_ptr<BlockRandomAccessDiagonalMatrix> m_;
 };
 
+// This version of the preconditioner is for use with CompressedRowSparseMatrix
+// Jacobians.
+class CERES_NO_EXPORT BlockCRSJacobiPreconditioner
+    : public CompressedRowSparseMatrixPreconditioner {
+ public:
+  // A must remain valid while the BlockJacobiPreconditioner is.
+  explicit BlockCRSJacobiPreconditioner(const CompressedRowSparseMatrix& A);
+  ~BlockCRSJacobiPreconditioner() override;
+  void RightMultiplyAndAccumulate(const double* x, double* y) const final {
+    m_->RightMultiplyAndAccumulate(x, y);
+  }
+  int num_rows() const final { return m_->num_rows(); }
+  int num_cols() const final { return m_->num_rows(); }
+  const CompressedRowSparseMatrix& matrix() const { return *m_; }
+
+ private:
+  bool UpdateImpl(const CompressedRowSparseMatrix& A, const double* D) final;
+
+  std::unique_ptr<CompressedRowSparseMatrix> m_;
+};
+
 }  // namespace ceres::internal
 
 #include "ceres/internal/reenable_warnings.h"

internal/ceres/block_jacobi_preconditioner_test.cc

@@ -41,58 +41,108 @@
 
 namespace ceres::internal {
 
-class BlockJacobiPreconditionerTest : public ::testing::Test {
- protected:
-  void SetUpFromProblemId(int problem_id) {
-    std::unique_ptr<LinearLeastSquaresProblem> problem =
-        CreateLinearLeastSquaresProblemFromId(problem_id);
+TEST(BlockSparseJacobiPreconditioner, _) {
+  constexpr int kNumtrials = 10;
+  BlockSparseMatrix::RandomMatrixOptions options;
+  options.num_col_blocks = 3;
+  options.min_col_block_size = 1;
+  options.max_col_block_size = 3;
 
-    CHECK(problem != nullptr);
-    A.reset(down_cast<BlockSparseMatrix*>(problem->A.release()));
-    D = std::move(problem->D);
+  options.num_row_blocks = 5;
+  options.min_row_block_size = 1;
+  options.max_row_block_size = 4;
+  options.block_density = 0.25;
+  std::mt19937 prng;
 
-    Matrix dense_a;
-    A->ToDenseMatrix(&dense_a);
-    dense_ata = dense_a.transpose() * dense_a;
-    dense_ata += VectorRef(D.get(), A->num_cols())
-                     .array()
-                     .square()
-                     .matrix()
-                     .asDiagonal();
-  }
+  for (int trial = 0; trial < kNumtrials; ++trial) {
+    auto jacobian = BlockSparseMatrix::CreateRandomMatrix(options, prng);
+    Vector diagonal = Vector::Ones(jacobian->num_cols());
+    Matrix dense_jacobian;
+    jacobian->ToDenseMatrix(&dense_jacobian);
+    Matrix hessian = dense_jacobian.transpose() * dense_jacobian;
+    hessian.diagonal() += diagonal.array().square().matrix();
 
-  void VerifyDiagonalBlocks(const int problem_id) {
-    SetUpFromProblemId(problem_id);
+    BlockSparseJacobiPreconditioner pre(*jacobian);
+    pre.Update(*jacobian, diagonal.data());
 
-    BlockJacobiPreconditioner pre(*A);
-    pre.Update(*A, D.get());
+    // The const_cast is needed to be able to call GetCell.
     auto* m = const_cast<BlockRandomAccessDiagonalMatrix*>(&pre.matrix());
-    EXPECT_EQ(m->num_rows(), A->num_cols());
-    EXPECT_EQ(m->num_cols(), A->num_cols());
+    EXPECT_EQ(m->num_rows(), jacobian->num_cols());
+    EXPECT_EQ(m->num_cols(), jacobian->num_cols());
 
-    const CompressedRowBlockStructure* bs = A->block_structure();
+    const CompressedRowBlockStructure* bs = jacobian->block_structure();
     for (int i = 0; i < bs->cols.size(); ++i) {
       const int block_size = bs->cols[i].size;
       int r, c, row_stride, col_stride;
       CellInfo* cell_info = m->GetCell(i, i, &r, &c, &row_stride, &col_stride);
-      MatrixRef m(cell_info->values, row_stride, col_stride);
-      Matrix actual_block_inverse = m.block(r, c, block_size, block_size);
-      Matrix expected_block = dense_ata.block(
+      Matrix actual_block_inverse =
+          MatrixRef(cell_info->values, row_stride, col_stride)
+              .block(r, c, block_size, block_size);
+      Matrix expected_block = hessian.block(
           bs->cols[i].position, bs->cols[i].position, block_size, block_size);
       const double residual = (actual_block_inverse * expected_block -
                                Matrix::Identity(block_size, block_size))
                                   .norm();
       EXPECT_NEAR(residual, 0.0, 1e-12) << "Block: " << i;
     }
+    options.num_col_blocks++;
+    options.num_row_blocks++;
   }
+}
 
-  std::unique_ptr<BlockSparseMatrix> A;
-  std::unique_ptr<double[]> D;
-  Matrix dense_ata;
-};
+TEST(CompressedRowSparseJacobiPreconditioner, _) {
+  constexpr int kNumtrials = 10;
+  CompressedRowSparseMatrix::RandomMatrixOptions options;
+  options.num_col_blocks = 3;
+  options.min_col_block_size = 1;
+  options.max_col_block_size = 3;
 
-TEST_F(BlockJacobiPreconditionerTest, SmallProblem) { VerifyDiagonalBlocks(2); }
+  options.num_row_blocks = 5;
+  options.min_row_block_size = 1;
+  options.max_row_block_size = 4;
+  options.block_density = 0.25;
+  std::mt19937 prng;
 
-TEST_F(BlockJacobiPreconditionerTest, LargeProblem) { VerifyDiagonalBlocks(3); }
+  for (int trial = 0; trial < kNumtrials; ++trial) {
+    auto jacobian =
+        CompressedRowSparseMatrix::CreateRandomMatrix(options, prng);
+    Vector diagonal = Vector::Ones(jacobian->num_cols());
+
+    Matrix dense_jacobian;
+    jacobian->ToDenseMatrix(&dense_jacobian);
+    Matrix hessian = dense_jacobian.transpose() * dense_jacobian;
+    hessian.diagonal() += diagonal.array().square().matrix();
+
+    BlockCRSJacobiPreconditioner pre(*jacobian);
+    pre.Update(*jacobian, diagonal.data());
+    auto& m = pre.matrix();
+
+    EXPECT_EQ(m.num_rows(), jacobian->num_cols());
+    EXPECT_EQ(m.num_cols(), jacobian->num_cols());
+
+    const auto& col_blocks = jacobian->col_blocks();
+    for (int i = 0, col = 0; i < col_blocks.size(); ++i) {
+      const int block_size = col_blocks[i];
+      int idx = m.rows()[col];
+      for (int j = 0; j < block_size; ++j) {
+        EXPECT_EQ(m.rows()[col + j + 1] - m.rows()[col + j], block_size);
+        for (int k = 0; k < block_size; ++k, ++idx) {
+          EXPECT_EQ(m.cols()[idx], col + k);
+        }
+      }
+
+      ConstMatrixRef actual_block_inverse(
+          m.values() + m.rows()[col], block_size, block_size);
+      Matrix expected_block = hessian.block(col, col, block_size, block_size);
+      const double residual = (actual_block_inverse * expected_block -
+                               Matrix::Identity(block_size, block_size))
+                                  .norm();
+      EXPECT_NEAR(residual, 0.0, 1e-12) << "Block: " << i;
+      col += block_size;
+    }
+    options.num_col_blocks++;
+    options.num_row_blocks++;
+  }
+}
 
 }  // namespace ceres::internal

internal/ceres/cgnr_solver.cc

@@ -137,7 +137,7 @@ LinearSolver::Summary CgnrSolver::SolveImpl(
   EventLogger event_logger("CgnrSolver::Solve");
   if (!preconditioner_) {
     if (options_.preconditioner_type == JACOBI) {
-      preconditioner_ = std::make_unique<BlockJacobiPreconditioner>(*A);
+      preconditioner_ = std::make_unique<BlockSparseJacobiPreconditioner>(*A);
     } else if (options_.preconditioner_type == SUBSET) {
       Preconditioner::Options preconditioner_options;
       preconditioner_options.type = SUBSET;