SymShiftInvert.h

// Copyright (C) 2020-2022 Yixuan Qiu <yixuan.qiu@cos.name>
//
// 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 https://mozilla.org/MPL/2.0/.
 
#ifndef SPECTRA_SYM_SHIFT_INVERT_H
#define SPECTRA_SYM_SHIFT_INVERT_H
 
#include <Eigen/Core>
#include <Eigen/SparseCore>
#include <Eigen/SparseLU>
#include <stdexcept>
#include <type_traits>  // std::conditional, std::is_same
 
#include "../LinAlg/BKLDLT.h"
#include "../Util/CompInfo.h"
 
namespace Spectra {
 
 
// Compute and factorize A-sigma*B without unnecessary copying
// Default case: A is sparse, B is sparse
template <bool AIsSparse, bool BIsSparse, int UploA, int UploB>
class SymShiftInvertHelper
{
public:
    template <typename Scalar, typename Fac, typename ArgA, typename ArgB>
    static bool factorize(Fac& fac, const ArgA& A, const ArgB& B, const Scalar& sigma)
    {
        using SpMat = typename ArgA::PlainObject;
        SpMat matA = A.template selfadjointView<UploA>();
        SpMat matB = B.template selfadjointView<UploB>();
        SpMat mat = matA - sigma * matB;
        // SparseLU solver
        fac.isSymmetric(true);
        fac.compute(mat);
        // Return true if successful
        return fac.info() == Eigen::Success;
    }
};
 
// A is dense, B is dense or sparse
template <bool BIsSparse, int UploA, int UploB>
class SymShiftInvertHelper<false, BIsSparse, UploA, UploB>
{
public:
    template <typename Scalar, typename Fac, typename ArgA, typename ArgB>
    static bool factorize(Fac& fac, const ArgA& A, const ArgB& B, const Scalar& sigma)
    {
        using Matrix = typename ArgA::PlainObject;
        // Make a copy of the <UploA> triangular part of A
        Matrix mat(A.rows(), A.cols());
        mat.template triangularView<UploA>() = A;
        // Update <UploA> triangular part of mat
        if (UploA == UploB)
            mat -= (B * sigma).template triangularView<UploA>();
        else
            mat -= (B * sigma).template triangularView<UploB>().transpose();
        // BKLDLT solver
        fac.compute(mat, UploA);
        // Return true if successful
        return fac.info() == CompInfo::Successful;
    }
};
 
// A is sparse, B is dense
template <int UploA, int UploB>
class SymShiftInvertHelper<true, false, UploA, UploB>
{
public:
    template <typename Scalar, typename Fac, typename ArgA, typename ArgB>
    static bool factorize(Fac& fac, const ArgA& A, const ArgB& B, const Scalar& sigma)
    {
        using Matrix = typename ArgB::PlainObject;
        // Construct the <UploB> triangular part of -sigma*B
        Matrix mat(B.rows(), B.cols());
        mat.template triangularView<UploB>() = -sigma * B;
        // Update <UploB> triangular part of mat
        if (UploA == UploB)
            mat += A.template triangularView<UploB>();
        else
            mat += A.template triangularView<UploA>().transpose();
        // BKLDLT solver
        fac.compute(mat, UploB);
        // Return true if successful
        return fac.info() == CompInfo::Successful;
    }
};
 
 
template <typename Scalar_, typename TypeA = Eigen::Sparse, typename TypeB = Eigen::Sparse,
          int UploA = Eigen::Lower, int UploB = Eigen::Lower,
          int FlagsA = Eigen::ColMajor, int FlagsB = Eigen::ColMajor,
          typename StorageIndexA = int, typename StorageIndexB = int>
class SymShiftInvert
{
public:
    using Scalar = Scalar_;
 
private:
    using Index = Eigen::Index;
 
    // Hypothetical type of the A matrix, either dense or sparse
    using DenseTypeA = Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic, FlagsA>;
    using SparseTypeA = Eigen::SparseMatrix<Scalar, FlagsA, StorageIndexA>;
    // Whether A is sparse
    using ASparse = std::is_same<TypeA, Eigen::Sparse>;
    // Actual type of the A matrix
    using MatrixA = typename std::conditional<ASparse::value, SparseTypeA, DenseTypeA>::type;
 
    // Hypothetical type of the B matrix, either dense or sparse
    using DenseTypeB = Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic, FlagsB>;
    using SparseTypeB = Eigen::SparseMatrix<Scalar, FlagsB, StorageIndexB>;
    // Whether B is sparse
    using BSparse = std::is_same<TypeB, Eigen::Sparse>;
    // Actual type of the B matrix
    using MatrixB = typename std::conditional<BSparse::value, SparseTypeB, DenseTypeB>::type;
 
    using Vector = Eigen::Matrix<Scalar, Eigen::Dynamic, 1>;
    using MapConstVec = Eigen::Map<const Vector>;
    using MapVec = Eigen::Map<Vector>;
 
    // The type of A-sigma*B if one of A and B is dense
    // DenseType = if (A is dense) MatrixA else MatrixB
    using DenseType = typename std::conditional<ASparse::value, MatrixB, MatrixA>::type;
    // The type of A-sigma*B
    // If both A and B are sparse, the result is MatrixA; otherwise the result is DenseType
    using ResType = typename std::conditional<ASparse::value && BSparse::value, MatrixA, DenseType>::type;
 
    // If both A and B are sparse, then the result A-sigma*B is sparse, so we use
    // sparseLU for factorization; otherwise A-sigma*B is dense, and we use BKLDLT
    using FacType = typename std::conditional<
        ASparse::value && BSparse::value,
        Eigen::SparseLU<ResType>,
        BKLDLT<Scalar>>::type;
 
    using ConstGenericMatrixA = const Eigen::Ref<const MatrixA>;
    using ConstGenericMatrixB = const Eigen::Ref<const MatrixB>;
 
    ConstGenericMatrixA m_matA;
    ConstGenericMatrixB m_matB;
    const Index m_n;
    FacType m_solver;
 
public:
    template <typename DerivedA, typename DerivedB>
    SymShiftInvert(const Eigen::EigenBase<DerivedA>& A, const Eigen::EigenBase<DerivedB>& B) :
        m_matA(A.derived()), m_matB(B.derived()), m_n(A.rows())
    {
        static_assert(
            static_cast<int>(DerivedA::PlainObject::IsRowMajor) == static_cast<int>(MatrixA::IsRowMajor),
            "SymShiftInvert: the \"FlagsA\" template parameter does not match the input matrix (Eigen::ColMajor/Eigen::RowMajor)");
 
        static_assert(
            static_cast<int>(DerivedB::PlainObject::IsRowMajor) == static_cast<int>(MatrixB::IsRowMajor),
            "SymShiftInvert: the \"FlagsB\" template parameter does not match the input matrix (Eigen::ColMajor/Eigen::RowMajor)");
 
        if (m_n != A.cols() || m_n != B.rows() || m_n != B.cols())
            throw std::invalid_argument("SymShiftInvert: A and B must be square matrices of the same size");
    }
 
    Index rows() const { return m_n; }
    Index cols() const { return m_n; }
 
    void set_shift(const Scalar& sigma)
    {
        constexpr bool AIsSparse = ASparse::value;
        constexpr bool BIsSparse = BSparse::value;
        using Helper = SymShiftInvertHelper<AIsSparse, BIsSparse, UploA, UploB>;
        const bool success = Helper::factorize(m_solver, m_matA, m_matB, sigma);
        if (!success)
            throw std::invalid_argument("SymShiftInvert: factorization failed with the given shift");
    }
 
    // y_out = inv(A - sigma * B) * x_in
    void perform_op(const Scalar* x_in, Scalar* y_out) const
    {
        MapConstVec x(x_in, m_n);
        MapVec y(y_out, m_n);
        y.noalias() = m_solver.solve(x);
    }
};
 
}  // namespace Spectra
 
#endif  // SPECTRA_SYM_SHIFT_INVERT_H