Mesh Oriented datABase  (version 5.5.1)
An array-based unstructured mesh library
MetisPartitioner.cpp
Go to the documentation of this file.
1 /**
2  * MOAB, a Mesh-Oriented datABase, is a software component for creating,
3  * storing and accessing finite element mesh data.
4  *
5  * Copyright 2004 Sandia Corporation. Under the terms of Contract
6  * DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government
7  * retains certain rights in this software.
8  *
9  * This library is free software; you can redistribute it and/or
10  * modify it under the terms of the GNU Lesser General Public
11  * License as published by the Free Software Foundation; either
12  * version 2.1 of the License, or (at your option) any later version.
13  *
14  */
15 
16 // Contributed by Lorenzo Alessio Botti (SpaFEDTe)
17 // This implementation is mostly borrowed from the mbzoltan MOAB partitioning tool
18 
19 #include <iostream>
20 #include <cassert>
21 #include <sstream>
22 #include <map>
23 #include <ctime>
24 
25 #include "moab/MetisPartitioner.hpp"
26 #include "moab/Interface.hpp"
27 #include "Internals.hpp"
28 #include "moab/Range.hpp"
29 #include "moab/WriteUtilIface.hpp"
30 #include "moab/MeshTopoUtil.hpp"
31 #include "moab/Skinner.hpp"
32 #include "MBTagConventions.hpp"
33 #include "moab/CN.hpp"
34 
35 using namespace moab;
36 
37 const bool debug = false;
38 
39 MetisPartitioner::MetisPartitioner( Interface* impl, const bool use_coords )
40  : PartitionerBase< idx_t >( impl, use_coords )
41 
42 {
43 }
44 
45 MetisPartitioner::~MetisPartitioner() {}
46 
47 ErrorCode MetisPartitioner::partition_mesh( const idx_t nparts,
48  const char* method,
49  const int part_dim,
50  const bool write_as_sets,
51  const bool write_as_tags,
52  const bool partition_tagged_sets,
53  const bool partition_tagged_ents,
54  const char* aggregating_tag,
55  const bool print_time )
56 {
57 #ifdef MOAB_HAVE_MPI
58  // should only be called in serial
59  if( mbpc->proc_config().proc_size() != 1 )
60  {
61  std::cout << "MetisPartitioner::partition_mesh_and_geometry must be called in serial." << std::endl;
62  return MB_FAILURE;
63  }
64 #endif
65 
66  if( NULL != method && strcmp( method, "ML_RB" ) != 0 && strcmp( method, "ML_KWAY" ) != 0 )
67  {
68  std::cout << "ERROR: Method must be "
69  << "ML_RB or ML_KWAY" << std::endl;
70  return MB_FAILURE;
71  }
72 
73  std::vector< double > pts; // x[0], y[0], z[0], ... from MOAB
74  std::vector< idx_t > ids; // poidx_t ids from MOAB
75  std::vector< idx_t > adjs, parts;
76  std::vector< idx_t > length;
77  Range elems;
78  // Get a mesh from MOAB and diide it across processors.
79 
80  clock_t t = clock();
81 
82  ErrorCode result;
83  if( !partition_tagged_sets && !partition_tagged_ents )
84  {
85  result = assemble_graph( part_dim, pts, ids, adjs, length, elems );MB_CHK_ERR( result );
86  }
87  else if( partition_tagged_sets )
88  {
89  result = assemble_taggedsets_graph( part_dim, pts, ids, adjs, length, elems, &( *aggregating_tag ) );MB_CHK_ERR( result );
90  }
91  else if( partition_tagged_ents )
92  {
93  result = assemble_taggedents_graph( part_dim, pts, ids, adjs, length, elems, &( *aggregating_tag ) );MB_CHK_ERR( result );
94  }
95  else
96  {
97  MB_SET_ERR( MB_FAILURE, "Either partition tags or sets for Metis partitoner" );
98  }
99 
100  if( print_time )
101  {
102  std::cout << " time to assemble graph: " << ( clock() - t ) / (double)CLOCKS_PER_SEC << "s. \n";
103  t = clock();
104  }
105 
106  std::cout << "Computing partition using " << method << " method for " << nparts << " processors..." << std::endl;
107 
108  idx_t nelems = length.size() - 1;
109  idx_t* assign_parts;
110  assign_parts = (idx_t*)malloc( sizeof( idx_t ) * nelems );
111  idx_t nconstraidx_ts = 1;
112  idx_t edgeCut = 0;
113  idx_t nOfPartitions = static_cast< idx_t >( nparts );
114  idx_t metis_RESULT;
115 
116  if( strcmp( method, "ML_KWAY" ) == 0 )
117  {
118  idx_t options[METIS_NOPTIONS];
119  METIS_SetDefaultOptions( options );
120  options[METIS_OPTION_CONTIG] = 1;
121  metis_RESULT = METIS_PartGraphKway( &nelems, &nconstraidx_ts, &length[0], &adjs[0], NULL, NULL, NULL,
122  &nOfPartitions, NULL, NULL, options, &edgeCut, assign_parts );
123  }
124  else if( strcmp( method, "ML_RB" ) == 0 )
125  {
126  idx_t options[METIS_NOPTIONS];
127  METIS_SetDefaultOptions( options );
128  options[METIS_OPTION_OBJTYPE] = METIS_OBJTYPE_CUT; // CUT
129  options[METIS_OPTION_IPTYPE] = METIS_IPTYPE_GROW; // GROW or RANDOM
130  options[METIS_OPTION_CTYPE] = METIS_CTYPE_RM; // RM or SHEM
131  options[METIS_OPTION_RTYPE] = METIS_RTYPE_FM; // FM
132  options[METIS_OPTION_NCUTS] = 10; // Number of different partitionings to compute, then
133  // chooses the best one, default = 1
134  options[METIS_OPTION_NITER] = 10; // Number of refinements steps, default = 10
135  options[METIS_OPTION_UFACTOR] = 30; // Imabalance, default = 1
136  options[METIS_OPTION_DBGLVL] = METIS_DBG_INFO;
137  metis_RESULT = METIS_PartGraphRecursive( &nelems, &nconstraidx_ts, &length[0], &adjs[0], NULL, NULL, NULL,
138  &nOfPartitions, NULL, NULL, options, &edgeCut, assign_parts );
139  }
140  else
141  MB_SET_ERR( MB_FAILURE, "Either ML_KWAY or ML_RB needs to be specified for Metis partitioner" );
142 
143  if( print_time )
144  {
145  std::cout << " time to partition: " << ( clock() - t ) / (double)CLOCKS_PER_SEC << "s. \n";
146  t = clock();
147  }
148 
149 #ifdef MOAB_HAVE_MPI
150  // assign global node ids, starting from one! TODO
151  if( assign_global_ids )
152  {
153  EntityHandle rootset = 0;
154  result = mbpc->assign_global_ids( rootset, part_dim, 1, true, false );MB_CHK_ERR( result );
155  }
156 #endif
157 
158  if( metis_RESULT != METIS_OK ) return MB_FAILURE;
159 
160  // take results & write onto MOAB partition sets
161  std::cout << "Saving partition information to MOAB..." << std::endl;
162  {
163  if( partition_tagged_sets || partition_tagged_ents )
164  {
165  result = write_aggregationtag_partition( nparts, elems, assign_parts, write_as_sets, write_as_tags );MB_CHK_ERR( result );
166  }
167  else
168  {
169  result = write_partition( nparts, elems, assign_parts, write_as_sets, write_as_tags );MB_CHK_ERR( result );
170  }
171  }
172 
173  if( print_time )
174  {
175  std::cout << " time to write partition in memory " << ( clock() - t ) / (double)CLOCKS_PER_SEC << "s. \n";
176  t = clock();
177  }
178  free( assign_parts );
179 
180  return MB_SUCCESS;
181 }
182 
183 ErrorCode MetisPartitioner::assemble_taggedents_graph( const int dimension,
184  std::vector< double >& coords,
185  std::vector< idx_t >& moab_ids,
186  std::vector< idx_t >& adjacencies,
187  std::vector< idx_t >& length,
188  Range& elems,
189  const char* aggregating_tag )
190 {
191  Tag partSetTag;
192  ErrorCode result = mbImpl->tag_get_handle( aggregating_tag, 1, MB_TYPE_INTEGER, partSetTag );
193  if( MB_SUCCESS != result ) return result;
194 
195  Range allSubElems;
196  result = mbImpl->get_entities_by_dimension( 0, dimension, allSubElems );
197  if( MB_SUCCESS != result || allSubElems.empty() ) return result;
198  idx_t partSet;
199  std::map< idx_t, Range > aggloElems;
200  for( Range::iterator rit = allSubElems.begin(); rit != allSubElems.end(); rit++ )
201  {
202  EntityHandle entity = *rit;
203  result = mbImpl->tag_get_data( partSetTag, &entity, 1, &partSet );
204  if( MB_SUCCESS != result ) return result;
205  if( partSet >= 0 ) aggloElems[partSet].insert( entity );
206  }
207  // clear aggregating tag data
208  TagType type;
209  result = mbImpl->tag_get_type( partSetTag, type );
210  if( type == MB_TAG_DENSE )
211  {
212  // clear tag on ents and sets
213  result = mbImpl->tag_delete( partSetTag );
214  if( MB_SUCCESS != result ) return result;
215  }
216  if( type == MB_TAG_SPARSE )
217  {
218  // clear tag on ents
219  result = mbImpl->tag_delete_data( partSetTag, allSubElems );
220  if( MB_SUCCESS != result ) return result;
221  // clear tag on sets
222  result = mbImpl->get_entities_by_type_and_tag( 0, MBENTITYSET, &partSetTag, 0, 1, elems );
223  if( MB_SUCCESS != result ) return result;
224  result = mbImpl->tag_delete_data( partSetTag, elems );
225  if( MB_SUCCESS != result ) return result;
226  elems.clear();
227  }
228  result =
229  mbImpl->tag_get_handle( "PARALLEL_PARTITION", 1, MB_TYPE_INTEGER, partSetTag, MB_TAG_SPARSE | MB_TAG_CREAT );
230  if( MB_SUCCESS != result ) return result;
231 
232  for( std::map< idx_t, Range >::iterator mit = aggloElems.begin(); mit != aggloElems.end(); mit++ )
233  {
234  EntityHandle new_set;
235  result = mbImpl->create_meshset( MESHSET_SET, new_set );
236  if( MB_SUCCESS != result ) return result;
237  result = mbImpl->add_entities( new_set, mit->second );
238  if( MB_SUCCESS != result ) return result;
239  result = mbImpl->tag_set_data( partSetTag, &new_set, 1, &mit->first );
240  if( MB_SUCCESS != result ) return result;
241  }
242 
243  result =
244  assemble_taggedsets_graph( dimension, coords, moab_ids, adjacencies, length, elems, &( *aggregating_tag ) );
245  return MB_SUCCESS;
246 }
247 
248 ErrorCode MetisPartitioner::assemble_taggedsets_graph( const int dimension,
249  std::vector< double >& coords,
250  std::vector< idx_t >& moab_ids,
251  std::vector< idx_t >& adjacencies,
252  std::vector< idx_t >& length,
253  Range& elems,
254  const char* aggregating_tag )
255 {
256  length.push_back( 0 );
257  // assemble a graph with vertices equal to elements of specified dimension, edges
258  // signified by list of other elements to which an element is connected
259 
260  // get the tagged elements
261  Tag partSetTag;
262  ErrorCode result = mbImpl->tag_get_handle( aggregating_tag, 1, MB_TYPE_INTEGER, partSetTag );MB_CHK_ERR( result );
263  // ErrorCode result = mbImpl->tag_get_handle("PARALLEL_PARTITION_SET", 1, MB_TYPE_INTEGER,
264  // partSetTag);MB_CHK_ERR(result);
265 
266  result = mbImpl->get_entities_by_type_and_tag( 0, MBENTITYSET, &partSetTag, 0, 1, elems );
267  if( MB_SUCCESS != result || elems.empty() ) return result;
268 
269  // assign globla ids to elem sets based on aggregating_tag data
270  Tag gid_tag;
271  idx_t zero1 = -1;
272  result =
273  mbImpl->tag_get_handle( "GLOBAL_ID_AGGLO", 1, MB_TYPE_INTEGER, gid_tag, MB_TAG_SPARSE | MB_TAG_CREAT, &zero1 );MB_CHK_ERR( result );
274  for( Range::iterator rit = elems.begin(); rit != elems.end(); rit++ )
275  {
276  idx_t partSet;
277  result = mbImpl->tag_get_data( partSetTag, &( *rit ), 1, &partSet );MB_CHK_ERR( result );
278  result = mbImpl->tag_set_data( gid_tag, &( *rit ), 1, &partSet );MB_CHK_ERR( result );
279  }
280  // clear aggregating tag data
281  TagType type;
282  result = mbImpl->tag_get_type( partSetTag, type );MB_CHK_ERR( result );
283  if( type == MB_TAG_DENSE )
284  {
285  result = mbImpl->tag_delete( partSetTag );MB_CHK_ERR( result );
286  }
287  if( type == MB_TAG_SPARSE )
288  {
289  result = mbImpl->tag_delete_data( partSetTag, elems );MB_CHK_ERR( result );
290  }
291 
292  // assemble the graph, using Skinner to get d-1 dimensional neighbors and then idx_tersecting to
293  // get adjacencies
294  std::vector< Range > skin_subFaces( elems.size() );
295  unsigned int i = 0;
296  for( Range::iterator rit = elems.begin(); rit != elems.end(); rit++ )
297  {
298  Range part_ents;
299  result = mbImpl->get_entities_by_handle( *rit, part_ents, false );
300  if( mbImpl->dimension_from_handle( *part_ents.rbegin() ) !=
301  mbImpl->dimension_from_handle( *part_ents.begin() ) )
302  {
303  Range::iterator lower = part_ents.lower_bound( CN::TypeDimensionMap[0].first ),
304  upper = part_ents.upper_bound( CN::TypeDimensionMap[dimension - 1].second );
305  part_ents.erase( lower, upper );
306  }
307  Skinner skinner( mbImpl );
308  result = skinner.find_skin( 0, part_ents, false, skin_subFaces[i], NULL, false, true, false );MB_CHK_ERR( result );
309  i++;
310  }
311  std::vector< EntityHandle > adjs;
312  std::vector< idx_t > neighbors;
313  double avg_position[3];
314  idx_t moab_id;
315  MeshTopoUtil mtu( mbImpl );
316  for( unsigned int k = 0; k < i; k++ )
317  {
318  // get bridge adjacencies for element k
319  adjs.clear();
320  for( unsigned int t = 0; t < i; t++ )
321  {
322  if( t != k )
323  {
324  Range subFaces = intersect( skin_subFaces[k], skin_subFaces[t] );
325  if( subFaces.size() > 0 ) adjs.push_back( elems[t] );
326  }
327  }
328  if( !adjs.empty() )
329  {
330  neighbors.resize( adjs.size() );
331  result = mbImpl->tag_get_data( gid_tag, &adjs[0], adjs.size(), &neighbors[0] );MB_CHK_ERR( result );
332  }
333  // copy those idx_to adjacencies vector
334  length.push_back( length.back() + (idx_t)adjs.size() );
335  std::copy( neighbors.begin(), neighbors.end(), std::back_inserter( adjacencies ) );
336  // get the graph vertex id for this element
337  const EntityHandle& setk = elems[k];
338  result = mbImpl->tag_get_data( gid_tag, &setk, 1, &moab_id );
339  moab_ids.push_back( moab_id );
340  // get average position of vertices
341  Range part_ents;
342  result = mbImpl->get_entities_by_handle( elems[k], part_ents, false );MB_CHK_ERR( result );
343  result = mtu.get_average_position( part_ents, avg_position );MB_CHK_ERR( result );
344  std::copy( avg_position, avg_position + 3, std::back_inserter( coords ) );
345  }
346  for( unsigned int k = 0; k < i; k++ )
347  {
348  for( unsigned int t = 0; t < k; t++ )
349  {
350  Range subFaces = intersect( skin_subFaces[k], skin_subFaces[t] );
351  if( subFaces.size() > 0 ) mbImpl->delete_entities( subFaces );
352  }
353  }
354 
355  if( debug )
356  {
357  std::cout << "Length vector: " << std::endl;
358  std::copy( length.begin(), length.end(), std::ostream_iterator< idx_t >( std::cout, ", " ) );
359  std::cout << std::endl;
360  std::cout << "Adjacencies vector: " << std::endl;
361  std::copy( adjacencies.begin(), adjacencies.end(), std::ostream_iterator< idx_t >( std::cout, ", " ) );
362  std::cout << std::endl;
363  std::cout << "Moab_ids vector: " << std::endl;
364  std::copy( moab_ids.begin(), moab_ids.end(), std::ostream_iterator< idx_t >( std::cout, ", " ) );
365  std::cout << std::endl;
366  std::cout << "Coords vector: " << std::endl;
367  std::copy( coords.begin(), coords.end(), std::ostream_iterator< double >( std::cout, ", " ) );
368  std::cout << std::endl;
369  }
370  return MB_SUCCESS;
371 }
372 
373 ErrorCode MetisPartitioner::assemble_graph( const int dimension,
374  std::vector< double >& coords,
375  std::vector< idx_t >& moab_ids,
376  std::vector< idx_t >& adjacencies,
377  std::vector< idx_t >& length,
378  Range& elems )
379 {
380  length.push_back( 0 );
381  // assemble a graph with vertices equal to elements of specified dimension, edges
382  // signified by list of other elements to which an element is connected
383 
384  // get the elements of that dimension
385  ErrorCode result = mbImpl->get_entities_by_dimension( 0, dimension, elems );
386  if( MB_SUCCESS != result || elems.empty() ) return result;
387 
388 #ifdef MOAB_HAVE_MPI
389  // assign global ids
390  if( assign_global_ids )
391  {
392  result = mbpc->assign_global_ids( 0, dimension, 0 );MB_CHK_ERR( result );
393  }
394 #endif
395 
396  // now assemble the graph, calling MeshTopoUtil to get bridge adjacencies through d-1
397  // dimensional neighbors
398  MeshTopoUtil mtu( mbImpl );
399  Range adjs;
400  // can use a fixed-size array 'cuz the number of lower-dimensional neighbors is limited
401  // by MBCN
402  int neighbors[5 * MAX_SUB_ENTITIES]; // these will be now indices in the elems range
403 
404  double avg_position[3];
405  int index_in_elems = 0;
406 
407  for( Range::iterator rit = elems.begin(); rit != elems.end(); rit++, index_in_elems++ )
408  {
409 
410  // get bridge adjacencies
411  adjs.clear();
412  result = mtu.get_bridge_adjacencies( *rit, ( dimension > 0 ? dimension - 1 : 3 ), dimension, adjs );MB_CHK_ERR( result );
413 
414  // get the indices in elems range of those
415  if( !adjs.empty() )
416  {
417  int i = 0;
418  assert( adjs.size() < 5 * MAX_SUB_ENTITIES );
419  for( Range::iterator ait = adjs.begin(); ait != adjs.end(); ait++, i++ )
420  {
421  EntityHandle adjEnt = *ait;
422  neighbors[i] = elems.index( adjEnt );
423  }
424  }
425 
426  // copy those idx_to adjacencies vector
427  length.push_back( length.back() + (idx_t)adjs.size() );
428  // conversion made to idx_t
429  std::copy( neighbors, neighbors + adjs.size(), std::back_inserter( adjacencies ) );
430 
431  // get average position of vertices
432  result = mtu.get_average_position( *rit, avg_position );MB_CHK_ERR( result );
433 
434  // get the graph vertex id for this element; it is now index in elems
435  moab_ids.push_back( index_in_elems ); // conversion made to idx_t
436 
437  // copy_to coords vector
438  std::copy( avg_position, avg_position + 3, std::back_inserter( coords ) );
439  }
440 
441  if( debug )
442  {
443  std::cout << "Length vector: " << std::endl;
444  std::copy( length.begin(), length.end(), std::ostream_iterator< idx_t >( std::cout, ", " ) );
445  std::cout << std::endl;
446  std::cout << "Adjacencies vector: " << std::endl;
447  std::copy( adjacencies.begin(), adjacencies.end(), std::ostream_iterator< idx_t >( std::cout, ", " ) );
448  std::cout << std::endl;
449  std::cout << "Moab_ids vector: " << std::endl;
450  std::copy( moab_ids.begin(), moab_ids.end(), std::ostream_iterator< idx_t >( std::cout, ", " ) );
451  std::cout << std::endl;
452  std::cout << "Coords vector: " << std::endl;
453  std::copy( coords.begin(), coords.end(), std::ostream_iterator< double >( std::cout, ", " ) );
454  std::cout << std::endl;
455  }
456 
457  return MB_SUCCESS;
458 }
459 
460 ErrorCode MetisPartitioner::write_aggregationtag_partition( const idx_t nparts,
461  Range& elems,
462  const idx_t* assignment,
463  const bool write_as_sets,
464  const bool write_as_tags )
465 {
466  ErrorCode result;
467 
468  // get the partition set tag
469  Tag part_set_tag;
470  result =
471  mbImpl->tag_get_handle( "PARALLEL_PARTITION", 1, MB_TYPE_INTEGER, part_set_tag, MB_TAG_SPARSE | MB_TAG_CREAT );MB_CHK_ERR( result );
472 
473  // get any sets already with this tag, and clear them
474  Range tagged_sets;
475  result =
476  mbImpl->get_entities_by_type_and_tag( 0, MBENTITYSET, &part_set_tag, NULL, 1, tagged_sets, Interface::UNION );MB_CHK_ERR( result );
477  if( !tagged_sets.empty() )
478  {
479  result = mbImpl->clear_meshset( tagged_sets );
480  if( !write_as_sets )
481  {
482  result = mbImpl->tag_delete_data( part_set_tag, tagged_sets );MB_CHK_ERR( result );
483  }
484  }
485 
486  if( write_as_sets )
487  {
488  // first, create partition sets and store in vector
489  partSets.clear();
490 
491  if( nparts > (idx_t)tagged_sets.size() )
492  {
493  // too few partition sets - create missing ones
494  idx_t num_new = nparts - tagged_sets.size();
495  for( idx_t i = 0; i < num_new; i++ )
496  {
497  EntityHandle new_set;
498  result = mbImpl->create_meshset( MESHSET_SET, new_set );MB_CHK_ERR( result );
499  tagged_sets.insert( new_set );
500  }
501  }
502  else if( nparts < (idx_t)tagged_sets.size() )
503  {
504  // too many partition sets - delete extras
505  idx_t num_del = tagged_sets.size() - nparts;
506  for( idx_t i = 0; i < num_del; i++ )
507  {
508  EntityHandle old_set = tagged_sets.pop_back();
509  result = mbImpl->delete_entities( &old_set, 1 );MB_CHK_ERR( result );
510  }
511  }
512 
513  // assign partition sets to vector
514  partSets.swap( tagged_sets );
515 
516  // write a tag to those sets denoting they're partition sets, with a value of the
517  // proc number
518  idx_t* dum_ids = new idx_t[nparts];
519  for( idx_t i = 0; i < nparts; i++ )
520  dum_ids[i] = i;
521 
522  result = mbImpl->tag_set_data( part_set_tag, partSets, dum_ids );MB_CHK_ERR( result );
523 
524  // assign entities to the relevant sets
525  std::vector< EntityHandle > tmp_part_sets;
526  std::copy( partSets.begin(), partSets.end(), std::back_inserter( tmp_part_sets ) );
527  Range::iterator rit;
528  unsigned j = 0;
529  for( rit = elems.begin(); rit != elems.end(); rit++, j++ )
530  {
531  result = mbImpl->add_entities( tmp_part_sets[assignment[j]], &( *rit ), 1 );MB_CHK_ERR( result );
532  }
533 
534  // check for empty sets, warn if there are any
535  Range empty_sets;
536  for( rit = partSets.begin(); rit != partSets.end(); rit++ )
537  {
538  int num_ents = 0;
539  result = mbImpl->get_number_entities_by_handle( *rit, num_ents );
540  if( MB_SUCCESS != result || !num_ents ) empty_sets.insert( *rit );
541  }
542  if( !empty_sets.empty() )
543  {
544  std::cout << "WARNING: " << empty_sets.size() << " empty sets in partition: ";
545  for( rit = empty_sets.begin(); rit != empty_sets.end(); rit++ )
546  std::cout << *rit << " ";
547  std::cout << std::endl;
548  }
549  }
550 
551  if( write_as_tags )
552  {
553  Tag gid_tag;
554  result = mbImpl->tag_get_handle( "GLOBAL_ID_AGGLO", 1, MB_TYPE_INTEGER, gid_tag, MB_TAG_SPARSE );MB_CHK_ERR( result );
555 
556  // allocate idx_teger-size partitions
557  unsigned int i = 0;
558  idx_t gid;
559  for( Range::iterator rit = elems.begin(); rit != elems.end(); rit++ )
560  {
561  result = mbImpl->tag_get_data( gid_tag, &( *rit ), 1, &gid );
562  Range part_ents;
563  // std::cout<<"part ents "<<part_ents.size()<<std::endl;
564  result = mbImpl->get_entities_by_handle( *rit, part_ents, false );MB_CHK_ERR( result );
565 
566  for( Range::iterator eit = part_ents.begin(); eit != part_ents.end(); eit++ )
567  {
568  result = mbImpl->tag_set_data( part_set_tag, &( *eit ), 1, &assignment[i] );MB_CHK_ERR( result );
569 
570  result = mbImpl->tag_set_data( gid_tag, &( *eit ), 1, &gid );MB_CHK_ERR( result );
571  }
572  i++;
573  }
574  }
575  return MB_SUCCESS;
576 }
577 
578 ErrorCode MetisPartitioner::write_partition( const idx_t nparts,
579  Range& elems,
580  const idx_t* assignment,
581  const bool write_as_sets,
582  const bool write_as_tags )
583 {
584  ErrorCode result;
585 
586  // get the partition set tag
587  Tag part_set_tag;
588  idx_t dum_id = -1, i;
589  result = mbImpl->tag_get_handle( "PARALLEL_PARTITION", 1, MB_TYPE_INTEGER, part_set_tag,
590  MB_TAG_SPARSE | MB_TAG_CREAT, &dum_id );MB_CHK_ERR( result );
591 
592  // get any sets already with this tag, and clear them
593  Range tagged_sets;
594  result =
595  mbImpl->get_entities_by_type_and_tag( 0, MBENTITYSET, &part_set_tag, NULL, 1, tagged_sets, Interface::UNION );MB_CHK_ERR( result );
596  if( !tagged_sets.empty() )
597  {
598  result = mbImpl->clear_meshset( tagged_sets );
599  if( !write_as_sets )
600  {
601  result = mbImpl->tag_delete_data( part_set_tag, tagged_sets );MB_CHK_ERR( result );
602  }
603  }
604 
605  if( write_as_sets )
606  {
607  // first, create partition sets and store in vector
608  partSets.clear();
609 
610  if( nparts > (int)tagged_sets.size() )
611  {
612  // too few partition sets - create missing ones
613  idx_t num_new = nparts - tagged_sets.size();
614  for( i = 0; i < num_new; i++ )
615  {
616  EntityHandle new_set;
617  result = mbImpl->create_meshset( MESHSET_SET, new_set );MB_CHK_ERR( result );
618  tagged_sets.insert( new_set );
619  }
620  }
621  else if( nparts < (idx_t)tagged_sets.size() )
622  {
623  // too many partition sets - delete extras
624  idx_t num_del = tagged_sets.size() - nparts;
625  for( i = 0; i < num_del; i++ )
626  {
627  EntityHandle old_set = tagged_sets.pop_back();
628  result = mbImpl->delete_entities( &old_set, 1 );MB_CHK_ERR( result );
629  }
630  }
631 
632  // assign partition sets to vector
633  partSets.swap( tagged_sets );
634 
635  // write a tag to those sets denoting they're partition sets, with a value of the
636  // proc number
637  int* dum_ids = new int[nparts]; // this remains integer
638  for( i = 0; i < nparts; i++ )
639  dum_ids[i] = i;
640 
641  result = mbImpl->tag_set_data( part_set_tag, partSets, dum_ids );
642  delete[] dum_ids;
643 
644  // assign entities to the relevant sets
645  std::vector< EntityHandle > tmp_part_sets;
646  std::copy( partSets.begin(), partSets.end(), std::back_inserter( tmp_part_sets ) );
647  Range::iterator rit;
648  for( i = 0, rit = elems.begin(); rit != elems.end(); rit++, i++ )
649  {
650  result = mbImpl->add_entities( tmp_part_sets[assignment[i]], &( *rit ), 1 );MB_CHK_ERR( result );
651  }
652 
653  // check for empty sets, warn if there are any
654  Range empty_sets;
655  for( rit = partSets.begin(); rit != partSets.end(); rit++ )
656  {
657  int num_ents = 0;
658  result = mbImpl->get_number_entities_by_handle( *rit, num_ents );
659  if( MB_SUCCESS != result || !num_ents ) empty_sets.insert( *rit );
660  }
661  if( !empty_sets.empty() )
662  {
663  std::cout << "WARNING: " << empty_sets.size() << " empty sets in partition: ";
664  for( rit = empty_sets.begin(); rit != empty_sets.end(); rit++ )
665  std::cout << *rit << " ";
666  std::cout << std::endl;
667  }
668  }
669 
670  if( write_as_tags )
671  {
672  if( sizeof( int ) != sizeof( idx_t ) )
673  {
674  // allocate idx_teger-size partitions
675  // first we have to copy to int, then assign
676  int* assg_int = new int[elems.size()];
677  for( int k = 0; k < (int)elems.size(); k++ )
678  assg_int[k] = assignment[k];
679  result = mbImpl->tag_set_data( part_set_tag, elems, assg_int );MB_CHK_ERR( result );
680  delete[] assg_int;
681  }
682  else
683  result = mbImpl->tag_set_data( part_set_tag, elems, assignment );MB_CHK_ERR( result );
684  }
685 
686  return MB_SUCCESS;
687 }