ParallelDomainProblem.cc

#include "ParallelDomainProblem.h"
#include "ProblemScal.h"
#include "ProblemInstat.h"
#include "ParMetisPartitioner.h"
#include "Mesh.h"
#include "Traverse.h"
#include "ElInfo.h"
#include "Element.h"
#include "MacroElement.h"
#include "PartitionElementData.h"
#include "DOFMatrix.h"
#include "DOFVector.h"
#include "VtkWriter.h"

#include "petscksp.h"

namespace AMDiS {

  ParallelDomainProblemBase::ParallelDomainProblemBase(const std::string& name,
						       ProblemIterationInterface *iIF,
						       ProblemTimeInterface *tIF,
						       FiniteElemSpace *fe)
    : iterationIF(iIF),
      timeIF(tIF),
      feSpace(fe),
      mesh(fe->getMesh()),
      initialPartitionMesh(true),
      nRankDOFs(0)
  {
    mpiRank = MPI::COMM_WORLD.Get_rank();
    mpiSize = MPI::COMM_WORLD.Get_size();
    mpiComm = MPI::COMM_WORLD;
    partitioner = new ParMetisPartitioner(mesh, &mpiComm);
  }

  void ParallelDomainProblemBase::initParallelization(AdaptInfo *adaptInfo)
  {
    if (mpiSize <= 1)
      return;

    // create an initial partitioning of the mesh
    partitioner->createPartitionData();
    // set the element weights, which are 1 at the very first begin
    setElemWeights(adaptInfo);
    // and now partition the mesh
    partitionMesh(adaptInfo);   


    /// === Determine to each dof the set of partitions the dof belongs to. ===

    std::map<const DegreeOfFreedom*, std::set<int> > partitionDofs;
    TraverseStack stack;
    ElInfo *elInfo = stack.traverseFirst(mesh, -1, Mesh::CALL_LEAF_EL);
    while (elInfo) {
      Element *element = elInfo->getElement();

      // Determine to each dof the partition(s) it corresponds to.
      for (int i = 0; i < 3; i++) 
	partitionDofs[element->getDOF(i)].insert(partitionVec[element->getIndex()]);
          
      elInfo = stack.traverseNext(elInfo);
    }

    /// === Determine the set of ranks dofs and the dofs ownership at the boundary. ===

    std::vector<const DegreeOfFreedom*> rankDofs;
    for (std::map<const DegreeOfFreedom*, std::set<int> >::iterator it = partitionDofs.begin();
	 it != partitionDofs.end();
	 ++it) {
      for (std::set<int>::iterator itpart1 = it->second.begin();
	   itpart1 != it->second.end();
	   ++itpart1) {
	if (*itpart1 == mpiRank) {
	  if (it->second.size() == 1) {
	    rankDofs.push_back(it->first);
	  } else {	    
	    // This dof is at the ranks boundary. It is owned by the rank only if
	    // the rank number is the highest of all ranks containing this dof.

	    bool insert = true;
	    int highestRank = mpiRank;
	    for (std::set<int>::iterator itpart2 = it->second.begin();
		 itpart2 != it->second.end();
		 ++itpart2) {
	      if (*itpart2 > mpiRank)
		insert = false;

	      if (*itpart2 > highestRank)
		highestRank = *itpart2;
	    }

	    if (insert)
	      rankDofs.push_back(it->first);

	    boundaryDofs[it->first] = highestRank;
	  }
	}
      }
    }


    // === Create interior boundary information ===

    elInfo = stack.traverseFirst(mesh, -1, Mesh::CALL_LEAF_EL | Mesh::FILL_NEIGH);
    while (elInfo) {
      Element *element = elInfo->getElement();

      // Hidde elements which are not part of ranks partition.
      PartitionElementData *partitionData = 
	dynamic_cast<PartitionElementData*>(element->getElementData(PARTITION_ED));   
      if (partitionData->getPartitionStatus() == IN) {
	for (int i = 0; i < 3; i++) {
	  if (!elInfo->getNeighbour(i))
	    continue;

	  PartitionElementData *neighbourPartitionData =
	    dynamic_cast<PartitionElementData*>(elInfo->getNeighbour(i)->getElementData(PARTITION_ED));
 	  if (neighbourPartitionData->getPartitionStatus() == OUT) {
 	    AtomicBoundary& bound = interiorBoundary.
	      getNewAtomicBoundary(partitionVec[elInfo->getNeighbour(i)->getIndex()]);
 	    bound.rankObject.el = element;
 	    bound.rankObject.subObjAtBoundary = EDGE;
 	    bound.rankObject.ithObjAtBoundary = i;
 	    bound.neighbourObject.el = elInfo->getNeighbour(i);
 	    bound.neighbourObject.subObjAtBoundary = EDGE;
 	    bound.neighbourObject.ithObjAtBoundary = -1;
 	  }
	}
      }

      elInfo = stack.traverseNext(elInfo);
    }


    // === Remove all macro elements that are not part of the rank partition. ===

    std::vector<MacroElement*> macrosToRemove;
    for (std::deque<MacroElement*>::iterator it = mesh->firstMacroElement();
	 it != mesh->endOfMacroElements();
	 ++it) {
      PartitionElementData *partitionData = 
	dynamic_cast<PartitionElementData*>
	((*it)->getElement()->getElementData(PARTITION_ED));
      if (partitionData->getPartitionStatus() != IN)
	macrosToRemove.push_back(*it);
    }

    mesh->removeMacroElements(macrosToRemove);

    // === Create local and global dofs ordering. ===

    int *gOrder = (int*)(malloc(sizeof(int) * rankDofs.size()));
    int *lOrder = (int*)(malloc(sizeof(int) * rankDofs.size()));

    for (std::vector<const DegreeOfFreedom*>::iterator it = rankDofs.begin();
	 it != rankDofs.end(); ++it)
      gOrder[nRankDOFs++] = (*it)[0];

    int rstart = 0;
    MPI_Scan(&nRankDOFs, &rstart, 1, MPI_INT, MPI_SUM, PETSC_COMM_WORLD);
    rstart -= nRankDOFs;
   
    for (int i = 0; i < nRankDOFs; i++)
      lOrder[i] = rstart + i;

    AOCreateBasic(PETSC_COMM_WORLD, nRankDOFs, gOrder, lOrder, &applicationOrdering);
    
    free(gOrder);
    free(lOrder);

    /// === Create information which dof indices must be send and which must be received. ===

    std::map<int, std::map<DegreeOfFreedom, DegreeOfFreedom> > sendNewDofs;
    std::map<int, std::vector<DegreeOfFreedom> > recvNewDofs;

    for (std::map<const DegreeOfFreedom*, int>::iterator it = boundaryDofs.begin();
	 it != boundaryDofs.end();
	 ++it) {

      if (it->second == mpiRank) {
	// If the boundary dof is a rank dof, it must be send to other ranks.

	// old global index
	int oldDofIndex = (it->first)[0];
	// search for new dof index in ranks partition for this boundary dof
	int newDofIndex = 0;
	for (int i = 0; i < static_cast<int>(rankDofs.size()); i++) {
	  if (rankDofs[i] == it->first) {
	    newDofIndex = rstart + i;
	    break;
	  }
	}

	// Search for all ranks that have this dof too.
	for (std::set<int>::iterator itRanks = partitionDofs[it->first].begin();
	     itRanks != partitionDofs[it->first].end();
	     ++itRanks) {
	  if (*itRanks != mpiRank) {
	    sendNewDofs[*itRanks][oldDofIndex] = newDofIndex;
	    sendDofs[*itRanks].push_back(newDofIndex);
	  }
	}
      } else {
	// If the boundary dof is not a rank dof, its new dof index, and later
	// also the dof values, must be received from another rank.
	recvNewDofs[it->second].push_back((it->first)[0]);
      }
    }

    /// === Send and receive the dof indices at boundary. ===

    std::vector<int*> sendBuffers(sendNewDofs.size());
    std::vector<int*> recvBuffers(recvNewDofs.size());
    
    int i = 0;
    for (std::map<int, std::map<DegreeOfFreedom, DegreeOfFreedom> >::iterator sendIt = sendNewDofs.begin();
	 sendIt != sendNewDofs.end();
	 ++sendIt, i++) {
      sendBuffers[i] = new int[sendIt->second.size() * 2];
      int c = 0;
      for (std::map<DegreeOfFreedom, DegreeOfFreedom>::iterator dofIt = sendIt->second.begin();
	   dofIt != sendIt->second.end();
	   ++dofIt, c += 2) {
	sendBuffers[i][c] = dofIt->first;
	sendBuffers[i][c + 1] = dofIt->second;
      }

      mpiComm.Isend(sendBuffers[i], sendIt->second.size() * 2, MPI_INT, sendIt->first, 0);
    }

    i = 0;
    for (std::map<int, std::vector<DegreeOfFreedom> >::iterator recvIt = recvNewDofs.begin();
	 recvIt != recvNewDofs.end();
	 ++recvIt, i++) {
      recvBuffers[i] = new int[recvIt->second.size() * 2];
      
      mpiComm.Irecv(recvBuffers[i], recvIt->second.size() * 2, MPI_INT, recvIt->first, 0);
    }


    mpiComm.Barrier();


    /// === Reset all DOFAdmins of the mesh. ===

    int nAdmins = mesh->getNumberOfDOFAdmin();
    for (int i = 0; i < nAdmins; i++) 
      for (int j = 0; j < mesh->getDOFAdmin(i).getSize(); j++)
	const_cast<DOFAdmin&>(mesh->getDOFAdmin(i)).setDOFFree(j, true);
    

    /// === Change dof indices at boundary from other ranks. ===

    i = 0;
    for (std::map<int, std::vector<DegreeOfFreedom> >::iterator recvIt = recvNewDofs.begin();
	 recvIt != recvNewDofs.end();
	 ++recvIt, i++) {

      for (int j = 0; j < static_cast<int>(recvIt->second.size()); j++) {
	for (std::map<const DegreeOfFreedom*, int>::iterator dofIt = boundaryDofs.begin();
	     dofIt != boundaryDofs.end();
	     ++dofIt) {
	  if ((dofIt->first)[0] == recvBuffers[i][j * 2]) {
	    int newDof = recvBuffers[i][j * 2 + 1];

	    recvDofs[recvIt->first].push_back(newDof);

	    const_cast<DegreeOfFreedom*>(dofIt->first)[0] = newDof;

	    for (int k = 0; k < nAdmins; k++) 
	      const_cast<DOFAdmin&>(mesh->getDOFAdmin(k)).setDOFFree(newDof, false);

	    break;
	  }
	}
      }

      delete [] recvBuffers[i];
    }

    i = 0;
    for (std::map<int, std::map<DegreeOfFreedom, DegreeOfFreedom> >::iterator sendIt = sendNewDofs.begin();
	 sendIt != sendNewDofs.end();
	 ++sendIt, i++) {
      delete [] sendBuffers[i];
    }

    /// === Change dof indices for rank partition. ===

    for (int i = 0; i < static_cast<int>(rankDofs.size()); i++) {
      const_cast<DegreeOfFreedom*>(rankDofs[i])[0] = rstart + i; 

      for (int k = 0; k < nAdmins; k++) 
	const_cast<DOFAdmin&>(mesh->getDOFAdmin(k)).setDOFFree(rstart + i, false);

    }

    /// === Create petsc matrix. ===
    int ierr;
    ierr = MatCreate(PETSC_COMM_WORLD, &petscMatrix);
    ierr = MatSetSizes(petscMatrix, rankDofs.size(), rankDofs.size(),
		       partitionDofs.size(), partitionDofs.size());
    ierr = MatSetType(petscMatrix, MATAIJ);

    ierr = VecCreate(PETSC_COMM_WORLD, &petscRhsVec);
    ierr = VecSetSizes(petscRhsVec, rankDofs.size(), partitionDofs.size());    
    ierr = VecSetType(petscRhsVec, VECMPI);

    ierr = VecCreate(PETSC_COMM_WORLD, &petscSolVec);
    ierr = VecSetSizes(petscSolVec, rankDofs.size(), partitionDofs.size());    
    ierr = VecSetType(petscSolVec, VECMPI);
  }

  void ParallelDomainProblemBase::exitParallelization(AdaptInfo *adaptInfo)
  {
    AODestroy(applicationOrdering);
  }

  void ParallelDomainProblemBase::fillPetscMatrix(DOFMatrix *mat, 
						  DOFVector<double> *vec)
  {
    using mtl::tag::major; using mtl::tag::nz; using mtl::begin; using mtl::end;
    namespace traits= mtl::traits;
    typedef DOFMatrix::base_matrix_type Matrix;

    traits::row<Matrix>::type row(mat->getBaseMatrix());
    traits::col<Matrix>::type col(mat->getBaseMatrix());
    traits::const_value<Matrix>::type value(mat->getBaseMatrix());

    typedef traits::range_generator<major, Matrix>::type cursor_type;
    typedef traits::range_generator<nz, cursor_type>::type icursor_type;

    std::cout.precision(10);
    for (cursor_type cursor = begin<major>(mat->getBaseMatrix()), cend = end<major>(mat->getBaseMatrix()); cursor != cend; ++cursor)
      for (icursor_type icursor = begin<nz>(cursor), icend = end<nz>(cursor); icursor != icend; ++icursor)
	if (value(*icursor) != 0.0) {
	  int r = row(*icursor);
	  int c = col(*icursor);
	  double v = value(*icursor);
	  MatSetValues(petscMatrix, 1, &r, 1, &c, &v, ADD_VALUES);
	}


    MatAssemblyBegin(petscMatrix, MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(petscMatrix, MAT_FINAL_ASSEMBLY);

    DOFVector<double>::Iterator dofIt(vec, USED_DOFS);
    for (dofIt.reset(); !dofIt.end(); ++dofIt) {
      int index = dofIt.getDOFIndex();
      double value = *dofIt;

      VecSetValues(petscRhsVec, 1, &index, &value, ADD_VALUES);
    }
  }

  void ParallelDomainProblemBase::solvePetscMatrix(DOFVector<double> *vec)
  {
    KSP ksp;
    PC pc;

    KSPCreate(PETSC_COMM_WORLD, &ksp);
    KSPSetOperators(ksp, petscMatrix, petscMatrix, DIFFERENT_NONZERO_PATTERN);
    KSPGetPC(ksp, &pc);
    PCSetType(pc, PCJACOBI);
    KSPSetTolerances(ksp, 1.e-7, PETSC_DEFAULT, PETSC_DEFAULT, PETSC_DEFAULT);
    KSPSetType(ksp, KSPBCGS);
    KSPMonitorSet(ksp, KSPMonitorDefault, PETSC_NULL, 0);
    KSPSolve(ksp, petscRhsVec, petscSolVec);

    PetscScalar *vecPointer;
    VecGetArray(petscSolVec, &vecPointer);

    DOFVector<double>::Iterator dofIt(vec, USED_DOFS);
    int counter = 0;
    for (dofIt.reset(); !dofIt.end(); ++dofIt)
      *dofIt = vecPointer[counter++];

    VecRestoreArray(petscSolVec, &vecPointer);

    std::vector<double*> sendBuffers(sendDofs.size());
    std::vector<double*> recvBuffers(recvDofs.size());
    
    int i = 0;
    for (std::map<int, std::vector<DegreeOfFreedom> >::iterator sendIt = sendDofs.begin();
	 sendIt != sendDofs.end();
	 ++sendIt, i++) {
      sendBuffers[i] = new double[sendIt->second.size()];

      for (int j = 0; j < sendIt->second.size(); j++)
	sendBuffers[i][j] = (*vec)[(sendIt->second)[j]];

      mpiComm.Isend(sendBuffers[i], sendIt->second.size(), MPI_DOUBLE, sendIt->first, 0);
    }

    i = 0;
    for (std::map<int, std::vector<DegreeOfFreedom> >::iterator recvIt = recvDofs.begin();
	 recvIt != recvDofs.end();
	 ++recvIt, i++) {
      recvBuffers[i] = new double[recvIt->second.size()];

      mpiComm.Irecv(recvBuffers[i], recvIt->second.size(), MPI_DOUBLE, recvIt->first, 0);
    }

    
    mpiComm.Barrier();
    
    i = 0;
    for (std::map<int, std::vector<DegreeOfFreedom> >::iterator recvIt = recvDofs.begin();
	 recvIt != recvDofs.end();
	 ++recvIt, i++) {
      for (int j = 0; j < recvIt->second.size(); j++) {
	if (MPI::COMM_WORLD.Get_rank() == 0)
	  std::cout << (recvIt->second)[j] << " = " << recvBuffers[i][j] << std::endl;      

	(*vec)[(recvIt->second)[j]] = recvBuffers[i][j];
      }

      delete [] recvBuffers[i];
    }
    
    for (int i = 0; i < sendBuffers.size(); i++)
      delete [] sendBuffers[i];    

    std::cout << "TEST" << std::endl;
    
    if (MPI::COMM_WORLD.Get_rank() == 0) {
      DOFVector<double>::Iterator it(vec, USED_DOFS);
      for (it.reset(); !it.end(); ++it) {
	std::cout << it.getDOFIndex() << " = " << *it << std::endl;
      }
    }
  }

  double ParallelDomainProblemBase::setElemWeights(AdaptInfo *adaptInfo) 
  {
    double localWeightSum = 0.0;
    int elNum = -1;

    elemWeights.clear();

    TraverseStack stack;
    ElInfo *elInfo = stack.traverseFirst(mesh, -1,
					 Mesh::CALL_EVERY_EL_PREORDER);
    while (elInfo) {
      Element *element = elInfo->getElement();

      // get partition data
      PartitionElementData *partitionData = dynamic_cast<PartitionElementData*>
	(element->getElementData(PARTITION_ED));

      if (partitionData && partitionData->getPartitionStatus() == IN) {
	if (partitionData->getLevel() == 0) {
	  elNum = element->getIndex();
	}
	TEST_EXIT(elNum != -1)("invalid element number\n");
	if (element->isLeaf()) {
	  elemWeights[elNum] += 1.0;
	  localWeightSum += 1.0;
	}
      }

      elInfo = stack.traverseNext(elInfo);
    }

    return localWeightSum;
  }

  void ParallelDomainProblemBase::partitionMesh(AdaptInfo *adaptInfo)
  {
    if (initialPartitionMesh) {
      initialPartitionMesh = false;
      partitioner->fillCoarsePartitionVec(&oldPartitionVec);
      partitioner->partition(&elemWeights, INITIAL);
    } else {
      oldPartitionVec = partitionVec;
      partitioner->partition(&elemWeights, ADAPTIVE_REPART, 100.0 /*0.000001*/);
    }    

    partitioner->fillCoarsePartitionVec(&partitionVec);
  }

  ParallelDomainProblemScal::ParallelDomainProblemScal(const std::string& name,
						       ProblemScal *problem,
						       ProblemInstatScal *problemInstat)
    : ParallelDomainProblemBase(name, problem, problemInstat, problem->getFESpace()),
      probScal(problem)
  {
  }

  Flag ParallelDomainProblemScal::oneIteration(AdaptInfo *adaptInfo, Flag toDo)
  {
    //    return iterationIF->oneIteration(adaptInfo, toDo);

    Flag flag = dynamic_cast<StandardProblemIteration*>(iterationIF)->buildAndAdapt(adaptInfo, toDo);

    fillPetscMatrix(probScal->getSystemMatrix(), probScal->getRHS());

    solvePetscMatrix(probScal->getSolution());

//     if (toDo.isSet(SOLVE))
//       iterationIF->getProblem()->solve(adaptInfo, false);

//     if (toDo.isSet(SOLVE_RHS))
//       iterationIF->getProblem()->solve(adaptInfo, true);

//     if (toDo.isSet(ESTIMATE)) 
//       iterationIF->getProblem()->estimate(adaptInfo);

    return flag;

  }

}