inline
size_t get_entities(stk::mesh::Part &part,
stk::mesh::EntityRank type,
const stk::mesh::BulkData &bulk,
stk::mesh::EntityVector &entities,
bool include_shared,
const stk::mesh::Selector *subset_selector)
{
stk::mesh::MetaData & meta = stk::mesh::MetaData::get(part);
stk::mesh::Selector own_share = meta.locally_owned_part();
if(include_shared)
own_share |= meta.globally_shared_part();
stk::mesh::Selector selector = part & own_share;
if(subset_selector)
selector &= *subset_selector;
get_selected_entities(selector, bulk.buckets(type), entities);
return entities.size();
}
STKUNIT_UNIT_TEST(UnitTestZoltanSimple, testUnit)
{
#ifdef STK_HAS_MPI
stk::ParallelMachine comm(MPI_COMM_WORLD);
#else
stk::ParallelMachine comm(0);
#endif
unsigned spatial_dimension = 2;
std::vector<std::string> rank_names = stk::mesh::fem::entity_rank_names(spatial_dimension);
const stk::mesh::EntityRank constraint_rank = rank_names.size();
rank_names.push_back("Constraint");
stk::mesh::fem::FEMMetaData fem_meta;
fem_meta.FEM_initialize( spatial_dimension, rank_names );
stk::mesh::MetaData & meta_data = stk::mesh::fem::FEMMetaData::get_meta_data(fem_meta);
stk::mesh::BulkData bulk_data( meta_data , comm , 100 );
const stk::mesh::EntityRank element_rank = fem_meta.element_rank();
stk::mesh::fem::CellTopology quad_top(shards::getCellTopologyData<shards::Quadrilateral<4> >());
stk::mesh::Part & quad_part( fem_meta.declare_part("quad", quad_top ) );
VectorField & coord_field( fem_meta.declare_field< VectorField >( "coordinates" ) );
ScalarField & weight_field( fem_meta.declare_field< ScalarField >( "element_weights" ) );
stk::mesh::put_field( coord_field , NODE_RANK , fem_meta.universal_part() );
stk::mesh::put_field(weight_field , element_rank , fem_meta.universal_part() );
fem_meta.commit();
const unsigned p_size = bulk_data.parallel_size();
const unsigned p_rank = bulk_data.parallel_rank();
bulk_data.modification_begin();
if ( p_rank == 0 ) {
std::vector<std::vector<stk::mesh::Entity*> > quads(nx);
for ( unsigned ix = 0 ; ix < nx ; ++ix ) quads[ix].resize(ny);
const unsigned nnx = nx + 1 ;
const unsigned nny = ny + 1 ;
for ( unsigned iy = 0 ; iy < ny ; ++iy ) {
for ( unsigned ix = 0 ; ix < nx ; ++ix ) {
stk::mesh::EntityId elem = 1 + ix + iy * nx ;
stk::mesh::EntityId nodes[4] ;
nodes[0] = 1 + ix + iy * nnx ;
nodes[1] = 2 + ix + iy * nnx ;
nodes[2] = 2 + ix + ( iy + 1 ) * nnx ;
nodes[3] = 1 + ix + ( iy + 1 ) * nnx ;
stk::mesh::Entity &q = stk::mesh::fem::declare_element( bulk_data , quad_part , elem , nodes );
quads[ix][iy] = &q;
}
}
for ( unsigned iy = 0 ; iy < ny ; ++iy ) {
for ( unsigned ix = 0 ; ix < nx ; ++ix ) {
stk::mesh::EntityId elem = 1 + ix + iy * nx ;
stk::mesh::Entity * e = bulk_data.get_entity( element_rank, elem );
double * const e_weight = stk::mesh::field_data( weight_field , *e );
*e_weight = 1.0;
}
}
for ( unsigned iy = 0 ; iy <= ny ; ++iy ) {
for ( unsigned ix = 0 ; ix <= nx ; ++ix ) {
stk::mesh::EntityId nid = 1 + ix + iy * nnx ;
stk::mesh::Entity * n = bulk_data.get_entity( NODE_RANK, nid );
double * const coord = stk::mesh::field_data( coord_field , *n );
coord[0] = .1*ix;
coord[1] = .1*iy;
coord[2] = 0;
}
}
{
const unsigned iy_left = 0;
const unsigned iy_right = ny;
stk::mesh::PartVector add(1, &fem_meta.locally_owned_part());
for ( unsigned ix = 0 ; ix <= nx ; ++ix ) {
stk::mesh::EntityId nid_left = 1 + ix + iy_left * nnx ;
stk::mesh::EntityId nid_right = 1 + ix + iy_right * nnx ;
stk::mesh::Entity * n_left = bulk_data.get_entity( NODE_RANK, nid_left );
stk::mesh::Entity * n_right = bulk_data.get_entity( NODE_RANK, nid_right );
const stk::mesh::EntityId constraint_entity_id = 1 + ix + nny * nnx;
stk::mesh::Entity & c = bulk_data.declare_entity( constraint_rank, constraint_entity_id, add );
bulk_data.declare_relation( c , *n_left , 0 );
bulk_data.declare_relation( c , *n_right , 1 );
}
}
}
// Only P0 has any nodes or elements
if ( p_rank == 0 ) {
STKUNIT_ASSERT( ! bulk_data.buckets( NODE_RANK ).empty() );
STKUNIT_ASSERT( ! bulk_data.buckets( element_rank ).empty() );
}
else {
STKUNIT_ASSERT( bulk_data.buckets( NODE_RANK ).empty() );
STKUNIT_ASSERT( bulk_data.buckets( element_rank ).empty() );
}
bulk_data.modification_end();
// create some sides and faces to rebalance.
stk::mesh::PartVector add_parts;
stk::mesh::create_adjacent_entities(bulk_data, add_parts);
// Zoltan partition is specialized fomm a virtual base class, stk::rebalance::Partition.
// Other specializations are possible.
Teuchos::ParameterList emptyList;
stk::rebalance::Zoltan zoltan_partition(comm, spatial_dimension, emptyList);
{
stk::mesh::Selector selector(fem_meta.universal_part());
stk::rebalance::rebalance(bulk_data, selector, &coord_field, &weight_field, zoltan_partition);
}
const double imbalance_threshold = stk::rebalance::check_balance(bulk_data, &weight_field, element_rank);
const bool do_rebal = 1.5 < imbalance_threshold;
// Check that we satisfy our threshhold
STKUNIT_ASSERT( !do_rebal );
if( (2 == p_size) || (4 == p_size) )
{
STKUNIT_ASSERT_NEAR(imbalance_threshold, 1.0, 1.e-8);
}
else
{
STKUNIT_ASSERT_LE(imbalance_threshold, 1.5);
}
// And verify that all dependent entities are on the same proc as their parent element
{
stk::mesh::EntityVector entities;
stk::mesh::Selector selector = fem_meta.locally_owned_part();
get_selected_entities(selector, bulk_data.buckets(NODE_RANK), entities);
bool result = stk::rebalance::verify_dependent_ownership(element_rank, entities);
//get_selected_entities(selector, bulk_data.buckets(constraint_rank), entities);
//result &= stk::rebalance::verify_dependent_ownership(element_rank, entities);
STKUNIT_ASSERT( result );
}
}
