Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use a quick union algorithm to speed up PFBlockAlgo (factor of 2-3x) #14138

Merged
merged 1 commit into from Apr 29, 2016
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
20 changes: 11 additions & 9 deletions RecoParticleFlow/PFProducer/interface/PFBlockAlgo.h
Expand Up @@ -62,6 +62,7 @@

#include <map>
#include <unordered_map>
#include <unordered_set>

namespace std {
template<>
Expand All @@ -76,7 +77,7 @@ namespace std {
struct equal_to<std::pair<size_t,size_t> > {
typedef std::pair<size_t,size_t> arg_type;
bool operator()(const arg_type& arg1, const arg_type& arg2) const {
return (arg1.first == arg2.first && arg1.second == arg2.second);
return ( (arg1.first == arg2.first) & (arg1.second == arg2.second) );
}
};
}
Expand All @@ -99,6 +100,8 @@ class PFBlockAlgo {
typedef ElementList::iterator IE;
typedef ElementList::const_iterator IEC;
typedef reco::PFBlockCollection::const_iterator IBC;
//for skipping ranges
typedef std::array<std::pair<size_t,size_t>,reco::PFBlockElement::kNBETypes> ElementRanges;

PFBlockAlgo();

Expand Down Expand Up @@ -132,12 +135,7 @@ class PFBlockAlgo {


private:
// flattened version of topological
// association of block elements
IE associate( ElementList& elems,
std::unordered_map<std::pair<size_t,size_t>,PFBlockLink>& links,
reco::PFBlock& );


/// compute missing links in the blocks
/// (the recursive procedure does not build all links)
void packLinks(reco::PFBlock& block,
Expand All @@ -153,11 +151,14 @@ class PFBlockAlgo {
PFBlockLink::Type& linktype,
reco::PFBlock::LinkTest& linktest,
double& dist) const;


std::auto_ptr< reco::PFBlockCollection > blocksNew_;
std::auto_ptr< reco::PFBlockCollection > blocks_;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

auto_ptr has been deprecated. When adding new code, wouldn't it be better to use shared_ptr or unique_ptr instead?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

blocksNew_ was from debugging and implementation anyway, I'll fix it.


// the test elements will be transferred to the blocks
ElementList elements_;
ElementList elements_;
std::vector<ElementList::value_type::pointer> bare_elements_;
ElementRanges ranges_;

/// if true, debug printouts activated
bool debug_;
Expand All @@ -170,6 +171,7 @@ class PFBlockAlgo {
const std::unordered_map<std::string,reco::PFBlockElement::Type>
_elementTypes;
std::vector<LinkTestPtr> _linkTests;
unsigned int _linkTestSquare[reco::PFBlockElement::kNBETypes][reco::PFBlockElement::kNBETypes];

std::vector<KDTreePtr> _kdtrees;
};
Expand Down
240 changes: 151 additions & 89 deletions RecoParticleFlow/PFProducer/src/PFBlockAlgo.cc
Expand Up @@ -8,13 +8,59 @@
#include "DataFormats/ParticleFlowReco/interface/PFRecHit.h"

#include <stdexcept>
#include <algorithm>
#include "TMath.h"

using namespace std;
using namespace reco;

#define INIT_ENTRY(name) {#name,name}

namespace {
class QuickUnion{
std::vector<unsigned> _id;
std::vector<unsigned> _size;
int _count;

public:
QuickUnion(const unsigned NBranches) {
_count = NBranches;
_id.resize(NBranches);
_size.resize(NBranches);
for( unsigned i = 0; i < NBranches; ++i ) {
_id[i] = i;
_size[i] = 1;
}
}

int count() const { return _count; }

unsigned find(unsigned p) {
while( p != _id[p] ) {
_id[p] = _id[_id[p]];
p = _id[p];
}
return p;
}

bool connected(unsigned p, unsigned q) { return find(p) == find(q); }

void unite(unsigned p, unsigned q) {
unsigned rootP = find(p);
unsigned rootQ = find(q);
_id[p] = q;

if(_size[rootP] < _size[rootQ] ) {
_id[rootP] = rootQ; _size[rootQ] += _size[rootP];
} else {
_id[rootQ] = rootP; _size[rootP] += _size[rootQ];
}
--_count;
}
};
}


//for debug only
//#define PFLOW_DEBUG

Expand All @@ -37,6 +83,12 @@ PFBlockAlgo::PFBlockAlgo() :

void PFBlockAlgo::setLinkers(const std::vector<edm::ParameterSet>& confs) {
constexpr unsigned rowsize = reco::PFBlockElement::kNBETypes;
for( unsigned i = 0; i < rowsize; ++i ) {
for( unsigned j = 0; j < rowsize; ++j ) {

_linkTestSquare[i][j] = 0;
}
}
_linkTests.resize(rowsize*rowsize);
const std::string prefix("PFBlockElement::");
const std::string pfx_kdtree("KDTree");
Expand All @@ -60,10 +112,12 @@ void PFBlockAlgo::setLinkers(const std::vector<edm::ParameterSet>& confs) {
}
const PFBlockElement::Type type1 = _elementTypes.at(link1);
const PFBlockElement::Type type2 = _elementTypes.at(link2);
const unsigned index = rowsize*std::max(type1,type2)+std::min(type1,type2);
const unsigned index = rowsize*std::max(type1,type2)+std::min(type1,type2);
BlockElementLinkerBase * linker =
BlockElementLinkerFactory::get()->create(linkerName,conf);
_linkTests[index].reset(linker);
_linkTestSquare[type1][type2] = index;
_linkTestSquare[type2][type1] = index;
// setup KDtree if requested
const bool useKDTree = conf.getParameter<bool>("useKDTree");
if( useKDTree ) {
Expand Down Expand Up @@ -96,8 +150,7 @@ PFBlockAlgo::~PFBlockAlgo() {
#endif
}

void
PFBlockAlgo::findBlocks() {
void PFBlockAlgo::findBlocks() {
// Glowinski & Gouzevitch
for( const auto& kdtree : _kdtrees ) {
kdtree->process();
Expand All @@ -107,93 +160,79 @@ PFBlockAlgo::findBlocks() {
if( blocks_.get() ) blocks_->clear();
else blocks_.reset( new reco::PFBlockCollection );
blocks_->reserve(elements_.size());
for(IE ie = elements_.begin();
ie != elements_.end();) {

#ifdef PFLOW_DEBUG
if(debug_) {
cout<<" PFBlockAlgo::findBlocks() ----------------------"<<endl;
cout<<" element "<<**ie<<endl;
cout<<" creating new block"<<endl;
}
#endif

blocks_->push_back( reco::PFBlock() );

std::unordered_map<std::pair<size_t,size_t>, PFBlockLink > links;
links.reserve(elements_.size());

ie = associate(elements_, links, blocks_->back());

packLinks( blocks_->back(), links );
QuickUnion qu(bare_elements_.size());
const auto elem_size = bare_elements_.size();
for( unsigned i = 0; i < elem_size; ++i ) {
for( unsigned j = i+1; j < elem_size; ++j ) {
if( !_linkTests[_linkTestSquare[bare_elements_[i]->type()][bare_elements_[j]->type()]] ) {
j = ranges_[bare_elements_[j]->type()].second;
continue;
}
auto p1(bare_elements_[i]), p2(bare_elements_[j]);
const PFBlockElement::Type type1 = p1->type();
const PFBlockElement::Type type2 = p2->type();
const unsigned index = _linkTestSquare[type1][type2];
if( _linkTests[index]->linkPrefilter(p1,p2) ) {
const double dist = _linkTests[index]->testLink(p1,p2);
// compute linking info if it is possible
if( dist > -0.5 ) {
qu.unite(i,j);
}
}
}
}

std::unordered_multimap<unsigned,unsigned> blocksmap(elements_.size());
std::vector<unsigned> keys;
keys.reserve(elements_.size());
for( unsigned i = 0; i < elements_.size(); ++i ) {
unsigned key = qu.find(i);
auto pos = std::lower_bound(keys.begin(),keys.end(),key);
if( *pos != key || 0 == keys.size() ) {
keys.insert(pos,key);
}
blocksmap.emplace(key,i);
}
//std::cout << "(new) Found " << blocks_->size() << " PFBlocks!" << std::endl;
}

// start from first element in elements_
// partition elements until block grows no further
// return the start of the new block
PFBlockAlgo::IE
PFBlockAlgo::associate( PFBlockAlgo::ElementList& elems,
std::unordered_map<std::pair<size_t,size_t>,PFBlockLink>& links,
reco::PFBlock& block) {
if( elems.size() == 0 ) return elems.begin();
ElementList::iterator scan_upper(elems.begin()), search_lower(elems.begin()),
scan_lower(elems.begin());
++scan_upper; ++search_lower;
double dist = -1;
PFBlockLink::Type linktype = PFBlockLink::NONE;
PFBlock::LinkTest linktest = PFBlock::LINKTEST_RECHIT;
block.addElement(scan_lower->get()); // seed the block
// the terminating condition of this loop is when the next range
// to scan has zero length (i.e. you have found no more nearest neighbours)
do {
scan_upper = search_lower;
// for each element added in the previous iteration we check to see what
// elements are linked to it
for( auto comp = scan_lower; comp != scan_upper; ++comp ) {
// group everything that's linked to the current element:
// std::partition moves all elements that return true for the
// function defined below (a.k.a. the linking function) to the
// front of the range provided
search_lower =
std::partition(search_lower,elems.end(),
[&](ElementList::value_type& a){
dist = -1.0;
// compute linking info if it is possible
if( linkPrefilter(comp->get(), a.get()) ) {
link( comp->get(), a.get(),
linktype, linktest, dist );
}
if( dist >= -0.5 ) {
const unsigned lidx = ((*comp)->type() < a->type() ?
(*comp)->index() :
a->index() );
const unsigned uidx = ((*comp)->type() >= a->type() ?
(*comp)->index() :
a->index() );
block.addElement( a.get() );
links.emplace( std::make_pair(lidx,uidx),
PFBlockLink(linktype, linktest, dist,
lidx, uidx ) );
return true;
} else {
return false;
}
});
for( auto key : keys ) {
blocks_->push_back( reco::PFBlock() );
auto range = blocksmap.equal_range(key);
auto& the_block = blocks_->back();
ElementList::value_type::pointer p1(bare_elements_[range.first->second]);
the_block.addElement(p1);
const unsigned block_size = blocksmap.count(key) + 1;
std::unordered_map<std::pair<size_t,size_t>, PFBlockLink > links(block_size*block_size);
auto itr = range.first;
++itr;
for( ; itr != range.second; ++itr ) {
ElementList::value_type::pointer p2(bare_elements_[itr->second]);
const PFBlockElement::Type type1 = p1->type();
const PFBlockElement::Type type2 = p2->type();
the_block.addElement(p2);
linktest = PFBlock::LINKTEST_RECHIT; //rechit by default
linktype = static_cast<PFBlockLink::Type>(1<<(type1-1)|1<<(type2-1));
const unsigned index = _linkTestSquare[type1][type2];
if( nullptr != _linkTests[index] ) {
const double dist = _linkTests[index]->testLink(p1,p2);
links.emplace( std::make_pair(p1->index(), p2->index()) ,
PFBlockLink( linktype, linktest, dist,
p1->index(), p2->index() ) );
}
}
// we then update the scan range lower boundary to point to the
// first element that we added in this round of association
scan_lower = scan_upper;
} while( search_lower != scan_upper );
// return the pointer to the first element not in the PFBlock we just made
return elems.erase(elems.begin(),scan_upper);
packLinks( the_block, links );
}

bare_elements_.clear();
elements_.clear();
}

void
PFBlockAlgo::packLinks( reco::PFBlock& block,
const std::unordered_map<std::pair<size_t,size_t>,PFBlockLink>& links ) const {

constexpr unsigned rowsize = reco::PFBlockElement::kNBETypes;

const edm::OwnVector< reco::PFBlockElement >& els = block.elements();

Expand Down Expand Up @@ -222,8 +261,12 @@ PFBlockAlgo::packLinks( reco::PFBlock& block,
}

if(!linked) {
const PFBlockElement::Type type1 = els[i1].type();
const PFBlockElement::Type type2 = els[i2].type();
const auto minmax = std::minmax(type1,type2);
const unsigned index = rowsize*minmax.second + minmax.first;
PFBlockLink::Type linktype = PFBlockLink::NONE;
bool bTestLink = linkPrefilter(&els[i1], &els[i2]);
bool bTestLink = ( nullptr == _linkTests[index] ? false : _linkTests[index]->linkPrefilter(&(els[i1]),&(els[i2])) );
if (bTestLink) link( & els[i1], & els[i2], linktype, linktest, dist);
}

Expand All @@ -247,13 +290,10 @@ inline bool
PFBlockAlgo::linkPrefilter(const reco::PFBlockElement* last,
const reco::PFBlockElement* next) const {
constexpr unsigned rowsize = reco::PFBlockElement::kNBETypes;
const PFBlockElement::Type& type1 = (last)->type();
const PFBlockElement::Type& type2 = (next)->type();
const PFBlockElement::Type type1 = (last)->type();
const PFBlockElement::Type type2 = (next)->type();
const unsigned index = rowsize*std::max(type1,type2) + std::min(type1,type2);
bool result = false;
if( index < _linkTests.size() && _linkTests[index] ) {
result = _linkTests[index]->linkPrefilter(last,next);
}
bool result = _linkTests[index]->linkPrefilter(last,next);
return result;
}

Expand All @@ -266,8 +306,8 @@ PFBlockAlgo::link( const reco::PFBlockElement* el1,
constexpr unsigned rowsize = reco::PFBlockElement::kNBETypes;
dist=-1.0;
linktest = PFBlock::LINKTEST_RECHIT; //rechit by default
PFBlockElement::Type type1 = el1->type();
PFBlockElement::Type type2 = el2->type();
const PFBlockElement::Type type1 = el1->type();
const PFBlockElement::Type type2 = el2->type();
linktype = static_cast<PFBlockLink::Type>(1<<(type1-1)|1<<(type2-1));
const unsigned index = rowsize*std::max(type1,type2) + std::min(type1,type2);
if(debug_ ) {
Expand All @@ -290,11 +330,33 @@ void PFBlockAlgo::updateEventSetup(const edm::EventSetup& es) {
// and kdtree preprocessors
void PFBlockAlgo::buildElements(const edm::Event& evt) {
// import block elements as defined in python configuration
ranges_.fill(std::make_pair(0,0));
elements_.clear();
for( const auto& importer : _importers ) {

importer->importToBlock(evt,elements_);
}

std::sort(elements_.begin(),elements_.end(),
[](const auto& a, const auto& b) { return a->type() < b->type(); } );

bare_elements_.resize(elements_.size());
for( unsigned i = 0; i < elements_.size(); ++i ) {
bare_elements_[i] = elements_[i].get();
}

// list is now partitioned, so mark the boundaries so we can efficiently skip chunks
unsigned current_type = ( elements_.size() ? elements_[0]->type() : 0 );
unsigned last_type = ( elements_.size() ? elements_.back()->type() : 0 );
ranges_[current_type].first = 0;
ranges_[last_type].second = elements_.size()-1;
for( size_t i = 0; i < elements_.size(); ++i ) {
const auto the_type = elements_[i]->type();
if( the_type != current_type ) {
ranges_[the_type].first = i;
ranges_[current_type].second = i-1;
current_type = the_type;
}
}
// -------------- Loop over block elements ---------------------

// Here we provide to all KDTree linkers the collections to link.
Expand Down