CROWN Changes for mumu to ee embedding

.gitignore

@@ -23,6 +23,7 @@ env.bak/
 venv.bak/
 .vscode
 .vscode/
+.codex
 
 # Jupyter Notebook
 .ipynb_checkpoints

code_generation/configuration.py

@@ -159,6 +159,18 @@
                 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 @@
             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

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,

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 <cmath>
 #include <nlohmann/json.hpp>
 #include <openssl/sha.h>
 
@@ -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<int>(emb_q_1)) {
+                return emb_p4_1;
+            }
+            if (reco_q == static_cast<int>(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<int>(reco_q == static_cast<int>(emb_q_1) ||
+                                reco_q == static_cast<int>(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 */
+#endif /* GUARD_EVENT_H */