PetscSolverSchur.cc

//
// Software License for AMDiS
//
// Copyright (c) 2010 Dresden University of Technology 
// All rights reserved.
// Authors: Simon Vey, Thomas Witkowski et al.
//
// This file is part of AMDiS
//
// See also license.opensource.txt in the distribution.


#include "parallel/PetscSolverSchur.h"
#include "parallel/StdMpi.h"
#include "parallel/MpiHelper.h"

namespace AMDiS {

  using namespace std;

  void PetscSolverSchur::updateDofData(int nComponents)
  {
    FUNCNAME("PetscSolverSchur::updateDofData()");

    TEST_EXIT_DBG(meshDistributor)("No mesh distributor object defined!\n");

    MPI::Intracomm& mpiComm = meshDistributor->getMpiComm();
    const FiniteElemSpace *feSpace = meshDistributor->getFeSpace(0);

    typedef map<int, DofContainer> RankToDofContainer;
    typedef map<DegreeOfFreedom, bool> DofIndexToBool;

    boundaryDofs.clear();
    std::set<DegreeOfFreedom> boundaryLocalDofs;
    for (DofComm::Iterator it(meshDistributor->getSendDofs(), feSpace);
	 !it.end(); it.nextRank())
      for (; !it.endDofIter(); it.nextDof()) {
	boundaryLocalDofs.insert(it.getDofIndex());	  
	boundaryDofs.insert(meshDistributor->mapDofToGlobal(feSpace, it.getDofIndex()));
      }

      
    nBoundaryDofs = boundaryDofs.size();
    mpi::getDofNumbering(mpiComm, nBoundaryDofs, 
			 rStartBoundaryDofs, nOverallBoundaryDofs);


    DofContainerSet& edgeDofs = 
      meshDistributor->getBoundaryDofInfo(feSpace).geoDofs[EDGE];
    DofContainerSet& vertexDofs = 
      meshDistributor->getBoundaryDofInfo(feSpace).geoDofs[VERTEX];
    int nEdgeDofs = edgeDofs.size();
    int nVertexDofs = vertexDofs.size();

    TEST_EXIT_DBG(nEdgeDofs + nVertexDofs == nBoundaryDofs)
      ("Should not happen!\n");

    int rStartEdgeDofs, nOverallEdgeDofs;
    mpi::getDofNumbering(mpiComm, nEdgeDofs, 
			 rStartEdgeDofs, nOverallEdgeDofs);

    int rStartVertexDofs, nOverallVertexDofs;
    mpi::getDofNumbering(mpiComm, nVertexDofs, 
			 rStartVertexDofs, nOverallVertexDofs);

    TEST_EXIT_DBG(nOverallEdgeDofs + nOverallVertexDofs == nOverallBoundaryDofs)
      ("Should not happen!\n");
    

    mapGlobalBoundaryDof.clear();
#if 1
    {
      int counter = rStartEdgeDofs;
      for (DofContainerSet::iterator it = edgeDofs.begin();
	   it != edgeDofs.end(); ++it)
	mapGlobalBoundaryDof[meshDistributor->mapDofToGlobal(feSpace, **it)] = 
	  counter++;
    }
    {
      int counter = nOverallEdgeDofs + rStartVertexDofs;
      for (DofContainerSet::iterator it = vertexDofs.begin();
	   it != vertexDofs.end(); ++it)
	mapGlobalBoundaryDof[meshDistributor->mapDofToGlobal(feSpace, **it)] = 
	  counter++;
    }
#else
 {
    int counter = rStartBoundaryDofs;
    for (std::set<DegreeOfFreedom>::iterator it = boundaryDofs.begin();
	 it != boundaryDofs.end(); ++it)
      mapGlobalBoundaryDof[*it] = counter++;
 }
#endif


    std::set<DegreeOfFreedom> otherBoundaryLocalDofs;
    for (DofComm::Iterator it(meshDistributor->getRecvDofs(), feSpace);
	 !it.end(); it.nextRank())
      for (; !it.endDofIter(); it.nextDof())
	otherBoundaryLocalDofs.insert(it.getDofIndex());
      
    interiorDofs.clear();
    DofIndexToBool& isRankDof = meshDistributor->getIsRankDof(feSpace);
    for (DofIndexToBool::iterator dofIt = isRankDof.begin(); 
	 dofIt != isRankDof.end(); ++dofIt) {
      if (dofIt->second && 
	  boundaryLocalDofs.count(dofIt->first) == 0 && 
	  otherBoundaryLocalDofs.count(dofIt->first) == 0)
	interiorDofs.insert(meshDistributor->mapDofToGlobal(feSpace, dofIt->first));
    }
      
    nInteriorDofs = interiorDofs.size();
    mpi::getDofNumbering(mpiComm, nInteriorDofs, 
			 rStartInteriorDofs, nOverallInteriorDofs);

    {
      int counter = rStartInteriorDofs;
      mapGlobalInteriorDof.clear();
      for (std::set<DegreeOfFreedom>::iterator it = interiorDofs.begin();
	   it != interiorDofs.end(); ++it)
	mapGlobalInteriorDof[*it] = counter++;
    }


    TEST_EXIT_DBG(nInteriorDofs > 0)("Should not happen!\n");


    StdMpi<vector<DegreeOfFreedom> > stdMpi(mpiComm);
    for (DofComm::Iterator it(meshDistributor->getSendDofs(), feSpace);
	 !it.end(); it.nextRank()) {
      stdMpi.getSendData(it.getRank()).resize(0);
      stdMpi.getSendData(it.getRank()).reserve(it.getDofs().size());

      for (; !it.endDofIter(); it.nextDof()) {
	int globalSendDof = 
	  meshDistributor->mapDofToGlobal(feSpace, it.getDofIndex());

	TEST_EXIT_DBG(mapGlobalBoundaryDof.count(globalSendDof))
	  ("No mapping for boundary DOF %d!\n", globalSendDof);

	stdMpi.getSendData(it.getRank()).push_back(mapGlobalBoundaryDof[globalSendDof]);
      }
    }

    stdMpi.updateSendDataSize();

    for (DofComm::Iterator it(meshDistributor->getRecvDofs(), feSpace);
	 !it.end(); it.nextRank())
      stdMpi.recv(it.getRank());

    stdMpi.startCommunication();

    for (DofComm::Iterator it(meshDistributor->getRecvDofs(), feSpace);
	 !it.end(); it.nextRank())
      for (; !it.endDofIter(); it.nextDof()) {
	int globalRecvDof = 
	  meshDistributor->mapDofToGlobal(feSpace, it.getDofIndex());
	mapGlobalBoundaryDof[globalRecvDof] = 
	  stdMpi.getRecvData(it.getRank())[it.getDofCounter()];
	boundaryDofs.insert(globalRecvDof);
      }


    // === Create PETSc IS structurs for interior and boundary DOFs. ===

    ISCreateStride(mpiComm, 
		   nInteriorDofs * nComponents,
		   (rStartInteriorDofs + rStartBoundaryDofs) * nComponents, 
		   1, &interiorIs);

    ISCreateStride(mpiComm,
		   nBoundaryDofs * nComponents,
		   (rStartInteriorDofs + rStartBoundaryDofs + nInteriorDofs) * nComponents, 
		   1, &boundaryIs);
  }


  void PetscSolverSchur::fillPetscMatrix(Matrix<DOFMatrix*> *mat)
  {
    FUNCNAME("PetscSolverSchur::fillPetscMatrix()");

    const FiniteElemSpace *feSpace = meshDistributor->getFeSpace(0);
    int nComponents = mat->getNumRows();
    updateDofData(nComponents);

    int nInteriorRows = nInteriorDofs * nComponents;
    int nOverallInteriorRows = nOverallInteriorDofs * nComponents;
    int nBoundaryRows = nBoundaryDofs * nComponents;
    int nOverallBoundaryRows = nOverallBoundaryDofs * nComponents;


    MatCreateMPIAIJ(mpiComm, 
		    nInteriorRows, nInteriorRows, 
		    nOverallInteriorRows, nOverallInteriorRows,
		    100, PETSC_NULL, 100, PETSC_NULL, &matA11);

    MatCreateMPIAIJ(mpiComm, 
		    nBoundaryRows, nBoundaryRows, 
		    nOverallBoundaryRows, nOverallBoundaryRows,
		    100, PETSC_NULL, 100, PETSC_NULL, &matA22);

    MatCreateMPIAIJ(mpiComm, 
		    nInteriorRows, nBoundaryRows, 
		    nOverallInteriorRows, nOverallBoundaryRows,
		    100, PETSC_NULL, 100, PETSC_NULL, &matA12);    

    MatCreateMPIAIJ(mpiComm, 
		    nBoundaryRows, nInteriorRows, 
		    nOverallBoundaryRows, nOverallInteriorRows,
		    100, PETSC_NULL, 100, PETSC_NULL, &matA21);


    for (int i = 0; i < nComponents; i++)
      for (int j = 0; j < nComponents; j++)
	if ((*mat)[i][j])
	  setDofMatrix((*mat)[i][j], nComponents, i, j);

    MatAssemblyBegin(matA11, MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(matA11, MAT_FINAL_ASSEMBLY);

    MatAssemblyBegin(matA12, MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(matA12, MAT_FINAL_ASSEMBLY);

    MatAssemblyBegin(matA21, MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(matA21, MAT_FINAL_ASSEMBLY);

    MatAssemblyBegin(matA22, MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(matA22, MAT_FINAL_ASSEMBLY);

    Mat tmpMat[2][2];
    tmpMat[0][0] = matA11;
    tmpMat[0][1] = matA12;
    tmpMat[1][0] = matA21;
    tmpMat[1][1] = matA22;

    IS tmpIS[2];
    tmpIS[0] = interiorIs;
    tmpIS[1] = boundaryIs;

    MatCreateNest(mpiComm, 2, &tmpIS[0], 2, &tmpIS[0], &tmpMat[0][0], &petscMatrix);
    MatNestSetVecType(petscMatrix, VECNEST);
    MatAssemblyBegin(petscMatrix, MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(petscMatrix, MAT_FINAL_ASSEMBLY);


    int nRankRows = meshDistributor->getNumberRankDofs(feSpace) * nComponents;
    int nOverallRows = meshDistributor->getNumberOverallDofs(feSpace) * nComponents;

    VecCreateMPI(mpiComm, nRankRows, nOverallRows, &petscSolVec);
    VecCreateMPI(mpiComm, nRankRows, nOverallRows, &petscTmpVec);
  }


  void PetscSolverSchur::fillPetscRhs(SystemVector *vec)
  {
    FUNCNAME("PetscSolverSchur::fillPetscRhs()");

    const FiniteElemSpace *feSpace = meshDistributor->getFeSpace(0);
    int nComponents = vec->getSize();
    int nRankRows = meshDistributor->getNumberRankDofs(feSpace) * nComponents;
    int nOverallRows = meshDistributor->getNumberOverallDofs(feSpace) * nComponents;

    VecCreateMPI(mpiComm, nRankRows, nOverallRows, &petscRhsVec);

    for (int i = 0; i < nComponents; i++)
      setDofVector(petscRhsVec, vec->getDOFVector(i), nComponents, i);

    VecAssemblyBegin(petscRhsVec);
    VecAssemblyEnd(petscRhsVec);
  }


  void PetscSolverSchur::solvePetscMatrix(SystemVector &vec, 
					  AdaptInfo *adaptInfo)
  {
    FUNCNAME("PetscSolverSchur::solvePetscMatrix()");

    const FiniteElemSpace *feSpace = meshDistributor->getFeSpace(0);
    int nComponents = vec.getSize();

    KSPCreate(mpiComm, &solver);

    KSPSetOperators(solver, petscMatrix, petscMatrix, SAME_NONZERO_PATTERN); 
    KSPSetTolerances(solver, 0.0, 1e-8, PETSC_DEFAULT, PETSC_DEFAULT);
    KSPSetFromOptions(solver);

    KSPGetPC(solver, &pc);
    PCSetType(pc, PCFIELDSPLIT);
    PCFieldSplitSetIS(pc, "interior", interiorIs);
    PCFieldSplitSetIS(pc, "boundary", boundaryIs);
    PCSetFromOptions(pc);


    KSPSolve(solver, petscRhsVec, petscSolVec);

    // === Transfere values from PETSc's solution vectors to AMDiS vectors. ===

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

    for (int i = 0; i < nComponents; i++) {
      DOFVector<double>::Iterator dofIt(vec.getDOFVector(i), USED_DOFS);
      for (dofIt.reset(); !dofIt.end(); ++dofIt) {
	DegreeOfFreedom globalRowDof = 
	  meshDistributor->mapDofToGlobal(feSpace, dofIt.getDOFIndex());
	if (boundaryDofs.count(globalRowDof)) {
	  int index =
	    (mapGlobalBoundaryDof[globalRowDof] - rStartBoundaryDofs + nInteriorDofs) * (i + 1);
	  *dofIt = vecPointer[index];
	} else {
	  int index = 
	    (mapGlobalInteriorDof[globalRowDof] - rStartInteriorDofs) * (i + 1);
	  *dofIt = vecPointer[index];
	}
      }
    }

    VecRestoreArray(petscSolVec, &vecPointer);


    // === Synchronize DOFs at common DOFs, i.e., DOFs that correspond to ===
    // === more than one partition.                                       ===
    meshDistributor->synchVector(vec);



    // === Print iteration counter and residual norm of the solution. ===
    printSolutionInfo(adaptInfo);


    // === Destroy PETSC's variables. ===

    VecDestroy(&petscRhsVec);
    VecDestroy(&petscSolVec);
    VecDestroy(&petscTmpVec);

    MatDestroy(&matA11);
    MatDestroy(&matA12);
    MatDestroy(&matA21);
    MatDestroy(&matA22);
    MatDestroy(&petscMatrix);

    KSPDestroy(&solver);
  }
  

  void PetscSolverSchur::setDofMatrix(DOFMatrix* mat, int dispMult, 
				      int dispAddRow, int dispAddCol)
  {
    FUNCNAME("PetscSolverSchur::setDofMatrix()");

    TEST_EXIT(mat)("No DOFMatrix!\n");

    const FiniteElemSpace* feSpace = meshDistributor->getFeSpace(0);
    using mtl::tag::row; using mtl::tag::nz; using mtl::begin; using mtl::end;
    namespace traits= mtl::traits;
    typedef DOFMatrix::base_matrix_type Matrix;

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

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

    vector<int> colsBoundary, colsInterior;
    vector<double> valuesBoundary, valuesInterior;
    colsBoundary.reserve(300);
    colsInterior.reserve(300);
    valuesBoundary.reserve(300);
    valuesInterior.reserve(300);

    for (cursor_type cursor = begin<row>(mat->getBaseMatrix()), 
	   cend = end<row>(mat->getBaseMatrix()); cursor != cend; ++cursor) {

      // Global index of the current row DOF.
      int globalRowDof = meshDistributor->mapDofToGlobal(feSpace, *cursor);
     
      colsBoundary.clear();
      colsInterior.clear();
      valuesBoundary.clear();
      valuesInterior.clear();
      
      for (icursor_type icursor = begin<nz>(cursor), icend = end<nz>(cursor); 
	   icursor != icend; ++icursor) {
	int globalColDof = meshDistributor->mapDofToGlobal(feSpace, col(*icursor));

	if (boundaryDofs.count(globalColDof)) {
	  TEST_EXIT_DBG(mapGlobalBoundaryDof.count(globalColDof))
	    ("Should not happen!\n");

	  int colIndex = 
	    mapGlobalBoundaryDof[globalColDof] * dispMult + dispAddCol;
	  
	  colsBoundary.push_back(colIndex);
	  valuesBoundary.push_back(value(*icursor));
	} else {
	  TEST_EXIT_DBG(mapGlobalInteriorDof.count(globalColDof))
	    ("Cannot find global interior mapping for global column DOF %d!\n", 
	     globalColDof);

	  int colIndex = 
	    mapGlobalInteriorDof[globalColDof] * dispMult + dispAddCol;

	  colsInterior.push_back(colIndex);
	  valuesInterior.push_back(value(*icursor));
	}
      }

      if (boundaryDofs.count(globalRowDof)) {
	TEST_EXIT_DBG(mapGlobalBoundaryDof.count(globalRowDof))
	  ("Should not happen!\n");

	int rowIndex = 
	  mapGlobalBoundaryDof[globalRowDof] * dispMult + dispAddRow;

 	MatSetValues(matA22, 1, &rowIndex, colsBoundary.size(), 
 		     &(colsBoundary[0]), &(valuesBoundary[0]), ADD_VALUES);
 	MatSetValues(matA21, 1, &rowIndex, colsInterior.size(), 
 		     &(colsInterior[0]), &(valuesInterior[0]), ADD_VALUES);
      } else {
	TEST_EXIT_DBG(mapGlobalInteriorDof.count(globalRowDof))
	  ("Cannot find global interior mapping for global row DOF %d!\n", 
	   globalRowDof);

	int rowIndex = 
	  mapGlobalInteriorDof[globalRowDof] * dispMult + dispAddRow;

  	MatSetValues(matA11, 1, &rowIndex, colsInterior.size(), 
  		     &(colsInterior[0]), &(valuesInterior[0]), ADD_VALUES);
  	MatSetValues(matA12, 1, &rowIndex, colsBoundary.size(), 
  		     &(colsBoundary[0]), &(valuesBoundary[0]), ADD_VALUES);
      }
    }
  }


  void PetscSolverSchur::setDofVector(Vec& petscVec, DOFVector<double>* vec,
				      int dispMult, int dispAdd, bool rankOnly)
  {
    const FiniteElemSpace *feSpace = meshDistributor->getFeSpace(0);

    DOFVector<double>::Iterator dofIt(vec, USED_DOFS);
    for (dofIt.reset(); !dofIt.end(); ++dofIt) {
      if (rankOnly && !meshDistributor->getIsRankDof(feSpace, dofIt.getDOFIndex()))
	continue;

      // Calculate global row index of the DOF.
      DegreeOfFreedom globalRowDof = 
	meshDistributor->mapDofToGlobal(feSpace, dofIt.getDOFIndex());

      double value = *dofIt;

      if (boundaryDofs.count(globalRowDof)) {
	TEST_EXIT_DBG(mapGlobalBoundaryDof.count(globalRowDof))
	  ("Should not happen!\n");

	int index = 
	  (rStartInteriorDofs + 
	   nInteriorDofs +
	   mapGlobalBoundaryDof[globalRowDof]) * dispMult + dispAdd;
	VecSetValues(petscRhsVec, 1, &index, &value, ADD_VALUES);
      } else {
	TEST_EXIT_DBG(mapGlobalInteriorDof.count(globalRowDof))
	  ("Should not happen!\n");
	
	int index = 
	  (rStartBoundaryDofs + 
	   mapGlobalInteriorDof[globalRowDof]) * dispMult + dispAdd;
	VecSetValues(petscRhsVec, 1, &index, &value, ADD_VALUES);
      }
    }
  }

}