coomatrix.h

#ifndef _LF_COOMATRIX_H
#define _LF_COOMATRIX_H
/***************************************************************************
 * LehrFEM++ - A simple C++ finite element libray for teaching
 * Developed from 2018 at the Seminar of Applied Mathematics of ETH Zurich,
 * lead developers Dr. R. Casagrande and Prof. R. Hiptmair
 ***************************************************************************/
 
#include "assembly_types.h"
// clang-format off
#include <Eigen/Sparse> // must be included after assembly_types.h
// clang-format on 
 
namespace lf::assemble {
template <typename SCALAR>
class COOMatrix {
 public:
  using Scalar = SCALAR;
  using Triplet = Eigen::Triplet<SCALAR>;
  using TripletVec = std::vector<Triplet>;
  using Index = Eigen::Index;
 
  COOMatrix(size_type num_rows, size_type num_cols)
      : rows_(num_rows), cols_(num_cols), triplets_() {}
 
  COOMatrix(const COOMatrix &) = default;
  COOMatrix(COOMatrix &&) noexcept = default;
  COOMatrix &operator=(const COOMatrix &) = default;
  COOMatrix &operator=(COOMatrix &&) noexcept = default;
  ~COOMatrix() = default;
 
  [[nodiscard]] Index rows() const { return rows_; }
  [[nodiscard]] Index cols() const { return cols_; }
  void AddToEntry(gdof_idx_t i, gdof_idx_t j, SCALAR increment) {
    rows_ = (i + 1 > rows_) ? i + 1 : rows_;
    cols_ = (j + 1 > cols_) ? j + 1 : cols_;
    triplets_.push_back(Eigen::Triplet<SCALAR>(i, j, increment));
  }
  void setZero() { triplets_.clear(); }
  template <typename PREDICATE>
  void setZero(PREDICATE &&pred) {
    auto new_last = std::remove_if(
        triplets_.begin(), triplets_.end(),
        [pred](Triplet &trp) { return (pred(trp.row(), trp.col())); });
    // Adjust size of triplet vector
    triplets_.erase(new_last, triplets_.end());
  }
 
  [[nodiscard]] TripletVec &triplets() { return triplets_; }
  [[nodiscard]] const TripletVec &triplets() const { return triplets_; }
 
  template <typename VECTOR>
  [[nodiscard]] Eigen::Matrix<SCALAR, Eigen::Dynamic, 1> MatVecMult(
      SCALAR alpha, const VECTOR &vec) const;
 
  template <typename VECTOR, typename RESULTVECTOR>
  void MatVecMult(SCALAR alpha, const VECTOR &vec, RESULTVECTOR &resvec) const;
 
  [[nodiscard]] Eigen::SparseMatrix<Scalar> makeSparse() const {
    Eigen::SparseMatrix<Scalar> result;
    LF_VERIFY_MSG(
        rows_ > 0 && cols_ > 0,
        "matrix has zero rows or columns, this is probably an error.");
    result.resize(rows_, cols_);
    result.setFromTriplets(triplets_.cbegin(), triplets_.cend());
    return result;
  }
  [[nodiscard]] Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>
  makeDense() const {
    Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic> mat(rows_, cols_);
    mat.setZero();
    for (const Triplet &trp : triplets_) {
      LF_ASSERT_MSG(
          ((trp.col() < cols_) && (trp.row() < rows_)),
          "Illegal triplet index (" << trp.col() << ',' << trp.row() << ")");
      mat(trp.row(), trp.col()) += trp.value();
    }
    return mat;
  }
 
  void PrintInfo(std::ostream &o) const {
    o << rows_ << " x " << cols_ << " COO matrix" << std::endl;
    for (const Triplet &trp : triplets_) {
      o << "(" << trp.row() << ',' << trp.col() << ") -> " << trp.value()
        << std::endl;
    }
  }
 
  template <typename SCALARTYPE>
  friend std::ostream &operator<<(std::ostream &o,  // NOLINT
                                  const COOMatrix<SCALARTYPE> &mat);
 
 private:
  size_type rows_, cols_; 
  TripletVec triplets_;   
};
 
// Implementation of output operator
template <typename SCALARTYPE>
std::ostream &operator<<(std::ostream &o, const COOMatrix<SCALARTYPE> &mat) {
  mat.PrintInfo();
  return o;
}
 
template <typename SCALAR>
template <typename VECTOR>
Eigen::Matrix<SCALAR, Eigen::Dynamic, 1> COOMatrix<SCALAR>::MatVecMult(
    SCALAR alpha, const VECTOR &vec) const {
  LF_ASSERT_MSG(vec.size() >= cols_,
                "size mismatch: " << cols_ << " <-> " << vec.size());
  Eigen::VectorXd result(rows_);
  result.setZero();
  for (const Triplet &trp : triplets_) {
    result[trp.row()] += trp.value() * (alpha * vec[trp.col()]);
  }
  return result;
}
 
template <typename SCALAR>
template <typename VECTOR, typename RESULTVECTOR>
void COOMatrix<SCALAR>::MatVecMult(SCALAR alpha, const VECTOR &vec,
                                   RESULTVECTOR &resvec) const {
  LF_ASSERT_MSG(vec.size() >= cols_,
                "Vector vec size mismatch: " << cols_ << " <-> " << vec.size());
  LF_ASSERT_MSG(
      resvec.size() >= rows_,
      "Vector result size mismatch: " << cols_ << " <-> " << resvec.size());
  for (const Triplet &trp : triplets_) {
    resvec[trp.row()] += trp.value() * (alpha * vec[trp.col()]);
  }
}
 
}  // namespace lf::assemble
 
#endif