VectorBackend.hpp

#pragma once

#include <petscvec.h>

#include <dune/common/ftraits.hh>
#include <dune/common/rangeutilities.hh>
#include <dune/common/shared_ptr.hh>
#include <dune/functions/common/multiindex.hh>

#include <amdis/Output.hpp>
#include <amdis/algorithm/ForEach.hpp>
#include <amdis/algorithm/InnerProduct.hpp>
#include <amdis/algorithm/Transform.hpp>
#include <amdis/common/FakeContainer.hpp>
#include <amdis/operations/Assigner.hpp>
#include <amdis/typetree/MultiIndex.hpp>

namespace AMDiS
{
  template <class DofMap>
  class PetscVector
  {
    template <class A>
    using AssignMode = std::conditional_t<std::is_same_v<A,Assigner::plus_assign>,
      std::integral_constant<::InsertMode, ADD_VALUES>,
      std::integral_constant<::InsertMode, INSERT_VALUES>>;

  public:
    using BaseVector = ::Vec;

    using value_type = PetscScalar;

    using size_type  = PetscInt;

  public:
    explicit PetscVector(DofMap const& dofMap)
      : dofMap_(dofMap)
    {}

    PetscVector(PetscVector&& other)
      : dofMap_(other.dofMap_)
    {
      swap(*this, other);
    }

    PetscVector(PetscVector const& other)
      : dofMap_(other.dofMap_)
      , initialized_(other.initialized_)
    {
      if (other.initialized_) {
        VecDuplicate(other.vector_, &vector_);
        VecCopy(other.vector_, vector_);
      }
    }

    ~PetscVector()
    {
      destroy();
    }

    PetscVector& operator=(PetscVector&& other)
    {
      swap(*this, other);
      return *this;
    }

    PetscVector& operator=(PetscVector const& other)
    {
      return *this = PetscVector(other);
    }

    BaseVector const& vector() const
    {
      return vector_;
    }

    BaseVector& vector()
    {
      return vector_;
    }

    std::size_t localSize() const
    {
      return dofMap_.localSize();
    }

    std::size_t globalSize() const
    {
      return dofMap_.globalSize();
    }

    // [[possibly collective]]
    template <class Basis>
    void init(Basis const&, bool clear)
    {
      Dune::Timer t;
      if (initialized_)
        destroy();

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

      if (clear)
        setZero();

      initialized_ = true;
    }

    // [[collective]]
    void finish()
    {
      Dune::Timer t;
      VecAssemblyBegin(vector_);
      VecAssemblyEnd(vector_);
      info(3, "  finish vector assembly needed {} seconds", t.elapsed());
    }

    // [[collective]]
    void synchronize()
    {
      Dune::Timer t;
      VecGhostUpdateBegin(vector_, INSERT_VALUES, SCATTER_FORWARD);
      VecGhostUpdateEnd(vector_, INSERT_VALUES, SCATTER_FORWARD);
      info(3, "  synchronize vector needed {} seconds", t.elapsed());
    }

    // [[not collective]]
    template <class MultiIndex>
    PetscScalar at(MultiIndex const& idx) const
    {
      PetscScalar y = 0;
      gather(std::array<MultiIndex,1>{idx}, &y);
      return y;
    }

    // [[not collective]]
    template <class MultiIndex, class Assign>
    void insert(MultiIndex const& dof, PetscScalar value, Assign)
    {
      assert(initialized_);
      VecSetValue(vector_, dofMap_.global(flatMultiIndex(dof)), value, AssignMode<Assign>::value);
    }

    // [[not collective]]
    template <class IndexRange, class OutputIterator>
    void gather(IndexRange const& localInd, OutputIterator buffer) const
    {
      forEach(localInd, [&buffer](auto&&, auto const& value) { *buffer = value; ++buffer; });
    }


    // [[not collective]]
    template <class IndexRange, class LocalValues, class MaskRange, class Assign>
    void scatter(IndexRange const& localInd, LocalValues const& localVal, MaskRange const& mask, Assign)
    {
      if constexpr(not std::is_same_v<typename LocalValues::value_type, PetscScalar>)
      {
        std::vector<PetscScalar> newLocalVal(localVal.begin(), localVal.end());
        scatter(localInd, newLocalVal, mask, Assign{});
        return;
      }

      assert(initialized_);
      assert(localInd.size() == localVal.size());

      // map local to global indices
      std::vector<PetscInt> globalInd;
      globalInd.reserve(localInd.size());

      auto ind_it = std::begin(localInd);
      auto mask_it = std::begin(mask);
      for (; ind_it != std::end(localInd); ++mask_it, ++ind_it)
        globalInd.push_back(*mask_it ? dofMap_.global(flatMultiIndex(*ind_it)) : -1);

      if (!std::is_same_v<MaskRange, FakeContainer<bool,true>>)
        VecSetOption(vector_, VEC_IGNORE_NEGATIVE_INDICES, PETSC_TRUE);

      VecSetValues(vector_, globalInd.size(), globalInd.data(), localVal.data(), AssignMode<Assign>::value);
      VecSetOption(vector_, VEC_IGNORE_NEGATIVE_INDICES, PETSC_FALSE);
    }

    // [[not collective]]
    template <class IndexRange, class Func>
    void forEach(IndexRange const& localInd, Func&& f) const
    {
      // Obtains the local ghosted representation of a parallel vector
      Vec localGhostedVec = nullptr;
      VecGhostGetLocalForm(vector_, &localGhostedVec);
      assert(localGhostedVec != nullptr);

      // A pointer to a contiguous array that contains this processor's portion of the vector data.
      PetscScalar* ptr;
      VecGetArray(localGhostedVec, &ptr);

      for (auto const& dof : localInd)
      {
        size_type i = flatMultiIndex(dof);
        if (dofMap_.owner(i)) {
          // the index is a purely local entry
          f(dof, ptr[dofMap_.dofToLocal(i)]);
        }
        else {
          // ghost entries are stored at the end of the array
          f(dof, ptr[dofMap_.localSize() + dofMap_.dofToGhost(i)]);
        }
      }

      VecRestoreArray(localGhostedVec, &ptr);
      VecGhostRestoreLocalForm(vector_, &localGhostedVec);
    }

    // [[not collective]]
    template <class Func>
    void forEachLocal(Func&& f) const
    {
      // A pointer to a contiguous array that contains this processor's portion of the vector data.
      PetscScalar const* ptr;
      VecGetArrayRead(vector_, &ptr);

      for (size_type i = 0; i < size_type(localSize()); ++i)
        f(ptr[dofMap_.dofToLocal(i)]);

      VecRestoreArrayRead(vector_, &ptr);
    }

    // [[mutable]]
    // [[not collective]]
    template <class Func>
    void forEachLocal(Func&& f)
    {
      // A pointer to a contiguous array that contains this processor's portion of the vector data.
      PetscScalar *ptr;
      VecGetArray(vector_, &ptr);

      for (size_type i = 0; i < size_type(localSize()); ++i)
        f(ptr[dofMap_.dofToLocal(i)]);

      VecRestoreArray(vector_, &ptr);
    }

    // [[not collective]]
    template <class Operation, class... VecIns>
    static void transformLocal(PetscVector& vecOut, Operation op, VecIns const&... vecIn)
    {
      PetscScalar *ptrOut;
      std::array<PetscScalar const*,sizeof...(VecIns)> ptrsIn{};

      // Obtain ptrs to internal local arrays
      VecGetArray(vecOut.vector_, &ptrOut);
      Ranges::forIndices<sizeof...(VecIns)>([&](auto ii) {
        VecGetArrayRead(std::get<ii>(std::tie(vecIn...)).vector_, &ptrsIn[ii]);
      });

      for (size_type i = 0; i < size_type(vecOut.localSize()); ++i) {
        auto idx = vecOut.dofMap_.dofToLocal(i);
        Ranges::applyIndices<sizeof...(VecIns)>([&](auto... ii) {
          ptrOut[idx] = op(ptrsIn[ii][idx]...);
        });
      }

      // Restore array obtained with VecGetArray[Read]()
      VecRestoreArray(vecOut.vector_, &ptrOut);
      Ranges::forIndices<sizeof...(VecIns)>([&](auto ii) {
        VecRestoreArrayRead(std::get<ii>(std::tie(vecIn...)).vector_, &ptrsIn[ii]);
      });
    }

    // [[not collective]]
    template <class T, class BinOp1, class BinOp2>
    static T innerProductLocal(PetscVector const& in1, PetscVector const& in2, T init, BinOp1 op1, BinOp2 op2)
    {
      PetscScalar const* ptr1;
      PetscScalar const* ptr2;

      // Obtain ptrs to internal local arrays
      VecGetArrayRead(in1.vector_, &ptr1);
      VecGetArrayRead(in2.vector_, &ptr2);

      for (size_type i = 0; i < size_type(in1.localSize()); ++i) {
        auto idx = in1.dofMap_.dofToLocal(i);
        init = op1(std::move(init), op2(ptr1[idx], ptr2[idx]));
      }

      // Restore array obtained with VecGetArrayRead()
      VecRestoreArrayRead(in1.vector_, &ptr1);
      VecRestoreArrayRead(in2.vector_, &ptr2);

      return init;
    }

    // [[not collective]]
    void set(PetscScalar value)
    {
      Vec localGhostedVec = nullptr;
      VecGhostGetLocalForm(vector_, &localGhostedVec);
      assert(localGhostedVec != nullptr);

      VecSet(localGhostedVec, value);
      VecGhostRestoreLocalForm(vector_, &localGhostedVec);
    }

    // [[not collective]]
    void setZero()
    {
      Vec localGhostedVec = nullptr;
      VecGhostGetLocalForm(vector_, &localGhostedVec);
      assert(localGhostedVec != nullptr);

      VecZeroEntries(localGhostedVec);
      VecGhostRestoreLocalForm(vector_, &localGhostedVec);
    }

    friend void swap(PetscVector& lhs, PetscVector& rhs)
    {
      assert(&lhs.dofMap_ == &rhs.dofMap_);
      std::swap(lhs.vector_, rhs.vector_);
      std::swap(lhs.initialized_, rhs.initialized_);
    }

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

  private:
    // Create a new Vector
    void create()
    {
      VecCreateGhost(comm(),
        dofMap_.localSize(), dofMap_.globalSize(), dofMap_.ghostSize(), dofMap_.ghostIndices().data(),
        &vector_);
    }

    // Destroy an old vector if created before
    void destroy()
    {
      if (initialized_)
        VecDestroy(&vector_);

      initialized_ = false;
    }

  private:
    // The local-to-global map
    DofMap const& dofMap_;

    mutable Vec vector_ = nullptr;
    bool initialized_ = false;
  };


  namespace Recursive
  {
    template <class DofMap>
    struct ForEach<PetscVector<DofMap>>
    {
      template <class Vec, class UnaryFunction>
      static void impl(Vec&& vec, UnaryFunction f)
      {
        vec.forEachLocal(f);
      }
    };

    template <class DofMap>
    struct Transform<PetscVector<DofMap>>
    {
      template <class Operation, class... VecIns>
      static void impl(PetscVector<DofMap>& vecOut, Operation op, VecIns const&... vecIn)
      {
        PetscVector<DofMap>::transformLocal(vecOut, op, vecIn...);
      }
    };

    template <class DofMap>
    struct InnerProduct<PetscVector<DofMap>>
    {
      template <class T, class BinOp1, class BinOp2>
      static T impl(PetscVector<DofMap> const& in1, PetscVector<DofMap> const& in2, T init, BinOp1 op1, BinOp2 op2)
      {
        return PetscVector<DofMap>::innerProductLocal(in1, in2, std::move(init), op1, op2);
      }
    };

  } // end namespace Recursive
} // end namespace AMDiS