Factorization of the JEC chain

include/event.hxx

@@ -264,6 +264,60 @@ inline ROOT::RDF::RNode Define(ROOT::RDF::RNode df,
     return df.Define(outputname, [value]() { return value; }, {});
 }
 
+
+/**
+ * @brief This function adds a new column to the dataframe, assigning it a
+ * vector of constant value for all entries. The length of the vector is
+ * determined by the column containing the number of objects in a collection,
+ * e.g., `nJet` for the `Jet` collection.
+ *
+ * @tparam T type of the value to be assigned
+ * @param df input dataframe
+ * @param outputname name of the new column
+ * @param number_column column containing the length of an object collection
+ * for each event
+ * @param value constant value to be assigned to the new column
+ *
+ * @return a dataframe with the new column
+ */
+template <typename T>
+inline ROOT::RDF::RNode Define(
+    ROOT::RDF::RNode df,
+    const std::string &outputname,
+    const std::string &number_column,
+    T const &value
+) {
+    return df.Define(
+        outputname,
+        [value] (const int &number_column) {
+            return ROOT::RVec<T>(number_column, value);
+        },
+        {number_column}
+    );
+}
+
+template <typename T>
+ROOT::RDF::RNode Concatenate(
+    ROOT::RDF::RNode df,
+    const std::string &outputname,
+    const std::string &inputname_1,
+    const std::string &inputname_2
+) {
+    auto cat_func = [] (
+        ROOT::RVec<T> input_1,
+        ROOT::RVec<T> input_2
+    ) {
+        return ROOT::VecOps::Concatenate(input_1, input_2);
+    };
+
+    return df.Define(
+        outputname,
+        cat_func,
+        {inputname_1, inputname_2}
+
+    );
+}
+
 /**
  * @brief This function defines a new column in the dataframe, where each
  * element is a randomly generated number. The random values are generated using
@@ -712,6 +766,35 @@ inline ROOT::RDF::RNode Sum(ROOT::RDF::RNode df, const std::string &outputname,
     return df.Define(outputname, sum_per_event, {quantity, index_vector});
 }
 
+/**
+ * @brief This function sums two vectors. The output vector contains the
+ * element-wise sum in each component.
+ *
+ * @tparam T type of the input column values
+ * @param df input dataframe
+ * @param outputname name of the new column containing the summed values
+ * @param quantity_1 name of the first vector column in the sum
+ * @param quantity_2 name of the second vector column in the sum
+ *
+ * @return a dataframe with the new column
+ */
+template <typename T>
+inline ROOT::RDF::RNode SumVectors(
+    ROOT::RDF::RNode df,
+    const std::string &outputname,
+    const std::string &quantity_1,
+    const std::string &quantity_2,
+    const T zero = T(0)
+) {
+    auto sum_func = [] (
+        const ROOT::RVec<T> &quantity_1,
+        const ROOT::RVec<T> &quantity_2
+    ) {
+        return quantity_1 + quantity_2;
+    };
+    return df.Define(outputname, sum_func, {quantity_1, quantity_2});
+}
+
 /**
  * @brief This function calculates the scalar sum of an arbitrary set of quantities
  * of type `float`.

include/jets.hxx

@@ -1,8 +1,273 @@
 #ifndef GUARD_JETS_H
 #define GUARD_JETS_H
 
+#include "../include/utility/CorrectionManager.hxx"
+#include "TRandom3.h"
+#include "correction.h"
+
+
 namespace physicsobject {
 namespace jet {
+namespace jec {
+typedef struct jec_result_t {
+    float jet_pt_l1;
+    float jet_pt_l2rel;
+    float jet_pt_l2l3res;
+    float jet_pt_syst;
+    float jet_pt_corr;
+} JECResult;
+const correction::Correction* load_nominal_jes_correction(
+    correctionManager::CorrectionManager &correction_manager,
+    const std::string &jec_file,
+    const std::string &jes_tag,
+    const std::string &type_tag,
+    const std::string &jes_level,
+    const std::string &jec_algo
+);
+const correction::Correction* load_shifted_jes_correction(
+    correctionManager::CorrectionManager &correction_manager,
+    const std::string &jec_file,
+    const std::string &jer_tag,
+    const std::string &type_tag,
+    const std::string &jes_shift,
+    const std::string &jec_algo
+);
+const correction::Correction* load_jer_correction(
+    correctionManager::CorrectionManager &correction_manager,
+    const std::string &jec_file,
+    const std::string &jer_tag,
+    const std::string &type_tag,
+    const std::string &jer_parameter,
+    const std::string &jec_algo
+);
+float apply_jes_l1 (
+    const float &jet_pt,
+    const float &jet_eta,
+    const float &jet_area,
+    const float &rho,
+    const correction::Correction *jes_l1_evaluator
+);
+float apply_jes_l2rel (
+    const float &jet_pt,
+    const float &jet_eta,
+    const float &jet_phi,
+    const std::string &era,
+    const correction::Correction *jes_l2rel_evaluator
+);
+float apply_jes_l2l3res (
+    const float &jet_pt,
+    const float &jet_eta,
+    const float &run,
+    const std::string &era,
+    const correction::Correction *jes_l2l3res_evaluator
+);
+float apply_jes_shifts (
+    const float &jet_pt,
+    const float &jet_eta,
+    const float &jet_phi,
+    const UChar_t &jet_id,
+    const std::vector<std::string> &jes_shift_sources,
+    const int &jes_shift_factor,
+    const std::vector<correction::Correction*> &jes_shift_evaluators
+);
+float apply_jer(
+    const float &jet_pt,
+    const float &jet_eta,
+    const float &jet_phi,
+    const float &rho,
+    const ROOT::RVec<float> &genjet_pt,
+    const ROOT::RVec<float> &genjet_eta,
+    const ROOT::RVec<float> &genjet_phi,
+    const correction::Correction *jer_resolution_evaluator,
+    const correction::Correction *jer_scalefactor_evaluator,
+    const std::string &jer_shift,
+    const float &jet_radius,
+    const std::string &era,
+    TRandom3 randgen
+);
+JECResult apply_full_jec_mc(
+    const float &jet_pt,
+    const float &jet_eta,
+    const float &jet_phi,
+    const UChar_t &jet_id,
+    const float &jet_area,
+    const float &rho,
+    const ROOT::RVec<float> &genjet_pt,
+    const ROOT::RVec<float> &genjet_eta,
+    const ROOT::RVec<float> &genjet_phi,
+    const std::vector<std::string> &jes_shift_sources,
+    const int &jes_shift_factor,
+    const std::string &jer_shift,
+    const float &jet_radius,
+    const std::string &era,
+    TRandom3 randgen,
+    const correction::Correction* jes_l1_evaluator,
+    const correction::Correction* jes_l2rel_evaluator,
+    const std::vector<correction::Correction*> &jes_shift_evaluators,
+    const correction::Correction *jer_resolution_evaluator,
+    const correction::Correction *jer_scalefactor_evaluator
+);
+JECResult apply_jes_shifts_and_jer_mc(
+    const float &jet_pt,
+    const float &jet_eta,
+    const float &jet_phi,
+    const UChar_t &jet_id,
+    const float &rho,
+    const ROOT::RVec<float> &genjet_pt,
+    const ROOT::RVec<float> &genjet_eta,
+    const ROOT::RVec<float> &genjet_phi,
+    const std::vector<std::string> &jes_shift_sources,
+    const int &jes_shift_factor,
+    const std::string &jer_shift,
+    const float &jet_radius,
+    const std::string &era,
+    TRandom3 randgen,
+    const std::vector<correction::Correction*> &jes_shift_evaluators,
+    const correction::Correction *jer_resolution_evaluator,
+    const correction::Correction *jer_scalefactor_evaluator
+);
+JECResult apply_full_jec_data(
+    const float &jet_pt,
+    const float &jet_eta,
+    const float &jet_phi,
+    const float &jet_area,
+    const float &rho,
+    const correction::Correction* jes_l1_evaluator,
+    const correction::Correction* jes_l2rel_evaluator,
+    const correction::Correction* jes_l2l3res_evaluator
+);
+ROOT::RDF::RNode Raw(
+    ROOT::RDF::RNode df,
+    const std::string &outputname,
+    const std::string &jet_quantity,
+    const std::string &jet_raw_factor
+);
+ROOT::RDF::RNode RawMuonSubtr(
+    ROOT::RDF::RNode df,
+    const std::string &outputname,
+    const std::string &jet_quantity,
+    const std::string &jet_raw_factor,
+    const std::string &jet_muon_subtr_factor
+);
+ROOT::RDF::RNode RawMuonSubtr(
+    ROOT::RDF::RNode df,
+    const std::string &outputname,
+    const std::string &jet_quantity,
+    const std::string &jet_muon_subtr_factor
+);
+ROOT::RDF::RNode PtCorrectionMC(
+    ROOT::RDF::RNode df,
+    correctionManager::CorrectionManager &correction_manager,
+    const std::string &output_jec_result,
+    const std::string &output_l1,
+    const std::string &output_l2rel,
+    const std::string &output_l2l3res,
+    const std::string &output_full,
+    const std::string &jet_pt_raw,
+    const std::string &jet_eta,
+    const std::string &jet_phi,
+    const std::string &jet_area,
+    const std::string &jet_id,
+    const std::string &genjet_pt,
+    const std::string &genjet_eta,
+    const std::string &genjet_phi,
+    const std::string &rho,
+    const std::string &jer_seed,
+    const std::string &jec_file,
+    const std::string &jec_algo,
+    const std::string &jes_tag,
+    const std::string &jer_tag,
+    const std::vector<std::string> &jes_shift_sources,
+    const int &jes_shift_factor,
+    const std::string &jer_shift,
+    const bool &reapply_jes,
+    const std::string &era
+);
+ROOT::RDF::RNode PtCorrectionData(
+    ROOT::RDF::RNode df,
+    correctionManager::CorrectionManager &correction_manager,
+    const std::string &output_jec_result,
+    const std::string &output_l1,
+    const std::string &output_l2rel,
+    const std::string &output_l2l3res,
+    const std::string &output_full,
+    const std::string &jet_pt_raw,
+    const std::string &jet_eta,
+    const std::string &jet_phi,
+    const std::string &jet_area,
+    const std::string &rho,
+    const std::string &run,
+    const std::string &jec_file,
+    const std::string &jec_algo,
+    const std::string &jes_tag,
+    const bool &reapply_jes,
+    const std::string &era
+);
+ROOT::RDF::RNode MassCorrectionFromPt(
+    ROOT::RDF::RNode df,
+    const std::string &outputname,
+    const std::string &jet_mass_raw,
+    const std::string &jet_pt_raw,
+    const std::string &jet_pt_corrected
+);
+}
+
+ROOT::RDF::RNode RawPt(ROOT::RDF::RNode df,
+                        const std::string &outputname,
+                        const std::string &pts,
+                        const std::string &jet_raw_factor);
+
+ROOT::RDF::RNode
+PtCorrectionL1(ROOT::RDF::RNode df,
+        correctionManager::CorrectionManager &correction_manager,
+        const std::string &outputname,
+        const std::string &jet_pt,
+        const std::string &jet_eta, 
+        const std::string &jet_phi,
+        const std::string &jet_area, 
+        const std::string &jet_raw_factor,
+        const std::string &jet_raw_muonfactor,
+        const std::string &corrjet_pt,
+        const std::string &corrjet_eta, 
+        const std::string &corrjet_phi,
+        const std::string &corrjet_area, 
+        const std::string &corrjet_raw_muonfactor,
+        const std::string &rho, 
+        const std::string &jec_file, 
+        const std::string &jec_algo,
+        const std::string &jes_tag_mc, 
+        const std::string &jes_tag_data, 
+        const std::string &era,
+        const bool &is_data,
+        const bool &is_embedding);
+
+ROOT::RDF::RNode
+PtCorrection(ROOT::RDF::RNode df,
+        correctionManager::CorrectionManager &correction_manager,
+        const std::string &outputname,
+        const std::string &jet_pts,
+        const std::string &jet_eta, 
+        const std::string &jet_phi,
+        const std::string &jet_area, 
+        const std::string &jet_id,
+        const std::string &corrjet_eta, 
+        const std::string &corrjet_phi,
+        const std::string &corrjet_area, 
+        const std::string &gen_jet_pt,
+        const std::string &gen_jet_eta, 
+        const std::string &gen_jet_phi,
+        const std::string &rho, 
+        const std::string &jer_seed,
+        const std::string &run, 
+        const std::string &jec_file, 
+        const std::string &jec_algo,
+        const std::string &jes_tag_mc, 
+        const std::string &jes_tag_data, 
+        const std::vector<std::string> &jes_shift_sources,
+        const std::string &jer_tag,
+        const int &jes_shift, const std::string &jer_shift,
+        const std::string &era, const bool &is_data,
+        const bool &is_embedding);
 
 ROOT::RDF::RNode
 PtCorrectionMC(ROOT::RDF::RNode df,

include/met.hxx

@@ -119,6 +119,19 @@ ROOT::RDF::RNode METPhiCorrection(ROOT::RDF::RNode df,
                         const bool is_mc,
                         const std::string &stat_variation,
                         const std::string &pileup_variation);
+ROOT::RDF::RNode Type1Correction(
+    ROOT::RDF::RNode df,
+    const std::string &outputname,
+    const std::string raw_met,
+    const std::string &t1jet_pt_l1corrected,
+    const std::string &t1jet_pt_corrected,
+    const std::string &t1jet_eta,
+    const std::string &t1jet_phi,
+    const std::string &t1jet_em_ef,
+    const float &t1jet_min_pt,
+    const float &t1jet_max_abs_eta,
+    const float &t1jet_max_em_ef
+);
 } // end namespace met
 
 namespace lorentzvector {
@@ -224,5 +237,6 @@ ROOT::RDF::RNode PropagateToMET(
     const std::string &pt, const std::string &eta,
     const std::string &phi, const std::string &mass,
     bool apply_propagation, float min_pt);
+
 } // end namespace physicsobject
 #endif /* GUARD_MET_H */