SparseInverse.h

// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2022 Julian Kent <jkflying@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

#ifndef EIGEN_SPARSEINVERSE_H
#define EIGEN_SPARSEINVERSE_H

// IWYU pragma: private
#include "./InternalHeaderCheck.h"

#include "../../../../Eigen/Sparse"
#include "../../../../Eigen/SparseLU"

namespace Eigen {

template <typename Scalar>
class KahanSum {
  // Straightforward Kahan summation for accurate accumulation of a sum of numbers
  Scalar _sum{};
  Scalar _correction{};

 public:
  Scalar value() { return _sum; }

  void operator+=(Scalar increment) {
    const Scalar correctedIncrement = increment + _correction;
    const Scalar previousSum = _sum;
    _sum += correctedIncrement;
    _correction = correctedIncrement - (_sum - previousSum);
  }
};
template <typename Scalar, Index Width = 16>
class FABSum {
  // https://epubs.siam.org/doi/pdf/10.1137/19M1257780
  // Fast and Accurate Blocked Summation
  // Uses naive summation for the fast sum, and Kahan summation for the accurate sum
  // Theoretically SIMD sum could be changed to a tree sum which would improve accuracy
  // over naive summation
  KahanSum<Scalar> _totalSum;
  Matrix<Scalar, Width, 1> _block;
  Index _blockUsed{};

 public:
  Scalar value() { return _block.topRows(_blockUsed).sum() + _totalSum.value(); }

  void operator+=(Scalar increment) {
    _block(_blockUsed++, 0) = increment;
    if (_blockUsed == Width) {
      _totalSum += _block.sum();
      _blockUsed = 0;
    }
  }
};

template <typename Derived, typename OtherDerived>
typename Derived::Scalar accurateDot(const SparseMatrixBase<Derived>& A, const SparseMatrixBase<OtherDerived>& other) {
  typedef typename Derived::Scalar Scalar;
  EIGEN_STATIC_ASSERT_VECTOR_ONLY(Derived)
  EIGEN_STATIC_ASSERT_VECTOR_ONLY(OtherDerived)
  EIGEN_STATIC_ASSERT_SAME_VECTOR_SIZE(Derived, OtherDerived)
  static_assert(internal::is_same<Scalar, typename OtherDerived::Scalar>::value, "mismatched types");

  internal::evaluator<Derived> thisEval(A.derived());
  typename Derived::ReverseInnerIterator i(thisEval, 0);

  internal::evaluator<OtherDerived> otherEval(other.derived());
  typename OtherDerived::ReverseInnerIterator j(otherEval, 0);

  FABSum<Scalar> res;
  while (i && j) {
    if (i.index() == j.index()) {
      res += numext::conj(i.value()) * j.value();
      --i;
      --j;
    } else if (i.index() > j.index())
      --i;
    else
      --j;
  }
  return res.value();
}

template <typename Scalar>
class SparseInverse {
 public:
  typedef SparseMatrix<Scalar, ColMajor> MatrixType;
  typedef SparseMatrix<Scalar, RowMajor> RowMatrixType;

  SparseInverse() {}

  SparseInverse(const SparseLU<MatrixType>& slu) { _result = computeInverse(slu); }

  SparseInverse& compute(const SparseMatrix<Scalar>& A) {
    SparseLU<MatrixType> slu;
    slu.compute(A);
    _result = computeInverse(slu);
    return *this;
  }

  const MatrixType& inverse() const { return _result; }

  static MatrixType computeInverse(const SparseLU<MatrixType>& slu) {
    if (slu.info() != Success) {
      return MatrixType(0, 0);
    }

    // Extract from SparseLU and decompose into L, inverse D and U terms
    Matrix<Scalar, Dynamic, 1> invD;
    RowMatrixType Upper;
    {
      RowMatrixType DU = slu.matrixU().toSparse();
      invD = DU.diagonal().cwiseInverse();
      Upper = (invD.asDiagonal() * DU).template triangularView<StrictlyUpper>();
    }
    MatrixType Lower = slu.matrixL().toSparse().template triangularView<StrictlyLower>();

    // Compute the inverse and reapply the permutation matrix from the LU decomposition
    return slu.colsPermutation().transpose() * computeInverse(Upper, invD, Lower) * slu.rowsPermutation();
  }

  static MatrixType computeInverse(const RowMatrixType& Upper, const Matrix<Scalar, Dynamic, 1>& inverseDiagonal,
                                   const MatrixType& Lower) {
    // Calculate the 'minimal set', which is the nonzeros of (L+U).transpose()
    // It could be zeroed, but we will overwrite all non-zeros anyways.
    MatrixType colInv = Lower.transpose().template triangularView<UnitUpper>();
    colInv += Upper.transpose();

    // We also need rowmajor representation in order to do efficient row-wise dot products
    RowMatrixType rowInv = Upper.transpose().template triangularView<UnitLower>();
    rowInv += Lower.transpose();

    // Use the Takahashi algorithm to build the supporting elements of the inverse
    // upwards and to the left, from the bottom right element, 1 col/row at a time
    for (Index recurseLevel = Upper.cols() - 1; recurseLevel >= 0; recurseLevel--) {
      const auto& col = Lower.col(recurseLevel);
      const auto& row = Upper.row(recurseLevel);

      // Calculate the inverse values for the nonzeros in this column
      typename MatrixType::ReverseInnerIterator colIter(colInv, recurseLevel);
      for (; recurseLevel < colIter.index(); --colIter) {
        const Scalar element = -accurateDot(col, rowInv.row(colIter.index()));
        colIter.valueRef() = element;
        rowInv.coeffRef(colIter.index(), recurseLevel) = element;
      }

      // Calculate the inverse values for the nonzeros in this row
      typename RowMatrixType::ReverseInnerIterator rowIter(rowInv, recurseLevel);
      for (; recurseLevel < rowIter.index(); --rowIter) {
        const Scalar element = -accurateDot(row, colInv.col(rowIter.index()));
        rowIter.valueRef() = element;
        colInv.coeffRef(recurseLevel, rowIter.index()) = element;
      }

      // And finally the diagonal, which corresponds to both row and col iterator now
      const Scalar diag = inverseDiagonal(recurseLevel) - accurateDot(row, colInv.col(recurseLevel));
      rowIter.valueRef() = diag;
      colIter.valueRef() = diag;
    }

    return colInv;
  }

 private:
  MatrixType _result;
};

}  // namespace Eigen
#endif