Forgotten to add the ZoltanPartitioner files to repository.

AMDiS/src/parallel/ZoltanPartitioner.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 <boost/lexical_cast.hpp>
+#include "parallel/ZoltanPartitioner.h"
+#include "Traverse.h"
+#include "ElInfo.h"
+
+namespace AMDiS {
+
+  ZoltanPartitioner::ZoltanPartitioner(MPI::Intracomm *comm)
+    : MeshPartitioner(comm),
+      zoltan(*comm)
+  {
+    zoltan.Set_Num_Obj_Fn(ZoltanFunctions::getNumObj, this);
+    zoltan.Set_Obj_List_Fn(ZoltanFunctions::getObjectList, this);
+    zoltan.Set_Num_Geom_Fn(ZoltanFunctions::getNumGeom, this);
+    zoltan.Set_Geom_Multi_Fn(ZoltanFunctions::getGeomMulti, this);
+  }
+
+
+  bool ZoltanPartitioner::partition(map<int, double> &elemWeights,
+				    PartitionMode mode)
+  {
+    FUNCNAME("ZoltanPartitioner::partition()");
+
+    using boost::lexical_cast;
+
+    int changes;
+    int nGid, nLid;
+
+    int nImportEls;   
+    ZOLTAN_ID_PTR import_global_ids, import_local_ids;
+    int *import_procs;
+    int *import_to_part;
+
+    int nExportEls;
+    ZOLTAN_ID_PTR export_global_ids, export_local_ids;
+    int *export_procs;
+    int *export_to_part;
+
+    zoltan.Set_Param("LB_METHOD", "RCB");
+    zoltan.Set_Param("NUM_GLOBAL_PARTS", 
+		     lexical_cast<string>(mpiComm->Get_size()).c_str());
+    int err = zoltan.LB_Partition(changes, nGid, nLid, 
+				  nImportEls,
+				  import_global_ids,
+				  import_local_ids,
+				  import_procs,
+				  import_to_part,
+				  nExportEls,
+				  export_global_ids,
+				  export_local_ids,
+				  export_procs,
+				  export_to_part);
+    
+    recvElements.clear();
+    sendElements.clear();
+
+    if (err == ZOLTAN_OK && changes != 0) {   
+      if (nImportEls > 0) {
+	for (int i = 0; i < nImportEls; i++) {
+	  int recvElIndex = import_global_ids[i];
+	  
+	  elementInRank[recvElIndex] = true;
+	  recvElements[import_procs[i]].push_back(recvElIndex);
+	}
+      }
+      
+      if (nExportEls > 0) {
+	for (int i = 0; i < nExportEls; i++) {
+	  int sendElIndex = export_global_ids[i];
+	  
+	  elementInRank[sendElIndex] = false;
+	  sendElements[export_procs[i]].push_back(sendElIndex);
+	}      
+      }
+    }
+
+    zoltan.LB_Free_Part(&import_global_ids, &import_local_ids, 
+			&import_procs, &import_to_part);
+    zoltan.LB_Free_Part(&export_global_ids, &export_local_ids, 
+			&export_procs, &export_to_part);      
+
+    if (err != ZOLTAN_OK)
+      return false;
+
+    return true;
+  }
+
+
+  void ZoltanPartitioner::getPartitionMap(map<int, int> &partitionMap)
+  {
+    FUNCNAME("ZoltanPartitioner::getPartitionMap()");
+
+    int mpiSize = mpiComm->Get_size();
+
+    vector<int> localElements;
+    for (map<int, bool>::iterator it = elementInRank.begin(); 
+	 it != elementInRank.end(); ++it)
+      if (it->second)
+	localElements.push_back(it->first);
+ 
+    int nLocalElements = localElements.size();
+    vector<int> nPartitionElements(mpiSize);
+    vector<int> elDist(mpiSize + 1);
+    mpiComm->Allgather(&nLocalElements, 1, MPI_INT, &(elDist[1]), 1, MPI_INT);
+
+    elDist[0] = 0;
+    nPartitionElements[0] = elDist[1];
+
+    for (int i = 2; i <= mpiSize; i++) {
+      nPartitionElements[i - 1] = elDist[i];
+      elDist[i] += elDist[i - 1];
+    }
+
+    int nOverallElements = elDist[mpiSize];
+
+    TEST_EXIT_DBG(nOverallElements == static_cast<int>(elementInRank.size()))
+      ("Number of elements differs: %d %d!\n", nOverallElements, elementInRank.size());
+
+    vector<int> partitionElements(nOverallElements);
+
+    // distribute partition elements
+    mpiComm->Allgatherv(&(localElements[0]),
+			nLocalElements, 
+			MPI_INT, 
+			&(partitionElements[0]), 
+			&(nPartitionElements[0]), 
+			&(elDist[0]),
+			MPI_INT);
+
+    // fill partitionMap
+    for (int i = 0; i < mpiSize; i++)
+      for (int j = 0; j < nPartitionElements[i]; j++)
+	partitionMap[partitionElements[elDist[i] + j]] = i;
+  }
+
+  
+
+  int ZoltanFunctions::getNumObj(void *data, int *ierr)
+  {
+    ZoltanPartitioner* zoltan = (ZoltanPartitioner*)data;
+    map<int, bool>& elInRank = zoltan->getElementInRank();
+
+    int nObjects = 0;
+    for (map<int, bool>::iterator it = elInRank.begin(); it != elInRank.end(); ++it)
+      if (it->second)
+	nObjects++;
+
+    *ierr = ZOLTAN_OK;
+
+    return nObjects;
+  }
+
+
+  void ZoltanFunctions::getObjectList(void *data, 
+				      int sizeGid, 
+				      int sizeLid,
+				      ZOLTAN_ID_PTR globalId, 
+				      ZOLTAN_ID_PTR localId,
+				      int wgt_dim, 
+				      float *obj_wgts, 
+				      int *ierr)
+  {
+    ZoltanPartitioner* zoltan = (ZoltanPartitioner*)data;
+    map<int, bool>& elInRank = zoltan->getElementInRank();
+    int localCounter = 0;
+
+    for (map<int, bool>::iterator it = elInRank.begin(); it != elInRank.end(); ++it) {
+      if (it->second) {
+	globalId[localCounter] = it->first;
+	localId[localCounter] = localCounter;
+	localCounter++;
+      }
+    }
+
+    *ierr = ZOLTAN_OK;
+  }
+
+
+  int ZoltanFunctions::getNumGeom(void *data, int *ierr)
+  {
+    ZoltanPartitioner* zoltan = (ZoltanPartitioner*)data;
+    *ierr = ZOLTAN_OK;
+
+    return zoltan->getMesh()->getGeo(WORLD);
+  }
+
+
+  void ZoltanFunctions::getGeomMulti(void *data, 
+				     int num_gid_entries, 
+				     int num_lid_entries, 
+				     int num_obj, 
+				     ZOLTAN_ID_PTR global_ids, 
+				     ZOLTAN_ID_PTR local_ids, 
+				     int num_dim, 
+				     double *geom_vec, 
+				     int *ierr)
+  {
+    FUNCNAME("ZoltanFunctions::getGeomMulti()");
+
+    ZoltanPartitioner* zoltan = (ZoltanPartitioner*)data;
+    Mesh *mesh = zoltan->getMesh();
+
+    TEST_EXIT_DBG(num_dim == mesh->getGeo(WORLD))("Should not happen!\n");
+
+    map<int, WorldVector<double> > elIndexToCoords;    
+    DimVec<double> bary(mesh->getDim(), DEFAULT_VALUE, 1.0 / mesh->getGeo(VERTEX));
+    WorldVector<double> coords;
+    TraverseStack stack;
+    ElInfo *elInfo = 
+      stack.traverseFirst(mesh, 0, Mesh::CALL_EL_LEVEL | Mesh::FILL_COORDS);
+    while (elInfo) {
+      elInfo->coordToWorld(bary, coords);
+      elIndexToCoords[elInfo->getElement()->getIndex()] = coords;
+
+      elInfo = stack.traverseNext(elInfo);
+    }
+
+    int c = 0;
+    for (int i = 0; i < num_obj; i++) {
+      int elIndex = global_ids[i];
+
+      TEST_EXIT_DBG(elIndexToCoords.count(elIndex))("Should not happen!\n");
+      
+      for (int j = 0; j < num_dim; j++)
+	geom_vec[c++] = elIndexToCoords[elIndex][j];
+    }
+
+    *ierr = ZOLTAN_OK;
+  }
+  
+}

AMDiS/src/parallel/ZoltanPartitioner.h

+// ============================================================================
+// ==                                                                        ==
+// == AMDiS - Adaptive multidimensional simulations                          ==
+// ==                                                                        ==
+// ==  http://www.amdis-fem.org                                              ==
+// ==                                                                        ==
+// ============================================================================
+//
+// 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.
+
+
+
+/** \file ZoltanPartitioner.h */
+
+#ifndef AMDIS_ZOLTAN_PARTITIONER_H
+#define AMDIS_ZOLTAN_PARTITIONER_H
+
+#include <zoltan_cpp.h>
+#include <set>
+#include "AMDiS_fwd.h"
+#include "parallel/MeshPartitioner.h"
+
+namespace AMDiS {
+
+  using namespace std;
+
+  class ZoltanPartitioner : public MeshPartitioner
+  {
+  public:
+    ZoltanPartitioner(MPI::Intracomm *comm);
+
+    ~ZoltanPartitioner() {}
+      
+    /// \ref MeshPartitioner::partition
+    bool partition(map<int, double> &elemWeights, PartitionMode mode = INITIAL);
+
+    /// \ref MeshPartitioner::getPartitionMap
+    void getPartitionMap(map<int, int> &partitionMap);
+
+  protected:
+    Zoltan zoltan;
+  };
+
+
+  class ZoltanFunctions
+  {
+  public:
+    static int getNumObj(void *data, int *ierr);   
+
+    static void getObjectList(void *data, 
+			      int sizeGid, 
+			      int sizeLid,
+			      ZOLTAN_ID_PTR globalId, 
+			      ZOLTAN_ID_PTR localId,
+			      int wgt_dim, 
+			      float *obj_wgts, 
+			      int *ierr);
+
+    static int getNumGeom(void *data, int *ierr);
+
+    static void getGeomMulti(void *data, 
+			     int num_gid_entries, 
+			     int num_lid_entries, 
+			     int num_obj, 
+			     ZOLTAN_ID_PTR global_ids, 
+			     ZOLTAN_ID_PTR local_ids, 
+			     int num_dim, 
+			     double *geom_vec, 
+			     int *ierr); 
+  };
+
+}
+
+#endif