PetscSolverFeti.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/PetscSolverFeti.h"
#include "parallel/PetscSolverFetiStructs.h"
#include "parallel/StdMpi.h"
#include "parallel/MpiHelper.h"
#include "io/VtkWriter.h"

namespace AMDiS {

  using namespace std;

  // y = mat * x
  int petscMultMatSchurPrimal(Mat mat, Vec x, Vec y)
  {
    // S_PiPi = K_PiPi - K_PiB inv(K_BB) K_BPi

    void *ctx;
    MatShellGetContext(mat, &ctx);
    SchurPrimalData* data = static_cast<SchurPrimalData*>(ctx);

    MatMult(*(data->mat_b_primal), x, data->tmp_vec_b);
    data->fetiSolver->solveLocalProblem(data->tmp_vec_b, data->tmp_vec_b);
    MatMult(*(data->mat_primal_b), data->tmp_vec_b, data->tmp_vec_primal);
    MatMult(*(data->mat_primal_primal), x, y);
    VecAXPBY(y, -1.0, 1.0, data->tmp_vec_primal);

    return 0;
  }


  // y = mat * x
  int petscMultMatFeti(Mat mat, Vec x, Vec y)
  {
    //    F = L inv(K_BB) trans(L) + L inv(K_BB) K_BPi inv(S_PiPi) K_PiB inv(K_BB) trans(L)
    // => F = L [I + inv(K_BB) K_BPi inv(S_PiPi) K_PiB] inv(K_BB) trans(L)

    void *ctx;
    MatShellGetContext(mat, &ctx);
    FetiData* data = static_cast<FetiData*>(ctx);

    MatMultTranspose(*(data->mat_lagrange), x, data->tmp_vec_b);
    data->fetiSolver->solveLocalProblem(data->tmp_vec_b, data->tmp_vec_b);
    MatMult(*(data->mat_lagrange), data->tmp_vec_b, data->tmp_vec_lagrange);

    MatMult(*(data->mat_primal_b), data->tmp_vec_b, data->tmp_vec_primal);
    KSPSolve(*(data->ksp_schur_primal), data->tmp_vec_primal, data->tmp_vec_primal);
    MatMult(*(data->mat_b_primal), data->tmp_vec_primal, data->tmp_vec_b);
    data->fetiSolver->solveLocalProblem(data->tmp_vec_b, data->tmp_vec_b);
    MatMult(*(data->mat_lagrange), data->tmp_vec_b, y);

    VecAXPBY(y, 1.0, 1.0, data->tmp_vec_lagrange);

    return 0;
  }


  // y = PC * x
  PetscErrorCode petscApplyFetiDirichletPrecon(PC pc, Vec x, Vec y)
  {
    // Get data for the preconditioner
    void *ctx;
    PCShellGetContext(pc, &ctx);
    FetiDirichletPreconData* data = static_cast<FetiDirichletPreconData*>(ctx);

    // Multiply with scaled Lagrange constraint matrix.
    MatMultTranspose(*(data->mat_lagrange_scaled), x, data->tmp_vec_b);


    // === Restriction of the B nodes to the boundary nodes. ===

    int nLocalB;
    int nLocalDuals;
    VecGetLocalSize(data->tmp_vec_b, &nLocalB);
    VecGetLocalSize(data->tmp_vec_duals0, &nLocalDuals);

    PetscScalar *local_b, *local_duals;
    VecGetArray(data->tmp_vec_b, &local_b);
    VecGetArray(data->tmp_vec_duals0, &local_duals);

    for (int i = nLocalB - nLocalDuals, j = 0; i < nLocalB; i++, j++)
      local_duals[j] = local_b[i];

    VecRestoreArray(data->tmp_vec_b, &local_b);
    VecRestoreArray(data->tmp_vec_duals0, &local_duals);


    // === K_DD - K_DI inv(K_II) K_ID ===

    MatMult(*(data->mat_duals_duals), data->tmp_vec_duals0, data->tmp_vec_duals1);

    MatMult(*(data->mat_interior_duals), data->tmp_vec_duals0, data->tmp_vec_interior);
    KSPSolve(*(data->ksp_interior), data->tmp_vec_interior, data->tmp_vec_interior);
    MatMult(*(data->mat_duals_interior), data->tmp_vec_interior, data->tmp_vec_duals0);

    VecAXPBY(data->tmp_vec_duals0, 1.0, -1.0, data->tmp_vec_duals1);


    // === Prolongation from local dual nodes to B nodes.

    VecGetArray(data->tmp_vec_b, &local_b);
    VecGetArray(data->tmp_vec_duals0, &local_duals);

    for (int i = nLocalB - nLocalDuals, j = 0; i < nLocalB; i++, j++)
      local_b[i] = local_duals[j];

    VecRestoreArray(data->tmp_vec_b, &local_b);
    VecRestoreArray(data->tmp_vec_duals0, &local_duals);


    // Multiply with scaled Lagrange constraint matrix.
    MatMult(*(data->mat_lagrange_scaled), data->tmp_vec_b, y);

    return 0;
  }


  // y = PC * x
  PetscErrorCode petscApplyFetiLumpedPrecon(PC pc, Vec x, Vec y)
  {
    // Get data for the preconditioner
    void *ctx;
    PCShellGetContext(pc, &ctx);
    FetiLumpedPreconData* data = static_cast<FetiLumpedPreconData*>(ctx);

    // Multiply with scaled Lagrange constraint matrix.
    MatMultTranspose(*(data->mat_lagrange_scaled), x, data->tmp_vec_b);


    // === Restriction of the B nodes to the boundary nodes. ===

    int nLocalB;
    int nLocalDuals;
    VecGetLocalSize(data->tmp_vec_b, &nLocalB);
    VecGetLocalSize(data->tmp_vec_duals0, &nLocalDuals);

    PetscScalar *local_b, *local_duals;
    VecGetArray(data->tmp_vec_b, &local_b);
    VecGetArray(data->tmp_vec_duals0, &local_duals);

    for (int i = nLocalB - nLocalDuals, j = 0; i < nLocalB; i++, j++)
      local_duals[j] = local_b[i];

    VecRestoreArray(data->tmp_vec_b, &local_b);
    VecRestoreArray(data->tmp_vec_duals0, &local_duals);


    // === K_DD ===

    MatMult(*(data->mat_duals_duals), data->tmp_vec_duals0, data->tmp_vec_duals1);


    // === Prolongation from local dual nodes to B nodes.

    VecGetArray(data->tmp_vec_b, &local_b);
    VecGetArray(data->tmp_vec_duals1, &local_duals);

    for (int i = nLocalB - nLocalDuals, j = 0; i < nLocalB; i++, j++)
      local_b[i] = local_duals[j];

    VecRestoreArray(data->tmp_vec_b, &local_b);
    VecRestoreArray(data->tmp_vec_duals0, &local_duals);


    // Multiply with scaled Lagrange constraint matrix.
    MatMult(*(data->mat_lagrange_scaled), data->tmp_vec_b, y);

    return 0;
  }


  PetscSolverFeti::PetscSolverFeti()
    : PetscSolver(),
      nComponents(-1),
      schurPrimalSolver(0)
  {
    FUNCNAME("PetscSolverFeti::PetscSolverFeti()");

    string preconditionerName = "";
    Parameters::get("parallel->solver->precon", preconditionerName);
    if (preconditionerName == "" || preconditionerName == "none") {
      MSG("Create FETI-DP solver with no preconditioner!\n");
      fetiPreconditioner = FETI_NONE;
    } else if (preconditionerName == "dirichlet") {
      MSG("Create FETI-DP solver with Dirichlet preconditioner!\n");
      fetiPreconditioner = FETI_DIRICHLET;
    } else if (preconditionerName == "lumped") {
      MSG("Create FETI-DP solver with lumped preconditioner!\n");
      fetiPreconditioner = FETI_LUMPED;
    } else {
      ERROR_EXIT("Preconditioner \"%s\" not available!\n", 
		 preconditionerName.c_str());
    }

    Parameters::get("parallel->feti->schur primal solver", schurPrimalSolver);
    TEST_EXIT(schurPrimalSolver == 0 || schurPrimalSolver == 1)
      ("Wrong solver \"%d\"for the Schur primal complement!\n", 
       schurPrimalSolver);
  }


  void PetscSolverFeti::updateDofData()
  {
    FUNCNAME("PetscSolverFeti::updateDofData()");

    TEST_EXIT(meshDistributor->getFeSpace()->getBasisFcts()->getDegree() == 1)
      ("Works for linear basis functions only!\n");
   
    for (unsigned int i = 0; i < meshDistributor->getFeSpaces().size(); i++) {
      const FiniteElemSpace *feSpace = meshDistributor->getFeSpace(i);

      createPrimals(feSpace);      
      createDuals(feSpace);
      createIndexB(feSpace);
    }

    createLagrange(meshDistributor->getFeSpaces());
  }


  void PetscSolverFeti::createPrimals(const FiniteElemSpace *feSpace)
  {
    FUNCNAME("PetscSolverFeti::createPrimals()");  

    // === Define all vertices on the interior boundaries of the macro mesh ===
    // === to be primal variables.                                          ===

    /// Set of DOF indices that are considered to be primal variables.
    
    DofIndexSet primals;
    DofContainerSet& vertices = 
      meshDistributor->getBoundaryDofInfo(feSpace).geoDofs[VERTEX];
    TEST_EXIT_DBG(vertices.size())("No primal vertices on this rank!\n");
    for (DofContainerSet::iterator it = vertices.begin(); 
	 it != vertices.end(); ++it)
      primals.insert(**it);


    // === Calculate the number of primals that are owned by the rank and ===
    // === create local indices of the primals starting at zero.          ===

    for (DofIndexSet::iterator it = primals.begin(); it != primals.end(); ++it)
      if (meshDistributor->getIsRankDof(feSpace, *it))
	primalDofMap[feSpace].insertRankDof(*it);

    primalDofMap[feSpace].update();

    MSG("nRankPrimals = %d   nOverallPrimals = %d\n", 
	primalDofMap[feSpace].nRankDofs, 
	primalDofMap[feSpace].nOverallDofs);


    // === Communicate primal's global index from ranks that own the     ===
    // === primals to ranks that contain this primals but are not owning ===
    // === them.                                                         ===
    
    typedef map<int, map<const FiniteElemSpace*, DofContainer> >::iterator it_type;

    StdMpi<vector<int> > stdMpi(meshDistributor->getMpiComm());
    for (DofComm::Iterator it(meshDistributor->getSendDofs(), feSpace);
	 !it.end(); it.nextRank())
      for (; !it.endDofIter(); it.nextDof())
	if (primalDofMap[feSpace].isSet(it.getDofIndex())) {
	  DegreeOfFreedom d = primalDofMap[feSpace][it.getDofIndex()];
	  stdMpi.getSendData(it.getRank()).push_back(d);
	}

    stdMpi.updateSendDataSize();

    for (DofComm::Iterator it(meshDistributor->getRecvDofs(), feSpace);
	 !it.end(); it.nextRank()) {
      bool recvFromRank = false;
      for (; !it.endDofIter(); it.nextDof()) {
	if (primals.count(it.getDofIndex()) && 
	    meshDistributor->getIsRankDof(feSpace, it.getDofIndex()) == false) {
	  recvFromRank = true;
	  break;
	}
      }

      if (recvFromRank) 
	stdMpi.recv(it.getRank());
    }

    stdMpi.startCommunication();

    for (DofComm::Iterator it(meshDistributor->getRecvDofs(), feSpace);
	 !it.end(); it.nextRank()) {
      int i = 0;
      for (; !it.endDofIter(); it.nextDof()) {
	if (primals.count(it.getDofIndex()) && 
	    meshDistributor->getIsRankDof(feSpace, it.getDofIndex()) == false) {
	  DegreeOfFreedom d = stdMpi.getRecvData(it.getRank())[i++];
	  primalDofMap[feSpace].insert(it.getDofIndex(), d);
	}
      }
    }

    TEST_EXIT_DBG(primals.size() == primalDofMap[feSpace].size())
      ("Number of primals %d, but number of global primals on this rank is %d!\n",
       primals.size(), primalDofMap[feSpace].size());
  }


  void PetscSolverFeti::createDuals(const FiniteElemSpace *feSpace)
  {
    FUNCNAME("PetscSolverFeti::createDuals()");

    // === Create for each dual node that is owned by the rank, the set ===
    // === of ranks that contain this node (denoted by W(x_j)).         ===

    boundaryDofRanks.clear();

    for (DofComm::Iterator it(meshDistributor->getSendDofs(), feSpace); 
	 !it.end(); it.nextRank())
      for (; !it.endDofIter(); it.nextDof()) {
	// If DOF is not primal, i.e., its a dual node
	if (primalDofMap[feSpace].isSet(it.getDofIndex()) == false) {
	  boundaryDofRanks[it.getDofIndex()].insert(mpiRank);
	  boundaryDofRanks[it.getDofIndex()].insert(it.getRank());
	}
      }
	

    // === Communicate these sets for all rank owned dual nodes to other ===
    // === ranks that also have this node.                               ===

    StdMpi<vector<std::set<int> > > stdMpi(meshDistributor->getMpiComm());

    for (DofComm::Iterator it(meshDistributor->getSendDofs(), feSpace);
	 !it.end(); it.nextRank())
      for (; !it.endDofIter(); it.nextDof())
	if (primalDofMap[feSpace].isSet(it.getDofIndex()) == false)
	  stdMpi.getSendData(it.getRank()).push_back(boundaryDofRanks[it.getDofIndex()]);

    stdMpi.updateSendDataSize();

    for (DofComm::Iterator it(meshDistributor->getRecvDofs(), feSpace); 
	 !it.end(); it.nextRank()) {
      bool recvFromRank = false;
      for (; !it.endDofIter(); it.nextDof()) {
	if (primalDofMap[feSpace].isSet(it.getDofIndex()) == false) {
	  recvFromRank = true;
	  break;
	}
      }

      if (recvFromRank)
	stdMpi.recv(it.getRank());
    }

    stdMpi.startCommunication();

    for (DofComm::Iterator it(meshDistributor->getRecvDofs(), feSpace); 
	 !it.end(); it.nextRank()) {
      int i = 0;
      for (; !it.endDofIter(); it.nextDof())
	if (primalDofMap[feSpace].isSet(it.getDofIndex()) == false)
	  boundaryDofRanks[it.getDofIndex()] = 
	    stdMpi.getRecvData(it.getRank())[i++];
    }

    // === Create global index of the dual nodes on each rank. ===

    DofContainer allBoundaryDofs;
    meshDistributor->getAllBoundaryDofs(feSpace, allBoundaryDofs);

    for (DofContainer::iterator it = allBoundaryDofs.begin();
	 it != allBoundaryDofs.end(); ++it)
      if (primalDofMap[feSpace].isSet(**it) == false)
	dualDofMap[feSpace].insertRankDof(**it);

    dualDofMap[feSpace].update(false);

    MSG("nRankDuals = %d   nOverallDuals = %d\n",
	dualDofMap[feSpace].nRankDofs, dualDofMap[feSpace].nOverallDofs);
  }

  
  void PetscSolverFeti::createLagrange(vector<const FiniteElemSpace*> &feSpaces)
  {
    FUNCNAME("PetscSolverFeti::createLagrange()");

    // === Reserve for each dual node, on the rank that owns this node, the ===
    // === appropriate number of Lagrange constraints.                      ===


#if 0    
    int nRankLagrange = 0;
    DofMapping& dualMap = dualDofMap[feSpace].getMap();
    for (DofMapping::iterator it = dualMap.begin(); it != dualMap.end(); ++it) {
      if (meshDistributor->getIsRankDof(feSpace, it->first)) {
	lagrangeMap[feSpace].insert(it->first, nRankLagrange);
	int degree = boundaryDofRanks[it->first].size();
	nRankLagrange += (degree * (degree - 1)) / 2;
      }
    }
    lagrangeMap[feSpace].nRankDofs = nRankLagrange;
    lagrangeMap[feSpace].update(false);
    
    MSG("nRankLagrange = %d  nOverallLagrange = %d\n",
	lagrangeMap[feSpace].nRankDofs, lagrangeMap[feSpace].nOverallDofs);


    // === Communicate lagrangeMap to all other ranks. ===

    StdMpi<vector<int> > stdMpi(meshDistributor->getMpiComm());

    for (DofComm::Iterator it(meshDistributor->getSendDofs(), feSpace);
	 !it.end(); it.nextRank()) {
      for (; !it.endDofIter(); it.nextDof())
	if (primalDofMap[feSpace].isSet(it.getDofIndex()) == false) {
	  DegreeOfFreedom d = lagrangeMap[feSpace][it.getDofIndex()];
	  stdMpi.getSendData(it.getRank()).push_back(d);
	}
    }

    stdMpi.updateSendDataSize();

    for (DofComm::Iterator it(meshDistributor->getRecvDofs(), feSpace);
	 !it.end(); it.nextRank()) {
      bool recvData = false;
      for (; !it.endDofIter(); it.nextDof())
	if (primalDofMap[feSpace].isSet(it.getDofIndex()) == false) {
	  recvData = true;
	  break;
	}
	  
      if (recvData)
	stdMpi.recv(it.getRank());
    }

    stdMpi.startCommunication();

    for (DofComm::Iterator it(meshDistributor->getRecvDofs(), feSpace);
	 !it.end(); it.nextRank()) {
      int counter = 0;
      for (; !it.endDofIter(); it.nextDof()) {
	if (primalDofMap[feSpace].isSet(it.getDofIndex()) == false) {
	  DegreeOfFreedom d = stdMpi.getRecvData(it.getRank())[counter++];
	  lagrangeMap[feSpace].insert(it.getDofIndex(), d); 
	}
      }
    }
#endif
  }


  void PetscSolverFeti::createIndexB(const FiniteElemSpace *feSpace)
  {
    FUNCNAME("PetscSolverFeti::createIndexB()");

    DOFAdmin* admin = feSpace->getAdmin();

    // === To ensure that all interior node on each rank are listen first in ===
    // === the global index of all B nodes, insert all interior nodes first, ===
    // === without defining a correct index.                                 ===

    nLocalInterior = 0;    
    for (int i = 0; i < admin->getUsedSize(); i++) {
      if (admin->isDofFree(i) == false && 
	  primalDofMap[feSpace].isSet(i) == false && 
	  dualDofMap[feSpace].isSet(i) == false) {
	localDofMap[feSpace].insertRankDof(i);
	nLocalInterior++;
      }
    }

    localDofMap[feSpace].update(false);

    TEST_EXIT_DBG(nLocalInterior + 
		  primalDofMap[feSpace].size() + 
		  dualDofMap[feSpace].size() == 
		  static_cast<unsigned int>(admin->getUsedDofs()))
      ("Should not happen!\n");

    // === And finally, add the global indicies of all dual nodes. ===

    for (DofMapping::iterator it = dualDofMap[feSpace].getMap().begin();
	 it != dualDofMap[feSpace].getMap().end(); ++it)
      localDofMap[feSpace].insertRankDof(it->first);

    localDofMap[feSpace].update(false);

#if 0
    dualDofMap[feSpace].addOffset(localDofMap[feSpace].rStartDofs + 
				  nLocalInterior);
#endif
  }


  void PetscSolverFeti::createMatLagrange(vector<const FiniteElemSpace*> &feSpaces)
  {
    FUNCNAME("PetscSolverFeti::createMatLagrange()");

    // === Create distributed matrix for Lagrange constraints. ===

    MatCreateMPIAIJ(PETSC_COMM_WORLD,
		    lagrangeMap.getRankDofs(feSpaces),
		    localDofMap.getRankDofs(),
		    lagrangeMap.getOverallDofs(feSpaces),
		    localDofMap.getOverallDofs(),	
		    2, PETSC_NULL, 2, PETSC_NULL,
		    &mat_lagrange);

    // === Create for all duals the corresponding Lagrange constraints. On ===
    // === each rank we traverse all pairs (n, m) of ranks, with n < m,    ===
    // === that contain this node. If the current rank number is r, and    ===
    // === n == r, the rank sets 1.0 for the corresponding constraint, if  ===
    // === m == r, than the rank sets -1.0 for the corresponding           ===
    // === constraint.                                                     ===

    for (unsigned int feIdx = 0; feIdx < feSpaces.size(); feIdx++) {
      const FiniteElemSpace *feSpace = feSpaces[feIdx];

      DofMapping &dualMap = dualDofMap[feSpace].getMap();
      for (DofMapping::iterator it = dualMap.begin(); it != dualMap.end(); ++it) {
	TEST_EXIT_DBG(boundaryDofRanks.count(it->first))
	  ("Should not happen!\n");

	// Global index of the first Lagrange constriant for this node.
	int constraintMatIndex = lagrangeMap.mapGlobal(it->first, feIdx);
	// Copy set of all ranks that contain this dual node.
	vector<int> W(boundaryDofRanks[it->first].begin(), 
		      boundaryDofRanks[it->first].end());
	// Number of ranks that contain this dual node.
	int degree = W.size();
	
	for (int i = 0; i < degree; i++) {
	  for (int j = i + 1; j < degree; j++) {
	    if (W[i] == mpiRank || W[j] == mpiRank) {
#if 1
	      // Get global index of the dual variable
	      int colIndex = localDofMap.mapGlobal(it->first, feIdx);
#else
	      int colIndex = it->second * nComponents + k;
#endif
	      double value = (W[i] == mpiRank ? 1.0 : -1.0);
	      MatSetValue(mat_lagrange, constraintMatIndex, colIndex, value, 
			  INSERT_VALUES);	      
	    }
	    
	    constraintMatIndex++;
	  }
	}
      }
    }

    MatAssemblyBegin(mat_lagrange, MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(mat_lagrange, MAT_FINAL_ASSEMBLY);
  }


  void PetscSolverFeti::createSchurPrimalKsp(vector<const FiniteElemSpace*> &feSpaces)
  {
    FUNCNAME("PetscSolverFeti::createSchurPrimal()");

    const FiniteElemSpace *feSpace = meshDistributor->getFeSpace(0);

    if (schurPrimalSolver == 0) {
      MSG("Create iterative schur primal solver!\n");

      schurPrimalData.mat_primal_primal = &mat_primal_primal;
      schurPrimalData.mat_primal_b = &mat_primal_b;
      schurPrimalData.mat_b_primal = &mat_b_primal;
      schurPrimalData.fetiSolver = this;
      
      VecCreateMPI(PETSC_COMM_WORLD, 
		   localDofMap.getRankDofs(), localDofMap.getOverallDofs(),
		   &(schurPrimalData.tmp_vec_b));
      VecCreateMPI(PETSC_COMM_WORLD, 
		   primalDofMap[feSpace].nRankDofs * nComponents, primalDofMap[feSpace].nOverallDofs * nComponents,
		   &(schurPrimalData.tmp_vec_primal));

      MatCreateShell(PETSC_COMM_WORLD,
		     primalDofMap[feSpace].nRankDofs * nComponents, primalDofMap[feSpace].nRankDofs * nComponents,
		     primalDofMap[feSpace].nOverallDofs * nComponents, primalDofMap[feSpace].nOverallDofs * nComponents,
		     &schurPrimalData, 
		     &mat_schur_primal);
      MatShellSetOperation(mat_schur_primal, MATOP_MULT, 
			   (void(*)(void))petscMultMatSchurPrimal);
      
      KSPCreate(PETSC_COMM_WORLD, &ksp_schur_primal);
      KSPSetOperators(ksp_schur_primal, mat_schur_primal, mat_schur_primal, SAME_NONZERO_PATTERN);
      KSPSetOptionsPrefix(ksp_schur_primal, "solver_sp_");
      KSPSetFromOptions(ksp_schur_primal);
    } else {
      MSG("Create direct schur primal solver!\n");

      TEST_EXIT(ksp_b)("No solver object for local problem!\n");

      double wtime = MPI::Wtime();

      int nRowsRankPrimal = primalDofMap[feSpace].nRankDofs * nComponents;
      int nRowsOverallPrimal = primalDofMap[feSpace].nOverallDofs * nComponents;
      int nRowsRankB = localDofMap.getRankDofs();
      int nRowsOverallB = localDofMap.getOverallDofs();

      Mat matBPi;
      MatCreateMPIAIJ(PETSC_COMM_WORLD,
		      nRowsRankB, nRowsRankPrimal, 
		      nRowsOverallB, nRowsOverallPrimal,
		      30, PETSC_NULL, 30, PETSC_NULL, &matBPi);
      Mat matPrimal;

      PetscInt nCols;
      const PetscInt *cols;
      const PetscScalar *values;

      map<int, vector<pair<int, double> > > mat_b_primal_cols;

      for (int i = 0; i < nRowsRankB; i++) {
	PetscInt row = localDofMap.getStartDofs() + i;
	MatGetRow(mat_b_primal, row, &nCols, &cols, &values);

	for (int j = 0; j < nCols; j++)
	  if (values[j] != 0.0)
	    mat_b_primal_cols[cols[j]].push_back(make_pair(i, values[j]));

	MatRestoreRow(mat_b_primal, row, &nCols, &cols, &values);
      }

      int maxLocalPrimal = mat_b_primal_cols.size();
      mpi::globalMax(maxLocalPrimal);

      TEST_EXIT(mat_b_primal_cols.size() ==
		primalDofMap[feSpace].size() * nComponents)("Should not happen!\n");
      for (map<int, vector<pair<int, double> > >::iterator it = mat_b_primal_cols.begin();
	   it != mat_b_primal_cols.end(); ++it) {
	Vec tmpVec;
	VecCreateSeq(PETSC_COMM_SELF, nRowsRankB, &tmpVec);

 	for (unsigned int i = 0; i < it->second.size(); i++) 
 	  VecSetValue(tmpVec, 
 		      it->second[i].first, it->second[i].second, INSERT_VALUES);

	VecAssemblyBegin(tmpVec);
	VecAssemblyEnd(tmpVec);

       	KSPSolve(ksp_b, tmpVec, tmpVec);

	PetscScalar *tmpValues;
	VecGetArray(tmpVec, &tmpValues);
	for (int i  = 0; i < nRowsRankB; i++)
	  MatSetValue(matBPi, 
		      localDofMap.getStartDofs() + i,
		      it->first,
		      tmpValues[i],
		      ADD_VALUES);
	VecRestoreArray(tmpVec, &tmpValues);

	VecDestroy(&tmpVec);
      }

      MatAssemblyBegin(matBPi, MAT_FINAL_ASSEMBLY);
      MatAssemblyEnd(matBPi, MAT_FINAL_ASSEMBLY);
      MatMatMult(mat_primal_b, matBPi, MAT_INITIAL_MATRIX, PETSC_DEFAULT, &matPrimal);
      MatAXPY(mat_primal_primal, -1.0, matPrimal, DIFFERENT_NONZERO_PATTERN);

      MatDestroy(&matPrimal);
      MatDestroy(&matBPi);

      MatInfo minfo;
      MatGetInfo(mat_primal_primal, MAT_GLOBAL_SUM, &minfo);
      MSG("Schur primal matrix nnz = %f\n", minfo.nz_used);

      KSPCreate(PETSC_COMM_WORLD, &ksp_schur_primal);
      KSPSetOperators(ksp_schur_primal, mat_primal_primal, 
		      mat_primal_primal, SAME_NONZERO_PATTERN);
      KSPSetOptionsPrefix(ksp_schur_primal, "solver_sp_");
      KSPSetFromOptions(ksp_schur_primal);

      MSG("Creating Schur primal matrix needed %.5f seconds.\n",
	  MPI::Wtime() - wtime);
    }
  }


  void PetscSolverFeti::destroySchurPrimalKsp()
  {
    FUNCNAME("PetscSolverFeti::destroySchurPrimal()");

    if (schurPrimalSolver == 0) {
      schurPrimalData.mat_primal_primal = PETSC_NULL;
      schurPrimalData.mat_primal_b = PETSC_NULL;
      schurPrimalData.mat_b_primal = PETSC_NULL;
      schurPrimalData.fetiSolver = NULL;

      VecDestroy(&schurPrimalData.tmp_vec_b);
      VecDestroy(&schurPrimalData.tmp_vec_primal);

      MatDestroy(&mat_schur_primal);
      KSPDestroy(&ksp_schur_primal);
    } else {
      KSPDestroy(&ksp_schur_primal);
    }
  }


  void PetscSolverFeti::createFetiKsp(vector<const FiniteElemSpace*> &feSpaces)
  {
    FUNCNAME("PetscSolverFeti::createFetiKsp()");

    const FiniteElemSpace *feSpace = meshDistributor->getFeSpace(0);

    // === Create FETI-DP solver object. ===

    fetiData.mat_primal_b = &mat_primal_b;
    fetiData.mat_b_primal = &mat_b_primal;
    fetiData.mat_lagrange = &mat_lagrange;
    fetiData.fetiSolver = this;
    fetiData.ksp_schur_primal = &ksp_schur_primal;

    VecCreateMPI(PETSC_COMM_WORLD, 
		 localDofMap.getRankDofs(), localDofMap.getOverallDofs(),
		 &(fetiData.tmp_vec_b));
    VecCreateMPI(PETSC_COMM_WORLD,
		 lagrangeMap[feSpace].nRankDofs * nComponents, 
		 lagrangeMap[feSpace].nOverallDofs * nComponents,
		 &(fetiData.tmp_vec_lagrange));
    VecCreateMPI(PETSC_COMM_WORLD, 
		 primalDofMap[feSpace].nRankDofs * nComponents, 
		 primalDofMap[feSpace].nOverallDofs * nComponents,
		 &(fetiData.tmp_vec_primal));

    MatCreateShell(PETSC_COMM_WORLD,
		   lagrangeMap[feSpace].nRankDofs * nComponents, 
		   lagrangeMap[feSpace].nRankDofs * nComponents,
		   lagrangeMap[feSpace].nOverallDofs * nComponents, 
		   lagrangeMap[feSpace].nOverallDofs * nComponents,
		   &fetiData, &mat_feti);
    MatShellSetOperation(mat_feti, MATOP_MULT, (void(*)(void))petscMultMatFeti);


    KSPCreate(PETSC_COMM_WORLD, &ksp_feti);
    KSPSetOperators(ksp_feti, mat_feti, mat_feti, SAME_NONZERO_PATTERN);
    KSPSetOptionsPrefix(ksp_feti, "solver_feti_");
    KSPSetFromOptions(ksp_feti);


    // === Create FETI-DP preconditioner object. ===

    if (fetiPreconditioner != FETI_NONE) {
      MatDuplicate(mat_lagrange, MAT_COPY_VALUES, &mat_lagrange_scaled);
      MatScale(mat_lagrange_scaled, 0.5);
    }

    switch (fetiPreconditioner) {
    case FETI_DIRICHLET:           
      KSPCreate(PETSC_COMM_SELF, &ksp_interior);
      KSPSetOperators(ksp_interior, mat_interior_interior, mat_interior_interior, 
		      SAME_NONZERO_PATTERN);
      KSPSetOptionsPrefix(ksp_interior, "solver_interior_");
      KSPSetFromOptions(ksp_interior);
            
      fetiDirichletPreconData.mat_lagrange_scaled = &mat_lagrange_scaled;
      fetiDirichletPreconData.mat_interior_interior = &mat_interior_interior;
      fetiDirichletPreconData.mat_duals_duals = &mat_duals_duals;
      fetiDirichletPreconData.mat_interior_duals = &mat_interior_duals;
      fetiDirichletPreconData.mat_duals_interior = &mat_duals_interior;
      fetiDirichletPreconData.ksp_interior = &ksp_interior;
      
      VecCreateMPI(PETSC_COMM_WORLD, 
		   localDofMap.getRankDofs(),localDofMap.getOverallDofs(),
		   &(fetiDirichletPreconData.tmp_vec_b));      
      MatGetVecs(mat_duals_duals, PETSC_NULL, 
		 &(fetiDirichletPreconData.tmp_vec_duals0));
      MatGetVecs(mat_duals_duals, PETSC_NULL, 
		 &(fetiDirichletPreconData.tmp_vec_duals1));
      MatGetVecs(mat_interior_interior, PETSC_NULL, 
		 &(fetiDirichletPreconData.tmp_vec_interior));
      
      KSPGetPC(ksp_feti, &precon_feti);
      PCSetType(precon_feti, PCSHELL);
      PCShellSetContext(precon_feti, static_cast<void*>(&fetiDirichletPreconData));
      PCShellSetApply(precon_feti, petscApplyFetiDirichletPrecon);
      
      break;

    case FETI_LUMPED:
      fetiLumpedPreconData.mat_lagrange_scaled = &mat_lagrange_scaled;
      fetiLumpedPreconData.mat_duals_duals = &mat_duals_duals;

      VecCreateMPI(PETSC_COMM_WORLD, 
		   localDofMap.getRankDofs(),
		   localDofMap.getOverallDofs(),
		   &(fetiLumpedPreconData.tmp_vec_b));
      MatGetVecs(mat_duals_duals, PETSC_NULL, 
		 &(fetiLumpedPreconData.tmp_vec_duals0));
      MatGetVecs(mat_duals_duals, PETSC_NULL, 
		 &(fetiLumpedPreconData.tmp_vec_duals1));

      KSPGetPC(ksp_feti, &precon_feti);
      PCSetType(precon_feti, PCSHELL);
      PCShellSetContext(precon_feti, static_cast<void*>(&fetiLumpedPreconData));
      PCShellSetApply(precon_feti, petscApplyFetiLumpedPrecon);
      
      break;
    default:
      break;
    }
  }
  

  void PetscSolverFeti::destroyFetiKsp()
  {
    FUNCNAME("PetscSolverFeti::destroyFetiKsp()");

    // === Destroy FETI-DP solver object. ===

    fetiData.mat_primal_b = PETSC_NULL;
    fetiData.mat_b_primal = PETSC_NULL;
    fetiData.mat_lagrange = PETSC_NULL;
    fetiData.fetiSolver = NULL;
    fetiData.ksp_schur_primal = PETSC_NULL;

    VecDestroy(&fetiData.tmp_vec_b);
    VecDestroy(&fetiData.tmp_vec_lagrange);
    VecDestroy(&fetiData.tmp_vec_primal);
    MatDestroy(&mat_feti);
    KSPDestroy(&ksp_feti);


    // === Destroy FETI-DP preconditioner object. ===

    switch (fetiPreconditioner) {
    case FETI_DIRICHLET:           
      KSPDestroy(&ksp_interior);

      fetiDirichletPreconData.mat_lagrange_scaled = NULL;
      fetiDirichletPreconData.mat_interior_interior = NULL;
      fetiDirichletPreconData.mat_duals_duals = NULL;
      fetiDirichletPreconData.mat_interior_duals = NULL;
      fetiDirichletPreconData.mat_duals_interior = NULL;
      fetiDirichletPreconData.ksp_interior = NULL;
      
      VecDestroy(&fetiDirichletPreconData.tmp_vec_b);
      VecDestroy(&fetiDirichletPreconData.tmp_vec_duals0);
      VecDestroy(&fetiDirichletPreconData.tmp_vec_duals1);
      VecDestroy(&fetiDirichletPreconData.tmp_vec_interior);
      MatDestroy(&mat_lagrange_scaled);
      break;

    case FETI_LUMPED:
      fetiLumpedPreconData.mat_lagrange_scaled = NULL;
      fetiLumpedPreconData.mat_duals_duals = NULL;

      VecDestroy(&fetiLumpedPreconData.tmp_vec_b);
      VecDestroy(&fetiLumpedPreconData.tmp_vec_duals0);
      VecDestroy(&fetiLumpedPreconData.tmp_vec_duals1);
      break;
    default:
      break;
    }
  }


  void PetscSolverFeti::recoverSolution(Vec &vec_sol_b,
					Vec &vec_sol_primal,
					SystemVector &vec)
  {
    FUNCNAME("PetscSolverFeti::recoverSolution()");

    const FiniteElemSpace *feSpace = meshDistributor->getFeSpace(0);

    // === Get local part of the solution for B variables. ===

    PetscScalar *localSolB;
    VecGetArray(vec_sol_b, &localSolB);


    // === Create scatter to get solutions of all primal nodes that are ===
    // === contained in rank's domain.                                  ===
    
    vector<PetscInt> globalIsIndex, localIsIndex;
    globalIsIndex.reserve(primalDofMap[feSpace].size() * nComponents);
    localIsIndex.reserve(primalDofMap[feSpace].size() * nComponents);

    {
      int counter = 0;
      for (DofMapping::iterator it = primalDofMap[feSpace].getMap().begin();
	   it != primalDofMap[feSpace].getMap().end(); ++it) {
	for (int i = 0; i < nComponents; i++) {
	  globalIsIndex.push_back(it->second * nComponents + i);
	  localIsIndex.push_back(counter++);
	}
      }
    }
    
    IS globalIs, localIs;
    ISCreateGeneral(PETSC_COMM_SELF, 
		    globalIsIndex.size(), 
		    &(globalIsIndex[0]),
		    PETSC_USE_POINTER,
		    &globalIs);

    ISCreateGeneral(PETSC_COMM_SELF, 
		    localIsIndex.size(), 
		    &(localIsIndex[0]),
		    PETSC_USE_POINTER,
		    &localIs);

    Vec local_sol_primal;
    VecCreateSeq(PETSC_COMM_SELF, localIsIndex.size(), &local_sol_primal);

    VecScatter primalScatter;
    VecScatterCreate(vec_sol_primal, globalIs, local_sol_primal, localIs, &primalScatter);
    VecScatterBegin(primalScatter, vec_sol_primal, local_sol_primal, 
		    INSERT_VALUES, SCATTER_FORWARD);
    VecScatterEnd(primalScatter, vec_sol_primal, local_sol_primal, 
		  INSERT_VALUES, SCATTER_FORWARD);

    ISDestroy(&globalIs);
    ISDestroy(&localIs);    
    VecScatterDestroy(&primalScatter);    

    PetscScalar *localSolPrimal;
    VecGetArray(local_sol_primal, &localSolPrimal);

    // === And copy from PETSc local vectors to the DOF vectors. ===

    for (int i = 0; i < nComponents; i++) {
      DOFVector<double>& dofVec = *(vec.getDOFVector(i));

      for (DofMapping::iterator it = localDofMap[feSpace].getMap().begin();
	   it != localDofMap[feSpace].getMap().end(); ++it) {

#if 1
	int petscIndex = localDofMap.mapLocal(it->first, i);
#else	  
	int petscIndex = it->second * nComponents + i;
#endif

	dofVec[it->first] = localSolB[petscIndex];
      }

      int counter = 0;
      for (DofMapping::iterator it = primalDofMap[feSpace].getMap().begin();
	   it != primalDofMap[feSpace].getMap().end(); ++it) {
	dofVec[it->first] = localSolPrimal[counter * nComponents + i];
	counter++;
      }
    }

    VecRestoreArray(vec_sol_b, &localSolB);
    VecRestoreArray(local_sol_primal, &localSolPrimal);
    VecDestroy(&local_sol_primal);
  }


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

    vector<const FiniteElemSpace*> feSpaces = getFeSpaces(mat);

    const FiniteElemSpace *feSpace = meshDistributor->getFeSpace(0);
    nComponents = mat->getSize();

    // === Create all sets and indices. ===

    primalDofMap.setMpiComm(mpiComm);
    dualDofMap.setMpiComm(mpiComm);
    lagrangeMap.setMpiComm(mpiComm);
    localDofMap.setMpiComm(mpiComm);

    primalDofMap.setFeSpaces(feSpaces);
    dualDofMap.setFeSpaces(feSpaces);
    lagrangeMap.setFeSpaces(feSpaces);
    localDofMap.setFeSpaces(feSpaces);

    updateDofData();

    primalDofMap.update();
    dualDofMap.update();
    lagrangeMap.update();
    localDofMap.update();

    // === Create matrices for the FETI-DP method. ===

    int nRowsRankB = localDofMap.getRankDofs();
    int nRowsOverallB = localDofMap.getOverallDofs();
    int nRowsRankPrimal = primalDofMap.getRankDofs(feSpaces);
    int nRowsOverallPrimal = primalDofMap.getOverallDofs(feSpaces);
    int nRowsDual = dualDofMap.getRankDofs(feSpaces);
    int nRowsInterior = nRowsRankB - nRowsDual;

    MatCreateSeqAIJ(PETSC_COMM_SELF, nRowsRankB, nRowsRankB, 30, PETSC_NULL,
		    &mat_b_b);

    MatCreateMPIAIJ(PETSC_COMM_WORLD,
		    nRowsRankPrimal, nRowsRankPrimal, 
		    nRowsOverallPrimal, nRowsOverallPrimal,
		    30, PETSC_NULL, 30, PETSC_NULL, &mat_primal_primal);

    MatCreateMPIAIJ(PETSC_COMM_WORLD,