MatrixBackend.hpp

#pragma once
 
#include <list>
#include <memory>
#include <string>
#include <vector>
 
#include <boost/numeric/mtl/matrix/compressed2D.hpp>
#include <boost/numeric/mtl/matrix/inserter.hpp>
#include <boost/numeric/mtl/utility/property_map.hpp>
#include <boost/numeric/mtl/utility/range_wrapper.hpp>
 
#include <amdis/Output.hpp>
#include <amdis/linearalgebra/SymmetryStructure.hpp>
#include <amdis/linearalgebra/mtl/SlotSize.hpp>
#include <amdis/linearalgebra/mtl/VectorBackend.hpp>
 
namespace AMDiS
{
  class DefaultIndexDistribution;
 
  template <class T>
  class MTLSparseMatrix
  {
  public:
    using BaseMatrix = mtl::compressed2D<T>;
 
    using value_type = typename BaseMatrix::value_type;
 
    using size_type = typename BaseMatrix::size_type;
 
  private:
    using Inserter = mtl::mat::inserter<BaseMatrix, mtl::operations::update_plus<value_type>>;
 
    using Pattern = SlotSize;
 
  public:
    MTLSparseMatrix(DefaultIndexDistribution const&, DefaultIndexDistribution const&)
    {}
 
    BaseMatrix& matrix()
    {
      assert( !inserter_ );
      return matrix_;
    }
 
    BaseMatrix const& matrix() const
    {
      assert( !inserter_ );
      return matrix_;
    }
 
    void init(Pattern const& pattern)
    {
      test_exit(!inserter_, "Matrix already in insertion mode!");
 
      std::size_t slotSize = nnz() > 0 ? 6*nnz() / (5*num_rows(matrix_))
                                       : pattern.rowSizeEstimate();
      matrix_.change_dim(pattern.rows(), pattern.cols());
      set_to_zero(matrix_);
 
      symmetry_ = pattern.symmetry();
      inserter_ = new Inserter(matrix_, slotSize);
    }
 
    void init()
    {
      test_exit(!inserter_, "Matrix already in insertion mode!");
 
      std::size_t slotSize = 6*nnz() / (5*num_rows(matrix_));
      set_to_zero(matrix_);
 
      inserter_ = new Inserter(matrix_, slotSize);
    }
 
 
    void finish()
    {
      delete inserter_;
      inserter_ = nullptr;
    }
 
 
    void insert(size_type r, size_type c, value_type const& value)
    {
      test_exit_dbg(inserter_, "Inserter not initialized!");
      test_exit_dbg(r < num_rows(matrix_) && c < num_cols(matrix_),
          "Indices out of range [0,{})x[0,{})", num_rows(matrix_), num_cols(matrix_));
      if (value != value_type(0) || r == c)
        (*inserter_)[r][c] += value;
    }
 
    template <class Ind, class LocalMat>
    void scatter(Ind const& idx, LocalMat const& mat)
    {
      scatter(idx, idx, mat);
    }
 
    template <class RowInd, class ColInd, class LocalMat>
    void scatter(RowInd const& rows, ColInd const& cols, LocalMat const& mat)
    {
      test_exit_dbg(inserter_, "Inserter not initialized!");
      for (size_type i = 0; i < size_type(rows.size()); ++i)
        for (size_type j = 0; j < size_type(cols.size()); ++j)
          if (mat[i][j] != value_type(0) || i == j)
            (*inserter_)[rows[i]][cols[j]] += mat[i][j];
    }
 
    template <class RowInd, class ColInd>
    void setUnitDiagonals(RowInd const& rowInd, ColInd const& colInd)
    {
      bool all_diagonals_found = true;
      auto const& indexer = matrix_.indexer;
      auto rowIt = rowInd.begin();
      auto rowEnd = rowInd.end();
      auto colIt = colInd.begin();
      for (; rowIt != rowEnd; ++rowIt, ++colIt) {
        auto pos = indexer(matrix_, *rowIt, *colIt);
        if (pos)
          matrix_.value_from_offset(pos.value()) = 1.0;
        else {
          all_diagonals_found = false;
          break;
        }
      }
 
      if (!all_diagonals_found) {
        std::size_t slotSize = 6*nnz() / (5*num_rows(matrix_));
        mtl::mat::inserter<BaseMatrix, mtl::update_store<value_type> > ins(matrix_, slotSize);
        auto rowIt = rowInd.begin();
        auto rowEnd = rowInd.end();
        auto colIt = colInd.begin();
        for (; rowIt != rowEnd; ++rowIt, ++colIt)
          ins(*rowIt,*colIt) << 1.0;
      }
    }
 
    template <class RowInd, class ColInd>
    void zeroRows(RowInd const& rowInd, ColInd const& colInd, bool diag)
    {
      for (auto i : rowInd) {
        size_type offset_begin = matrix_.ref_major()[i];
        size_type offset_end = matrix_.ref_major()[i+1];
        std::fill(matrix_.address_data()+offset_begin, matrix_.address_data()+offset_end, value_type(0));
      }
 
      if (diag)
        setUnitDiagonals(rowInd,colInd);
    }
 
    template <class RowInd>
    void zeroRows(RowInd const& rowInd, bool diag)
    {
      zeroRows(rowInd,rowInd,diag);
    }
 
    template <class RowInd, class ColInd>
    void zeroRowsColumnsImpl(RowInd const& rowInd, ColInd const& colInd, bool diag, MTLVector<T> const* x = nullptr, MTLVector<T>* b = nullptr)
    {
      assert(rowInd.size() == colInd.size());
 
      std::vector<bool> colBitVec(num_cols(matrix_), false);
      for (std::size_t i = 0; i < colInd.size(); ++i)
        colBitVec[colInd[i]] = true;
 
      zeroRows(rowInd, false);
 
      // Define the property maps
      auto row   = mtl::mat::row_map(matrix_);
      auto col   = mtl::mat::col_map(matrix_);
      auto value = mtl::mat::value_map(matrix_);
 
      // iterate over the matrix
      std::size_t slotSize = 0;
      for (auto r : mtl::rows_of(matrix_)) {   // rows or columns
        std::size_t rowSize = 0;
        for (auto i : mtl::nz_of(r)) {          // non-zeros within
          ++rowSize;
          if (colBitVec[col(i)]) {
            if (diag && x && b)
              b->vector()[row(i)] -= value(i) * x->vector()[col(i)];
            value(i, 0);
          }
        }
        slotSize = std::max(slotSize, rowSize);
      }
 
      // set diagonal entry
      if (diag)
        setUnitDiagonals(rowInd,colInd);
    }
 
    template <class RowInd>
    void zeroRowsColumns(RowInd const& rowInd, bool diag)
    {
      zeroRowsColumnsImpl(rowInd,rowInd,diag,nullptr,nullptr);
    }
 
    template <class RowInd>
    void zeroRowsColumns(RowInd const& rowInd, bool diag, MTLVector<T> const& x, MTLVector<T>& b)
    {
      zeroRowsColumnsImpl(rowInd,rowInd,diag,&x,&b);
    }
 
    template <class RowInd, class ColInd>
    void zeroRowsColumns(RowInd const& rowInd, ColInd const& colInd, bool diag)
    {
      zeroRowsColumnsImpl(rowInd,colInd,diag,nullptr,nullptr);
    }
 
    template <class RowInd, class ColInd>
    void zeroRowsColumns(RowInd const& rowInd, ColInd const& colInd, bool diag, MTLVector<T> const& x, MTLVector<T>& b)
    {
      zeroRowsColumnsImpl(rowInd,colInd,diag,&x,&b);
    }
 
    std::size_t nnz() const
    {
      return matrix_.nnz();
    }
 
    SymmetryStructure symmetry() const
    {
      return symmetry_;
    }
 
  private:
    BaseMatrix matrix_;
 
    Inserter* inserter_ = nullptr;
 
    SymmetryStructure symmetry_ = SymmetryStructure::unknown;
  };
 
} // end namespace AMDiS