Additional functions for hadronic recoil quantities

include/basefunctions.hxx

@@ -141,6 +141,100 @@ inline ROOT::RDF::RNode DefineQuantity(ROOT::RDF::RNode df,
                                        T const &value) {
     return df.Define(outputname, [value]() { return value; }, {});
 }
+
+/// Function to sum all elements of the column with name `quantity` with `ROOT::VecOps::RVec<T>`
+/// objects.
+///
+/// This function is a template implementation, i.e. call SumPerEvent<T> if the column
+/// `quantity` contains ROOT::VecOps::RVec<T> objects.
+///
+/// Elements of the `ROOT::VecOps::RVec<T>`, which should enter the sum, can be selected with
+/// index lists from the column `collection_index` as `ROOT::VecOps::RVec<int>` objects
+/// per entry.
+///
+/// Internally, `ROOT::VecOps::Sum` is used to calculate the sum. A custom zero element, which is
+/// a second optional argument of `ROOT::VecOps::Sum`, can be passed to this function with setting
+/// the parameter `zero`. Its default value is `T(0)`. For instance, when dealing with
+/// `ROOT::Math::PtEtaPhiMVector` objects, the `zero` parameter must be set to
+/// `ROOT::Math::PtEtaPhiMVector(0., 0., 0., 0)` in order to enable summation with this function.
+///
+/// \param df Input dataframe
+/// \param outputname name of the output column
+/// \param quantity column name of the vector variable which is summed per entry
+/// \param collection_index column name for index lists of the elements which are going to be summed up
+/// \param zero zero element passed as second argument to the `ROOT::VecOps::Sum` function
+///
+/// \returns a dataframe with the new column
+template <typename T>
+inline ROOT::RDF::RNode SumPerEvent(ROOT::RDF::RNode df,
+                                    const std::string &outputname,
+                                    const std::string &quantity,
+                                    const std::string &collection_index,
+                                    const T zero = T(0)) {
+    auto sum_per_event = [zero](const ROOT::RVec<T> &quantity, const ROOT::RVec<int> &collection_index) {
+        Logger::get("SumPerEvent")->debug(
+                    "sum values {} at indices {}", quantity, collection_index);
+        T sum = ROOT::VecOps::Sum(ROOT::VecOps::Take(quantity, collection_index), zero);
+        Logger::get("SumPerEvent")->debug(
+                    "sum {}", sum);
+        return sum;
+    };
+    return df.Define(outputname, sum_per_event, {quantity, collection_index});
+}
+
+/// Function to sum all elements of the column with name `quantity` with `ROOT::VecOps::RVec<T>`
+/// objects.
+///
+/// This function is a template implementation, i.e. call SumPerEvent<T> if the column
+/// `quantity` contains ROOT::VecOps::RVec<T> objects.
+///
+/// Internally, `ROOT::VecOps::Sum` is used to calculate the sum. A custom zero element, which is
+/// a second optional argument of `ROOT::VecOps::Sum`, can be passed to this function with setting
+/// the parameter `zero`. Its default value is `T(0)`.
+///
+/// \param df Input dataframe
+/// \param outputname name of the output column
+/// \param quantity column name of the vector variable which is summed per entry
+/// \param zero zero element passed as second argument to the `ROOT::VecOps::Sum` function
+///
+/// \returns a dataframe with the new column
+template <typename T>
+inline ROOT::RDF::RNode SumPerEvent(ROOT::RDF::RNode df,
+                                    const std::string &outputname,
+                                    const std::string &quantity,
+                                    const T zero = T(0)) {
+    auto sum_per_event = [zero](const ROOT::RVec<T> &quantity) {
+        Logger::get("SumPerEvent")->debug(
+                    "sum values {}", quantity);
+        T sum = ROOT::VecOps::Sum(quantity, zero);
+        Logger::get("SumPerEvent")->debug(
+                    "sum {}", sum);
+        return sum;
+    };
+    return df.Define(outputname, sum_per_event, {quantity});
+}
+
+/// This function creates a new column `outputname` with the negatives of type of the values
+/// in the column `inputname`.
+///
+/// Note that this function is implemented as a template, so specify the type `T` of the objects
+/// in the input column when calling this function with `Negative<T>(...)`.
+///
+/// \param df Input dataframe
+/// \param outputname name of the output column
+/// \param inputname column name of the input column
+///
+/// \returns a dataframe with the new column
+template <typename T>
+inline ROOT::RDF::RNode Negative(ROOT::RDF::RNode df,
+                                 const std::string &outputname,
+                                 const std::string &inputname) {
+    auto negative = [](const T &input) {
+        return -input;
+    };
+    return df.Define(outputname, negative, {inputname});
+}
+
 /// This function filters events, where neither of the input flags is true.
 /// This is used to filter events which do not pass an underlying requirement in
 /// any systematic variation.

include/lorentzvectors.hxx

@@ -13,6 +13,19 @@ ROOT::RDF::RNode buildparticle(ROOT::RDF::RNode df,
                                const std::vector<std::string> &quantities,
                                const std::string &outputname,
                                const int &position);
+ROOT::RDF::RNode BuildP4Collection(ROOT::RDF::RNode df,
+                                   const std::string &outputname,
+                                   const std::string &pts,
+                                   const std::string &etas,
+                                   const std::string &phis,
+                                   const std::string &masses,
+                                   const std::string &collection_index);
+ROOT::RDF::RNode BuildP4Collection(ROOT::RDF::RNode df,
+                                          const std::string &outputname,
+                                          const std::string &pts,
+                                          const std::string &etas,
+                                          const std::string &phis,
+                                          const std::string &masses);
 ROOT::RDF::RNode build(ROOT::RDF::RNode df,
                        const std::vector<std::string> &obj_quantities,
                        const int pairindex, const std::string &obj_p4_name);

src/lorentzvectors.cxx

@@ -69,6 +69,85 @@ ROOT::RDF::RNode buildparticle(ROOT::RDF::RNode df,
     return df1;
 }
 
+/**
+ * @brief Build a new column with a collection of Lorentz vectors per dataframe entry, which are
+ * created from the pt, eta, phi and mass columns of a collection.
+ *
+ * For instance, this can be used to create four-vector objects from the four-vector component
+ * columns of a NanoAOD collection.
+ *
+ * The function expects pt, eta phi and mass columns, which contain `ROOT::VecOps::RVec<float>`
+ * objects. In addition, the argument `collection_index` must point to a column which contains
+ * index lists of type `ROOT::VecOps::RVec<int>`. These index lists contain the indices of
+ * elements, for which four-vectors are built. The output column contains collections of
+ * four-vectors `ROOT::VecOps::RVec<PtEtaPhiMVector>`. The output collection only contains
+ * the elements, which have been selected in the `collection_index` list.
+ *
+ * @param df Input dataframe
+ * @param outputname name of the output column
+ * @param pts column name with the pt values of the collection's objects
+ * @param etas column name with the eta values of the collection's objects
+ * @param phis column name with the phi values of the collection's objects
+ * @param masses column name with the mass values of the collection's objects
+ * @param collection_index column name for index lists of the elements for which four-vectors are built
+ *
+ * @return a dataframe with the new column
+ */
+ROOT::RDF::RNode BuildP4Collection(ROOT::RDF::RNode df,
+                                   const std::string &outputname,
+                                   const std::string &pts,
+                                   const std::string &etas,
+                                   const std::string &phis,
+                                   const std::string &masses,
+                                   const std::string &collection_index) {
+    auto build_collection_p4 = [](const ROOT::RVec<float> &pts,
+                                  const ROOT::RVec<float> &etas,
+                                  const ROOT::RVec<float> &phis,
+                                  const ROOT::RVec<float> &masses,
+                                  const ROOT::RVec<int> &collection_index) {
+        auto pts_ = ROOT::VecOps::Take(pts, collection_index);
+        auto etas_ = ROOT::VecOps::Take(etas, collection_index);
+        auto phis_ = ROOT::VecOps::Take(phis, collection_index);
+        auto masses_ = ROOT::VecOps::Take(masses, collection_index);
+        return ROOT::VecOps::Construct<ROOT::Math::PtEtaPhiMVector>(pts_, etas_, phis_, masses_);
+    };
+    return df.Define(outputname, build_collection_p4, {pts, etas, phis, masses, collection_index});
+}
+
+/** 
+ * @brief Build a new column with a collection of Lorentz vectors per dataframe entry, which are
+ * created from the pt, eta, phi and mass columns of a collection.
+ *
+ * For instance, this can be used to create four-vector objects from the four-vector component
+ * columns of a NanoAOD collection.
+ *
+ * The function expects pt, eta phi and mass columns, which contain `ROOT::VecOps::RVec<float>`
+ * objects. The output column contains four-vectors `ROOT::VecOps::RVec<PtEtaPhiMVector>`.
+ *
+ * @param df Input dataframe
+ * @param outputname name of the output column
+ * @param pts column name with the pt values of the collection's objects
+ * @param etas column name with the eta values of the collection's objects
+ * @param phis column name with the phi values of the collection's objects
+ * @param masses column name with the mass values of the collection's objects
+ *
+ * @return a dataframe with the new column
+ */
+ROOT::RDF::RNode BuildP4Collection(ROOT::RDF::RNode df,
+                                          const std::string &outputname,
+                                          const std::string &pts,
+                                          const std::string &etas,
+                                          const std::string &phis,
+                                          const std::string &masses) {
+    auto build_collection_p4 = [](const ROOT::RVec<float> &pts,
+                                  const ROOT::RVec<float> &etas,
+                                  const ROOT::RVec<float> &phis,
+                                  const ROOT::RVec<float> &masses) {
+        return ROOT::VecOps::Construct<ROOT::Math::PtEtaPhiMVector>(pts, etas, phis, masses);
+    };
+    return df.Define(outputname, build_collection_p4, {pts, etas, phis, masses});
+}
+
 /**
  * @brief Function used to construct a 4-vector for a pair particle.
  *