petscmatrixbackend.hh

Go to the documentation of this file.

// SPDX-FileCopyrightText: Copyright © DUNE Project contributors, see file AUTHORS.md
// SPDX-License-Identifier: LicenseRef-GPL-2.0-only-with-DUNE-exception OR LGPL-3.0-or-later
 
#ifndef DUNE_ASSEMBLER_BACKENDS_PETSCMATRIXBACKEND_HH
#define DUNE_ASSEMBLER_BACKENDS_PETSCMATRIXBACKEND_HH
 
#include <algorithm>
#include <cassert>
#include <numeric>
#include <vector>
 
#include <dune/assembler/backends/flatmatrix.hh>
#include <dune/assembler/backends/flatset.hh>
#include <dune/assembler/backends/localpattern.hh>
#include <dune/common/std/extents.hh>
 
#include <petscmat.h>
 
namespace Dune::Assembler {
 
 
class PETScLocalPattern
{
  using size_type = PetscInt;
 
public:
  PETScLocalPattern (size_type rows = 0, size_type cols = 0)
    : rows_(rows,0)
    , cols_(cols,0)
  {}
 
  template <class S, std::size_t r, std::size_t c>
  PETScLocalPattern (Std::extents<S,r,c> const& extents)
    : PETScLocalPattern(extents.extent(0), extents.extent(1))
  {}
 
  void resize (size_type rows, size_type cols)
  {
    rows_.resize(rows);
    cols_.resize(cols);
  }
 
  void add (size_type row, size_type col)
  {
    rows_[row] = true;
    cols_[col] = true;
  }
 
  void addAll ()
  {
    containsAll_ = true;
  }
 
  [[nodiscard]] bool row (size_type i) const
  {
    return rows_[i];
  }
 
  [[nodiscard]] bool col (size_type j) const
  {
    return cols_[j];
  }
 
  [[nodiscard]] bool containsAll () const
  {
    return containsAll_;
  }
 
  void clear ()
  {
    for (std::size_t i = 0; i < rows_.size(); ++i)
      rows_[i] = 0;
    for (std::size_t i = 0; i < cols_.size(); ++i)
      cols_[i] = 0;
    containsAll_ = false;
  }
 
private:
  std::vector<char> rows_;
  std::vector<char> cols_;
  bool containsAll_ = false;
};
 
 
class PETScMatrixBackend
{
public:
  using Matrix = ::Mat;
  using value_type = PetscScalar;
  using size_type = PetscInt;
  using LocalMatrix = FlatMatrix<value_type>;
  using LocalPattern = PETScLocalPattern;
 
public:
  explicit PETScMatrixBackend (Matrix& matrix) noexcept
    : matrix_(&matrix)
  {}
 
  void setZero ()
  {
    MatZeroEntries(*matrix_);
  }
 
  void finish (MatAssemblyType type = MAT_FINAL_ASSEMBLY)
  {
    assert(matrix_);
    MatAssemblyBegin(*matrix_, type);
    MatAssemblyEnd(*matrix_, type);
  }
 
  template <class LocalView>
  void scatter (const LocalView& localView, const LocalMatrix& localMatrix, const LocalPattern& localPattern)
  {
    assert(matrix_);
 
    rowIndices_.resize(localView.size());
    for (std::size_t i = 0; i < localView.size(); ++i)
      rowIndices_[i] = localView.index(i);
 
    // Use negative indices to indicate ignored rows/columns
    if (!localPattern.containsAll())
    {
      for (std::size_t i = 0; i < localView.size(); ++i)
        if (!localPattern.row(i))
          rowIndices_[i] = -1;
    }
 
    MatSetValues(*matrix_, rowIndices_.size(), rowIndices_.data(), rowIndices_.size(), rowIndices_.data(),
                 localMatrix.container_data(), ADD_VALUES);
  }
 
 
  template <class RowLocalView, class ColLocalView>
  void scatter (const RowLocalView& rowLocalView, const ColLocalView& colLocalView,
                const LocalMatrix& localMatrix, const LocalPattern& localPattern)
  {
    assert(matrix_);
    // specialization for the case that both localviews are identical
    if ((void*)(&rowLocalView) == (void*)(&colLocalView))
      return scatter(rowLocalView, localMatrix, localPattern);
 
    rowIndices_.resize(rowLocalView.size());
    colIndices_.resize(colLocalView.size());
 
    for (std::size_t i = 0; i < rowLocalView.size(); ++i)
      rowIndices_[i] = rowLocalView.index(i);
    for (std::size_t j = 0; j < colLocalView.size(); ++j)
      colIndices_[j] = colLocalView.index(j);
 
    // Use negative indices to indicate ignored rows/columns
    if (!localPattern.containsAll())
    {
      for (std::size_t i = 0; i < rowLocalView.size(); ++i)
        if (!localPattern.row(i))
          rowIndices_[i] = -1;
      for (std::size_t j = 0; j < colLocalView.size(); ++j)
        if (!localPattern.col(j))
          colIndices_[j] = -1;
    }
 
    MatSetValues(*matrix_, rowIndices_.size(), rowIndices_.data(), colIndices_.size(), colIndices_.data(),
                 localMatrix.container_data(), ADD_VALUES);
  }
 
 
  struct PatternBuilder
  {
  public:
    using Matrix = ::Mat;
    using LocalPattern = PETScLocalPattern;
 
  public:
    explicit PatternBuilder (Matrix& matrix)
      : matrix_(&matrix)
    {}
 
    void setupMatrix (bool setColumnIndices = false)
    {
      assert(matrix_);
 
      if (*matrix_ == nullptr)
        MatCreate(PETSC_COMM_SELF, matrix_);
 
      MatSetSizes(*matrix_, rowSize_, colSize_, rowSize_, colSize_);
      MatSetType(*matrix_, MATSEQAIJ);
 
      // number of nonzeros per row
      std::vector<size_type> nnzPerRow(nz_.size());
      std::transform(nz_.begin(), nz_.end(), nnzPerRow.begin(), [](auto const& s) {
        return s.size();
      });
 
      // create a MATSEQAIJ matrix with provided number of nonzeros per row
      MatSeqAIJSetPreallocation(*matrix_, {}, nnzPerRow.data());
 
      if (setColumnIndices)
      {
        // total number of nonzeros
        size_type nnz = std::accumulate(nnzPerRow.begin(), nnzPerRow.end(), size_type(0));
 
        // insert the column indices from the sparsity pattern
        std::vector<size_type> columnIndices(nnz);
        auto it = columnIndices.begin();
        for (const auto& row : nz_)
          it = std::copy(row.begin(), row.end(), it);
        MatSeqAIJSetColumnIndices(*matrix_, columnIndices.data());
 
        // set all entries to zero
        MatZeroEntries(*matrix_);
      }
    }
 
    template <class RowBasis, class ColBasis>
    void resize (const RowBasis& rowBasis, const ColBasis& colBasis)
    {
      rowSize_ = rowBasis.dimension();
      colSize_ = colBasis.dimension();
      nz_.resize(rowSize_);
    }
 
    // use a given pattern
    template <class RowLocalView, class ColLocalView>
    void scatter (const RowLocalView& rowLocalView, const ColLocalView& colLocalView,
                  const LocalPattern& localPattern)
    {
      if (localPattern.containsAll())
      {
        for (std::size_t i = 0; i < rowLocalView.size(); ++i)
          for (std::size_t j = 0; j < colLocalView.size(); ++j)
            nz_[rowLocalView.index(i)].insert(colLocalView.index(j));
      }
      else
      {
        for (std::size_t i = 0; i < rowLocalView.size(); ++i)
          if (localPattern.row(i))
            for (std::size_t j = 0; j < colLocalView.size(); ++j)
              if (localPattern.col(j))
                nz_[rowLocalView.index(i)].insert(colLocalView.index(j));
      }
    }
 
    // use the full pattern
    template <class RowLocalView, class ColLocalView>
    void scatter (const RowLocalView& rowLocalView, const ColLocalView& colLocalView)
    {
      for (std::size_t i = 0; i < rowLocalView.size(); ++i)
        for (std::size_t j = 0; j < colLocalView.size(); ++j)
          nz_[rowLocalView.index(i)].insert(colLocalView.index(j));
    }
 
  private:
    Matrix* matrix_ = nullptr;
    size_type rowSize_ = 0;
    size_type colSize_ = 0;
 
    std::vector<FlatSet<size_type>> nz_;
  };
 
  PatternBuilder patternBuilder ()
  {
    return PatternBuilder{*matrix_};
  }
 
  std::size_t nnz () const
  {
    MatInfo info;
    MatGetInfo(*matrix_, MAT_LOCAL, &info);
    return std::size_t(info.nz_used);
  }
 
private:
  Matrix* matrix_ = nullptr;
 
  std::vector<size_type> rowIndices_{};
  std::vector<size_type> colIndices_{};
};
 
} // end namespace Dune::Assembler
 
#endif // DUNE_ASSEMBLER_BACKENDS_PETSCMATRIXBACKEND_HH