diff --git a/include/basefunctions.hxx b/include/basefunctions.hxx index 5ec94511..bd76dfb5 100644 --- a/include/basefunctions.hxx +++ b/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` +/// objects. +/// +/// This function is a template implementation, i.e. call SumPerEvent if the column +/// `quantity` contains ROOT::VecOps::RVec objects. +/// +/// Elements of the `ROOT::VecOps::RVec`, which should enter the sum, can be selected with +/// index lists from the column `collection_index` as `ROOT::VecOps::RVec` 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 +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 &quantity, const ROOT::RVec &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` +/// objects. +/// +/// This function is a template implementation, i.e. call SumPerEvent if the column +/// `quantity` contains ROOT::VecOps::RVec 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 +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 &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(...)`. +/// +/// \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 +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. diff --git a/include/lorentzvectors.hxx b/include/lorentzvectors.hxx index 18544021..aae12f5b 100644 --- a/include/lorentzvectors.hxx +++ b/include/lorentzvectors.hxx @@ -13,6 +13,19 @@ ROOT::RDF::RNode buildparticle(ROOT::RDF::RNode df, const std::vector &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 &obj_quantities, const int pairindex, const std::string &obj_p4_name); diff --git a/src/lorentzvectors.cxx b/src/lorentzvectors.cxx index 7c2dc991..526bb8c0 100644 --- a/src/lorentzvectors.cxx +++ b/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` + * objects. In addition, the argument `collection_index` must point to a column which contains + * index lists of type `ROOT::VecOps::RVec`. 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`. 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 &pts, + const ROOT::RVec &etas, + const ROOT::RVec &phis, + const ROOT::RVec &masses, + const ROOT::RVec &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(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` + * objects. The output column contains four-vectors `ROOT::VecOps::RVec`. + * + * @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 &pts, + const ROOT::RVec &etas, + const ROOT::RVec &phis, + const ROOT::RVec &masses) { + return ROOT::VecOps::Construct(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. *