File indexing completed on 2023-03-17 11:21:14
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) \
0018 { #name, name }
0019
0020 namespace {
0021 class QuickUnion {
0022 std::vector<unsigned> id_;
0023 std::vector<unsigned> size_;
0024 int count_;
0025
0026 public:
0027 QuickUnion(const unsigned NBranches) {
0028 count_ = NBranches;
0029 id_.resize(NBranches);
0030 size_.resize(NBranches);
0031 for (unsigned i = 0; i < NBranches; ++i) {
0032 id_[i] = i;
0033 size_[i] = 1;
0034 }
0035 }
0036
0037 int count() const { return count_; }
0038
0039 unsigned find(unsigned p) {
0040 while (p != id_[p]) {
0041 id_[p] = id_[id_[p]];
0042 p = id_[p];
0043 }
0044 return p;
0045 }
0046
0047 bool connected(unsigned p, unsigned q) { return find(p) == find(q); }
0048
0049 void unite(unsigned p, unsigned q) {
0050 unsigned rootP = find(p);
0051 unsigned rootQ = find(q);
0052 id_[p] = q;
0053
0054 if (size_[rootP] < size_[rootQ]) {
0055 id_[rootP] = rootQ;
0056 size_[rootQ] += size_[rootP];
0057 } else {
0058 id_[rootQ] = rootP;
0059 size_[rootP] += size_[rootQ];
0060 }
0061 --count_;
0062 }
0063 };
0064 }
0065
0066
0067
0068
0069 PFBlockAlgo::PFBlockAlgo()
0070 : debug_(false),
0071 elementTypes_({INIT_ENTRY(PFBlockElement::TRACK),
0072 INIT_ENTRY(PFBlockElement::PS1),
0073 INIT_ENTRY(PFBlockElement::PS2),
0074 INIT_ENTRY(PFBlockElement::ECAL),
0075 INIT_ENTRY(PFBlockElement::HCAL),
0076 INIT_ENTRY(PFBlockElement::GSF),
0077 INIT_ENTRY(PFBlockElement::BREM),
0078 INIT_ENTRY(PFBlockElement::HFEM),
0079 INIT_ENTRY(PFBlockElement::HFHAD),
0080 INIT_ENTRY(PFBlockElement::SC),
0081 INIT_ENTRY(PFBlockElement::HO),
0082 INIT_ENTRY(PFBlockElement::HGCAL)}) {}
0083
0084 void PFBlockAlgo::setLinkers(const std::vector<edm::ParameterSet>& confs) {
0085 constexpr unsigned rowsize = reco::PFBlockElement::kNBETypes;
0086 for (unsigned i = 0; i < rowsize; ++i) {
0087 for (unsigned j = 0; j < rowsize; ++j) {
0088 linkTestSquare_[i][j] = 0;
0089 }
0090 }
0091 linkTests_.resize(rowsize * rowsize);
0092 const std::string prefix("PFBlockElement::");
0093 const std::string pfx_kdtree("KDTree");
0094 for (const auto& conf : confs) {
0095 const std::string& linkerName = conf.getParameter<std::string>("linkerName");
0096 const std::string& linkTypeStr = conf.getParameter<std::string>("linkType");
0097 size_t split = linkTypeStr.find(':');
0098 if (split == std::string::npos) {
0099 throw cms::Exception("MalformedLinkType") << "\"" << linkTypeStr << "\" is not a valid link type definition."
0100 << " This string should have the form \"linkFrom:linkTo\"";
0101 }
0102 std::string link1(prefix + linkTypeStr.substr(0, split));
0103 std::string link2(prefix + linkTypeStr.substr(split + 1, std::string::npos));
0104 if (!(elementTypes_.count(link1) && elementTypes_.count(link2))) {
0105 throw cms::Exception("InvalidBlockElementType")
0106 << "One of \"" << link1 << "\" or \"" << link2 << "\" are invalid block element types!";
0107 }
0108 const PFBlockElement::Type type1 = elementTypes_.at(link1);
0109 const PFBlockElement::Type type2 = elementTypes_.at(link2);
0110 const unsigned index = rowsize * std::max(type1, type2) + std::min(type1, type2);
0111 linkTests_[index] = BlockElementLinkerFactory::get()->create(linkerName, conf);
0112 linkTestSquare_[type1][type2] = index;
0113 linkTestSquare_[type2][type1] = index;
0114
0115 const bool useKDTree = conf.getParameter<bool>("useKDTree");
0116 if (useKDTree) {
0117 kdtrees_.emplace_back(KDTreeLinkerFactory::get()->create(pfx_kdtree + linkerName, conf));
0118 kdtrees_.back()->setTargetType(std::min(type1, type2));
0119 kdtrees_.back()->setFieldType(std::max(type1, type2));
0120 }
0121 }
0122 }
0123
0124 void PFBlockAlgo::setImporters(const std::vector<edm::ParameterSet>& confs, edm::ConsumesCollector& cc) {
0125 importers_.reserve(confs.size());
0126 for (const auto& conf : confs) {
0127 const std::string& importerName = conf.getParameter<std::string>("importerName");
0128 importers_.emplace_back(BlockElementImporterFactory::get()->create(importerName, conf, cc));
0129 }
0130 }
0131
0132 PFBlockAlgo::~PFBlockAlgo() {
0133 #ifdef PFLOW_DEBUG
0134 if (debug_)
0135 cout << "~PFBlockAlgo - number of remaining elements: " << elements_.size() << endl;
0136 #endif
0137 }
0138
0139 reco::PFBlockCollection PFBlockAlgo::findBlocks() {
0140
0141 for (const auto& kdtree : kdtrees_) {
0142 kdtree->process();
0143 }
0144
0145 reco::PFBlockCollection blocks;
0146
0147 blocks.reserve(elements_.size());
0148
0149 QuickUnion qu(elements_.size());
0150 const auto elem_size = elements_.size();
0151 for (unsigned i = 0; i < elem_size; ++i) {
0152 for (unsigned j = i + 1; j < elem_size; ++j) {
0153 if (qu.connected(i, j))
0154 continue;
0155 if (!linkTests_[linkTestSquare_[elements_[i]->type()][elements_[j]->type()]]) {
0156 j = ranges_[elements_[j]->type()].second;
0157 continue;
0158 }
0159 auto p1(elements_[i].get()), p2(elements_[j].get());
0160 const PFBlockElement::Type type1 = p1->type();
0161 const PFBlockElement::Type type2 = p2->type();
0162 const unsigned index = linkTestSquare_[type1][type2];
0163 if (linkTests_[index]->linkPrefilter(p1, p2)) {
0164 const double dist = linkTests_[index]->testLink(p1, p2);
0165
0166 if (dist > -0.5) {
0167 qu.unite(i, j);
0168 }
0169 }
0170 }
0171 }
0172
0173 std::unordered_multimap<unsigned, unsigned> blocksmap(elements_.size());
0174 std::vector<unsigned> keys;
0175 keys.reserve(elements_.size());
0176 for (unsigned i = 0; i < elements_.size(); ++i) {
0177 unsigned key = i;
0178 while (key != qu.find(key))
0179 key = qu.find(key);
0180 auto pos = std::lower_bound(keys.begin(), keys.end(), key);
0181 if (pos == keys.end() || *pos != key) {
0182 keys.insert(pos, key);
0183 }
0184 blocksmap.emplace(key, i);
0185 }
0186
0187 for (auto key : keys) {
0188 blocks.push_back(reco::PFBlock());
0189 auto range = blocksmap.equal_range(key);
0190 auto& the_block = blocks.back();
0191 ElementList::value_type::pointer p1(elements_[range.first->second].get());
0192 the_block.addElement(p1);
0193 const unsigned block_size = blocksmap.count(key) + 1;
0194
0195 std::unordered_map<std::pair<unsigned int, unsigned int>, double> links(min(1000000u, block_size * block_size));
0196 auto itr = range.first;
0197 ++itr;
0198 for (; itr != range.second; ++itr) {
0199 ElementList::value_type::pointer p2(elements_[itr->second].get());
0200 const PFBlockElement::Type type1 = p1->type();
0201 const PFBlockElement::Type type2 = p2->type();
0202 the_block.addElement(p2);
0203 const unsigned index = linkTestSquare_[type1][type2];
0204 if (nullptr != linkTests_[index]) {
0205 const double dist = linkTests_[index]->testLink(p1, p2);
0206 links.emplace(std::make_pair(p1->index(), p2->index()), dist);
0207 }
0208 }
0209 packLinks(the_block, links);
0210 }
0211
0212 elements_.clear();
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