add SVFIT producer

KIT-CMS · Mar 16, 2023 · 8777d03 · 8777d03
1 parent 3bc06ab
commit 8777d03
Show file tree

Hide file tree

Showing 2 changed files with 114 additions and 0 deletions.
diff --git a/include/quantities.hxx b/include/quantities.hxx
@@ -26,6 +26,23 @@ ROOT::RDF::RNode charge(ROOT::RDF::RNode df, const std::string &outputname,
                         const std::string &chargecolumn);
 ROOT::RDF::RNode m_vis(ROOT::RDF::RNode df, const std::string &outputname,
                        const std::vector<std::string> &inputvectors);
+ROOT::RDF::RNode p4_fastmtt(ROOT::RDF::RNode df, const std::string &outputname,
+                        const std::string &pt_1,
+                        const std::string &pt_2,
+                        const std::string &eta_1,
+                        const std::string &eta_2,
+                        const std::string &phi_1,
+                        const std::string &phi_2,
+                        const std::string &mass_1,
+                        const std::string &mass_2,
+                        const std::string &met_pt,
+                        const std::string &met_phi,
+                        const std::string &met_cov_xx,
+                        const std::string &met_cov_xy,
+                        const std::string &met_cov_yy,
+                        const std::string &decay_mode_1,
+                        const std::string &decay_mode_2,
+                        const std::string &finalstate);
 ROOT::RDF::RNode pt_vis(ROOT::RDF::RNode df, const std::string &outputname,
                         const std::vector<std::string> &inputvectors);
 ROOT::RDF::RNode pzetamissvis(ROOT::RDF::RNode df,

diff --git a/src/quantities.cxx b/src/quantities.cxx
@@ -1,6 +1,8 @@
 #ifndef GUARD_QUANTITIES_H
 #define GUARD_QUANTITIES_H
 
+#include "../include/SVFit/FastMTT.hxx"
+#include "../include/SVFit/MeasuredTauLepton.hxx"
 #include "../include/basefunctions.hxx"
 #include "../include/defaults.hxx"
 #include "../include/utility/Logger.hxx"
@@ -176,6 +178,101 @@ ROOT::RDF::RNode m_vis(ROOT::RDF::RNode df, const std::string &outputname,
         },
         inputvectors);
 }
+
+
+/**
+ * @brief Function used to calculate the FastMTT p4 from the given inputs. The implementation is based on https://github.com/SVfit/ClassicSVfit/tree/fastMTT_19_02_2019
+ * 
+ * @param df The dataframe to add the quantity to
+ * @param outputname name of the new column containing the lorentz vector value
+ * @param pt_1 the name of the column containing the pt of the first particle
+ * @param pt_2 the name of the column containing the pt of the second particle
+ * @param eta_1  the name of the column containing the eta of the first particle
+ * @param eta_2 the name of the column containing the eta of the second particle
+ * @param phi_1 the name of the column containing the phi of the first particle
+ * @param phi_2 the name of the column containing the phi of the second particle
+ * @param mass_1 the name of the column containing the mass of the first particle
+ * @param mass_2 the name of the column containing the mass of the second particle
+ * @param met_pt the name of the column containing the met pt
+ * @param met_phi the name of the column containing the met phi
+ * @param met_cov_xx the name of the column containing the met covariance xx
+ * @param met_cov_xy the name of the column containing the met covariance xy
+ * @param met_cov_yy the name of the column containing the met covariance yy
+ * @param decay_mode_1 the name of the column containing the decay mode of the first particle
+ * @param decay_mode_2 the name of the column containing the decay mode of the second particle
+ * @param finalstate the final state of the ditaudecay. Supported are "mt", "et", "tt", "em"
+ * @return ROOT::RDF::RNode
+ */
+ROOT::RDF::RNode
+p4_fastmtt(ROOT::RDF::RNode df, const std::string &outputname,
+           const std::string &pt_1, const std::string &pt_2,
+           const std::string &eta_1, const std::string &eta_2,
+           const std::string &phi_1, const std::string &phi_2,
+           const std::string &mass_1, const std::string &mass_2,
+           const std::string &met_pt, const std::string &met_phi,
+           const std::string &met_cov_xx, const std::string &met_cov_xy,
+           const std::string &met_cov_yy, const std::string &decay_mode_1,
+           const std::string &decay_mode_2, const std::string &finalstate)
+    {
+    auto calculate_fast_mtt =
+        [finalstate](const float &pt_1, const float &pt_2, const float &eta_1,
+           const float &eta_2, const float &phi_1, const float &phi_2,
+           const float &mass_1, const float &mass_2, const float &met_pt,
+           const float &met_phi, const float &met_cov_xx,
+           const float &met_cov_xy, const float &met_cov_yy,
+           const int &decay_mode_1, const int &decay_mode_2) {
+            std::vector<fastmtt::MeasuredTauLepton> measuredTauLeptons;
+            TMatrixD covMET(2, 2);
+            covMET[0][0] = met_cov_xx;
+            covMET[1][0] = met_cov_xy;
+            covMET[0][1] = met_cov_xy;
+            covMET[1][1] = met_cov_yy;
+            // build the met lorentz vector
+            ROOT::Math::PtEtaPhiMVector met(met_pt, 0.0, met_phi, 0.0);
+            // set the decay modes according to the final state
+            auto decay_obj_1 = fastmtt::MeasuredTauLepton::kTauToHadDecay;
+            auto decay_obj_2 = fastmtt::MeasuredTauLepton::kTauToHadDecay;
+            int dm_1, dm_2;
+            if (finalstate == "mt") {
+                dm_1 = -1;
+                dm_2 = decay_mode_2;
+                auto decay_obj_1 = fastmtt::MeasuredTauLepton::kTauToMuDecay;
+                auto decay_obj_2 = fastmtt::MeasuredTauLepton::kTauToHadDecay;
+            } else if (finalstate == "et") {
+                dm_1 = -1;
+                dm_2 = decay_mode_2;
+                auto decay_obj_1 = fastmtt::MeasuredTauLepton::kTauToElecDecay;
+                auto decay_obj_2 = fastmtt::MeasuredTauLepton::kTauToHadDecay;
+            } else if (finalstate == "tt") {
+                dm_1 = decay_mode_1;
+                dm_2 = decay_mode_2;
+            } else if (finalstate == "em") {
+                dm_1 = -1;
+                dm_2 = -1;
+                auto decay_obj_1 = fastmtt::MeasuredTauLepton::kTauToElecDecay;
+                auto decay_obj_2 = fastmtt::MeasuredTauLepton::kTauToMuDecay;
+            } else {
+                Logger::get("FastMTT")->error(
+                    "Final state {} not supported by FastMTT", finalstate);
+                return (ROOT::Math::PtEtaPhiMVector)LorentzVector();
+            }
+            measuredTauLeptons.push_back(fastmtt::MeasuredTauLepton(
+                decay_obj_1, pt_1, eta_1, phi_1, mass_1, dm_1));
+            measuredTauLeptons.push_back(fastmtt::MeasuredTauLepton(
+                decay_obj_2, pt_2, eta_2, phi_2, mass_2, dm_2));
+            FastMTT FastMTTAlgo;
+            FastMTTAlgo.run(measuredTauLeptons, met.X(), met.Y(), covMET);
+            LorentzVector result = FastMTTAlgo.getBestP4();
+            // ROOT::Math::PtEtaPhiMVector result(_result.Pt(), _result.Eta(),
+            //                                    _result.Phi(), _result.M());
+            Logger::get("FastMTT")->debug("FastMTT result: {}", result.M());
+            return (ROOT::Math::PtEtaPhiMVector)result;
+        };
+    return df.Define(outputname, calculate_fast_mtt,
+                     {pt_1, pt_2, eta_1, eta_2, phi_1, phi_2, mass_1, mass_2,
+                      met_pt, met_phi, met_cov_xx, met_cov_xy, met_cov_yy,
+                      decay_mode_1, decay_mode_2});
+}
 /// Function to calculate the visible pt from a pair of lorentz vectors and
 /// add it to the dataframe. The visible pt is calculated as the pt of the
 /// lorentz vector of the dilepton system.