MatrixBackend.hpp

#pragma once

#include <algorithm>
#include <memory>
#include <vector>

#include <petscmat.h>

#include <dune/common/timer.hh>

#include <amdis/Output.hpp>
#include <amdis/common/FakeContainer.hpp>
#include <amdis/linearalgebra/petsc/MatrixNnzStructure.hpp>
#include <amdis/linearalgebra/petsc/VectorBackend.hpp>
#include <amdis/linearalgebra/SymmetryStructure.hpp>

namespace AMDiS
{
  template <class> class Constraints;

  template <class DofMap>
  class PetscMatrix
  {
    template <class> friend class Constraints;

  public:
    using BaseMatrix = ::Mat;

    using value_type = PetscScalar;

    using size_type = PetscInt;

  public:
    PetscMatrix(DofMap const& rowDofMap, DofMap const& colDofMap)
      : rowDofMap_(rowDofMap)
      , colDofMap_(colDofMap)
    {}

    // disable copy and move operations
    PetscMatrix(PetscMatrix const&) = delete;
    PetscMatrix(PetscMatrix&&) = delete;
    PetscMatrix& operator=(PetscMatrix const&) = delete;
    PetscMatrix& operator=(PetscMatrix&&) = delete;

    ~PetscMatrix()
    {
      destroy();
    }

    BaseMatrix& matrix()
    {
      return matrix_;
    }

    BaseMatrix const& matrix() const
    {
      return matrix_;
    }

    template <class RowIndex, class ColIndex>
    void insert(RowIndex const& r, ColIndex const& c, PetscScalar value)
    {
      PetscInt row = rowDofMap_.global(r);
      PetscInt col = colDofMap_.global(c);
      MatSetValue(matrix_, row, col, value, ADD_VALUES);
    }

    template <class LocalInd, class LocalMatrix>
    void scatter(LocalInd const& localInd, LocalMatrix const& localMat)
    {
      if (&rowDofMap_ == &colDofMap_) {
        thread_local std::vector<PetscInt> idx;
        idx.resize(localInd.size());

        // create a vector of global indices from the local indices using the local-to-global map
        std::transform(localInd.begin(), localInd.end(), idx.begin(),
          [this](auto const& mi) { return rowDofMap_.global(mi); });

        MatSetValues(matrix_, idx.size(), idx.data(), idx.size(), idx.data(), localMat.data(), ADD_VALUES);
      } else {
        scatter(localInd, localInd, localMat);
      }
    }

    template <class RowLocalInd, class ColLocalInd, class LocalMatrix>
    void scatter(RowLocalInd const& rowLocalInd, ColLocalInd const& colLocalInd, LocalMatrix const& localMat)
    {
      thread_local std::vector<PetscInt> ri;
      thread_local std::vector<PetscInt> ci;
      ri.resize(rowLocalInd.size());
      ci.resize(colLocalInd.size());

      // create vectors of global indices from the local indices using the local-to-global map
      std::transform(rowLocalInd.begin(), rowLocalInd.end(), ri.begin(),
        [this](auto const& mi) { return rowDofMap_.global(mi); });
      std::transform(colLocalInd.begin(), colLocalInd.end(), ci.begin(),
        [this](auto const& mi) { return colDofMap_.global(mi); });

      MatSetValues(matrix_, ri.size(), ri.data(), ci.size(), ci.data(), localMat.data(), ADD_VALUES);
    }

    template <class RowInd>
    void zeroRows(std::vector<RowInd> const& rowLocalInd, bool diag)
    {
      thread_local std::vector<PetscInt> ri;
      ri.resize(rowLocalInd.size());
      std::transform(rowLocalInd.begin(), rowLocalInd.end(), ri.begin(),
        [this](auto const& mi) { return rowDofMap_.global(mi); });

      MatZeroRows(matrix_, ri.size(), ri.data(), diag ? 1.0 : 0.0, PETSC_NULL, PETSC_NULL);
    }

    template <class RowInd, class ColInd>
    void zeroRows(std::vector<RowInd> const& rowLocalInd, std::vector<ColInd> const& colLocalInd, bool diag)
    {
      thread_local std::vector<PetscInt> ri;
      ri.resize(rowLocalInd.size());
      std::transform(rowLocalInd.begin(), rowLocalInd.end(), ri.begin(),
        [this](auto const& mi) { return rowDofMap_.global(mi); });
      MatZeroRows(matrix_, ri.size(), ri.data(), 0.0, PETSC_NULL, PETSC_NULL);

      if (diag) {
        thread_local std::vector<PetscInt> ci;
        ci.resize(colLocalInd.size());
        std::transform(colLocalInd.begin(), colLocalInd.end(), ci.begin(),
          [this](auto const& mi) { return colDofMap_.global(mi); });
        PetscBool mat_new_nonzero_allocation_err;
        MatGetOption(matrix_, MAT_NEW_NONZERO_ALLOCATION_ERR, &mat_new_nonzero_allocation_err);
        MatSetOption(matrix_, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_FALSE);
        for (std::size_t i = 0; i < ri.size(); ++i)
          MatSetValue(matrix_, ri[i], ci[i], 1.0, INSERT_VALUES);
        MatAssemblyBegin(matrix_, MAT_FINAL_ASSEMBLY);
        MatAssemblyEnd(matrix_, MAT_FINAL_ASSEMBLY);
        MatSetOption(matrix_, MAT_NEW_NONZERO_ALLOCATION_ERR, mat_new_nonzero_allocation_err);
      }
    }

    template <class RowInd>
    void zeroRowsColumns(std::vector<RowInd> const& rowLocalInd, bool diag)
    {
      thread_local std::vector<PetscInt> ri;
      ri.resize(rowLocalInd.size());
      std::transform(rowLocalInd.begin(), rowLocalInd.end(), ri.begin(),
        [this](auto const& mi) { return rowDofMap_.global(mi); });
      MatZeroRowsColumns(matrix_, ri.size(), ri.data(), diag ? 1.0 : 0.0, PETSC_NULL, PETSC_NULL);
    }

    template <class RowInd>
    void zeroRowsColumns(std::vector<RowInd> const& rowLocalInd, bool diag, PetscVector<DofMap> const& x, PetscVector<DofMap>& b)
    {
      thread_local std::vector<PetscInt> ri;
      ri.resize(rowLocalInd.size());
      std::transform(rowLocalInd.begin(), rowLocalInd.end(), ri.begin(),
        [this](auto const& mi) { return rowDofMap_.global(mi); });
      MatZeroRowsColumns(matrix_, ri.size(), ri.data(), diag ? 1.0 : 0.0, x.vector(), b.vector());
    }

    template <class RowInd, class ColInd>
    void zeroRowsColumns(std::vector<RowInd> const& rowLocalInd, std::vector<ColInd> const& colLocalInd, bool diag)
    {
      error_exit("Different row and column indices not implemented for PetscMatrix::zeroRowsColumns. Either pass only one index vector, if the row and column indices are the same, or use PetscMatrix::zeroRows that does not eliminate the columns.");
    }

    template <class RowInd, class ColInd>
    void zeroRowsColumns(std::vector<RowInd> const& rowLocalInd, std::vector<ColInd> const& colLocalInd, bool diag, PetscVector<DofMap> const& x, PetscVector<DofMap>& b)
    {
      error_exit("Different row and column indices not implemented for PetscMatrix::zeroRowsColumns. Either pass only one index vector, if the row and column indices are the same, or use PetscMatrix::zeroRows that does not eliminate the columns.");
    }

    void init(MatrixNnzStructure const& pattern)
    {
      Dune::Timer t;

      // destroy an old matrix if created before
      destroy();
      info(3, "  destroy old matrix needed {} seconds", t.elapsed());
      t.reset();

      // create a MATAIJ or MATSEQAIJ matrix with provided sparsity pattern
      MatCreateAIJ(comm(),
        rowDofMap_.localSize(), colDofMap_.localSize(),
        rowDofMap_.globalSize(), colDofMap_.globalSize(),
        0, pattern.d_nnz().data(), 0, pattern.o_nnz().data(), &matrix_);

      // keep sparsity pattern even if we delete a row / column with e.g. MatZeroRows
      MatSetOption(matrix_, MAT_KEEP_NONZERO_PATTERN, PETSC_TRUE);

      // set symmetry properties of the matrix
      switch (pattern.symmetry()) {
        case SymmetryStructure::spd:
          MatSetOption(matrix_, MAT_SPD, PETSC_TRUE);
          break;
        case SymmetryStructure::symmetric:
          MatSetOption(matrix_, MAT_SYMMETRIC, PETSC_TRUE);
          break;
        case SymmetryStructure::hermitian:
          MatSetOption(matrix_, MAT_HERMITIAN, PETSC_TRUE);
          break;
        case SymmetryStructure::structurally_symmetric:
          MatSetOption(matrix_, MAT_STRUCTURALLY_SYMMETRIC, PETSC_TRUE);
          break;
        default:
          /* do nothing */
          break;
      }

      info(3, "  create new matrix needed {} seconds", t.elapsed());
      t.reset();

      initialized_ = true;
    }

    void init()
    {
      MatZeroEntries(matrix_);
      initialized_ = true;
    }

    void finish()
    {
      Dune::Timer t;
      MatAssemblyBegin(matrix_, MAT_FINAL_ASSEMBLY);
      MatAssemblyEnd(matrix_, MAT_FINAL_ASSEMBLY);
      info(3, "  finish matrix assembling needed {} seconds", t.elapsed());
    }

    std::size_t nnz() const
    {
      MatInfo info;
      MatGetInfo(matrix_, MAT_LOCAL, &info);
      return std::size_t(info.nz_used);
    }

    // An MPI Communicator or PETSC_COMM_SELF
    MPI_Comm comm() const
    {
      return rowDofMap_.comm();
    }

  private:
    // Destroy the matrix if initialized before.
    void destroy()
    {
      if (initialized_)
        MatDestroy(&matrix_);
    }

  private:
    // The local-to-global maps
    DofMap const& rowDofMap_;
    DofMap const& colDofMap_;

    Mat matrix_;

    bool initialized_ = false;
  };

} // end namespace AMDiS