Merge pull request #46856 from MohitS704/backport_filter

backporting MCMultiParticleMassFilter

GeneratorInterface/GenFilters/plugins/MCMultiParticleMassFilter.cc

@@ -0,0 +1,234 @@
+#include "GeneratorInterface/GenFilters/plugins/MCFilterZboostHelper.h"
+#include "DataFormats/Common/interface/Handle.h"
+#include "FWCore/Framework/interface/global/EDFilter.h"
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/Framework/interface/Frameworkfwd.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+#include "FWCore/MessageLogger/interface/MessageLogger.h"
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+#include "FWCore/Utilities/interface/EDGetToken.h"
+#include "FWCore/Utilities/interface/InputTag.h"
+#include "SimDataFormats/GeneratorProducts/interface/HepMCProduct.h"
+
+#include <vector>
+#include <iostream>
+#include <unordered_map>
+#include <tuple>
+
+class MCMultiParticleMassFilter : public edm::global::EDFilter<> {
+public:
+  explicit MCMultiParticleMassFilter(const edm::ParameterSet&);
+  ~MCMultiParticleMassFilter() override;
+
+private:
+  bool filter(edm::StreamID, edm::Event& iEvent, const edm::EventSetup&) const override;
+  bool recurseLoop(std::vector<HepMC::GenParticle*>& particlesThatPassCuts, std::vector<int> indices, int depth) const;
+
+  /* Member data */
+  const edm::EDGetTokenT<edm::HepMCProduct> token_;
+  const std::vector<int> particleID;
+  std::vector<double> ptMin;
+  std::vector<double> etaMax;
+  std::vector<int> status;
+
+  //Maps each particle ID provided to its required pt, max eta, and status
+  std::unordered_map<int, std::tuple<double, double, int>> cutPerParticle;
+
+  const double minTotalMass;
+  const double maxTotalMass;
+
+  double minTotalMassSq;
+  double maxTotalMassSq;
+  int nParticles;
+
+  int particleSumTo;
+  int particleProdTo;
+};
+
+using namespace edm;
+using namespace std;
+
+MCMultiParticleMassFilter::MCMultiParticleMassFilter(const edm::ParameterSet& iConfig)
+    : token_(consumes<edm::HepMCProduct>(
+          iConfig.getUntrackedParameter<edm::InputTag>("src", edm::InputTag("generator", "unsmeared")))),
+      particleID(iConfig.getParameter<std::vector<int>>("ParticleID")),
+      ptMin(iConfig.getParameter<std::vector<double>>("PtMin")),
+      etaMax(iConfig.getParameter<std::vector<double>>("EtaMax")),
+      status(iConfig.getParameter<std::vector<int>>("Status")),
+      minTotalMass(iConfig.getParameter<double>("minTotalMass")),
+      maxTotalMass(iConfig.getParameter<double>("maxTotalMass")) {
+  nParticles = particleID.size();
+  minTotalMassSq = minTotalMass * minTotalMass;
+  maxTotalMassSq = maxTotalMass * maxTotalMass;
+
+  //These two values dictate what particles it accepts as combinations
+  particleSumTo = 0;
+  particleProdTo = 1;
+  for (const int i : particleID) {
+    particleSumTo += i;
+    particleProdTo *= i;
+  }
+
+  // if any of the vectors are of size 1, take that to mean it is a new default
+  double defaultPtMin = 1.8;
+  if ((int)ptMin.size() == 1) {
+    defaultPtMin = ptMin[0];
+  }
+
+  if ((int)ptMin.size() < nParticles) {
+    edm::LogWarning("MCMultiParticleMassFilter") << "Warning: Given pT value size"
+                                                    "< size of the number of particle IDs."
+                                                    " Filling remaining values with "
+                                                 << defaultPtMin << endl;
+    while ((int)ptMin.size() < nParticles) {
+      ptMin.push_back(defaultPtMin);
+    }
+  } else if ((int)ptMin.size() > nParticles) {
+    edm::LogWarning("MCMultiParticleMassFilter") << "Warning: Given pT value size"
+                                                    "> size of the number of particle IDs."
+                                                    " Ignoring extraneous values "
+                                                 << endl;
+    ptMin.resize(nParticles);
+  }
+
+  double defaultEtaMax = 2.7;
+  if ((int)etaMax.size() == 1) {
+    defaultEtaMax = etaMax[0];
+  }
+  if ((int)etaMax.size() < nParticles) {
+    edm::LogWarning("MCMultiParticleMassFilter") << "Warning: Given eta value size"
+                                                    "< size of the number of particle IDs."
+                                                    " Filling remaining values with "
+                                                 << defaultEtaMax << endl;
+    while ((int)etaMax.size() < nParticles) {
+      etaMax.push_back(defaultEtaMax);
+    }
+  } else if ((int)etaMax.size() > nParticles) {
+    edm::LogWarning("MCMultiParticleMassFilter") << "Warning: Given eta value size"
+                                                    "> size of the number of particle IDs."
+                                                    " Ignoring extraneous values "
+                                                 << endl;
+    etaMax.resize(nParticles);
+  }
+
+  int defaultStatus = 1;
+  if ((int)status.size() == 1) {
+    defaultStatus = status[0];
+  }
+  if ((int)status.size() < nParticles) {
+    edm::LogWarning("MCMultiParticleMassFilter") << "Warning: Given status value size"
+                                                    "< size of the number of particle IDs."
+                                                    " Filling remaining values with "
+                                                 << defaultStatus << endl;
+    while ((int)status.size() < nParticles) {
+      status.push_back(defaultStatus);
+    }
+  } else if ((int)status.size() > nParticles) {
+    edm::LogWarning("MCMultiParticleMassFilter") << "Warning: Given status value size"
+                                                    "> size of the number of particle IDs."
+                                                    " Ignoring extraneous values "
+                                                 << endl;
+    status.resize(nParticles);
+  }
+
+  for (int i = 0; i < nParticles; i++) {
+    std::tuple<double, double, int> cutForParticle = std::make_tuple(ptMin[i], etaMax[i], status[i]);
+    cutPerParticle[particleID[i]] = cutForParticle;
+  }  //assign the set of cuts you decided upon matched to each ID value in order
+}
+
+MCMultiParticleMassFilter::~MCMultiParticleMassFilter() {}
+
+bool MCMultiParticleMassFilter::filter(edm::StreamID, edm::Event& iEvent, const edm::EventSetup&) const {
+  edm::Handle<edm::HepMCProduct> evt;
+  iEvent.getByToken(token_, evt);
+  const HepMC::GenEvent* myGenEvent = evt->GetEvent();
+
+  std::vector<HepMC::GenParticle*> particlesThatPassCuts = std::vector<HepMC::GenParticle*>();
+  for (HepMC::GenEvent::particle_const_iterator p = myGenEvent->particles_begin(); p != myGenEvent->particles_end();
+       ++p) {
+    for (const int i : particleID) {
+      if (i == (*p)->pdg_id()) {
+        //if the particle ID is one of the ones you specified, check for cuts per ID
+        const auto cuts = cutPerParticle.at(i);
+        if (((*p)->status() == get<2>(cuts)) && ((*p)->momentum().perp() >= get<0>(cuts)) &&
+            (std::fabs((*p)->momentum().eta()) <= get<1>(cuts))) {
+          particlesThatPassCuts.push_back(*p);
+          break;
+        }
+      }
+    }
+  }
+  int nIterables = particlesThatPassCuts.size();
+  if (nIterables < nParticles) {
+    return false;
+  } else {
+    int i = 0;
+    //only iterate while there are enough particles that pass cuts
+    while ((nIterables - i) >= nParticles) {
+      vector<int> indices;
+      //start recursing from index 0, 1, 2, ...
+      indices.push_back(i);
+      bool success = recurseLoop(particlesThatPassCuts, indices, 1);
+      if (success) {
+        return true;
+      }
+      i++;
+    }
+  }
+  return false;
+}
+
+bool MCMultiParticleMassFilter::recurseLoop(std::vector<HepMC::GenParticle*>& particlesThatPassCuts,
+                                            std::vector<int> indices,
+                                            int depth) const {
+  int lastIndex = indices.back();
+  int nIterables = (int)(particlesThatPassCuts.size());
+  if (lastIndex >= nIterables) {
+    return false;
+  } else if (depth == nParticles) {
+    //you have the right number of particles!
+    int tempSum = 0;
+    int tempProd = 1;
+
+    double px, py, pz, e;
+    px = py = pz = e = 0;
+    for (const int i : indices) {
+      int id = particlesThatPassCuts[i]->pdg_id();
+      tempSum += id;
+      tempProd *= id;
+      HepMC::FourVector tempVec = particlesThatPassCuts[i]->momentum();
+      px += tempVec.px();
+      py += tempVec.py();
+      pz += tempVec.pz();
+      e += tempVec.e();
+    }
+    if ((tempSum != particleSumTo) || (tempProd != particleProdTo)) {
+      return false;  //check if the ids are what you want
+    }
+    double invMassSq = e * e - px * px - py * py - pz * pz;
+    //taking the root is computationally expensive, so use the squared term
+    if ((invMassSq >= minTotalMassSq) && (invMassSq <= maxTotalMassSq)) {
+      return true;
+    }
+    return false;
+  } else {
+    vector<bool> recursionResult;
+    //propagate recursion across all combinations of remaining indices
+    for (int i = 1; i < nIterables - lastIndex; i++) {
+      vector<int> newIndices = indices;
+      newIndices.push_back(lastIndex + i);
+      //always up the depth by 1 to make sure there is no infinite recursion
+      recursionResult.push_back(recurseLoop(particlesThatPassCuts, newIndices, depth + 1));
+    }
+    //search the results to look for one successful combination
+    for (bool r : recursionResult) {
+      if (r) {
+        return true;
+      }
+    }
+    return false;
+  }
+}
+DEFINE_FWK_MODULE(MCMultiParticleMassFilter);
+