function to select quantities on exact match (#250)

KIT-CMS · Feb 13, 2024 · ac72633 · ac72633
1 parent 61dbbf3
commit ac72633
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 0 deletions.
diff --git a/include/basefunctions.hxx b/include/basefunctions.hxx
@@ -349,6 +349,19 @@ inline auto FilterAbsMin(const float &cut) {
     };
 }
 
+/// Function to apply an exact filter requirement to an integer quantity.
+/// Returns true if the value is equal to the given value
+///
+/// \param cut The value of the filter
+///
+/// \returns a lambda function to be used in RDF Define
+inline auto FilterEqualInt(const int &cut) {
+    return [cut](const ROOT::RVec<int> &values) {
+        ROOT::RVec<int> mask = values == cut;
+        return mask;
+    };
+}
+
 /// Function to combine two RVec Masks by multiplying the two RVec elementwise
 ///
 /// \param mask_1 The first mask

diff --git a/include/physicsobjects.hxx b/include/physicsobjects.hxx
@@ -79,6 +79,8 @@ ROOT::RDF::RNode CheckForDiLeptonPairs(
     const std::string &leptons_phi, const std::string &leptons_mass,
     const std::string &leptons_charge, const std::string &leptons_mask,
     const float dR_cut);
+ROOT::RDF::RNode SelectInt(ROOT::RDF::RNode df, const std::string &maskname,
+                           const std::string &nameID, const int &IDvalue);
 namespace muon {
 ROOT::RDF::RNode CutID(ROOT::RDF::RNode df, const std::string &maskname,
                        const std::string &nameID);

diff --git a/src/physicsobjects.cxx b/src/physicsobjects.cxx
@@ -406,6 +406,21 @@ ROOT::RDF::RNode CheckForDiLeptonPairs(
                           leptons_charge, leptons_mask});
     return df1;
 }
+/// Function to select objects based on matching a specific integer value
+///
+/// \param[in] df the input dataframe
+/// \param[out] maskname the name of the new mask to be added as column to
+/// the dataframe
+/// \param[in] nameID name of the ID column in the NanoAOD
+/// \param[in] IDvalue value that has to match
+///
+/// \return a dataframe containing the new mask
+ROOT::RDF::RNode SelectInt(ROOT::RDF::RNode df, const std::string &maskname,
+                         const std::string &nameID, const int &IDvalue) {
+    auto df1 =
+        df.Define(maskname, basefunctions::FilterEqualInt(IDvalue), {nameID});
+    return df1;
+}
 
 /// Muon specific functions
 namespace muon {