MatrixNnzStructure.inc.hpp

#pragma once

#include <map>
#include <set>

#include <dune/common/timer.hh>
#include <dune/grid/common/partitionset.hh>

#include <amdis/AdaptiveGrid.hpp>
#include <amdis/Output.hpp>
#if HAVE_MPI
#include <amdis/common/parallel/Communicator.hpp>
#include <amdis/common/parallel/RequestOperations.hpp>
#endif
#include <amdis/linearalgebra/AttributeSet.hpp>

namespace AMDiS {

template <class RowBasis, class ColBasis>
void MatrixNnzStructure::init(RowBasis const& rowBasis, ColBasis const& /*colBasis*/)
{
  Dune::Timer t;

  auto const& basis = rowBasis; // TODO: generalize to row != col basis
  auto const& dofMap = basis.communicator().dofMap();
  std::size_t localSize = dofMap.localSize();
  test_exit(localSize > 0, "Local domain has no owned DOFs!");

  std::vector<std::set<PetscInt>> d_nz(localSize), o_nz(localSize);

  // column indices of row DOFs belonging to remote rank
  // global-row -> {global-col}
  std::map<PetscInt, std::set<PetscInt>> remote_nz;

  // build up local nz structure
  auto localView = basis.localView();
  for (const auto& element : elements(basis.gridView(), Dune::Partitions::interior)) {
    localView.bind(element);
    std::size_t size = localView.size();

    for (std::size_t i = 0; i < size; ++i) {
      auto row = localView.index(i);
      auto global_row = dofMap.global(row);
      if (dofMap.owner(row)) {
        auto local_row = dofMap.globalToLocal(global_row);
        assert(local_row < localSize);
        for (std::size_t j = 0; j < size; ++j) {
          auto col = localView.index(j);
          if (dofMap.owner(col))
            d_nz[local_row].insert(dofMap.global(col));
          else
            o_nz[local_row].insert(dofMap.global(col));
        }
      } else {
        auto& remote_row = remote_nz[global_row];
        for (std::size_t j = 0; j < size; ++j) {
          auto col = localView.index(j);
          remote_row.insert(dofMap.global(col));
        }
      }
    }

    localView.unbind();
  }

  // insert size of sets into diagonal block nnz and off-diagonal block nnz
  dnnz_.resize(localSize, 0);
  onnz_.resize(localSize, 0);

  std::transform(d_nz.begin(), d_nz.end(), dnnz_.begin(), [](auto const& s) { return s.size(); });
  std::transform(o_nz.begin(), o_nz.end(), onnz_.begin(), [](auto const& s) { return s.size(); });

#if HAVE_MPI
  Mpi::Communicator world(Environment::comm());

  // exchange remote rows
  using Attribute = DefaultAttributeSet::Type;
  using RemoteNz = std::vector<PetscInt>;
  std::vector<RemoteNz> sendList(world.size());
  std::vector<std::size_t> receiveList(world.size(), 0);

  for (auto const& rim : basis.communicator().remoteIndices()) {
    int p = rim.first;

    auto* sourceRemoteIndexList = rim.second.first;
    auto* targetRemoteIndexList = rim.second.second;

    // receive from overlap
    for (auto const& ri : *sourceRemoteIndexList) {
      auto const& lip = ri.localIndexPair();
      Attribute remoteAttr = ri.attribute();
      Attribute myAttr = lip.local().attribute();
      if (myAttr == Attribute::owner && remoteAttr == Attribute::overlap) {
        assert(dofMap.owner(lip.local().local()));
        receiveList[p]++;
      }
    }

    // send to owner
    for (auto const& ri : *targetRemoteIndexList) {
      auto const& lip = ri.localIndexPair();
      Attribute remoteAttr = ri.attribute();
      Attribute myAttr = lip.local().attribute();
      if (myAttr == Attribute::overlap && remoteAttr == Attribute::owner) {
        assert(!dofMap.owner(lip.local().local()));

        PetscInt global_row = dofMap.global(lip.local().local());
        assert(dofMap.globalOwner(p, global_row));

        PetscInt remoteDnnz = 0;
        PetscInt remoteOnnz = 0;
        for (PetscInt global_col : remote_nz[global_row]) {
          if (dofMap.globalOwner(p, global_col))
            remoteDnnz++;
          else
            remoteOnnz++;
        }

        auto& data = sendList[p];
        data.push_back(global_row);
        data.push_back(remoteDnnz);
        data.push_back(remoteOnnz);
      }
    }
  }

  // 1. send remote nz structure to owner
  std::vector<Mpi::Request> sendRequests;
  for (int p = 0; p < world.size(); ++p) {
    if (!sendList[p].empty()) {
      sendRequests.emplace_back( world.isend(sendList[p], p, tag_) );
    }
  }

  // 2. receive nz structure belonging to own DOFs from remote node
  std::vector<Mpi::Request> recvRequests;
  std::vector<RemoteNz> recvData(world.size());
  for (int p = 0; p < world.size(); ++p) {
    if (receiveList[p] > 0)
      recvRequests.emplace_back( world.irecv(recvData[p], p, tag_) );
  }

  Mpi::wait_all(recvRequests.begin(), recvRequests.end());

  // 3. insert all remote nz data into my local nnz sets
  for (int p = 0; p < world.size(); ++p) {
    auto const& data = recvData[p];
    assert(data.size() == 3*receiveList[p]);
    for (std::size_t i = 0; i < receiveList[p]; ++i) {
      PetscInt global_row = data[3*i];
      assert(dofMap.globalOwner(global_row));

      PetscInt row_dnnz = data[3*i+1];
      assert(row_dnnz < PetscInt(localSize));

      PetscInt row_onnz = data[3*i+2];
      assert(row_onnz < PetscInt(dofMap.globalSize()));

      auto local_row = dofMap.globalToLocal(global_row);
      assert(local_row < localSize);

      dnnz_[local_row] += row_dnnz;
      onnz_[local_row] += row_onnz;
    }
  }

  // 4. check that the sum of diagonal and off-diagonal entries does not exceed the size of the matrix
  // This may happen for very small size matrices
  for (std::size_t i = 0; i < localSize; ++i) {
    dnnz_[i] = std::min(dnnz_[i], PetscInt(localSize));
    if (onnz_[i] + dnnz_[i] > PetscInt(dofMap.globalSize())) {
      PetscInt diff = onnz_[i] + dnnz_[i] - dofMap.globalSize();
      onnz_[i] -= diff;
    }
  }

  Mpi::wait_all(sendRequests.begin(), sendRequests.end());

#endif // HAVE_MPI

  info(2,"update MatrixNnzStructure need {} sec", t.elapsed());
}

} // end namespace AMDiS