From 352ede96e4c331daae4e1be9a5f3f50fff951b8d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Antonio=20S=C3=A1nchez?= <cantonios@google.com>
Date: Fri, 8 Mar 2024 19:18:10 +0000
Subject: [PATCH] Fix incomplete cholesky.

---
 .../IncompleteCholesky.h                      | 35 +++++++++++--------
 Eigen/src/SparseCore/SparseMatrix.h           | 31 +++++++++++++---
 test/incomplete_cholesky.cpp                  | 20 ++++++++++-
 3 files changed, 67 insertions(+), 19 deletions(-)

diff --git a/Eigen/src/IterativeLinearSolvers/IncompleteCholesky.h b/Eigen/src/IterativeLinearSolvers/IncompleteCholesky.h
index f456c96bd..a97b9054c 100644
--- a/Eigen/src/IterativeLinearSolvers/IncompleteCholesky.h
+++ b/Eigen/src/IterativeLinearSolvers/IncompleteCholesky.h
@@ -32,17 +32,19 @@ namespace Eigen {
  *
  * \implsparsesolverconcept
  *
- * It performs the following incomplete factorization: \f$ S P A P' S \approx L L' \f$
- * where L is a lower triangular factor, S is a diagonal scaling matrix, and P is a
- * fill-in reducing permutation as computed by the ordering method.
+ * It performs the following incomplete factorization: \f$ S P A P' S + \sigma I \approx L L' \f$
+ * where L is a lower triangular factor, S is a diagonal scaling matrix, P is a
+ * fill-in reducing permutation as computed by the ordering method, and \f$ \sigma \f$ is a shift
+ * for ensuring the decomposed matrix is positive definite.
  *
  * \b Shifting \b strategy: Let \f$ B = S P A P' S \f$  be the scaled matrix on which the factorization is carried out,
  * and \f$ \beta \f$ be the minimum value of the diagonal. If \f$ \beta > 0 \f$ then, the factorization is directly
- * performed on the matrix B. Otherwise, the factorization is performed on the shifted matrix \f$ B + (\sigma+|\beta| I
- * \f$ where \f$ \sigma \f$ is the initial shift value as returned and set by setInitialShift() method. The default
- * value is \f$ \sigma = 10^{-3} \f$. If the factorization fails, then the shift in doubled until it succeed or a
- * maximum of ten attempts. If it still fails, as returned by the info() method, then you can either increase the
- * initial shift, or better use another preconditioning technique.
+ * performed on the matrix B, and \sigma = 0. Otherwise, the factorization is performed on the shifted matrix \f$ B +
+ * \sigma I \f$ for a shifting factor  \f$ \sigma \f$.  We start with \f$ \sigma = \sigma_0 - \beta \f$, where \f$
+ * \sigma_0 \f$ is the initial shift value as returned and set by setInitialShift() method. The default value is \f$
+ * \sigma_0 = 10^{-3} \f$. If the factorization fails, then the shift in doubled until it succeed or a maximum of ten
+ * attempts. If it still fails, as returned by the info() method, then you can either increase the initial shift, or
+ * better use another preconditioning technique.
  *
  */
 template <typename Scalar, int UpLo_ = Lower, typename OrderingType_ = AMDOrdering<int> >
@@ -176,6 +178,9 @@ class IncompleteCholesky : public SparseSolverBase<IncompleteCholesky<Scalar, Up
     return m_perm;
   }
 
+  /** \returns the final shift parameter from the computation */
+  RealScalar shift() const { return m_shift; }
+
  protected:
   FactorType m_L;             // The lower part stored in CSC
   VectorRx m_scale;           // The vector for scaling the matrix
@@ -184,6 +189,7 @@ class IncompleteCholesky : public SparseSolverBase<IncompleteCholesky<Scalar, Up
   bool m_factorizationIsOk;
   ComputationInfo m_info;
   PermutationType m_perm;
+  RealScalar m_shift;  // The final shift parameter.
 
  private:
   inline void updateList(Ref<const VectorIx> colPtr, Ref<VectorIx> rowIdx, Ref<VectorSx> vals, const Index& col,
@@ -220,8 +226,9 @@ void IncompleteCholesky<Scalar, UpLo_, OrderingType>::factorize(const MatrixType
   // matrix has a zero on the diagonal.
   bool modified = false;
   for (Index i = 0; i < mat.cols(); ++i) {
-    if (numext::is_exactly_zero(m_L.coeff(i, i))) {
-      m_L.insert(i, i) = Scalar(0);
+    bool inserted = false;
+    m_L.findOrInsertCoeff(i, i, &inserted);
+    if (inserted) {
       modified = true;
     }
   }
@@ -270,8 +277,8 @@ void IncompleteCholesky<Scalar, UpLo_, OrderingType>::factorize(const MatrixType
 
   FactorType L_save = m_L;
 
-  RealScalar shift = 0;
-  if (mindiag <= RealScalar(0.)) shift = m_initialShift - mindiag;
+  m_shift = RealScalar(0);
+  if (mindiag <= RealScalar(0.)) m_shift = m_initialShift - mindiag;
 
   m_info = NumericalIssue;
 
@@ -279,7 +286,7 @@ void IncompleteCholesky<Scalar, UpLo_, OrderingType>::factorize(const MatrixType
   int iter = 0;
   do {
     // Apply the shift to the diagonal elements of the matrix
-    for (Index j = 0; j < n; j++) vals[colPtr[j]] += shift;
+    for (Index j = 0; j < n; j++) vals[colPtr[j]] += m_shift;
 
     // jki version of the Cholesky factorization
     Index j = 0;
@@ -323,7 +330,7 @@ void IncompleteCholesky<Scalar, UpLo_, OrderingType>::factorize(const MatrixType
         if (++iter >= 10) return;
 
         // increase shift
-        shift = numext::maxi(m_initialShift, RealScalar(2) * shift);
+        m_shift = numext::maxi(m_initialShift, RealScalar(2) * m_shift);
         // restore m_L, col_pattern, and listCol
         vals = Map<const VectorSx>(L_save.valuePtr(), nnz);
         rowIdx = Map<const VectorIx>(L_save.innerIndexPtr(), nnz);
diff --git a/Eigen/src/SparseCore/SparseMatrix.h b/Eigen/src/SparseCore/SparseMatrix.h
index 81b0a11e6..ecf406fef 100644
--- a/Eigen/src/SparseCore/SparseMatrix.h
+++ b/Eigen/src/SparseCore/SparseMatrix.h
@@ -217,15 +217,18 @@ class SparseMatrix : public SparseCompressedBase<SparseMatrix<Scalar_, Options_,
     return m_data.atInRange(m_outerIndex[outer], end, inner);
   }
 
-  /** \returns a non-const reference to the value of the matrix at position \a i, \a j
+  /** \returns a non-const reference to the value of the matrix at position \a i, \a j.
    *
    * If the element does not exist then it is inserted via the insert(Index,Index) function
    * which itself turns the matrix into a non compressed form if that was not the case.
+   * The output parameter `inserted` is set to true.
+   *
+   * Otherwise, if the element does exist, `inserted` will be set to false.
    *
    * This is a O(log(nnz_j)) operation (binary search) plus the cost of insert(Index,Index)
    * function if the element does not already exist.
    */
-  inline Scalar& coeffRef(Index row, Index col) {
+  inline Scalar& findOrInsertCoeff(Index row, Index col, bool* inserted) {
     eigen_assert(row >= 0 && row < rows() && col >= 0 && col < cols());
     const Index outer = IsRowMajor ? row : col;
     const Index inner = IsRowMajor ? col : row;
@@ -240,17 +243,37 @@ class SparseMatrix : public SparseCompressedBase<SparseMatrix<Scalar_, Options_,
         m_innerNonZeros[outer]++;
         m_data.index(end) = StorageIndex(inner);
         m_data.value(end) = Scalar(0);
+        if (inserted != nullptr) {
+          *inserted = true;
+        }
         return m_data.value(end);
       }
     }
-    if ((dst < end) && (m_data.index(dst) == inner))
+    if ((dst < end) && (m_data.index(dst) == inner)) {
       // this coefficient exists, return a refernece to it
+      if (inserted != nullptr) {
+        *inserted = false;
+      }
       return m_data.value(dst);
-    else
+    } else {
+      if (inserted != nullptr) {
+        *inserted = true;
+      }
       // insertion will require reconfiguring the buffer
       return insertAtByOuterInner(outer, inner, dst);
+    }
   }
 
+  /** \returns a non-const reference to the value of the matrix at position \a i, \a j
+   *
+   * If the element does not exist then it is inserted via the insert(Index,Index) function
+   * which itself turns the matrix into a non compressed form if that was not the case.
+   *
+   * This is a O(log(nnz_j)) operation (binary search) plus the cost of insert(Index,Index)
+   * function if the element does not already exist.
+   */
+  inline Scalar& coeffRef(Index row, Index col) { return findOrInsertCoeff(row, col, nullptr); }
+
   /** \returns a reference to a novel non zero coefficient with coordinates \a row x \a col.
    * The non zero coefficient must \b not already exist.
    *
diff --git a/test/incomplete_cholesky.cpp b/test/incomplete_cholesky.cpp
index 2b03ba197..fccc20778 100644
--- a/test/incomplete_cholesky.cpp
+++ b/test/incomplete_cholesky.cpp
@@ -54,10 +54,28 @@ void bug1150() {
   }
 }
 
+void test_non_spd() {
+  Eigen::SparseMatrix<double> A(2, 2);
+  A.insert(0, 0) = 0;
+  A.insert(1, 1) = 3;
+
+  Eigen::IncompleteCholesky<double> solver(A);
+
+  // Recover original matrix.
+  Eigen::MatrixXd M = solver.permutationP().transpose() *
+                      (solver.scalingS().asDiagonal().inverse() *
+                       (solver.matrixL() * solver.matrixL().transpose() -
+                        solver.shift() * Eigen::MatrixXd::Identity(A.rows(), A.cols())) *
+                       solver.scalingS().asDiagonal().inverse()) *
+                      solver.permutationP();
+  VERIFY_IS_APPROX(A.toDense(), M);
+}
+
 EIGEN_DECLARE_TEST(incomplete_cholesky) {
   CALL_SUBTEST_1((test_incomplete_cholesky_T<double, int>()));
   CALL_SUBTEST_2((test_incomplete_cholesky_T<std::complex<double>, int>()));
   CALL_SUBTEST_3((test_incomplete_cholesky_T<double, long int>()));
 
-  CALL_SUBTEST_1((bug1150<0>()));
+  CALL_SUBTEST_4((bug1150<0>()));
+  CALL_SUBTEST_4(test_non_spd());
 }