GenFilters/plugins/MCMultiParticleMassFilter.cc

0001 #include "GeneratorInterface/GenFilters/plugins/MCFilterZboostHelper.h"
0002 #include "DataFormats/Common/interface/Handle.h"
0003 #include "FWCore/Framework/interface/global/EDFilter.h"
0004 #include "FWCore/Framework/interface/Event.h"
0005 #include "FWCore/Framework/interface/Frameworkfwd.h"
0006 #include "FWCore/Framework/interface/MakerMacros.h"
0007 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0008 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0009 #include "FWCore/Utilities/interface/EDGetToken.h"
0010 #include "FWCore/Utilities/interface/InputTag.h"
0011 #include "SimDataFormats/GeneratorProducts/interface/HepMCProduct.h"
0012
0013 #include <vector>
0014 #include <iostream>
0015 #include <unordered_map>
0016 #include <tuple>
0017
0018 class MCMultiParticleMassFilter : public edm::global::EDFilter<> {
0019 public:
0020   explicit MCMultiParticleMassFilter(const edm::ParameterSet&);
0021   ~MCMultiParticleMassFilter() override;
0022
0023 private:
0024   bool filter(edm::StreamID, edm::Event& iEvent, const edm::EventSetup&) const override;
0025   bool recurseLoop(std::vector<HepMC::GenParticle*>& particlesThatPassCuts, std::vector<int> indices, int depth) const;
0026
0027   /* Member data */
0028   const edm::EDGetTokenT<edm::HepMCProduct> token_;
0029   const std::vector<int> particleID;
0030   std::vector<double> ptMin;
0031   std::vector<double> etaMax;
0032   std::vector<int> status;
0033
0034   //Maps each particle ID provided to its required pt, max eta, and status
0035   std::unordered_map<int, std::tuple<double, double, int>> cutPerParticle;
0036
0037   const double minTotalMass;
0038   const double maxTotalMass;
0039
0040   double minTotalMassSq;
0041   double maxTotalMassSq;
0042   int nParticles;
0043
0044   int particleSumTo;
0045   int particleProdTo;
0046 };
0047
0048 using namespace edm;
0049 using namespace std;
0050
0051 MCMultiParticleMassFilter::MCMultiParticleMassFilter(const edm::ParameterSet& iConfig)
0052     : token_(consumes<edm::HepMCProduct>(
0053           iConfig.getUntrackedParameter<edm::InputTag>("src", edm::InputTag("generator", "unsmeared")))),
0054       particleID(iConfig.getParameter<std::vector<int>>("ParticleID")),
0055       ptMin(iConfig.getParameter<std::vector<double>>("PtMin")),
0056       etaMax(iConfig.getParameter<std::vector<double>>("EtaMax")),
0057       status(iConfig.getParameter<std::vector<int>>("Status")),
0058       minTotalMass(iConfig.getParameter<double>("minTotalMass")),
0059       maxTotalMass(iConfig.getParameter<double>("maxTotalMass")) {
0060   nParticles = particleID.size();
0061   minTotalMassSq = minTotalMass * minTotalMass;
0062   maxTotalMassSq = maxTotalMass * maxTotalMass;
0063
0064   //These two values dictate what particles it accepts as combinations
0065   particleSumTo = 0;
0066   particleProdTo = 1;
0067   for (const int i : particleID) {
0068     particleSumTo += i;
0069     particleProdTo *= i;
0070   }
0071
0072   // if any of the vectors are of size 1, take that to mean it is a new default
0073   double defaultPtMin = 1.8;
0074   if ((int)ptMin.size() == 1) {
0075     defaultPtMin = ptMin[0];
0076   }
0077
0078   if ((int)ptMin.size() < nParticles) {
0079     edm::LogWarning("MCMultiParticleMassFilter") << "Warning: Given pT value size"
0080                                                     "< size of the number of particle IDs."
0081                                                     " Filling remaining values with "
0082                                                  << defaultPtMin << endl;
0083     while ((int)ptMin.size() < nParticles) {
0084       ptMin.push_back(defaultPtMin);
0085     }
0086   } else if ((int)ptMin.size() > nParticles) {
0087     edm::LogWarning("MCMultiParticleMassFilter") << "Warning: Given pT value size"
0088                                                     "> size of the number of particle IDs."
0089                                                     " Ignoring extraneous values "
0090                                                  << endl;
0091     ptMin.resize(nParticles);
0092   }
0093
0094   double defaultEtaMax = 2.7;
0095   if ((int)etaMax.size() == 1) {
0096     defaultEtaMax = etaMax[0];
0097   }
0098   if ((int)etaMax.size() < nParticles) {
0099     edm::LogWarning("MCMultiParticleMassFilter") << "Warning: Given eta value size"
0100                                                     "< size of the number of particle IDs."
0101                                                     " Filling remaining values with "
0102                                                  << defaultEtaMax << endl;
0103     while ((int)etaMax.size() < nParticles) {
0104       etaMax.push_back(defaultEtaMax);
0105     }
0106   } else if ((int)etaMax.size() > nParticles) {
0107     edm::LogWarning("MCMultiParticleMassFilter") << "Warning: Given eta value size"
0108                                                     "> size of the number of particle IDs."
0109                                                     " Ignoring extraneous values "
0110                                                  << endl;
0111     etaMax.resize(nParticles);
0112   }
0113
0114   int defaultStatus = 1;
0115   if ((int)status.size() == 1) {
0116     defaultStatus = status[0];
0117   }
0118   if ((int)status.size() < nParticles) {
0119     edm::LogWarning("MCMultiParticleMassFilter") << "Warning: Given status value size"
0120                                                     "< size of the number of particle IDs."
0121                                                     " Filling remaining values with "
0122                                                  << defaultStatus << endl;
0123     while ((int)status.size() < nParticles) {
0124       status.push_back(defaultStatus);
0125     }
0126   } else if ((int)status.size() > nParticles) {
0127     edm::LogWarning("MCMultiParticleMassFilter") << "Warning: Given status value size"
0128                                                     "> size of the number of particle IDs."
0129                                                     " Ignoring extraneous values "
0130                                                  << endl;
0131     status.resize(nParticles);
0132   }
0133
0134   for (int i = 0; i < nParticles; i++) {
0135     std::tuple<double, double, int> cutForParticle = std::make_tuple(ptMin[i], etaMax[i], status[i]);
0136     cutPerParticle[particleID[i]] = cutForParticle;
0137   }  //assign the set of cuts you decided upon matched to each ID value in order
0138 }
0139
0140 MCMultiParticleMassFilter::~MCMultiParticleMassFilter() {}
0141
0142 bool MCMultiParticleMassFilter::filter(edm::StreamID, edm::Event& iEvent, const edm::EventSetup&) const {
0143   edm::Handle<edm::HepMCProduct> evt;
0144   iEvent.getByToken(token_, evt);
0145   const HepMC::GenEvent* myGenEvent = evt->GetEvent();
0146
0147   std::vector<HepMC::GenParticle*> particlesThatPassCuts = std::vector<HepMC::GenParticle*>();
0148   for (HepMC::GenEvent::particle_const_iterator p = myGenEvent->particles_begin(); p != myGenEvent->particles_end();
0149        ++p) {
0150     for (const int i : particleID) {
0151       if (i == (*p)->pdg_id()) {
0152         //if the particle ID is one of the ones you specified, check for cuts per ID
0153         const auto cuts = cutPerParticle.at(i);
0154         if (((*p)->status() == get<2>(cuts)) && ((*p)->momentum().perp() >= get<0>(cuts)) &&
0155             (std::fabs((*p)->momentum().eta()) <= get<1>(cuts))) {
0156           particlesThatPassCuts.push_back(*p);
0157           break;
0158         }
0159       }
0160     }
0161   }
0162   int nIterables = particlesThatPassCuts.size();
0163   if (nIterables < nParticles) {
0164     return false;
0165   } else {
0166     int i = 0;
0167     //only iterate while there are enough particles that pass cuts
0168     while ((nIterables - i) >= nParticles) {
0169       vector<int> indices;
0170       //start recursing from index 0, 1, 2, ...
0171       indices.push_back(i);
0172       bool success = recurseLoop(particlesThatPassCuts, indices, 1);
0173       if (success) {
0174         return true;
0175       }
0176       i++;
0177     }
0178   }
0179   return false;
0180 }
0181
0182 bool MCMultiParticleMassFilter::recurseLoop(std::vector<HepMC::GenParticle*>& particlesThatPassCuts,
0183                                             std::vector<int> indices,
0184                                             int depth) const {
0185   int lastIndex = indices.back();
0186   int nIterables = (int)(particlesThatPassCuts.size());
0187   if (lastIndex >= nIterables) {
0188     return false;
0189   } else if (depth == nParticles) {
0190     //you have the right number of particles!
0191     int tempSum = 0;
0192     int tempProd = 1;
0193
0194     double px, py, pz, e;
0195     px = py = pz = e = 0;
0196     for (const int i : indices) {
0197       int id = particlesThatPassCuts[i]->pdg_id();
0198       tempSum += id;
0199       tempProd *= id;
0200       HepMC::FourVector tempVec = particlesThatPassCuts[i]->momentum();
0201       px += tempVec.px();
0202       py += tempVec.py();
0203       pz += tempVec.pz();
0204       e += tempVec.e();
0205     }
0206     if ((tempSum != particleSumTo) || (tempProd != particleProdTo)) {
0207       return false;  //check if the ids are what you want
0208     }
0209     double invMassSq = e * e - px * px - py * py - pz * pz;
0210     //taking the root is computationally expensive, so use the squared term
0211     if ((invMassSq >= minTotalMassSq) && (invMassSq <= maxTotalMassSq)) {
0212       return true;
0213     }
0214     return false;
0215   } else {
0216     vector<bool> recursionResult;
0217     //propagate recursion across all combinations of remaining indices
0218     for (int i = 1; i < nIterables - lastIndex; i++) {
0219       vector<int> newIndices = indices;
0220       newIndices.push_back(lastIndex + i);
0221       //always up the depth by 1 to make sure there is no infinite recursion
0222       recursionResult.push_back(recurseLoop(particlesThatPassCuts, newIndices, depth + 1));
0223     }
0224     //search the results to look for one successful combination
0225     for (bool r : recursionResult) {
0226       if (r) {
0227         return true;
0228       }
0229     }
0230     return false;
0231   }
0232 }
0233 DEFINE_FWK_MODULE(MCMultiParticleMassFilter);