MetisPartitioner.cpp

Go to the documentation of this file.

/**
 * MOAB, a Mesh-Oriented datABase, is a software component for creating,
 * storing and accessing finite element mesh data.
 *
 * Copyright 2004 Sandia Corporation.  Under the terms of Contract
 * DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government
 * retains certain rights in this software.
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 */

// Contributed by Lorenzo Alessio Botti (SpaFEDTe)
// This implementation is mostly borrowed from the mbzoltan MOAB partitioning tool

#include <iostream>
#include <cassert>
#include <sstream>
#include <map>
#include <ctime>

#include "moab/MetisPartitioner.hpp"
#include "moab/Interface.hpp"
#include "Internals.hpp"
#include "moab/Range.hpp"
#include "moab/WriteUtilIface.hpp"
#include "moab/MeshTopoUtil.hpp"
#include "moab/Skinner.hpp"
#include "MBTagConventions.hpp"
#include "moab/CN.hpp"

using namespace moab;

const bool debug = false;

MetisPartitioner::MetisPartitioner( Interface* impl, const bool use_coords )
    : PartitionerBase< idx_t >( impl, use_coords )

{
}

MetisPartitioner::~MetisPartitioner() {}

ErrorCode MetisPartitioner::partition_mesh( const idx_t nparts, const char* method, const int part_dim,
                                            const bool write_as_sets, const bool write_as_tags,
                                            const bool partition_tagged_sets, const bool partition_tagged_ents,
                                            const char* aggregating_tag, const bool print_time )
{
#ifdef MOAB_HAVE_MPI
    // should only be called in serial
    if( mbpc->proc_config().proc_size() != 1 )
    {
        std::cout << "MetisPartitioner::partition_mesh_and_geometry must be called in serial." << std::endl;
        return MB_FAILURE;
    }
#endif

    if( NULL != method && strcmp( method, "ML_RB" ) != 0 && strcmp( method, "ML_KWAY" ) != 0 )
    {
        std::cout << "ERROR: Method must be "
                  << "ML_RB or ML_KWAY" << std::endl;
        return MB_FAILURE;
    }

    std::vector< double > pts;  // x[0], y[0], z[0], ... from MOAB
    std::vector< idx_t > ids;   // poidx_t ids from MOAB
    std::vector< idx_t > adjs, parts;
    std::vector< idx_t > length;
    Range elems;
    // Get a mesh from MOAB and diide it across processors.

    clock_t t = clock();

    ErrorCode result;
    if( !partition_tagged_sets && !partition_tagged_ents )
    {
        result = assemble_graph( part_dim, pts, ids, adjs, length, elems );MB_CHK_ERR( result );
    }
    else if( partition_tagged_sets )
    {
        result = assemble_taggedsets_graph( part_dim, pts, ids, adjs, length, elems, &( *aggregating_tag ) );MB_CHK_ERR( result );
    }
    else if( partition_tagged_ents )
    {
        result = assemble_taggedents_graph( part_dim, pts, ids, adjs, length, elems, &( *aggregating_tag ) );MB_CHK_ERR( result );
    }
    else
    {
        MB_SET_ERR( MB_FAILURE, "Either partition tags or sets for Metis partitoner" );
    }

    if( print_time )
    {
        std::cout << " time to assemble graph: " << ( clock() - t ) / (double)CLOCKS_PER_SEC << "s. \n";
        t = clock();
    }

    std::cout << "Computing partition using " << method << " method for " << nparts << " processors..." << std::endl;

    idx_t nelems = length.size() - 1;
    idx_t* assign_parts;
    assign_parts         = (idx_t*)malloc( sizeof( idx_t ) * nelems );
    idx_t nconstraidx_ts = 1;
    idx_t edgeCut        = 0;
    idx_t nOfPartitions  = static_cast< idx_t >( nparts );
    idx_t metis_RESULT;

    if( strcmp( method, "ML_KWAY" ) == 0 )
    {
        idx_t options[METIS_NOPTIONS];
        METIS_SetDefaultOptions( options );
        options[METIS_OPTION_CONTIG] = 1;
        metis_RESULT = METIS_PartGraphKway( &nelems, &nconstraidx_ts, &length[0], &adjs[0], NULL, NULL, NULL,
                                            &nOfPartitions, NULL, NULL, options, &edgeCut, assign_parts );
    }
    else if( strcmp( method, "ML_RB" ) == 0 )
    {
        idx_t options[METIS_NOPTIONS];
        METIS_SetDefaultOptions( options );
        options[METIS_OPTION_OBJTYPE] = METIS_OBJTYPE_CUT;  // CUT
        options[METIS_OPTION_IPTYPE]  = METIS_IPTYPE_GROW;  // GROW or RANDOM
        options[METIS_OPTION_CTYPE]   = METIS_CTYPE_RM;     // RM or SHEM
        options[METIS_OPTION_RTYPE]   = METIS_RTYPE_FM;     // FM
        options[METIS_OPTION_NCUTS]   = 10;                 // Number of different partitionings to compute, then
                                                            // chooses the best one, default = 1
        options[METIS_OPTION_NITER]   = 10;                 // Number of refinements steps, default = 10
        options[METIS_OPTION_UFACTOR] = 30;                 // Imabalance, default = 1
        options[METIS_OPTION_DBGLVL]  = METIS_DBG_INFO;
        metis_RESULT = METIS_PartGraphRecursive( &nelems, &nconstraidx_ts, &length[0], &adjs[0], NULL, NULL, NULL,
                                                 &nOfPartitions, NULL, NULL, options, &edgeCut, assign_parts );
    }
    else
        MB_SET_ERR( MB_FAILURE, "Either ML_KWAY or ML_RB needs to be specified for Metis partitioner" );

    if( print_time )
    {
        std::cout << " time to partition: " << ( clock() - t ) / (double)CLOCKS_PER_SEC << "s. \n";
        t = clock();
    }

#ifdef MOAB_HAVE_MPI
    // assign global node ids, starting from one! TODO
    if( assign_global_ids )
    {
        EntityHandle rootset = 0;
        result               = mbpc->assign_global_ids( rootset, part_dim, 1, true, false );MB_CHK_ERR( result );
    }
#endif

    if( metis_RESULT != METIS_OK ) return MB_FAILURE;

    // take results & write onto MOAB partition sets
    std::cout << "Saving partition information to MOAB..." << std::endl;
    {
        if( partition_tagged_sets || partition_tagged_ents )
        {
            result = write_aggregationtag_partition( nparts, elems, assign_parts, write_as_sets, write_as_tags );MB_CHK_ERR( result );
        }
        else
        {
            result = write_partition( nparts, elems, assign_parts, write_as_sets, write_as_tags );MB_CHK_ERR( result );
        }
    }

    if( print_time )
    {
        std::cout << " time to write partition in memory " << ( clock() - t ) / (double)CLOCKS_PER_SEC << "s. \n";
        t = clock();
    }
    free( assign_parts );

    return MB_SUCCESS;
}

ErrorCode MetisPartitioner::assemble_taggedents_graph( const int dimension, std::vector< double >& coords,
                                                       std::vector< idx_t >& moab_ids,
                                                       std::vector< idx_t >& adjacencies, std::vector< idx_t >& length,
                                                       Range& elems, const char* aggregating_tag )
{
    Tag partSetTag;
    ErrorCode result = mbImpl->tag_get_handle( aggregating_tag, 1, MB_TYPE_INTEGER, partSetTag );
    if( MB_SUCCESS != result ) return result;

    Range allSubElems;
    result = mbImpl->get_entities_by_dimension( 0, dimension, allSubElems );
    if( MB_SUCCESS != result || allSubElems.empty() ) return result;
    idx_t partSet;
    std::map< idx_t, Range > aggloElems;
    for( Range::iterator rit = allSubElems.begin(); rit != allSubElems.end(); rit++ )
    {
        EntityHandle entity = *rit;
        result              = mbImpl->tag_get_data( partSetTag, &entity, 1, &partSet );
        if( MB_SUCCESS != result ) return result;
        if( partSet >= 0 ) aggloElems[partSet].insert( entity );
    }
    // clear aggregating tag data
    TagType type;
    result = mbImpl->tag_get_type( partSetTag, type );
    if( type == MB_TAG_DENSE )
    {
        // clear tag on ents and sets
        result = mbImpl->tag_delete( partSetTag );
        if( MB_SUCCESS != result ) return result;
    }
    if( type == MB_TAG_SPARSE )
    {
        // clear tag on ents
        result = mbImpl->tag_delete_data( partSetTag, allSubElems );
        if( MB_SUCCESS != result ) return result;
        // clear tag on sets
        result = mbImpl->get_entities_by_type_and_tag( 0, MBENTITYSET, &partSetTag, 0, 1, elems );
        if( MB_SUCCESS != result ) return result;
        result = mbImpl->tag_delete_data( partSetTag, elems );
        if( MB_SUCCESS != result ) return result;
        elems.clear();
    }
    result =
        mbImpl->tag_get_handle( "PARALLEL_PARTITION", 1, MB_TYPE_INTEGER, partSetTag, MB_TAG_SPARSE | MB_TAG_CREAT );
    if( MB_SUCCESS != result ) return result;

    for( std::map< idx_t, Range >::iterator mit = aggloElems.begin(); mit != aggloElems.end(); mit++ )
    {
        EntityHandle new_set;
        result = mbImpl->create_meshset( MESHSET_SET, new_set );
        if( MB_SUCCESS != result ) return result;
        result = mbImpl->add_entities( new_set, mit->second );
        if( MB_SUCCESS != result ) return result;
        result = mbImpl->tag_set_data( partSetTag, &new_set, 1, &mit->first );
        if( MB_SUCCESS != result ) return result;
    }

    result =
        assemble_taggedsets_graph( dimension, coords, moab_ids, adjacencies, length, elems, &( *aggregating_tag ) );
    return MB_SUCCESS;
}

ErrorCode MetisPartitioner::assemble_taggedsets_graph( const int dimension, std::vector< double >& coords,
                                                       std::vector< idx_t >& moab_ids,
                                                       std::vector< idx_t >& adjacencies, std::vector< idx_t >& length,
                                                       Range& elems, const char* aggregating_tag )
{
    length.push_back( 0 );
    // assemble a graph with vertices equal to elements of specified dimension, edges
    // signified by list of other elements to which an element is connected

    // get the tagged elements
    Tag partSetTag;
    ErrorCode result = mbImpl->tag_get_handle( aggregating_tag, 1, MB_TYPE_INTEGER, partSetTag );MB_CHK_ERR( result );
    // ErrorCode result = mbImpl->tag_get_handle("PARALLEL_PARTITION_SET", 1, MB_TYPE_INTEGER,
    // partSetTag);MB_CHK_ERR(result);

    result = mbImpl->get_entities_by_type_and_tag( 0, MBENTITYSET, &partSetTag, 0, 1, elems );
    if( MB_SUCCESS != result || elems.empty() ) return result;

    // assign globla ids to elem sets based on aggregating_tag data
    Tag gid_tag;
    idx_t zero1 = -1;
    result =
        mbImpl->tag_get_handle( "GLOBAL_ID_AGGLO", 1, MB_TYPE_INTEGER, gid_tag, MB_TAG_SPARSE | MB_TAG_CREAT, &zero1 );MB_CHK_ERR( result );
    for( Range::iterator rit = elems.begin(); rit != elems.end(); rit++ )
    {
        idx_t partSet;
        result = mbImpl->tag_get_data( partSetTag, &( *rit ), 1, &partSet );MB_CHK_ERR( result );
        result = mbImpl->tag_set_data( gid_tag, &( *rit ), 1, &partSet );MB_CHK_ERR( result );
    }
    // clear aggregating tag data
    TagType type;
    result = mbImpl->tag_get_type( partSetTag, type );MB_CHK_ERR( result );
    if( type == MB_TAG_DENSE )
    {
        result = mbImpl->tag_delete( partSetTag );MB_CHK_ERR( result );
    }
    if( type == MB_TAG_SPARSE )
    {
        result = mbImpl->tag_delete_data( partSetTag, elems );MB_CHK_ERR( result );
    }

    // assemble the graph, using Skinner to get d-1 dimensional neighbors and then idx_tersecting to
    // get adjacencies
    std::vector< Range > skin_subFaces( elems.size() );
    unsigned int i = 0;
    for( Range::iterator rit = elems.begin(); rit != elems.end(); rit++ )
    {
        Range part_ents;
        result = mbImpl->get_entities_by_handle( *rit, part_ents, false );
        if( mbImpl->dimension_from_handle( *part_ents.rbegin() ) !=
            mbImpl->dimension_from_handle( *part_ents.begin() ) )
        {
            Range::iterator lower = part_ents.lower_bound( CN::TypeDimensionMap[0].first ),
                            upper = part_ents.upper_bound( CN::TypeDimensionMap[dimension - 1].second );
            part_ents.erase( lower, upper );
        }
        Skinner skinner( mbImpl );
        result = skinner.find_skin( 0, part_ents, false, skin_subFaces[i], NULL, false, true, false );MB_CHK_ERR( result );
        i++;
    }
    std::vector< EntityHandle > adjs;
    std::vector< idx_t > neighbors;
    double avg_position[3];
    idx_t moab_id;
    MeshTopoUtil mtu( mbImpl );
    for( unsigned int k = 0; k < i; k++ )
    {
        // get bridge adjacencies for element k
        adjs.clear();
        for( unsigned int t = 0; t < i; t++ )
        {
            if( t != k )
            {
                Range subFaces = intersect( skin_subFaces[k], skin_subFaces[t] );
                if( subFaces.size() > 0 ) adjs.push_back( elems[t] );
            }
        }
        if( !adjs.empty() )
        {
            neighbors.resize( adjs.size() );
            result = mbImpl->tag_get_data( gid_tag, &adjs[0], adjs.size(), &neighbors[0] );MB_CHK_ERR( result );
        }
        // copy those idx_to adjacencies vector
        length.push_back( length.back() + (idx_t)adjs.size() );
        std::copy( neighbors.begin(), neighbors.end(), std::back_inserter( adjacencies ) );
        // get the graph vertex id for this element
        const EntityHandle& setk = elems[k];
        result                   = mbImpl->tag_get_data( gid_tag, &setk, 1, &moab_id );
        moab_ids.push_back( moab_id );
        // get average position of vertices
        Range part_ents;
        result = mbImpl->get_entities_by_handle( elems[k], part_ents, false );MB_CHK_ERR( result );
        result = mtu.get_average_position( part_ents, avg_position );MB_CHK_ERR( result );
        std::copy( avg_position, avg_position + 3, std::back_inserter( coords ) );
    }
    for( unsigned int k = 0; k < i; k++ )
    {
        for( unsigned int t = 0; t < k; t++ )
        {
            Range subFaces = intersect( skin_subFaces[k], skin_subFaces[t] );
            if( subFaces.size() > 0 ) mbImpl->delete_entities( subFaces );
        }
    }

    if( debug )
    {
        std::cout << "Length vector: " << std::endl;
        std::copy( length.begin(), length.end(), std::ostream_iterator< idx_t >( std::cout, ", " ) );
        std::cout << std::endl;
        std::cout << "Adjacencies vector: " << std::endl;
        std::copy( adjacencies.begin(), adjacencies.end(), std::ostream_iterator< idx_t >( std::cout, ", " ) );
        std::cout << std::endl;
        std::cout << "Moab_ids vector: " << std::endl;
        std::copy( moab_ids.begin(), moab_ids.end(), std::ostream_iterator< idx_t >( std::cout, ", " ) );
        std::cout << std::endl;
        std::cout << "Coords vector: " << std::endl;
        std::copy( coords.begin(), coords.end(), std::ostream_iterator< double >( std::cout, ", " ) );
        std::cout << std::endl;
    }
    return MB_SUCCESS;
}

ErrorCode MetisPartitioner::assemble_graph( const int dimension, std::vector< double >& coords,
                                            std::vector< idx_t >& moab_ids, std::vector< idx_t >& adjacencies,
                                            std::vector< idx_t >& length, Range& elems )
{
    length.push_back( 0 );
    // assemble a graph with vertices equal to elements of specified dimension, edges
    // signified by list of other elements to which an element is connected

    // get the elements of that dimension
    ErrorCode result = mbImpl->get_entities_by_dimension( 0, dimension, elems );
    if( MB_SUCCESS != result || elems.empty() ) return result;

#ifdef MOAB_HAVE_MPI
    // assign global ids
    if( assign_global_ids )
    {
        result = mbpc->assign_global_ids( 0, dimension, 0 );MB_CHK_ERR( result );
    }
#endif

    // now assemble the graph, calling MeshTopoUtil to get bridge adjacencies through d-1
    // dimensional neighbors
    MeshTopoUtil mtu( mbImpl );
    Range adjs;
    // can use a fixed-size array 'cuz the number of lower-dimensional neighbors is limited
    // by MBCN
    int neighbors[5 * MAX_SUB_ENTITIES];  // these will be now indices in the elems range

    double avg_position[3];
    int index_in_elems = 0;

    for( Range::iterator rit = elems.begin(); rit != elems.end(); rit++, index_in_elems++ )
    {

        // get bridge adjacencies
        adjs.clear();
        result = mtu.get_bridge_adjacencies( *rit, ( dimension > 0 ? dimension - 1 : 3 ), dimension, adjs );MB_CHK_ERR( result );

        // get the indices in elems range of those
        if( !adjs.empty() )
        {
            int i = 0;
            assert( adjs.size() < 5 * MAX_SUB_ENTITIES );
            for( Range::iterator ait = adjs.begin(); ait != adjs.end(); ait++, i++ )
            {
                EntityHandle adjEnt = *ait;
                neighbors[i]        = elems.index( adjEnt );
            }
        }

        // copy those idx_to adjacencies vector
        length.push_back( length.back() + (idx_t)adjs.size() );
        // conversion made to idx_t
        std::copy( neighbors, neighbors + adjs.size(), std::back_inserter( adjacencies ) );

        // get average position of vertices
        result = mtu.get_average_position( *rit, avg_position );MB_CHK_ERR( result );

        // get the graph vertex id for this element; it is now index in elems
        moab_ids.push_back( index_in_elems );  // conversion made to idx_t

        // copy_to coords vector
        std::copy( avg_position, avg_position + 3, std::back_inserter( coords ) );
    }

    if( debug )
    {
        std::cout << "Length vector: " << std::endl;
        std::copy( length.begin(), length.end(), std::ostream_iterator< idx_t >( std::cout, ", " ) );
        std::cout << std::endl;
        std::cout << "Adjacencies vector: " << std::endl;
        std::copy( adjacencies.begin(), adjacencies.end(), std::ostream_iterator< idx_t >( std::cout, ", " ) );
        std::cout << std::endl;
        std::cout << "Moab_ids vector: " << std::endl;
        std::copy( moab_ids.begin(), moab_ids.end(), std::ostream_iterator< idx_t >( std::cout, ", " ) );
        std::cout << std::endl;
        std::cout << "Coords vector: " << std::endl;
        std::copy( coords.begin(), coords.end(), std::ostream_iterator< double >( std::cout, ", " ) );
        std::cout << std::endl;
    }

    return MB_SUCCESS;
}

ErrorCode MetisPartitioner::write_aggregationtag_partition( const idx_t nparts, Range& elems, const idx_t* assignment,
                                                            const bool write_as_sets, const bool write_as_tags )
{
    ErrorCode result;

    // get the partition set tag
    Tag part_set_tag;
    result =
        mbImpl->tag_get_handle( "PARALLEL_PARTITION", 1, MB_TYPE_INTEGER, part_set_tag, MB_TAG_SPARSE | MB_TAG_CREAT );MB_CHK_ERR( result );

    // get any sets already with this tag, and clear them
    Range tagged_sets;
    result =
        mbImpl->get_entities_by_type_and_tag( 0, MBENTITYSET, &part_set_tag, NULL, 1, tagged_sets, Interface::UNION );MB_CHK_ERR( result );
    if( !tagged_sets.empty() )
    {
        result = mbImpl->clear_meshset( tagged_sets );
        if( !write_as_sets )
        {
            result = mbImpl->tag_delete_data( part_set_tag, tagged_sets );MB_CHK_ERR( result );
        }
    }

    if( write_as_sets )
    {
        // first, create partition sets and store in vector
        partSets.clear();

        if( nparts > (idx_t)tagged_sets.size() )
        {
            // too few partition sets - create missing ones
            idx_t num_new = nparts - tagged_sets.size();
            for( idx_t i = 0; i < num_new; i++ )
            {
                EntityHandle new_set;
                result = mbImpl->create_meshset( MESHSET_SET, new_set );MB_CHK_ERR( result );
                tagged_sets.insert( new_set );
            }
        }
        else if( nparts < (idx_t)tagged_sets.size() )
        {
            // too many partition sets - delete extras
            idx_t num_del = tagged_sets.size() - nparts;
            for( idx_t i = 0; i < num_del; i++ )
            {
                EntityHandle old_set = tagged_sets.pop_back();
                result               = mbImpl->delete_entities( &old_set, 1 );MB_CHK_ERR( result );
            }
        }

        // assign partition sets to vector
        partSets.swap( tagged_sets );

        // write a tag to those sets denoting they're partition sets, with a value of the
        // proc number
        idx_t* dum_ids = new idx_t[nparts];
        for( idx_t i = 0; i < nparts; i++ )
            dum_ids[i] = i;

        result = mbImpl->tag_set_data( part_set_tag, partSets, dum_ids );MB_CHK_ERR( result );

        // assign entities to the relevant sets
        std::vector< EntityHandle > tmp_part_sets;
        std::copy( partSets.begin(), partSets.end(), std::back_inserter( tmp_part_sets ) );
        Range::iterator rit;
        unsigned j = 0;
        for( rit = elems.begin(); rit != elems.end(); rit++, j++ )
        {
            result = mbImpl->add_entities( tmp_part_sets[assignment[j]], &( *rit ), 1 );MB_CHK_ERR( result );
        }

        // check for empty sets, warn if there are any
        Range empty_sets;
        for( rit = partSets.begin(); rit != partSets.end(); rit++ )
        {
            int num_ents = 0;
            result       = mbImpl->get_number_entities_by_handle( *rit, num_ents );
            if( MB_SUCCESS != result || !num_ents ) empty_sets.insert( *rit );
        }
        if( !empty_sets.empty() )
        {
            std::cout << "WARNING: " << empty_sets.size() << " empty sets in partition: ";
            for( rit = empty_sets.begin(); rit != empty_sets.end(); rit++ )
                std::cout << *rit << " ";
            std::cout << std::endl;
        }
    }

    if( write_as_tags )
    {
        Tag gid_tag;
        result = mbImpl->tag_get_handle( "GLOBAL_ID_AGGLO", 1, MB_TYPE_INTEGER, gid_tag, MB_TAG_SPARSE );MB_CHK_ERR( result );

        // allocate idx_teger-size partitions
        unsigned int i = 0;
        idx_t gid;
        for( Range::iterator rit = elems.begin(); rit != elems.end(); rit++ )
        {
            result = mbImpl->tag_get_data( gid_tag, &( *rit ), 1, &gid );
            Range part_ents;
            //      std::cout<<"part ents "<<part_ents.size()<<std::endl;
            result = mbImpl->get_entities_by_handle( *rit, part_ents, false );MB_CHK_ERR( result );

            for( Range::iterator eit = part_ents.begin(); eit != part_ents.end(); eit++ )
            {
                result = mbImpl->tag_set_data( part_set_tag, &( *eit ), 1, &assignment[i] );MB_CHK_ERR( result );

                result = mbImpl->tag_set_data( gid_tag, &( *eit ), 1, &gid );MB_CHK_ERR( result );
            }
            i++;
        }
    }
    return MB_SUCCESS;
}

ErrorCode MetisPartitioner::write_partition( const idx_t nparts, Range& elems, const idx_t* assignment,
                                             const bool write_as_sets, const bool write_as_tags )
{
    ErrorCode result;

    // get the partition set tag
    Tag part_set_tag;
    idx_t dum_id = -1, i;
    result       = mbImpl->tag_get_handle( "PARALLEL_PARTITION", 1, MB_TYPE_INTEGER, part_set_tag,
                                     MB_TAG_SPARSE | MB_TAG_CREAT, &dum_id );MB_CHK_ERR( result );

    // get any sets already with this tag, and clear them
    Range tagged_sets;
    result =
        mbImpl->get_entities_by_type_and_tag( 0, MBENTITYSET, &part_set_tag, NULL, 1, tagged_sets, Interface::UNION );MB_CHK_ERR( result );
    if( !tagged_sets.empty() )
    {
        result = mbImpl->clear_meshset( tagged_sets );
        if( !write_as_sets )
        {
            result = mbImpl->tag_delete_data( part_set_tag, tagged_sets );MB_CHK_ERR( result );
        }
    }

    if( write_as_sets )
    {
        // first, create partition sets and store in vector
        partSets.clear();

        if( nparts > (int)tagged_sets.size() )
        {
            // too few partition sets - create missing ones
            idx_t num_new = nparts - tagged_sets.size();
            for( i = 0; i < num_new; i++ )
            {
                EntityHandle new_set;
                result = mbImpl->create_meshset( MESHSET_SET, new_set );MB_CHK_ERR( result );
                tagged_sets.insert( new_set );
            }
        }
        else if( nparts < (idx_t)tagged_sets.size() )
        {
            // too many partition sets - delete extras
            idx_t num_del = tagged_sets.size() - nparts;
            for( i = 0; i < num_del; i++ )
            {
                EntityHandle old_set = tagged_sets.pop_back();
                result               = mbImpl->delete_entities( &old_set, 1 );MB_CHK_ERR( result );
            }
        }

        // assign partition sets to vector
        partSets.swap( tagged_sets );

        // write a tag to those sets denoting they're partition sets, with a value of the
        // proc number
        int* dum_ids = new int[nparts];  // this remains integer
        for( i = 0; i < nparts; i++ )
            dum_ids[i] = i;

        result = mbImpl->tag_set_data( part_set_tag, partSets, dum_ids );
        delete[] dum_ids;

        // assign entities to the relevant sets
        std::vector< EntityHandle > tmp_part_sets;
        std::copy( partSets.begin(), partSets.end(), std::back_inserter( tmp_part_sets ) );
        Range::iterator rit;
        for( i = 0, rit = elems.begin(); rit != elems.end(); rit++, i++ )
        {
            result = mbImpl->add_entities( tmp_part_sets[assignment[i]], &( *rit ), 1 );MB_CHK_ERR( result );
        }

        // check for empty sets, warn if there are any
        Range empty_sets;
        for( rit = partSets.begin(); rit != partSets.end(); rit++ )
        {
            int num_ents = 0;
            result       = mbImpl->get_number_entities_by_handle( *rit, num_ents );
            if( MB_SUCCESS != result || !num_ents ) empty_sets.insert( *rit );
        }
        if( !empty_sets.empty() )
        {
            std::cout << "WARNING: " << empty_sets.size() << " empty sets in partition: ";
            for( rit = empty_sets.begin(); rit != empty_sets.end(); rit++ )
                std::cout << *rit << " ";
            std::cout << std::endl;
        }
    }

    if( write_as_tags )
    {
        if( sizeof( int ) != sizeof( idx_t ) )
        {
            // allocate idx_teger-size partitions
            // first we have to copy to int, then assign
            int* assg_int = new int[elems.size()];
            for( int k = 0; k < (int)elems.size(); k++ )
                assg_int[k] = assignment[k];
            result = mbImpl->tag_set_data( part_set_tag, elems, assg_int );MB_CHK_ERR( result );
            delete[] assg_int;
        }
        else
            result = mbImpl->tag_set_data( part_set_tag, elems, assignment );MB_CHK_ERR( result );
    }

    return MB_SUCCESS;
}