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