Files
Aurora/thirdparty/include/Spectra/MatOp/SymShiftInvert.h
2023-06-02 10:49:02 +08:00

246 lines
9.9 KiB
C++

// 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 {
/// \cond
// 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;
}
};
/// \endcond
///
/// \ingroup MatOp
///
/// This class defines matrix operations required by the generalized eigen solver
/// in the shift-and-invert mode. Given two symmetric matrices \f$A\f$ and \f$B\f$,
/// it solves the linear equation \f$y=(A-\sigma B)^{-1}x\f$, where \f$\sigma\f$ is a real shift.
/// Each of \f$A\f$ and \f$B\f$ can be a dense or sparse matrix.
///
/// This class is intended to be used with the SymGEigsShiftSolver generalized eigen solver.
///
/// \tparam Scalar_ The element type of the matrices.
/// Currently supported types are `float`, `double`, and `long double`.
/// \tparam TypeA The type of the \f$A\f$ matrix, indicating whether \f$A\f$ is
/// dense or sparse. Possible values are `Eigen::Dense` and `Eigen::Sparse`.
/// \tparam TypeB The type of the \f$B\f$ matrix, indicating whether \f$B\f$ is
/// dense or sparse. Possible values are `Eigen::Dense` and `Eigen::Sparse`.
/// \tparam UploA Whether the lower or upper triangular part of \f$A\f$ should be used.
/// Possible values are `Eigen::Lower` and `Eigen::Upper`.
/// \tparam UploB Whether the lower or upper triangular part of \f$B\f$ should be used.
/// Possible values are `Eigen::Lower` and `Eigen::Upper`.
/// \tparam FlagsA Additional flags for the matrix class of \f$A\f$.
/// Possible values are `Eigen::ColMajor` and `Eigen::RowMajor`.
/// \tparam FlagsB Additional flags for the matrix class of \f$B\f$.
/// Possible values are `Eigen::ColMajor` and `Eigen::RowMajor`.
/// \tparam StorageIndexA The storage index type of the \f$A\f$ matrix, only used when \f$A\f$
/// is a sparse matrix.
/// \tparam StorageIndexB The storage index type of the \f$B\f$ matrix, only used when \f$B\f$
/// is a sparse matrix.
///
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:
///
/// Element type of the matrix.
///
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:
///
/// Constructor to create the matrix operation object.
///
/// \param A A dense or sparse matrix object, whose type can be `Eigen::Matrix<...>`,
/// `Eigen::SparseMatrix<...>`, `Eigen::Map<Eigen::Matrix<...>>`,
/// `Eigen::Map<Eigen::SparseMatrix<...>>`, `Eigen::Ref<Eigen::Matrix<...>>`,
/// `Eigen::Ref<Eigen::SparseMatrix<...>>`, etc.
/// \param B A dense or sparse matrix object.
///
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");
}
///
/// Return the number of rows of the underlying matrix.
///
Index rows() const { return m_n; }
///
/// Return the number of columns of the underlying matrix.
///
Index cols() const { return m_n; }
///
/// Set the real shift \f$\sigma\f$.
///
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");
}
///
/// Perform the shift-invert operation \f$y=(A-\sigma B)^{-1}x\f$.
///
/// \param x_in Pointer to the \f$x\f$ vector.
/// \param y_out Pointer to the \f$y\f$ vector.
///
// 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