BlockMTLMatrix.hpp

/** \file BlockMTLMatrix.hpp */

#pragma once

#include <array>

#include <boost/numeric/mtl/matrices.hpp>

#include <dune/amdis/Basic.hpp>
#include <dune/amdis/Loops.hpp>
#include <dune/amdis/linear_algebra/LinearAlgebraBase.hpp>

namespace AMDiS
{
  /// A wrapper for AMDiS::SolverMatrix to be used in MTL/ITL solvers
  template <class MTLMatrix, size_t _N, size_t _M>
  class BlockMTLMatrix
    : public std::array<std::array<MTLMatrix, _M>, _N>
  {
    using Self = BlockMTLMatrix;
    
  public:
    /// The index/size - type
    using size_type  = typename MTLMatrix::size_type;
    
    /// The type of the elements of the MTLMatrix
    using value_type = typename MTLMatrix::value_type;
    
    /// The underlying mtl matrix type
    using BaseMatrix = MTLMatrix;

  public:
    /// Return the (R,C)'th matrix block
    template <size_t R, size_t C>
    auto& operator()(const index_<R>, const index_<C>)
    {
      static_assert(R < N() && C < M(), "Indices out of range [0,N)x[0,M)");
      return std::get<C>(std::get<R>(*this));
    }
    
    /// Return the (R,C)'th matrix block
    template <size_t R, size_t C>
    auto const& operator()(const index_<R>, const index_<C>) const
    {
      static_assert(R < N() && C < M(), "Indices out of range [0,N)x[0,M)");
      return std::get<C>(std::get<R>(*this));
    }
    
    /// Return the number of row blocks
    static constexpr size_t N() { return _N; }
    
    /// Return the number of column blocks
    static constexpr size_t M() { return _M; }

    /// perform blockwise multiplication A*b -> x
    template <class VectorIn, class VectorOut, class Assign>
    void mult(VectorIn const& b, VectorOut& x, Assign) const
    {
      // create iranges to access array blocks
      std::array<mtl::irange, _N> r_rows;
      std::array<mtl::irange, _M> r_cols;
      getRanges(r_rows, r_cols);
      
      For<0, _N>::loop([&](const auto _i) {
        bool first = true;
        
        // a reference to the i'th block of x
        VectorOut x_i(x[r_rows[_i]]);
        For<0, _M>::loop([&](const auto _j) {
          auto const& A_ij = this->operator()(_i, _j);
          if (num_rows(A_ij) > 0) {
            // a reference to the j'th block of b
            const VectorIn b_j(b[r_cols[_j]]);
            if (first) {
              Assign::first_update(x_i, A_ij * b_j);
              first = false;
            }
            else {
              Assign::update(x_i, A_ij * b_j);
            }
          }
        });
      });
    }

    /// A Multiplication operator returns a multiplication-expresssion.
    /// Calls \ref mult internally.
    template <class VectorIn>
    mtl::vec::mat_cvec_multiplier<Self, VectorIn> 
    operator*(VectorIn const& v) const
    {
      return {*this, v};
    }
    
    /// Fill an array of irange corresponding to the row-sizes, used 
    /// to access sub-vectors
    void getRowRanges(std::array<mtl::irange, _N>& r_rows) const
    {      
      size_t start = 0;
      For<0, _N>::loop([&](const auto _r) {
        size_t finish = start + num_rows((*this)(_r,index_<0>()));
        r_rows[_r].set(start, finish);
        start = finish;
      });
    }
    
    /// Fill an array of irange corresponding to the column-sizes, used 
    /// to access sub-vectors
    void getColRanges(std::array<mtl::irange, _M>& r_cols) const
    {      
      size_t start = 0;
      For<0, _M>::loop([&](const auto _c) {
        size_t finish = start + num_cols((*this)(index_<0>(),_c));
        r_cols[_c].set(start, finish);
        start = finish;
      });
    }
    
    /// Fill two arrays of irange corresponding to row and column sizes.
    /// \see getRowRanges() and \see getColRanges()
    void getRanges(std::array<mtl::irange, _N>& r_rows, 
                   std::array<mtl::irange, _M>& r_cols) const
    {      
      getRowRanges(r_rows);
      getColRanges(r_cols);
    }
  };
  
  
  namespace Impl
  {
    /// Specialization of Impl::MTLMatrix from \file LinearAlgebraBase.hpp
    template <class MTLMatrix, size_t _N, size_t _M>
    struct BaseMatrix<BlockMTLMatrix<MTLMatrix, _N, _M>>
    {
      using type = MTLMatrix;
    };
  }

  /// Return the number of overall rows of a BlockMTLMatrix
  template <class MTLMatrix, size_t _N, size_t _M>
  inline size_t num_rows(BlockMTLMatrix<MTLMatrix, _N, _M> const& A)
  {
    size_t nRows = 0;
    For<0, _N>::loop([&](const auto _r) {
      nRows += num_rows(A(_r,index_<0>()));
    });
    return nRows;
  }

  /// Return the number of overall columns of a BlockMTLMatrix
  template <class MTLMatrix, size_t _N, size_t _M>
  inline size_t num_cols(BlockMTLMatrix<MTLMatrix, _N, _M> const& A)
  {
    size_t nCols = 0;
    For<0, _M>::loop([&](const auto _c) {
      nCols += num_cols(A(index_<0>(),_c));
    });
    return nCols;
  }

  /// Return the size, i.e. rows*columns of a BlockMTLMatrix
  template <class MTLMatrix, size_t _N, size_t _M>
  inline size_t size(BlockMTLMatrix<MTLMatrix, _N, _M> const& A)
  {
    return num_rows(A) * num_cols(A);
  }

  /// Nullify a BlockMTLMatrix, i.e. nullify each block.
  template <class MTLMatrix, size_t _N, size_t _M>
  inline void set_to_zero(BlockMTLMatrix<MTLMatrix, _N, _M>& A) 
  {
    For<0, _N>::loop([&](const auto _r) {
      For<0, _M>::loop([&](const auto _c) {
        set_to_zero(A(_r,_c));
      });
    });
  }

} // end namespace AMDiS


/// \cond HIDDEN_SYMBOLS
namespace mtl
{
  template <class MTLMatrix, size_t _N, size_t _M>
  struct Collection<AMDiS::BlockMTLMatrix<MTLMatrix, _N, _M>>
  {
    using value_type = typename MTLMatrix::value_type;
    using size_type  = typename MTLMatrix::size_type;
  };

  namespace ashape
  {
    template <class MTLMatrix, size_t _N, size_t _M>
    struct ashape_aux<AMDiS::BlockMTLMatrix<MTLMatrix, _N, _M>>
    {
      using type = nonscal;
    };

  } // end namespace ashape

} // end namespace mtl
/// \endcond