Merge pull request #15287 from makortel/backport80_prs_14138_14441_14…

…754_14844

Backport of #14138: Speed up PFBlockAlgo (including the fixes from #14441 #14754 #14844)

RecoParticleFlow/PFProducer/interface/PFBlockAlgo.h

@@ -62,6 +62,7 @@
 
 #include <map>
 #include <unordered_map>
+#include <unordered_set>
 
 namespace std {
   template<>
@@ -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) );
     }
   };
 }
@@ -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();
 
@@ -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, 
@@ -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_;
 
   // 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_;
@@ -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;
 };

RecoParticleFlow/PFProducer/src/PFBlockAlgo.cc

@@ -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
 
@@ -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");
@@ -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 ) {
@@ -96,8 +150,7 @@ PFBlockAlgo::~PFBlockAlgo() {
 #endif  
 }
 
-void 
-PFBlockAlgo::findBlocks() {
+void PFBlockAlgo::findBlocks() {
   // Glowinski & Gouzevitch
   for( const auto& kdtree : _kdtrees ) {
     kdtree->process();
@@ -107,93 +160,81 @@ 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 = 0; j < elem_size; ++j ) {
+      if( qu.connected(i,j) || j == i ) continue;
+      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 = i; 
+    while( key != qu.find(key) ) key = qu.find(key); // make sure we always find the root node...
+    auto pos  = std::lower_bound(keys.begin(),keys.end(),key);
+    if( pos == keys.end() || *pos != key ) {
+      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();
 
@@ -222,8 +263,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);
       }
 
@@ -247,13 +292,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;  
 }
 
@@ -266,8 +308,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_ ) { 
@@ -290,11 +332,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.