Constraints.hpp

#pragma once

#include <vector>

#include <amdis/linearalgebra/Constraints.hpp>
#include <amdis/linearalgebra/petsc/MatrixBackend.hpp>
#include <amdis/linearalgebra/petsc/VectorBackend.hpp>

namespace AMDiS
{
  template <class DofMap>
  struct Constraints<PetscMatrix<DofMap>>
  {
    using Matrix = PetscMatrix<DofMap>;
    using Vector = PetscVector<DofMap>;


    template <class BitVector>
    static void dirichletBC(Matrix& mat, Vector& x, Vector& b, BitVector const& nodes, bool setDiagonal = true)
    {
      thread_local std::vector<PetscInt> rows;
      rows.clear();
      auto const& dofMap = mat.dofMap_;
      for (std::size_t i = 0; i < nodes.size(); ++i)
        if (nodes[i])
          rows.push_back(dofMap.global(i));

      // test for symmetry of the matrix
      PetscBool set, flg;
      MatIsSymmetricKnown(mat.matrix(), &set, &flg);

      if (set == PETSC_TRUE && flg == PETSC_TRUE)
        MatZeroRowsColumns(mat.matrix(), rows.size(), rows.data(), setDiagonal ? 1.0 : 0.0, x.vector(), b.vector());
      else
        MatZeroRows(mat.matrix(), rows.size(), rows.data(), setDiagonal ? 1.0 : 0.0, x.vector(), b.vector());
    }



    template <class BitVector, class Associations>
    static void periodicBC(Matrix& mat, Vector& x, Vector& b, BitVector const& left, Associations const& left2right, bool setDiagonal = true)
    {
      thread_local std::vector<PetscInt> rows, associated_rows;
      rows.clear();
      associated_rows.clear();
      auto const& dofMap = mat.dofMap_;
      for (PetscInt i = 0; i < PetscInt(left.size()); ++i) {
        if (left[i]) {
          // get global row index
          PetscInt row_local[2] = {i, at(left2right,i)};
          PetscInt row[2] = {dofMap.global(row_local[0]), dofMap.global(row_local[1])};

          rows.push_back(row[0]);
          associated_rows.push_back(row[1]);

          // copy row to associated row
          PetscInt ncols;
          PetscInt const* cols;
          PetscScalar const* vals;
          MatGetRow(mat.matrix(), row[0], &ncols, &cols, &vals);
          MatSetValues(mat.matrix(), 1, &row[1], ncols, cols, vals, ADD_VALUES);
          MatRestoreRow(mat.matrix(), row[0], &ncols, &cols, &vals);
        }
      }

      MatAssemblyBegin(mat.matrix(), MAT_FLUSH_ASSEMBLY); // Maybe MAT_FLUSH_ASSEMBLY
      MatAssemblyEnd(mat.matrix(), MAT_FLUSH_ASSEMBLY);

      // clear the periodic rows and add a 1 on the diagonal
      MatSetOption(mat.matrix(), MAT_NO_OFF_PROC_ZERO_ROWS, PETSC_TRUE);
      MatZeroRows(mat.matrix(), rows.size(), rows.data(), setDiagonal ? 1 : 0, PETSC_NULL, PETSC_NULL);

      // Insert the -1 value at the off-diagonal associated columns
      // NOTE: the following is probably very slow
      for (std::size_t i = 0; i < rows.size(); ++i)
        MatSetValue(mat.matrix(), rows[i], associated_rows[i], -1, ADD_VALUES);
      MatAssemblyBegin(mat.matrix(), MAT_FINAL_ASSEMBLY); // Maybe MAT_FLUSH_ASSEMBLY
      MatAssemblyEnd(mat.matrix(), MAT_FINAL_ASSEMBLY);

      std::vector<PetscScalar> data(rows.size());

      // copy periodic rhs values to associated rows
      VecGetValues(b.vector(), rows.size(), rows.data(), data.data());
      VecSetValues(b.vector(), rows.size(), associated_rows.data(), data.data(), ADD_VALUES);
      VecAssemblyBegin(b.vector());
      VecAssemblyEnd(b.vector());

      // set the periodic rhs components to 0
      data.assign(rows.size(), 0);
      VecSetValues(b.vector(), rows.size(), rows.data(), data.data(), INSERT_VALUES);
      VecAssemblyBegin(b.vector());
      VecAssemblyEnd(b.vector());

      // symmetrize the solution vector
      VecGetValues(x.vector(), rows.size(), rows.data(), data.data());
      VecSetValues(x.vector(), rows.size(), associated_rows.data(), data.data(), INSERT_VALUES);
      VecAssemblyBegin(x.vector());
      VecAssemblyEnd(x.vector());
    }

  private:
    template <class Associations>
    static PetscInt at(Associations const& m, std::size_t idx)
    {
      auto it = m.find(idx);
      assert(it != m.end());
      return it->second;
    }
  };

} // end namespace AMDiS