b2spliss_solver.H

//------------------------------------------------------------------------
// b2spliss_solver.H --
//
//
// written Neda Ebrahimi Pour <neda.ebrahimipour@dlr.de>
//
// (c) 2022-2023 Deutsches Zentrum für Luft- und Raumfahrt (DLR) e.V.
//               Linder Höhe, 51147 Köln
//
// All Rights Reserved.  Proprietary source code.  The contents of
// this file may not be disclosed to third parties, copied or
// duplicated in any form, in whole or in part, without the prior
// written permission of SMR.
//------------------------------------------------------------------------
#ifndef B2SPLISS_SOLVER_H_
#define B2SPLISS_SOLVER_H_
 
#include <string>
#include <vector>
 
#include "b2ppconfig.h"
#include "b2sparse_solver.H"
#include "spliss/spliss.h"
 
namespace b2000::b2linalg {
 
// Struct for all configurable variables for Spliss
typedef struct {
    double residuum;
    size_t iteration;
    bool verbose;
    std::string precond_type;
    double relaxation;
} Spliss_config;
 
// LDLT solver for Spliss
template <typename T>
class Spliss_LDLt_sparse_solver : public LDLt_sparse_solver<T> {
public:
    Spliss_LDLt_sparse_solver() : updated_value(false) {}
 
    ~Spliss_LDLt_sparse_solver() {}
 
    void init(
          size_t s, size_t nnz, const size_t* colind, const size_t* rowind, const T* value,
          const int connectivity, const Dictionary& dictionary) {
        // No distributed communication for now, a function local assignement
        auto communicator = std::make_shared<spliss::LocalCommunicator>();
 
        // B2000 only stores the lower triangular matrix (symmetric matrix pattern).
        // Spliss requires both, the lower and upper part of
        // the matrix, thus we need to fill the upper part, too (if condition)!
        std::vector<std::pair<spliss::LocalIndexStoreT, spliss::GlobalIndexStoreT>>
              vector_of_positions;
        size_t r_ptr = colind[0];
        for (size_t j = 0; j != s; ++j) {
            for (size_t r_end = colind[j + 1]; r_ptr != r_end; ++r_ptr) {
                vector_of_positions.emplace_back(rowind[r_ptr], j);
                // Store input indices for potential update of the matrix values,
                // used in the resolve function.
                indices.emplace_back(rowind[r_ptr], j);
 
                if (rowind[r_ptr] != j) { vector_of_positions.emplace_back(j, rowind[r_ptr]); }
            }
        }
        // The number of local blocks is "s", the size per block is 1!
        family = std::make_shared<const spliss::VectorFamily>(communicator, s, 1);
 
        // First two entries (family) provide the size of colums and rows (s x s matrix)
        // Last the matrix pattern (colum major).
        auto connect = std::make_shared<MatrixConnectivityT>(
              family, family, CreateCSRPatternFromVectorOfPositions(*family, vector_of_positions));
 
        auto connectivityDiag = std::make_shared<MatrixConnectivityT>(
              family, family, spliss::CreateDiagonalSparsityPattern(*family));
 
        // Coloring for partitioning (parallel execution)
        auto colors = spliss::ComputeColors(*connect);
 
        // Reference to nonzero values
        // Required later in resolve function, in the case
        // the values change during runtime!
        A_reference = value;
 
        A = std::make_shared<MatrixT>(family, family, colors, connect, connectivityDiag);
 
        A->GetAccess().SetZero();
 
        conf.residuum = dictionary.get_double("SPLISS_RESIDUUM", 1.0E-8);
        conf.iteration = dictionary.get_int("SPLISS_MAX_ITER", 1000000);
        conf.verbose = dictionary.get_bool("SPLISS_VERBOSE", false);
        conf.precond_type = dictionary.get_string("SPLISS_PRECOND", "LU_JACOBI_CG");
        conf.relaxation = dictionary.get_double("SPLISS_RELAXATION", 0.8);
 
        using VectorT = typename MatrixT::SourceVectorT;
        VectorT vec_format(family);
        auto criterion = std::make_shared<spliss::ConservativeStopCriterion<T>>(
              conf.residuum, spliss::ResidualCheckMode::Relative, conf.iteration, conf.verbose);
 
        auto&& stack = spliss::InitSolverStack(A);
 
        if (conf.precond_type == "ILU_CG") {
            solver = stack.template Append<SimpleILUDecomposition>(A->GetOffDiagonalData())
                           .template Append<spliss::CG>(criterion)
                           .GetSolver();
        } else if (conf.precond_type == "LU_CG") {
            solver = stack.template Append<spliss::LUDecomposition>()
                           .template Append<spliss::CG>(criterion)
                           .GetSolver();
        } else if (conf.precond_type == "CG") {
            solver = stack.template Append<spliss::CG>(criterion).GetSolver();
        } else if (conf.precond_type == "JACOBI") {
            solver = stack.template Append<spliss::LUDecomposition>()
                           .template Append<spliss::Jacobi>(criterion, conf.relaxation)
                           .GetSolver();
        } else {
            auto criterionJacobi = std::make_shared<spliss::ConfidentStopCriterion<T>>(5u);
            solver = stack.template Append<spliss::LUDecomposition>()
                           .template Append<spliss::Jacobi>(criterionJacobi, conf.relaxation)
                           .template Append<spliss::CG>(criterion)
                           .GetSolver();
        }
 
        updated_value = true;
    }
 
    void update_value() { updated_value = true; }
 
    void resolve(
          size_t s, size_t nrhs, const T* b, size_t ldb, T* x, size_t ldx,
          char left_or_right = ' ') {
        logging::Logger logger = logging::get_logger("linear_algebra.sparse_solver.spliss");
        if (s == 0) { return; }
 
        using VectorT = typename MatrixT::SourceVectorT;
        VectorT b_value(family), x_value(family);
 
        if (updated_value) {
            auto a = A->GetAccess();
 
            // Fill the matrix with nonzero values.
            // Here the lower and upper part needs to be filled.
            size_t r_ptr = 0;
            for (const auto& i : indices) {
                a(i.first, i.second)(0, 0) = A_reference[r_ptr];
                r_ptr++;
                if (i.first != i.second) {
                    a(i.second, i.first)(0, 0) = a(i.first, i.second)(0, 0);
                }
            }
 
            solver->Prepare();
            updated_value = false;
        }
 
        for (size_t i = 0; i != nrhs; ++i) {
            x_value.SetZero();
 
            {
                auto b_access = b_value.GetAccess();
                for (size_t j = 0; j != s; ++j) { b_access[j][0] = b[j + i * ldb]; }
            }
            auto status = solver->Apply(b_value, x_value);
 
            std::cout << "Total number of iterations: " << status->LastIterationNumber()
                      << std::endl;
 
#ifdef SPLISS_DEBUG_OUTPUT
            auto convergedData = status->GetConvergenceData();
 
            for (size_t k = 0; k != convergedData.size(); k++) {
                std::cout << "Obtained solution: " << k << ", " << convergedData[k] << std::endl;
            }
#endif  // SPLISS_DEBUG_OUTPUT
 
            if (status->GetReason() != spliss::SolverStoppingReason::RelativeReductionSuccess) {
                Exception e;
                e << "Solver stopped without success!" << THROW;
            }
 
            {
                auto x_access = x_value.GetAccess();
                for (size_t j = 0; j != s; ++j) { x[j + i * ldx] = x_access[j][0]; }
            }
        }
    }
 
protected:
    bool updated_value;
    std::shared_ptr<const spliss::VectorFamily> family;
    // Matrix storage typ, only CSR is supported
    using MatrixConnectivityT = spliss::MatrixConnectivity<spliss::CSRPattern>;
    // Blockstorage, here we use 1 Block per nonzero entry in the matrix
    // Template Parameters: DataScalarType, ConnectivityType, RowsPerBlock, ColumnsPerBlock
    using StorageT = spliss::CompactBlockStorage<T, MatrixConnectivityT, 1, 1>;
    // DataBasedMatrix< ScalarType, ConnectivityType, RowsPerBlock, ColumnsPerBlock,
    // DataScalarType, StorageType
    using MatrixT = spliss::DataBasedMatrix<T, MatrixConnectivityT, 1, 1, T, StorageT>;
    std::shared_ptr<MatrixT> A;
    // Setting for ILU decomposition
    template <typename ScalarType, typename StorageType>
    using SimpleILUDecomposition = spliss::ILUDecomposition<ScalarType, StorageType, StorageType>;
    // Spliss configuration struct
    Spliss_config conf;
    // Reference to the Matrix
    const T* A_reference;
    // Vector holding the indices of the nonzero rows and colums
    std::vector<std::pair<size_t, size_t>> indices;
    // Pointer for solve Spliss
    std::shared_ptr<spliss::SolverInterface<T>> solver;
};
 
template <typename T>
class Spliss_LDLt_extension_sparse_solver : public LDLt_extension_sparse_solver<T>,
                                            public Spliss_LDLt_sparse_solver<T> {
public:
    Spliss_LDLt_extension_sparse_solver() : Spliss_LDLt_sparse_solver<T>(), div(0) {}
 
    void init(
          size_t size_, size_t nnz_, const size_t* colind_, const size_t* rowind_, const T* value_,
          size_t size_ext_, const int connectivity, const Dictionary& dictionary) {
        if (size_ext_ != 1) { UnimplementedError() << THROW; }
 
        Spliss_LDLt_sparse_solver<T>::init(
              size_, nnz_, colind_, rowind_, value_, connectivity, dictionary);
        m_ab.resize(size_ * 2);
    }
 
    void update_value() { Spliss_LDLt_sparse_solver<T>::update_value(); }
 
    void resolve(
          size_t s, size_t nrhs, const T* b, size_t ldb, T* x, size_t ldx, const T* ma_ = 0,
          const T* mb_ = 0, const T* mc_ = 0, char left_or_right = ' ') {
        if (mc_ != 0) {
            std::copy(ma_, ma_ + s - 1, &m_ab[0]);
            std::copy(mb_, mb_ + s - 1, &m_ab[s - 1]);
            Spliss_LDLt_sparse_solver<T>::resolve(s - 1, 2, &m_ab[0], s - 1, &m_ab[0], s - 1);
            div = 1 / (*mc_ - blas::dot(s - 1, ma_, 1, &m_ab[s - 1], 1));
        }
 
        for (size_t i = 0; i != nrhs; ++i) {
            const T x2 = x[ldx * i + s - 1] =
                  div * (b[ldb * i + s - 1] - blas::dot(s - 1, b + ldb * i, 1, &m_ab[s - 1], 1));
            Spliss_LDLt_sparse_solver<T>::resolve(s - 1, 1, b + ldb * i, ldb, x + ldx * i, ldx);
            blas::axpy(s - 1, -x2, &m_ab[0], 1, x + ldx * i, 1);
        }
    }
 
protected:
    std::vector<T> m_ab;
    T div;
};
 
template <>
class Spliss_LDLt_sparse_solver<csda<double>> : public LDLt_sparse_solver<csda<double>> {
public:
    Spliss_LDLt_sparse_solver() {
        UnimplementedError() << "The Spliss solver is only enabled for doubles." << THROW;
    }
 
    void init(
          size_t s, size_t nnz, const size_t* colind, const size_t* rowind,
          const csda<double>* value, const int connectivity, const Dictionary& dictionary) {}
 
    void update_value() {}
 
    void resolve(
          size_t s, size_t nrhs, const csda<double>* b, size_t ldb, csda<double>* x, size_t ldx,
          char left_or_right = ' ') {
        UnimplementedError() << THROW;
    }
};
 
template <>
class Spliss_LDLt_extension_sparse_solver<csda<double>>
    : public LDLt_extension_sparse_solver<csda<double>> {
public:
    Spliss_LDLt_extension_sparse_solver() {
        UnimplementedError() << "The Spliss solver is only enabled for doubles." << THROW;
    }
 
    void init(
          size_t s, size_t nnz, const size_t* colind, const size_t* rowind,
          const csda<double>* value, size_t, const int connectivity, const Dictionary& dictionary) {
    }
 
    void update_value() {}
 
    void resolve(
          size_t s, size_t nrhs, const csda<double>* b, size_t ldb, csda<double>* x, size_t ldx,
          const csda<double>* ma = 0, const csda<double>* mb = 0, const csda<double>* mc = 0,
          char left_or_right = ' ') {
        UnimplementedError() << THROW;
    }
};
 
}  // namespace b2000::b2linalg
 
#endif  // B2SPLISS_SOLVER_H_