BlockMatrix.hpp

#pragma once
 
#include <cassert>
#include <cmath>
#include <vector>
 
#include <boost/numeric/mtl/matrix/inserter.hpp>
#include <boost/numeric/mtl/utility/property_map.hpp>
#include <boost/numeric/mtl/utility/range_wrapper.hpp>
 
namespace AMDiS
{
  template <class Matrix>
  class BlockMatrix
      : public mtl::dense2D<Matrix>
  {
    using Inserter = mtl::mat::inserter<Matrix, mtl::operations::update_plus<typename Matrix::value_type>>;
    template <class, class, class, class> friend class BlockPreconditioner;
 
  public:
    BlockMatrix() = default;
 
    // Mapping should be an instance of BlockedBasisMapping
    template <class Mapping>
    BlockMatrix(Matrix const& matrix, Mapping const& mapping, bool diagonalsOnly = false)
    {
      copyFrom(matrix, mapping, diagonalsOnly);
    }
 
    // Copy the non-blocked `matrix` into this blocked matrix, using the provided `mapping`.
    template <class Mapping>
    void copyFrom(Matrix const& matrix, Mapping const& mapping, bool diagonalsOnly = false)
    {
      initSizes(mapping.rowSizeInfo(), rowSizes_);
      initSizes(mapping.colSizeInfo(), colSizes_);
 
      this->change_dim(rows(), cols());
 
      std::size_t avgRowSize = std::round((1.2 * (matrix.nnz() / num_rows(matrix))) / rows());
 
      // start insertion
      mtl::dense2D<Inserter*> inserter(rows(), cols());
      for (std::size_t i = 0; i < rows(); ++i) {
        for (std::size_t j = 0; j < cols(); ++j) {
          (*this)[i][j].change_dim(rows(i), cols(j));
          set_to_zero((*this)[i][j]);
          inserter[i][j] = new Inserter((*this)[i][j], avgRowSize);
        }
      }
 
      // assign non-zeros to blocks
      auto row   = mtl::mat::row_map(matrix);
      auto col   = mtl::mat::col_map(matrix);
      auto value = mtl::mat::const_value_map(matrix);
 
      for (auto row_it : mtl::rows_of(matrix)) {
        for (auto it : mtl::nz_of(row_it)) {
          auto r = mapping.row(row(it));
          auto c = mapping.col(col(it));
 
          assert(r[0] < rows() && c[0] < cols());
          auto& ins = *inserter[r[0]][c[0]];
          if (!diagonalsOnly || r[0] == c[0])
            ins[r[1]][c[1]] << value(it);
        }
      }
 
      // finish insertion
      for (std::size_t i = 0; i < rows(); ++i)
        for (std::size_t j = 0; j < cols(); ++j)
          delete inserter[i][j];
    }
 
    std::size_t rows() const
    {
      return rowSizes_.size();
    }
 
    std::size_t rows(std::size_t i) const
    {
      return rowSizes_[i];
    }
 
    std::size_t cols() const
    {
      return colSizes_.size();
    }
 
    std::size_t cols(std::size_t i) const
    {
      return colSizes_[i];
    }
 
  private:
 
    // extract the hierarchic size into rowSizes or colSizes
    template <class SizeInfo>
    static void initSizes(SizeInfo const& sizeInfo, std::vector<std::size_t>& sizes)
    {
      using SizePrefix = typename SizeInfo::SizePrefix;
      SizePrefix sizePrefix{};
 
      std::size_t numBlocks = sizeInfo.size(sizePrefix);
      sizes.resize(numBlocks);
 
      sizePrefix.push_back(0);
      for (std::size_t i = 0; i < numBlocks; ++i) {
        sizePrefix.back() = i;
        sizes[i] = sizeInfo.size(sizePrefix);
      }
    }
 
  private:
    std::vector<std::size_t> rowSizes_, colSizes_;
  };
 
} // end namespace AMDiS