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 "AMDiS.h"
#include "MatrixVector.h"
#include "parallel/PetscHelper.h"
#include "parallel/PetscSolverFeti.h"
#include "parallel/PetscSolverFetiDebug.h"
#include "parallel/PetscSolverFetiMonitor.h"
#include "parallel/PetscSolverFetiStructs.h"
#include "parallel/PetscSolverFetiOperators.h"
#include "parallel/PetscSolverFetiTimings.h"
#include "parallel/StdMpi.h"
#include "parallel/MpiHelper.h"
#include "parallel/PetscSolverGlobalMatrix.h"
#include "io/VtkWriter.h"

namespace AMDiS {

  using namespace std;

  PetscSolverFeti::PetscSolverFeti()
    : PetscSolver(),
      primalDofMap(COMPONENT_WISE),
      dualDofMap(COMPONENT_WISE),
      interfaceDofMap(COMPONENT_WISE),
      localDofMap(COMPONENT_WISE),
      lagrangeMap(COMPONENT_WISE),
      interiorDofMap(COMPONENT_WISE),
      schurPrimalSolver(0),
      multiLevelTest(false),
      subdomain(NULL),
      massMatrixSolver(NULL),
      meshLevel(0),
      rStartInterior(0),
      nGlobalOverallInterior(0),
      printTimings(false),
      stokesMode(false),
      augmentedLagrange(false),
      pressureComponent(-1)
  {
    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);

    Parameters::get("parallel->multi level test", multiLevelTest);
    if (multiLevelTest)
      meshLevel = 1;

    Parameters::get("parallel->print timings", printTimings);

    Parameters::get("parallel->feti->augmented lagrange", augmentedLagrange);

    Parameters::get("parallel->feti->symmetric", isSymmetric);
  }


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

    TEST_EXIT_DBG(meshLevel + 1 == meshDistributor->getMeshLevelData().getLevelNumber())
      ("Mesh hierarchy does not contain %d levels!\n", meshLevel + 1);

    MeshLevelData& levelData = meshDistributor->getMeshLevelData();

    stokesMode = false;
    Parameters::get("parallel->feti->stokes mode", stokesMode);
    if (stokesMode) {
      Parameters::get("parallel->feti->pressure component", pressureComponent);
      TEST_EXIT(pressureComponent >= 0)
	("FETI-DP in Stokes mode, no pressure component defined!\n");
    }
			   
    if (subdomain == NULL) {
      subdomain = new PetscSolverGlobalMatrix();
      subdomain->setSymmetric(isSymmetric);
      subdomain->setHandleDirichletRows(false);

      if (meshLevel == 0) {
	subdomain->setMeshDistributor(meshDistributor, 
				      mpiCommGlobal, mpiCommLocal);
      } else {
	subdomain->setMeshDistributor(meshDistributor, 
				      levelData.getMpiComm(meshLevel - 1),
				      levelData.getMpiComm(meshLevel));
	subdomain->setLevel(meshLevel);
      }
    }

    primalDofMap.init(levelData, componentSpaces, feSpaces);   
    dualDofMap.init(levelData, componentSpaces, feSpaces, false);
    localDofMap.init(levelData, componentSpaces, feSpaces, meshLevel != 0);
    lagrangeMap.init(levelData, componentSpaces, feSpaces);

    if (stokesMode)
      interfaceDofMap.init(levelData, componentSpaces, feSpaces);

    if (fetiPreconditioner == FETI_DIRICHLET) {
      TEST_EXIT(meshLevel == 0)
	("Dirichlet preconditioner not yet implemented for multilevel FETI-DP\n");

      interiorDofMap.init(levelData, componentSpaces, feSpaces, false);
    }
  }


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

    int nComponents = mat.getSize();
    for (int component = 0; component < nComponents; component++) {
      DOFMatrix* dofMat = mat[component][component];
      if (!dofMat)
	continue;

      dirichletRows[component] = dofMat->getDirichletRows();
    }
  }


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

    double timeCounter = MPI::Wtime();

    MeshLevelData& levelData = meshDistributor->getMeshLevelData();

    primalDofMap.clear();
    dualDofMap.clear();
    lagrangeMap.clear();
    localDofMap.clear();
    if (fetiPreconditioner == FETI_DIRICHLET)
      interiorDofMap.clear();

    primalDofMap.setDofComm(meshDistributor->getDofComm());
    lagrangeMap.setDofComm(meshDistributor->getDofComm());

    primalDofMap.setMpiComm(levelData.getMpiComm(0), 0);
    dualDofMap.setMpiComm(levelData.getMpiComm(0), 0);
    lagrangeMap.setMpiComm(levelData.getMpiComm(0), 0);
    localDofMap.setMpiComm(levelData.getMpiComm(meshLevel), meshLevel);
    if (fetiPreconditioner == FETI_DIRICHLET)
      interiorDofMap.setMpiComm(levelData.getMpiComm(meshLevel), meshLevel);

    if (meshLevel == 0)
      localDofMap.setDofComm(meshDistributor->getDofComm());
    else
      localDofMap.setDofComm(meshDistributor->getDofCommSd());

    if (stokesMode) {
      interfaceDofMap.clear();
      interfaceDofMap.setDofComm(meshDistributor->getDofComm());
      interfaceDofMap.setMpiComm(levelData.getMpiComm(0), 0);
    }

    int nComponents = componentSpaces.size();
    for (int component = 0; component < nComponents; component++) {
      createPrimals(component);  
      createDuals(component);
      createInterfaceNodes(component);
      createIndexB(component);
    }

    primalDofMap.update();
    dualDofMap.update();
    localDofMap.update();
    if (fetiPreconditioner == FETI_DIRICHLET)
      interiorDofMap.update();

    if (stokesMode)
      interfaceDofMap.update();

    for (int component = 0; component < nComponents; component++) {
      createLagrange(component);
      createAugmentedLagrange(component);
    }

    lagrangeMap.update();


    // === ===

    if (meshLevel == 0) {
      rStartInterior = 0;
      nGlobalOverallInterior = localDofMap.getOverallDofs();
    } else {
      MeshLevelData& levelData = meshDistributor->getMeshLevelData();

      int groupRowsInterior = 0;
      if (levelData.getMpiComm(1).Get_rank() == 0)
	groupRowsInterior = localDofMap.getOverallDofs();

      mpi::getDofNumbering(mpiCommGlobal, groupRowsInterior,
			   rStartInterior, nGlobalOverallInterior);

      int tmp = 0;
      if (levelData.getMpiComm(1).Get_rank() == 0)
	tmp = rStartInterior;

      levelData.getMpiComm(1).Allreduce(&tmp, &rStartInterior, 1, MPI_INT, MPI_SUM);
    }

    MSG("FETI-DP data created on mesh level %d\n", meshLevel);
    for (unsigned int i = 0; i < componentSpaces.size(); i++) {
      const FiniteElemSpace *feSpace = componentSpaces[i];

      MSG("FETI-DP data for %d-ith component (FE space %p):\n", i, feSpace);

      if (i == pressureComponent) {
	MSG("  nRankInterface = %d  nOverallInterface = %d\n",
	    interfaceDofMap[i].nRankDofs, 
	    interfaceDofMap[i].nOverallDofs);
      } else {
	MSG("  nRankPrimals = %d   nLocalPrimals = %d  nOverallPrimals = %d\n", 
	    primalDofMap[i].nRankDofs, 
	    primalDofMap[i].nLocalDofs,
	    primalDofMap[i].nOverallDofs);
	
	MSG("  nRankDuals = %d  nOverallDuals = %d\n",
	    dualDofMap[i].nRankDofs, 
	    dualDofMap[i].nOverallDofs);
	
	MSG("  nRankLagrange = %d  nOverallLagrange = %d\n",
	    lagrangeMap[i].nRankDofs, 
	    lagrangeMap[i].nOverallDofs);
      }
    }

    subdomain->setDofMapping(&localDofMap);
    subdomain->setCoarseSpaceDofMapping(&primalDofMap); 
    if (stokesMode)
      subdomain->setCoarseSpaceDofMapping(&interfaceDofMap, pressureComponent);

    if (printTimings) {
      MPI::COMM_WORLD.Barrier();
      timeCounter = MPI::Wtime() - timeCounter;
      MSG("FETI-DP timing 01: %.5f seconds (creation of basic data structures)\n", 
	  timeCounter);
    }
  }


  void PetscSolverFeti::createPrimals(int component)
  {
    FUNCNAME("PetscSolverFeti::createPrimals()");  

    if (component == pressureComponent)
      return;

    const FiniteElemSpace *feSpace = componentSpaces[component];

    // === 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.
    DofContainerSet& vertices = 
      meshDistributor->getBoundaryDofInfo(feSpace, meshLevel).geoDofs[VERTEX];

    DofIndexSet primals;
    for (DofContainerSet::iterator it = vertices.begin(); 
	 it != vertices.end(); ++it) {
      
      if (dirichletRows[component].count(**it))
	continue;

      if (meshLevel == 0) {
	primals.insert(**it);
      } else {
	double e = 1e-8;
	WorldVector<double> c;
	feSpace->getMesh()->getDofIndexCoords(*it, feSpace, c);

	if (fabs(c[0]) < e || fabs(c[1]) < e ||
	    fabs(c[0] - 25.0) < e || fabs(c[1] - 25.0) < e ||
	    (fabs(c[0] - 12.5) < e && fabs(c[1] - 12.5) < e)) {
	  MSG("PRIMAL COORD %f %f\n", c[0], c[1]);
	  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 (dofMap[feSpace].isRankDof(*it)) {
	primalDofMap[component].insertRankDof(*it);
      } else
  	primalDofMap[component].insertNonRankDof(*it);
  }


  void PetscSolverFeti::createDuals(int component)
  {
    FUNCNAME("PetscSolverFeti::createDuals()");

    if (component == pressureComponent)
      return;

    const FiniteElemSpace *feSpace = componentSpaces[component];

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

    DofContainer allBoundaryDofs;
    meshDistributor->getAllBoundaryDofs(feSpace, meshLevel, allBoundaryDofs);
    
    for (DofContainer::iterator it = allBoundaryDofs.begin();
	 it != allBoundaryDofs.end(); ++it) {
      if (dirichletRows[component].count(**it))
	continue;

      if (isPrimal(component, **it))
	continue;

      if (meshLevel == 0) {
	dualDofMap[component].insertRankDof(**it);
      } else {
	if (dofMapSd[feSpace].isRankDof(**it))
	  dualDofMap[component].insertRankDof(**it);
      }	  
    }
  }

  
  void PetscSolverFeti::createInterfaceNodes(int component)
  {
    FUNCNAME("PetscSolverFeti::createInterfaceNodes()");

    if (component != pressureComponent)
      return;

    const FiniteElemSpace *feSpace = componentSpaces[component];
    DofContainer allBoundaryDofs;
    meshDistributor->getAllBoundaryDofs(feSpace, meshLevel, allBoundaryDofs);

    for (DofContainer::iterator it = allBoundaryDofs.begin();
	 it != allBoundaryDofs.end(); ++it) {
      if (dirichletRows[component].count(**it))
	continue;      
      
      if (dofMap[feSpace].isRankDof(**it))
	interfaceDofMap[component].insertRankDof(**it);
      else
	interfaceDofMap[component].insertNonRankDof(**it);      
    }
  }


  void PetscSolverFeti::createLagrange(int component)
  {
    FUNCNAME("PetscSolverFeti::createLagrange()");

    if (component == pressureComponent)
      return;

    const FiniteElemSpace *feSpace = componentSpaces[component];
    boundaryDofRanks[feSpace].clear();

    // Stores for all rank owned communication DOFs, if the counterpart is
    // a rank owned DOF in its subdomain. Thus, the following map stores to
    // each rank number all DOFs that fulfill this requirenment.
    map<int, std::set<DegreeOfFreedom> > sdRankDofs;

    if (meshLevel > 0) {
      StdMpi<vector<int> > stdMpi(mpiCommGlobal);

      for (DofComm::Iterator it(meshDistributor->getDofComm().getRecvDofs(), 
				meshLevel, feSpace);
	   !it.end(); it.nextRank()) {

	vector<int> subdomainRankDofs;
	subdomainRankDofs.reserve(it.getDofs().size());

	for (; !it.endDofIter(); it.nextDof()) {
	  if (dofMapSd[feSpace].isRankDof(it.getDofIndex()))
	    subdomainRankDofs.push_back(1);
	  else
	    subdomainRankDofs.push_back(0);
	}

	stdMpi.send(it.getRank(), subdomainRankDofs);
      }	     

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

      stdMpi.startCommunication();

      for (DofComm::Iterator it(meshDistributor->getDofComm().getSendDofs(), 
				meshLevel, feSpace); 
	   !it.end(); it.nextRank())
	for (; !it.endDofIter(); it.nextDof())
	  if (!isPrimal(component, it.getDofIndex()))
	    if (stdMpi.getRecvData(it.getRank())[it.getDofCounter()] == 1)
	      sdRankDofs[it.getRank()].insert(it.getDofIndex());
    }

    if (dualDofMap[component].nLocalDofs == 0)
      return;


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

    int mpiRank = meshDistributor->getMpiRank();
    for (DofComm::Iterator it(meshDistributor->getDofComm().getSendDofs(), 
			      meshLevel, feSpace); 
	 !it.end(); it.nextRank()) {
      for (; !it.endDofIter(); it.nextDof()) {
	if (!isPrimal(component, it.getDofIndex())) {
	  boundaryDofRanks[feSpace][it.getDofIndex()].insert(mpiRank);

 	  if (meshLevel == 0 ||
 	      (meshLevel > 0 && sdRankDofs[it.getRank()].count(it.getDofIndex())))
	    boundaryDofRanks[feSpace][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->getDofComm().getSendDofs(), meshLevel, feSpace);
	 !it.end(); it.nextRank())
      for (; !it.endDofIter(); it.nextDof())
	if (!isPrimal(component, it.getDofIndex()))
 	  if (meshLevel == 0 ||
 	      (meshLevel > 0 && sdRankDofs[it.getRank()].count(it.getDofIndex())))
	    stdMpi.getSendData(it.getRank()).push_back(boundaryDofRanks[feSpace][it.getDofIndex()]);

    stdMpi.updateSendDataSize();

    for (DofComm::Iterator it(meshDistributor->getDofComm().getRecvDofs(), meshLevel, feSpace); 
	 !it.end(); it.nextRank()) {
      bool recvFromRank = false;
      for (; !it.endDofIter(); it.nextDof()) {
	if (!isPrimal(component, it.getDofIndex())) {
 	  if (meshLevel == 0 ||
 	      (meshLevel > 0 && 
	       dofMapSd[feSpace].isRankDof(it.getDofIndex()))) {
	    recvFromRank = true;
	    break;
	  }
	}
      }

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

    stdMpi.startCommunication();

    for (DofComm::Iterator it(meshDistributor->getDofComm().getRecvDofs(), meshLevel, feSpace); 
	 !it.end(); it.nextRank()) {
      int i = 0;
      for (; !it.endDofIter(); it.nextDof())
	if (!isPrimal(component, it.getDofIndex()))
 	  if (meshLevel == 0 ||
 	      (meshLevel > 0 && 
	       dofMapSd[feSpace].isRankDof(it.getDofIndex())))	    
	    boundaryDofRanks[feSpace][it.getDofIndex()] = 
	      stdMpi.getRecvData(it.getRank())[i++];
	  else
	    lagrangeMap[component].insertNonRankDof(it.getDofIndex());
    }


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

    int nRankLagrange = 0;
    DofMap& dualMap = dualDofMap[component].getMap();
    for (DofMap::iterator it = dualMap.begin(); it != dualMap.end(); ++it) {
      if (dofMap[feSpace].isRankDof(it->first)) {
	lagrangeMap[component].insertRankDof(it->first, nRankLagrange);
	int degree = boundaryDofRanks[feSpace][it->first].size();
	nRankLagrange += (degree * (degree - 1)) / 2;
      } else {
	lagrangeMap[component].insertNonRankDof(it->first);
      }
    }
    lagrangeMap[component].nRankDofs = nRankLagrange;
  }


  void PetscSolverFeti::createAugmentedLagrange(int component)
  {
    FUNCNAME("PetscSolverFeti::createAugmentedLagrange()");

    if (!augmentedLagrange)
      return;
  }


  void PetscSolverFeti::createIndexB(int component)
  {
    FUNCNAME("PetscSolverFeti::createIndexB()");


    const FiniteElemSpace *feSpace = componentSpaces[component];
    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.                                 ===

    int nLocalInterior = 0;    
    for (int i = 0; i < admin->getUsedSize(); i++) {
      if (admin->isDofFree(i) ||
	  isPrimal(component, i) ||
	  isDual(component, i) ||
	  isInterface(component, i) ||
	  dirichletRows[component].count(i))
	continue;      

      if (meshLevel == 0) {
	localDofMap[component].insertRankDof(i, nLocalInterior);
	
	if (fetiPreconditioner == FETI_DIRICHLET)
	  interiorDofMap[component].insertRankDof(i, nLocalInterior);
	
	nLocalInterior++;	
      } else {
	if (dofMapSd[feSpace].isRankDof(i))
	  localDofMap[component].insertRankDof(i);
	else
	  localDofMap[component].insertNonRankDof(i);
	
	TEST_EXIT_DBG(fetiPreconditioner == FETI_NONE)
	  ("Not yet implemnted!\n");	
      }
    }
    
    // === And finally, add the global indicies of all dual nodes. ===

    for (DofMap::iterator it = dualDofMap[component].getMap().begin();
	 it != dualDofMap[component].getMap().end(); ++it) {
      if (meshLevel == 0) {
	localDofMap[component].insertRankDof(it->first);
      } else {
	if (dofMapSd[feSpace].isRankDof(it->first))
	  localDofMap[component].insertRankDof(it->first);
	else 
	  localDofMap[component].insertNonRankDof(it->first);
      }
    }
  }


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

    double wtime = MPI::Wtime();
    int mpiRank = meshDistributor->getMpiRank();

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

    MatCreateAIJ(mpiCommGlobal,
		 lagrangeMap.getRankDofs(), localDofMap.getRankDofs(),
		 lagrangeMap.getOverallDofs(), nGlobalOverallInterior,
		 2, PETSC_NULL, 2, PETSC_NULL,
		 &mat_lagrange);
    MatSetOption(mat_lagrange, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_FALSE);

    Vec vec_scale_lagrange;
    lagrangeMap.createVec(vec_scale_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 component = 0; component < componentSpaces.size(); component++) {
      DofMap &dualMap = dualDofMap[component].getMap();

      for (DofMap::iterator it = dualMap.begin(); it != dualMap.end(); ++it) {
	TEST_EXIT_DBG(boundaryDofRanks[componentSpaces[component]].count(it->first))
	  ("Should not happen!\n");
	
	// Global index of the first Lagrange constriant for this node.
	int index = lagrangeMap.getMatIndex(component, it->first);

	// Copy set of all ranks that contain this dual node.
	vector<int> W(boundaryDofRanks[componentSpaces[component]][it->first].begin(), 
		      boundaryDofRanks[componentSpaces[component]][it->first].end());
	// Number of ranks that contain this dual node.
	int degree = W.size();

	TEST_EXIT_DBG(degree > 1)("Should not happen!\n");
	
	int counter = 0;
	for (int i = 0; i < degree; i++) {
	  for (int j = i + 1; j < degree; j++) {
	    if (W[i] == mpiRank || W[j] == mpiRank) {
	      MatSetValue(mat_lagrange, 
			  index + counter, 
			  localDofMap.getMatIndex(component, it->first) + rStartInterior,
			  (W[i] == mpiRank ? 1.0 : -1.0),
			  INSERT_VALUES);
	    }
	    counter++;
	  }
	}

	// === Create scaling factors for scaling the lagrange matrix, which ===
	// === is required for FETI-DP preconditioners.                      ===
	
	if (dofMap[componentSpaces[component]].isRankDof(it->first)) {
	  int nConstraints = (degree * (degree - 1)) / 2;
	  for (int i = 0; i < nConstraints; i++) {
	    VecSetValue(vec_scale_lagrange,
			index + i,
			1.0 / static_cast<double>(degree),
			INSERT_VALUES);
	  }
	}
      }
    }

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


    int nZeroRows = PetscSolverFetiDebug::testZeroRows(mat_lagrange);
    int m,n;
    MatGetSize(mat_lagrange, &m ,&n);
    MSG("Lagrange matrix has %d zero rows and global size of %d %d!\n", nZeroRows, m, n);

    PetscViewer petscView;
    PetscViewerBinaryOpen(PETSC_COMM_WORLD, "lagrange.mat", 
			  FILE_MODE_WRITE, &petscView);
    MatView(mat_lagrange, petscView);
    PetscViewerDestroy(&petscView);


    // === If required, create \ref mat_lagrange_scaled ===

    VecAssemblyBegin(vec_scale_lagrange);
    VecAssemblyEnd(vec_scale_lagrange);

    if (fetiPreconditioner != FETI_NONE || stokesMode) {
      MatDuplicate(mat_lagrange, MAT_COPY_VALUES, &mat_lagrange_scaled);
      MatDiagonalScale(mat_lagrange_scaled, vec_scale_lagrange, PETSC_NULL);
    }

    VecDestroy(&vec_scale_lagrange);


    // === Print final timings. ===

    if (printTimings) {
      MPI::COMM_WORLD.Barrier();
      MSG("FETI-DP timing 05: %.5f seconds (creation of lagrange constraint matrix)\n", 
	  MPI::Wtime() - wtime);
    }
  }


  bool PetscSolverFeti::testWirebasketEdge(BoundaryObject &edge, const FiniteElemSpace *feSpace)
  {
    Element *el = edge.el;
    int i0 = el->getVertexOfEdge(edge.ithObj, 0);
    int i1 = el->getVertexOfEdge(edge.ithObj, 1);
    DegreeOfFreedom d0 = el->getDof(i0, 0);
    DegreeOfFreedom d1 = el->getDof(i1, 0);
    WorldVector<double> c0, c1;
    el->getMesh()->getDofIndexCoords(d0, feSpace, c0);
    el->getMesh()->getDofIndexCoords(d1, feSpace, c1);
    bool xe = fabs(c0[0] - c1[0]) < 1e-8;
    bool ye = fabs(c0[1] - c1[1]) < 1e-8;
    bool ze = fabs(c0[2] - c1[2]) < 1e-8;
    int counter = static_cast<int>(xe) + static_cast<int>(ye) + static_cast<int>(ze);
    return (counter == 2);
  }


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

    if (!augmentedLagrange)
      return;

    double wtime = MPI::Wtime();

    nOverallEdges = 0;
    InteriorBoundary &intBound = meshDistributor->getIntBoundary();
    std::set<BoundaryObject> allEdges;
    for (InteriorBoundary::iterator it(intBound.getOwn()); !it.end(); ++it) {
      if ((it->rankObj.subObj == FACE ||
	   (it->rankObj.subObj == EDGE && 
	    testWirebasketEdge(it->rankObj, feSpaces[0]))) && 
	  allEdges.count(it->rankObj) == 0) {
	bool dirichletOnlyEdge = true;

	DofContainer edgeDofs;
	it->rankObj.el->getAllDofs(feSpaces[0], it->rankObj, edgeDofs);
	
	for (DofContainer::iterator dit = edgeDofs.begin();
	     dit != edgeDofs.end(); ++dit) {
	  if (dirichletRows[0].count(**dit) == 0) {
	    dirichletOnlyEdge = false;
	    break;
	  }
	}

	if (!dirichletOnlyEdge)
	  allEdges.insert(it->rankObj);
      }
    }

    nRankEdges = allEdges.size();    
    int rStartEdges = 0;
    mpi::getDofNumbering(mpiCommGlobal, nRankEdges, rStartEdges, nOverallEdges);

    MSG("nRankEdges = %d, nOverallEdges = %d\n", nRankEdges, nOverallEdges);
    
    nRankEdges *= componentSpaces.size();
    rStartEdges *= componentSpaces.size();
    nOverallEdges *= componentSpaces.size();

    MatCreateAIJ(mpiCommGlobal,
		 nRankEdges, lagrangeMap.getRankDofs(),
		 nOverallEdges, lagrangeMap.getOverallDofs(),
		 2, PETSC_NULL, 2, PETSC_NULL, 
		 &mat_augmented_lagrange);
    MatSetOption(mat_augmented_lagrange, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_FALSE);

    int rowCounter = rStartEdges;
    for (std::set<BoundaryObject>::iterator edgeIt = allEdges.begin(); 
	 edgeIt != allEdges.end(); ++edgeIt) {
      for (int component = 0; component < componentSpaces.size(); component++) {
	DofContainer edgeDofs;
	edgeIt->el->getAllDofs(componentSpaces[component], *edgeIt, edgeDofs);
	
	for (DofContainer::iterator it = edgeDofs.begin();
	     it != edgeDofs.end(); ++it) {
	  TEST_EXIT_DBG(isPrimal(component, **it) == false)
	    ("Should not be primal!\n");
	  
	  if (dirichletRows[component].count(**it))
	    continue;
	  
	  int col = lagrangeMap.getMatIndex(component, **it);
	  double value = 1.0;
	  MatSetValue(mat_augmented_lagrange, rowCounter, col, value, INSERT_VALUES);
	}
	
	rowCounter++;
      }      
    }

    MatAssemblyBegin(mat_augmented_lagrange, MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(mat_augmented_lagrange, MAT_FINAL_ASSEMBLY);

    int nZeroRows = PetscSolverFetiDebug::testZeroRows(mat_augmented_lagrange);
    int m,n;
    MatGetSize(mat_augmented_lagrange, &m ,&n);
    MSG("Augmented lagrange matrix has %d zero rows and global size of %d %d!\n", nZeroRows, m, n);

    if (printTimings) {
      MPI::COMM_WORLD.Barrier();
      MSG("FETI-DP timing 05a: %.5f seconds (creation of augmented lagrange constraint matrix)\n", 
	  MPI::Wtime() - wtime);
    }
  }


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

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

      if (augmentedLagrange == false) {
	schurPrimalData.subSolver = subdomain;
	
	localDofMap.createVec(schurPrimalData.tmp_vec_b, nGlobalOverallInterior);
	primalDofMap.createVec(schurPrimalData.tmp_vec_primal);
	
	MatCreateShell(mpiCommGlobal,
		       primalDofMap.getRankDofs(), 
		       primalDofMap.getRankDofs(), 
		       primalDofMap.getOverallDofs(), 
		       primalDofMap.getOverallDofs(),
		       &schurPrimalData, 
		       &mat_schur_primal);
	MatShellSetOperation(mat_schur_primal, MATOP_MULT, 
			     (void(*)(void))petscMultMatSchurPrimal);	
      } else {
	schurPrimalAugmentedData.subSolver = subdomain;
	schurPrimalAugmentedData.nestedVec = true;

	localDofMap.createVec(schurPrimalAugmentedData.tmp_vec_b0, nGlobalOverallInterior);
	localDofMap.createVec(schurPrimalAugmentedData.tmp_vec_b1, nGlobalOverallInterior);
	primalDofMap.createVec(schurPrimalAugmentedData.tmp_vec_primal);
	lagrangeMap.createVec(schurPrimalAugmentedData.tmp_vec_lagrange);

	schurPrimalAugmentedData.mat_lagrange = &mat_lagrange;
	schurPrimalAugmentedData.mat_augmented_lagrange = &mat_augmented_lagrange;

	MatCreateShell(mpiCommGlobal,
		       primalDofMap.getRankDofs() + nRankEdges, 
		       primalDofMap.getRankDofs() + nRankEdges, 
		       primalDofMap.getOverallDofs() + nOverallEdges, 
		       primalDofMap.getOverallDofs() + nOverallEdges,
		       &schurPrimalAugmentedData, 
		       &mat_schur_primal);
	MatShellSetOperation(mat_schur_primal, MATOP_MULT, 
			     (void(*)(void))petscMultMatSchurPrimalAugmented);
      }

      KSPCreate(mpiCommGlobal, &ksp_schur_primal);
      KSPSetOperators(ksp_schur_primal, mat_schur_primal, mat_schur_primal, SAME_NONZERO_PATTERN);
      KSPSetOptionsPrefix(ksp_schur_primal, "schur_primal_");
      KSPSetType(ksp_schur_primal, KSPGMRES);
      KSPSetFromOptions(ksp_schur_primal);
    } else {
      MSG("Create direct schur primal solver!\n");

      double wtime = MPI::Wtime();

      TEST_EXIT_DBG(meshLevel == 0)
	("Does not support for multilevel, check usage of localDofMap.\n");


      // === Create explicit matrix representation of the Schur primal system. ===

      if (!augmentedLagrange)
	createMatExplicitSchurPrimal();
      else
	createMatExplicitAugmentedSchurPrimal();


      // === Create KSP solver object and set appropriate solver options. ====

      KSPCreate(mpiCommGlobal, &ksp_schur_primal);
      KSPSetOperators(ksp_schur_primal, mat_schur_primal, mat_schur_primal,
		      SAME_NONZERO_PATTERN);
      KSPSetOptionsPrefix(ksp_schur_primal, "schur_primal_");
      KSPSetType(ksp_schur_primal, KSPPREONLY);
      PC pc_schur_primal;      
      KSPGetPC(ksp_schur_primal, &pc_schur_primal);
      PCSetType(pc_schur_primal, PCLU);
      PCFactorSetMatSolverPackage(pc_schur_primal, MATSOLVERMUMPS);
      KSPSetFromOptions(ksp_schur_primal);


      // === And finally print timings, if required. ===