diff --git a/.gitignore b/.gitignore index f0a97d00..cb045cb6 100644 --- a/.gitignore +++ b/.gitignore @@ -23,6 +23,7 @@ env.bak/ venv.bak/ .vscode .vscode/ +.codex # Jupyter Notebook .ipynb_checkpoints diff --git a/code_generation/configuration.py b/code_generation/configuration.py index 1ada2826..90c78e80 100644 --- a/code_generation/configuration.py +++ b/code_generation/configuration.py @@ -159,6 +159,18 @@ def _set_sample_parameters(self) -> None: sample_parameters["is_{}".format(sampletype)] = True else: sample_parameters["is_{}".format(sampletype)] = False + if "data" in self.sample: + sample_parameters["is_data"] = True + else: + sample_parameters["is_data"] = False + if "dyjets" in self.sample: + sample_parameters["is_dyjets"] = True + else: + sample_parameters["is_dyjets"] = False + if "wjets" in self.sample: + sample_parameters["is_wjets"] = True + else: + sample_parameters["is_wjets"] = False for scope in self.scopes: self.config_parameters[scope].update(sample_parameters) @@ -636,7 +648,7 @@ def _remove_empty_configkeys(self, config) -> None: Returns: None """ - for key in config: + for key in list(config.keys()): if isinstance(config[key], dict): self._remove_empty_configkeys(config[key]) # special case for extended vector producers, here we can have a list, that contains empty dicts diff --git a/include/event.hxx b/include/event.hxx index 5b834f98..8c3871e6 100644 --- a/include/event.hxx +++ b/include/event.hxx @@ -327,6 +327,31 @@ GenerateSeed( const UInt_t &master_seed = 42 ); +// Select the embedding-muon p4 that has the same charge as the reco object. +ROOT::RDF::RNode MatchEmbeddingMuonP4ByCharge( + ROOT::RDF::RNode df, + const std::string &outputname, + const std::string &reco_q, + const std::string &emb_q_1, + const std::string &emb_q_2, + const std::string &emb_p4_1, + const std::string &emb_p4_2); + +// Keep the old logic explicit: require that such a charge match exists. +ROOT::RDF::RNode HasEmbeddingMuonChargeMatch( + ROOT::RDF::RNode df, + const std::string &outputname, + const std::string &reco_q, + const std::string &emb_q_1, + const std::string &emb_q_2); + +// Require opposite charge for a reconstructed dilepton pair. +ROOT::RDF::RNode HasOppositeCharge( + ROOT::RDF::RNode df, + const std::string &outputname, + const std::string &q_1, + const std::string &q_2); + /** * @brief This function creates a new column in the dataframe by applying * element-wise negation to an existing `quantity` column. @@ -927,7 +952,6 @@ Quantity(ROOT::RDF::RNode df, const std::string &filtername, }, {quantity}, filtername); } - ROOT::RDF::RNode GoldenJSON(ROOT::RDF::RNode df, correctionManager::CorrectionManager &correctionManager, diff --git a/src/event.cxx b/src/event.cxx index 5d3a306d..fe0006d6 100644 --- a/src/event.cxx +++ b/src/event.cxx @@ -1,10 +1,12 @@ #ifndef GUARD_EVENT_H #define GUARD_EVENT_H +#include "../include/defaults.hxx" #include "../include/utility/CorrectionManager.hxx" #include "../include/utility/Logger.hxx" #include "ROOT/RDataFrame.hxx" #include "TRandom3.h" +#include #include #include @@ -68,10 +70,117 @@ GenerateSeed( ); } +/** + * @brief This function selects the embedding muon Lorentz vector that matches + * the charge of the reconstructed particle. + * + * The reconstructed particle charge is compared to the charges of the two + * embedding muons. If the charge matches the first embedding muon, its Lorentz + * vector is returned. If it matches the second embedding muon, the second + * Lorentz vector is returned. If no charge match is found, a default Lorentz + * vector is returned and the corresponding `HasEmbeddingMuonChargeMatch` + * helper is expected to veto the event later. + * + * @param df input dataframe + * @param outputname name of the new column containing the matched embedding muon Lorentz vector + * @param reco_q name of the reconstructed particle charge column + * @param emb_q_1 name of the leading embedding muon charge column + * @param emb_q_2 name of the trailing embedding muon charge column + * @param emb_p4_1 name of the leading embedding muon Lorentz vector column + * @param emb_p4_2 name of the trailing embedding muon Lorentz vector column + * + * @return a dataframe with the new Lorentz vector column + */ +ROOT::RDF::RNode MatchEmbeddingMuonP4ByCharge( + ROOT::RDF::RNode df, + const std::string &outputname, + const std::string &reco_q, + const std::string &emb_q_1, + const std::string &emb_q_2, + const std::string &emb_p4_1, + const std::string &emb_p4_2) { + auto match_embedding_muon = + [](const int &reco_q, + const float &emb_q_1, + const float &emb_q_2, + const ROOT::Math::PtEtaPhiMVector &emb_p4_1, + const ROOT::Math::PtEtaPhiMVector &emb_p4_2) { + if (reco_q == static_cast(emb_q_1)) { + return emb_p4_1; + } + if (reco_q == static_cast(emb_q_2)) { + return emb_p4_2; + } + return default_lorentzvector; + }; + + return df.Define(outputname, + match_embedding_muon, + {reco_q, emb_q_1, emb_q_2, emb_p4_1, emb_p4_2}); +} + +/** + * @brief This function checks whether the reconstructed particle charge matches + * the charge of at least one embedding muon. + * + * The reconstructed particle charge is compared to the charges of the two + * embedding muons. The returned flag is `1` if a charge match is found and `0` + * otherwise. This keeps the old selection logic explicit before the `DeltaR` + * cut is applied. + * + * @param df input dataframe + * @param outputname name of the new column containing the charge-match flag + * @param reco_q name of the reconstructed particle charge column + * @param emb_q_1 name of the leading embedding muon charge column + * @param emb_q_2 name of the trailing embedding muon charge column + * + * @return a dataframe with the new flag column + */ +ROOT::RDF::RNode HasEmbeddingMuonChargeMatch( + ROOT::RDF::RNode df, + const std::string &outputname, + const std::string &reco_q, + const std::string &emb_q_1, + const std::string &emb_q_2) { + auto has_match = [](const int &reco_q, + const float &emb_q_1, + const float &emb_q_2) { + return static_cast(reco_q == static_cast(emb_q_1) || + reco_q == static_cast(emb_q_2)); + }; + + return df.Define(outputname, has_match, {reco_q, emb_q_1, emb_q_2}); +} + +/** + * @brief This function checks whether two reconstructed particles have + * opposite charge. + * + * The returned flag is `true` if the product of the two charges is `-1`, + * i.e. one particle is positively charged and the other negatively charged. + * + * @param df input dataframe + * @param outputname name of the new column containing the opposite-sign flag + * @param q_1 name of the first reconstructed charge column + * @param q_2 name of the second reconstructed charge column + * + * @return a dataframe with the new boolean flag column + */ +ROOT::RDF::RNode HasOppositeCharge( + ROOT::RDF::RNode df, + const std::string &outputname, + const std::string &q_1, + const std::string &q_2) { + auto has_opposite_charge = [](const int &q_1, const int &q_2) { + return q_1 * q_2 == -1; + }; + + return df.Define(outputname, has_opposite_charge, {q_1, q_2}); +} + } // end namespace quantity namespace filter { - /** * @brief This function applies a filter to the input dataframe using a Golden * JSON file, which contains a mapping of valid run-luminosity pairs. The @@ -129,4 +238,4 @@ GoldenJSON(ROOT::RDF::RNode df, } // end namespace filter } // end namespace event -#endif /* GUARD_EVENT_H */ \ No newline at end of file +#endif /* GUARD_EVENT_H */