PFProducer/src/PFBlockAlgo.cc

0001 #include "RecoParticleFlow/PFProducer/interface/PFBlockAlgo.h"
0002 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0003 #include "FWCore/PluginManager/interface/PluginFactory.h"
0004 #include "FWCore/Utilities/interface/Exception.h"
0005
0006 #include <algorithm>
0007 #include <iostream>
0008 #include <array>
0009 #include <iterator>
0010 #include <sstream>
0011 #include <type_traits>
0012 #include <utility>
0013
0014 using namespace std;
0015 using namespace reco;
0016
0017 #define INIT_ENTRY(name) {#name, name}
0018
0019 namespace {
0020   class QuickUnion {
0021     std::vector<unsigned> id_;
0022     std::vector<unsigned> size_;
0023     int count_;
0024
0025   public:
0026     QuickUnion(const unsigned NBranches) {
0027       count_ = NBranches;
0028       id_.resize(NBranches);
0029       size_.resize(NBranches);
0030       for (unsigned i = 0; i < NBranches; ++i) {
0031         id_[i] = i;
0032         size_[i] = 1;
0033       }
0034     }
0035
0036     int count() const { return count_; }
0037
0038     unsigned find(unsigned p) {
0039       while (p != id_[p]) {
0040         id_[p] = id_[id_[p]];
0041         p = id_[p];
0042       }
0043       return p;
0044     }
0045
0046     bool connected(unsigned p, unsigned q) { return find(p) == find(q); }
0047
0048     void unite(unsigned p, unsigned q) {
0049       unsigned rootP = find(p);
0050       unsigned rootQ = find(q);
0051       id_[p] = q;
0052
0053       if (size_[rootP] < size_[rootQ]) {
0054         id_[rootP] = rootQ;
0055         size_[rootQ] += size_[rootP];
0056       } else {
0057         id_[rootQ] = rootP;
0058         size_[rootP] += size_[rootQ];
0059       }
0060       --count_;
0061     }
0062   };
0063 }  // namespace
0064
0065 //for debug only
0066 //#define PFLOW_DEBUG
0067
0068 PFBlockAlgo::PFBlockAlgo()
0069     : debug_(false),
0070       elementTypes_({INIT_ENTRY(PFBlockElement::TRACK),
0071                      INIT_ENTRY(PFBlockElement::PS1),
0072                      INIT_ENTRY(PFBlockElement::PS2),
0073                      INIT_ENTRY(PFBlockElement::ECAL),
0074                      INIT_ENTRY(PFBlockElement::HCAL),
0075                      INIT_ENTRY(PFBlockElement::GSF),
0076                      INIT_ENTRY(PFBlockElement::BREM),
0077                      INIT_ENTRY(PFBlockElement::HFEM),
0078                      INIT_ENTRY(PFBlockElement::HFHAD),
0079                      INIT_ENTRY(PFBlockElement::SC),
0080                      INIT_ENTRY(PFBlockElement::HO),
0081                      INIT_ENTRY(PFBlockElement::HGCAL)}) {}
0082
0083 void PFBlockAlgo::setLinkers(const std::vector<edm::ParameterSet>& confs) {
0084   constexpr unsigned rowsize = reco::PFBlockElement::kNBETypes;
0085   for (unsigned i = 0; i < rowsize; ++i) {
0086     for (unsigned j = 0; j < rowsize; ++j) {
0087       linkTestSquare_[i][j] = 0;
0088     }
0089   }
0090   linkTests_.resize(rowsize * rowsize);
0091   const std::string prefix("PFBlockElement::");
0092   const std::string pfx_kdtree("KDTree");
0093   for (const auto& conf : confs) {
0094     const std::string& linkerName = conf.getParameter<std::string>("linkerName");
0095     const std::string& linkTypeStr = conf.getParameter<std::string>("linkType");
0096     size_t split = linkTypeStr.find(':');
0097     if (split == std::string::npos) {
0098       throw cms::Exception("MalformedLinkType") << "\"" << linkTypeStr << "\" is not a valid link type definition."
0099                                                 << " This string should have the form \"linkFrom:linkTo\"";
0100     }
0101     std::string link1(prefix + linkTypeStr.substr(0, split));
0102     std::string link2(prefix + linkTypeStr.substr(split + 1, std::string::npos));
0103     if (!(elementTypes_.count(link1) && elementTypes_.count(link2))) {
0104       throw cms::Exception("InvalidBlockElementType")
0105           << "One of \"" << link1 << "\" or \"" << link2 << "\" are invalid block element types!";
0106     }
0107     const PFBlockElement::Type type1 = elementTypes_.at(link1);
0108     const PFBlockElement::Type type2 = elementTypes_.at(link2);
0109     const unsigned index = rowsize * std::max(type1, type2) + std::min(type1, type2);
0110     linkTests_[index] = BlockElementLinkerFactory::get()->create(linkerName, conf);
0111     linkTestSquare_[type1][type2] = index;
0112     linkTestSquare_[type2][type1] = index;
0113     // setup KDtree if requested
0114     const bool useKDTree = conf.getParameter<bool>("useKDTree");
0115     if (useKDTree) {
0116       kdtrees_.emplace_back(KDTreeLinkerFactory::get()->create(pfx_kdtree + linkerName, conf));
0117       kdtrees_.back()->setTargetType(std::min(type1, type2));
0118       kdtrees_.back()->setFieldType(std::max(type1, type2));
0119     }  // useKDTree
0120   }  // loop over confs
0121 }
0122
0123 void PFBlockAlgo::setImporters(const std::vector<edm::ParameterSet>& confs, edm::ConsumesCollector& cc) {
0124   importers_.reserve(confs.size());
0125   for (const auto& conf : confs) {
0126     const std::string& importerName = conf.getParameter<std::string>("importerName");
0127     importers_.emplace_back(BlockElementImporterFactory::get()->create(importerName, conf, cc));
0128   }
0129 }
0130
0131 PFBlockAlgo::~PFBlockAlgo() {
0132 #ifdef PFLOW_DEBUG
0133   if (debug_)
0134     cout << "~PFBlockAlgo - number of remaining elements: " << elements_.size() << endl;
0135 #endif
0136 }
0137
0138 reco::PFBlockCollection PFBlockAlgo::findBlocks() {
0139   // Glowinski & Gouzevitch
0140   for (const auto& kdtree : kdtrees_) {
0141     kdtree->process();
0142   }
0143   // !Glowinski & Gouzevitch
0144   reco::PFBlockCollection blocks;
0145   // the blocks have not been passed to the event, and need to be cleared
0146   blocks.reserve(elements_.size());
0147
0148   QuickUnion qu(elements_.size());
0149   const auto elem_size = elements_.size();
0150   for (unsigned i = 0; i < elem_size; ++i) {
0151     for (unsigned j = i + 1; j < elem_size; ++j) {
0152       if (qu.connected(i, j))
0153         continue;
0154       if (!linkTests_[linkTestSquare_[elements_[i]->type()][elements_[j]->type()]]) {
0155         j = ranges_[elements_[j]->type()].second;
0156         continue;
0157       }
0158       auto p1(elements_[i].get()), p2(elements_[j].get());
0159       const PFBlockElement::Type type1 = p1->type();
0160       const PFBlockElement::Type type2 = p2->type();
0161       const unsigned index = linkTestSquare_[type1][type2];
0162       if (linkTests_[index]->linkPrefilter(p1, p2)) {
0163         const double dist = linkTests_[index]->testLink(p1, p2);
0164         // compute linking info if it is possible
0165         if (dist > -0.5) {
0166           qu.unite(i, j);
0167         }
0168       }
0169     }
0170   }
0171
0172   std::unordered_multimap<unsigned, unsigned> blocksmap(elements_.size());
0173   std::vector<unsigned> keys;
0174   keys.reserve(elements_.size());
0175   for (unsigned i = 0; i < elements_.size(); ++i) {
0176     unsigned key = i;
0177     while (key != qu.find(key))
0178       key = qu.find(key);  // make sure we always find the root node...
0179     auto pos = std::lower_bound(keys.begin(), keys.end(), key);
0180     if (pos == keys.end() || *pos != key) {
0181       keys.insert(pos, key);
0182     }
0183     blocksmap.emplace(key, i);
0184   }
0185
0186   for (auto key : keys) {
0187     blocks.push_back(reco::PFBlock());
0188     auto range = blocksmap.equal_range(key);
0189     auto& the_block = blocks.back();
0190     ElementList::value_type::pointer p1(elements_[range.first->second].get());
0191     the_block.addElement(p1);
0192     const unsigned block_size = blocksmap.count(key) + 1;
0193     //reserve up to 1M or 8MB; pay rehash cost for more
0194     std::unordered_map<std::pair<unsigned int, unsigned int>, double> links(min(1000000u, block_size * block_size));
0195     auto itr = range.first;
0196     ++itr;
0197     for (; itr != range.second; ++itr) {
0198       ElementList::value_type::pointer p2(elements_[itr->second].get());
0199       const PFBlockElement::Type type1 = p1->type();
0200       const PFBlockElement::Type type2 = p2->type();
0201       the_block.addElement(p2);
0202       const unsigned index = linkTestSquare_[type1][type2];
0203       if (nullptr != linkTests_[index]) {
0204         const double dist = linkTests_[index]->testLink(p1, p2);
0205         links.emplace(std::make_pair(p1->index(), p2->index()), dist);
0206       }
0207     }
0208     packLinks(the_block, links);
0209   }
0210
0211   elements_.clear();
0212   elements_.shrink_to_fit();
0213
0214   return blocks;
0215 }
0216
0217 void PFBlockAlgo::packLinks(reco::PFBlock& block,
0218                             const std::unordered_map<std::pair<unsigned int, unsigned int>, double>& links) const {
0219   constexpr unsigned rowsize = reco::PFBlockElement::kNBETypes;
0220
0221   const edm::OwnVector<reco::PFBlockElement>& els = block.elements();
0222
0223   block.bookLinkData();
0224   unsigned elsize = els.size();
0225   //First Loop: update all link data
0226   for (unsigned i1 = 0; i1 < elsize; ++i1) {
0227     for (unsigned i2 = 0; i2 < i1; ++i2) {
0228       // no reflexive link
0229       //if( i1==i2 ) continue;
0230
0231       double dist = -1;
0232
0233       bool linked = false;
0234
0235       // are these elements already linked ?
0236       // this can be optimized
0237       const auto link_itr = links.find(std::make_pair(i2, i1));
0238       if (link_itr != links.end()) {
0239         dist = link_itr->second;
0240         linked = true;
0241       }
0242
0243       if (!linked) {
0244         const PFBlockElement::Type type1 = els[i1].type();
0245         const PFBlockElement::Type type2 = els[i2].type();
0246         const auto minmax = std::minmax(type1, type2);
0247         const unsigned index = rowsize * minmax.second + minmax.first;
0248         bool bTestLink =
0249             (nullptr == linkTests_[index] ? false : linkTests_[index]->linkPrefilter(&(els[i1]), &(els[i2])));
0250         if (bTestLink)
0251           link(&els[i1], &els[i2], dist);
0252       }
0253
0254       //loading link data according to link test used: RECHIT
0255       //block.setLink( i1, i2, chi2, block.linkData() );
0256       block.setLink(i1, i2, dist, block.linkData());
0257     }
0258   }
0259 }
0260
0261 inline void PFBlockAlgo::link(const reco::PFBlockElement* el1, const reco::PFBlockElement* el2, double& dist) const {
0262   constexpr unsigned rowsize = reco::PFBlockElement::kNBETypes;
0263   dist = -1.0;
0264   const PFBlockElement::Type type1 = el1->type();
0265   const PFBlockElement::Type type2 = el2->type();
0266   const unsigned index = rowsize * std::max(type1, type2) + std::min(type1, type2);
0267   if (debug_) {
0268     std::cout << " PFBlockAlgo links type1 " << type1 << " type2 " << type2 << std::endl;
0269   }
0270
0271   // index is always checked in the preFilter above, no need to check here
0272   dist = linkTests_[index]->testLink(el1, el2);
0273 }
0274
0275 void PFBlockAlgo::updateEventSetup(const edm::EventSetup& es) {
0276   for (auto& importer : importers_) {
0277     importer->updateEventSetup(es);
0278   }
0279 }
0280
0281 // see plugins/importers and plugins/kdtrees
0282 // for the definitions of available block element importers
0283 // and kdtree preprocessors
0284 void PFBlockAlgo::buildElements(const edm::Event& evt) {
0285   // import block elements as defined in python configuration
0286   ranges_.fill(std::make_pair(0, 0));
0287   elements_.clear();
0288   for (const auto& importer : importers_) {
0289     importer->importToBlock(evt, elements_);
0290   }
0291
0292   std::sort(elements_.begin(), elements_.end(), [](const auto& a, const auto& b) { return a->type() < b->type(); });
0293
0294   // list is now partitioned, so mark the boundaries so we can efficiently skip chunks
0295   unsigned current_type = (!elements_.empty() ? elements_[0]->type() : 0);
0296   unsigned last_type = (!elements_.empty() ? elements_.back()->type() : 0);
0297   ranges_[current_type].first = 0;
0298   ranges_[last_type].second = elements_.size() - 1;
0299   for (size_t i = 0; i < elements_.size(); ++i) {
0300     const auto the_type = elements_[i]->type();
0301     if (the_type != current_type) {
0302       ranges_[the_type].first = i;
0303       ranges_[current_type].second = i - 1;
0304       current_type = the_type;
0305     }
0306   }
0307   // -------------- Loop over block elements ---------------------
0308
0309   // Here we provide to all KDTree linkers the collections to link.
0310   // Glowinski & Gouzevitch
0311
0312   for (ElementList::iterator it = elements_.begin(); it != elements_.end(); ++it) {
0313     for (const auto& kdtree : kdtrees_) {
0314       if ((*it)->type() == kdtree->targetType()) {
0315         kdtree->insertTargetElt(it->get());
0316       }
0317       if ((*it)->type() == kdtree->fieldType()) {
0318         kdtree->insertFieldClusterElt(it->get());
0319       }
0320     }
0321   }
0322   //std::cout << "(new) imported: " << elements_.size() << " elements!" << std::endl;
0323 }
0324
0325 std::ostream& operator<<(std::ostream& out, const PFBlockAlgo& a) {
0326   if (!out)
0327     return out;
0328
0329   out << "====== Particle Flow Block Algorithm ======= ";
0330   out << endl;
0331   out << "number of unassociated elements : " << a.elements_.size() << endl;
0332   out << endl;
0333
0334   for (auto const& element : a.elements_) {
0335     out << "\t" << *element << endl;
0336   }
0337
0338   return out;
0339 }
0340
0341 // a little history, ideas we may want to keep around for later
0342 /*
0343   // Links between the two preshower layers are not used for now - disable
0344   case PFBlockLink::PS1andPS2:
0345     {
0346       PFClusterRef  ps1ref = lowEl->clusterRef();
0347       PFClusterRef  ps2ref = highEl->clusterRef();
0348       assert( !ps1ref.isNull() );
0349       assert( !ps2ref.isNull() );
0350       // PJ - 14-May-09 : A link by rechit is needed here !
0351       // dist = testPS1AndPS2( *ps1ref, *ps2ref );
0352       dist = -1.;
0353       linktest = PFBlock::LINKTEST_RECHIT;
0354       break;
0355     }
0356   case PFBlockLink::TRACKandPS1:
0357   case PFBlockLink::TRACKandPS2:
0358     {
0359       //cout<<"TRACKandPS"<<endl;
0360       PFRecTrackRef trackref = lowEl->trackRefPF();
0361       PFClusterRef  clusterref = highEl->clusterRef();
0362       assert( !trackref.isNull() );
0363       assert( !clusterref.isNull() );
0364       // PJ - 14-May-09 : A link by rechit is needed here !
0365       dist = testTrackAndPS( *trackref, *clusterref );
0366       linktest = PFBlock::LINKTEST_RECHIT;
0367       break;
0368     }
0369     // GSF Track/Brem Track and preshower cluster links are not used for now - disable
0370   case PFBlockLink::PS1andGSF:
0371   case PFBlockLink::PS2andGSF:
0372     {
0373       PFClusterRef  psref = lowEl->clusterRef();
0374       assert( !psref.isNull() );
0375       const reco::PFBlockElementGsfTrack *  GsfEl =  dynamic_cast<const reco::PFBlockElementGsfTrack*>(highEl);
0376       const PFRecTrack * myTrack =  &(GsfEl->GsftrackPF());
0377       // PJ - 14-May-09 : A link by rechit is needed here !
0378       dist = testTrackAndPS( *myTrack, *psref );
0379       linktest = PFBlock::LINKTEST_RECHIT;
0380       break;
0381     }
0382   case PFBlockLink::PS1andBREM:
0383   case PFBlockLink::PS2andBREM:
0384     {
0385       PFClusterRef  psref = lowEl->clusterRef();
0386       assert( !psref.isNull() );
0387       const reco::PFBlockElementBrem * BremEl =  dynamic_cast<const reco::PFBlockElementBrem*>(highEl);
0388       const PFRecTrack * myTrack = &(BremEl->trackPF());
0389       // PJ - 14-May-09 : A link by rechit is needed here !
0390       dist = testTrackAndPS( *myTrack, *psref );
0391       linktest = PFBlock::LINKTEST_RECHIT;
0392       break;
0393     }
0394     */