File indexing completed on 2024-11-19 23:20:19
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 }
0064
0065
0066
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
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 }
0120 }
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
0140 for (const auto& kdtree : kdtrees_) {
0141 kdtree->process();
0142 }
0143
0144 reco::PFBlockCollection blocks;
0145
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
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);
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
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
0226 for (unsigned i1 = 0; i1 < elsize; ++i1) {
0227 for (unsigned i2 = 0; i2 < i1; ++i2) {
0228
0229
0230
0231 double dist = -1;
0232
0233 bool linked = false;
0234
0235
0236
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
0255
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
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
0282
0283
0284 void PFBlockAlgo::buildElements(const edm::Event& evt) {
0285
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
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
0308
0309
0310
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
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
0342
0343
0344
0345
0346
0347
0348
0349
0350
0351
0352
0353
0354
0355
0356
0357
0358
0359
0360
0361
0362
0363
0364
0365
0366
0367
0368
0369
0370
0371
0372
0373
0374
0375
0376
0377
0378
0379
0380
0381
0382
0383
0384
0385
0386
0387
0388
0389
0390
0391
0392
0393
0394