MatrixBackend.hpp

#pragma once

#include <list>
#include <memory>
#include <string>
#include <vector>

#include <Eigen/SparseCore>

#include <dune/common/timer.hh>

#include <amdis/Output.hpp>
#include <amdis/linearalgebra/eigen/MatrixSize.hpp>
#include <amdis/linearalgebra/eigen/VectorBackend.hpp>

namespace AMDiS
{
  class DefaultIndexDistribution;

  template <class T, int Orientation = Eigen::RowMajor>
  class EigenSparseMatrix
  {
  public:
    using BaseMatrix = Eigen::SparseMatrix<T, Orientation>;

    using value_type = typename BaseMatrix::Scalar;

    using size_type = typename BaseMatrix::Index;

  public:
    EigenSparseMatrix(DefaultIndexDistribution const&, DefaultIndexDistribution const&)
    {}

    BaseMatrix& matrix()
    {
      return matrix_;
    }

    BaseMatrix const& matrix() const
    {
      return matrix_;
    }


    void insert(size_type r, size_type c, value_type const& value)
    {
      test_exit_dbg(r < matrix_.rows() && c < matrix_.cols(),
          "Indices out of range [0,{})x[0,{})", matrix_.rows(), matrix_.cols());
      tripletList_.emplace_back(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)
    {
      //tripletList_.reserve(tripletList_.size() + rows.size()*cols.size());
      for (size_type i = 0; i < size_type(rows.size()); ++i)
        for (size_type j = 0; j < size_type(cols.size()); ++j)
          tripletList_.emplace_back(rows[i], cols[j], mat[i][j]);
    }

    template <class RowInd, class ColInd>
    void zeroRows(RowInd const& rowInd, ColInd const& colInd, bool diag)
    {
      if constexpr(Orientation == Eigen::RowMajor) {
        // loop over the matrix rows
        auto rowIt = rowInd.begin();
        auto rowEnd = rowInd.end();
        auto colIt = colInd.begin();
        for (; rowIt != rowEnd; ++rowIt, ++colIt) {
          for (typename BaseMatrix::InnerIterator it(matrix_, *rowIt); it; ++it)
            it.valueRef() = (diag && *colIt == it.col() ? T(1) : T(0));
        }
      } else {
        error_exit("EigenMatrix::zeroRows with row and col indices is only implemented for RowMajor matrices.");
      }
    }

    template <class RowInd>
    void zeroRows(RowInd const& rowInd, bool diag)
    {
      if constexpr(Orientation == Eigen::ColMajor) {
        std::vector<bool> rowBitVec(matrix_.rows());
        for (std::size_t i = 0; i < rowInd.size(); ++i)
          rowBitVec[rowInd[i]] = true;

        // loop over the matrix columns
        for (size_type c = 0; c < matrix_.outerSize(); ++c) {
          for (typename BaseMatrix::InnerIterator it(matrix_, c); it; ++it) {
            if (rowBitVec[it.row()])
              it.valueRef() = (diag && it.row() == it.col() ? T(1) : T(0));
          }
        }
      } else if constexpr(Orientation == Eigen::RowMajor) {
        // loop over the matrix rows
        for (size_type r : rowInd) {
          for (typename BaseMatrix::InnerIterator it(matrix_, r); it; ++it)
            it.valueRef() = (diag && it.row() == it.col() ? T(1) : T(0));
        }
      }
    }

    template <class RowInd, class ColInd>
    void zeroRowsColumnsImpl(RowInd const& rowInd, ColInd const& colInd, bool diag, EigenVector<T> const* x = nullptr, EigenVector<T>* b = nullptr)
    {
      static_assert(Orientation == Eigen::RowMajor, "Only implemented for RowMajor matrices.");
      assert(rowInd.size() == colInd.size());

      std::vector<bool> colBitVec(matrix_.cols());
      for (std::size_t i = 0; i < rowInd.size(); ++i) {
        colBitVec[colInd[i]] = true;
      }

      // eliminate rows
      zeroRows(rowInd,false);

      // eliminate columns
      for (size_type r = 0; r < matrix_.outerSize(); ++r) {
        for (typename BaseMatrix::InnerIterator it(matrix_, r); it; ++it) {
          if (colBitVec[it.col()]) {
            if (diag && x && b)
              b->vector()[r] -= it.value() * x->vector()[it.col()];
            it.valueRef() = 0;
          }
        }
      }

      // set diagonal entry
      if (diag) {
        for (std::size_t i = 0; i < rowInd.size(); ++i)
          matrix_.coeffRef(rowInd[i],colInd[i]) = 1;
        if (!matrix_.isCompressed())
          matrix_.makeCompressed();
      }
    }
    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, EigenVector<T>  const& x, EigenVector<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, EigenVector<T>  const& x, EigenVector<T> & b)
    {
      zeroRowsColumnsImpl(rowInd,colInd,diag,&x,&b);
    }

    void init(MatrixSize const& sizes)
    {
      matrix_.resize(sizes.rows(), sizes.cols());
      matrix_.setZero();
    }

    void init()
    {
      matrix_.setZero();
    }

    void finish()
    {
      matrix_.setFromTriplets(tripletList_.begin(), tripletList_.end());
      matrix_.makeCompressed();

      tripletList_.clear(); // NOTE: this does not free the memory
    }

    std::size_t nnz() const
    {
      return matrix_.nonZeros();
    }

  private:
    BaseMatrix matrix_;

    std::vector<Eigen::Triplet<value_type, Eigen::Index>> tripletList_;
  };

} // end namespace AMDiS