BlockPreconditioner.hpp

#pragma once
 
#include <type_traits>
#include <vector>
 
#include <boost/numeric/mtl/utility/irange.hpp>
#include <boost/numeric/mtl/utility/iset.hpp>
#include <boost/numeric/mtl/utility/srange.hpp>
#include <boost/numeric/mtl/vector/strided_vector_ref.hpp>
 
#include <dune/functions/functionspacebases/basistags.hh>
 
#include <amdis/functions/BlockedBasisMapping.hpp>
#include <amdis/linearalgebra/mtl/BlockMatrix.hpp>
#include <amdis/linearalgebra/mtl/PreconditionerInterface.hpp>
 
namespace AMDiS
{
  namespace Impl
  {
    void initMtlIndexSet(std::vector<std::size_t> const& sizes, std::vector<mtl::irange>& indexSet)
    {
      indexSet.resize(sizes.size());
      std::size_t start = 0;
      for (std::size_t i = 0; i < sizes.size(); ++i) {
        indexSet[i].set(start, start + sizes[i]);
        start += sizes[i];
      }
    }
 
    void initMtlIndexSet(std::vector<std::size_t> const& sizes, std::vector<mtl::srange>& indexSet)
    {
      // NOTE: Only for PowerPreBasis
      std::size_t blocks = sizes.size();
      std::size_t dimension = sizes[0];
 
      // all blocks have the same size
      std::for_each(sizes.begin(), sizes.end(), [&](std::size_t s) { assert(s == dimension); });
 
      for (std::size_t i = 0; i < blocks; ++i)
        indexSet.emplace_back(i, blocks*dimension, blocks);
    }
 
    void initMtlIndexSet(std::vector<std::size_t> const& sizes, std::vector<mtl::iset>& indexSet)
    {
      // NOTE: Only for PowerPreBasis
      std::size_t blocks = sizes.size();
      std::size_t dimension = sizes[0];
 
      // all blocks have the same size
      std::for_each(sizes.begin(), sizes.end(), [&](std::size_t s) { assert(s == dimension); });
 
      indexSet.resize(blocks);
      for (std::size_t i = 0; i < dimension; ++i) {
        for (std::size_t b = 0; b < blocks; ++b)
          indexSet[b].push_back(i*blocks + b);
      }
    }
 
 
    using IMS_FlatLexicographic = Dune::Functions::BasisFactory::FlatLexicographic;
    using IMS_FlatInterleaved = Dune::Functions::BasisFactory::FlatInterleaved;
 
 
    template <class RowBasis, class ColBasis>
    class BlockPreconditionerImpl
    {
      using RowIMS = typename RowBasis::PreBasis::IndexMergingStrategy;
      using RowIndexSet = std::conditional_t<
        std::is_same<RowIMS, IMS_FlatLexicographic>::value, mtl::irange, std::conditional_t<
        std::is_same<RowIMS, IMS_FlatInterleaved>::value,   mtl::srange, void>>;
 
      using ColIMS = typename ColBasis::PreBasis::IndexMergingStrategy;
      using ColIndexSet = std::conditional_t<
        std::is_same<ColIMS, IMS_FlatLexicographic>::value, mtl::irange, std::conditional_t<
        std::is_same<ColIMS, IMS_FlatInterleaved>::value,   mtl::srange, void>>;
 
    public:
      // extract row and column ranges from block sizes
      void initIndexSets(std::vector<std::size_t> const& rowSizes, std::vector<std::size_t> const& colSizes)
      {
        initMtlIndexSet(rowSizes, rowIndexSet_);
        initMtlIndexSet(colSizes, colIndexSet_);
      }
 
      RowIndexSet const& rowIndices(std::size_t i) const
      {
        return rowIndexSet_[i];
      }
 
      ColIndexSet const& colIndices(std::size_t i) const
      {
        return colIndexSet_[i];
      }
 
    private:
 
      std::vector<RowIndexSet> rowIndexSet_;
      std::vector<ColIndexSet> colIndexSet_;
    };
 
 
    // specialization of row and col basis are the same
    template <class Basis>
    class BlockPreconditionerImpl<Basis, Basis>
    {
      using IMS = typename Basis::PreBasis::IndexMergingStrategy;
      using IndexSet = std::conditional_t<
        std::is_same<IMS, IMS_FlatLexicographic>::value, mtl::irange, std::conditional_t<
        std::is_same<IMS, IMS_FlatInterleaved>::value,   mtl::srange, void>>;
 
    public:
      // extract row and column ranges from block sizes
      void initIndexSets(std::vector<std::size_t> const& rowSizes, std::vector<std::size_t> const& /*colSizes*/)
      {
        initMtlIndexSet(rowSizes, indexSet_);
      }
 
      IndexSet const& rowIndices(std::size_t i) const
      {
        return indexSet_[i];
      }
 
      IndexSet const& colIndices(std::size_t i) const
      {
        return indexSet_[i];
      }
 
    private:
 
      std::vector<IndexSet> indexSet_;
    };
 
  } // end namespace Impl
 
 
  template <class M, class V, class RowBasis, class ColBasis = RowBasis>
  class BlockPreconditioner
      : public Impl::BlockPreconditionerImpl<RowBasis, ColBasis>
      , public PreconditionerInterface<M,V,V>
  {
  public:
 
    BlockPreconditioner(std::shared_ptr<RowBasis> rowBasis, std::shared_ptr<ColBasis> colBasis)
      : blockMapping_(std::move(rowBasis), std::move(colBasis))
    {}
 
    BlockPreconditioner(std::shared_ptr<RowBasis> rowBasis)
      : blockMapping_(std::move(rowBasis))
    {}
 
    void init(M const& A) override
    {
      initBlocks(A);
    }
 
    void initBlocks(M const& A, bool diagonalsOnly = false)
    {
      subMatrix_.copyFrom(A, blockMapping_, diagonalsOnly);
      this->initIndexSets(subMatrix_.rowSizes_, subMatrix_.colSizes_);
    }
 
  public:
 
    M const& subMatrix(std::size_t i, std::size_t j) const
    {
      return subMatrix_[i][j];
    }
 
    M& subMatrix(std::size_t i, std::size_t j)
    {
      return subMatrix_[i][j];
    }
 
  public:
 
    template <class Vector>
    auto subVector(Vector&& vec, mtl::irange const& ir) const
    {
      return mtl::vec::sub_vector(FWD(vec), ir.start(), ir.finish());
    }
 
    using StridedVector = mtl::vec::strided_vector_ref<typename V::value_type>;
    StridedVector subVector(V& vec, mtl::srange const& sr) const
    {
      return StridedVector{sr.size()/sr.stride(), vec.address_data() + sr.start(), sr.stride()};
    }
 
    StridedVector const subVector(V const& vec, mtl::srange const& sr) const
    {
      return subVector(const_cast<V&>(vec), sr);
    }
 
  public:
 
    std::size_t rows() const
    {
      return subMatrix_.rows();
    }
 
    std::size_t cols() const
    {
      return subMatrix_.cols();
    }
 
  protected:
    BlockedBasisMapping<RowBasis, ColBasis> blockMapping_;
    BlockMatrix<M> subMatrix_;
  };
 
} // end namespace AMDiS