PWGJE: adding utility task for HF jet tagging (#3589)

Co-authored-by: Nima Zardoshti <[email protected]>

PWGJE/Core/CMakeLists.txt

@@ -18,5 +18,6 @@ o2physics_target_root_dictionary(PWGJECore
               HEADERS JetFinder.h
                       JetUtilities.h
                       FastJetUtilities.h
+                      JetTaggingUtilities.h
               LINKDEF PWGJECoreLinkDef.h)
 endif()

PWGJE/Core/JetTaggingUtilities.h

@@ -0,0 +1,280 @@
+// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
+// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
+// All rights not expressly granted are reserved.
+//
+// This software is distributed under the terms of the GNU General Public
+// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
+//
+// In applying this license CERN does not waive the privileges and immunities
+// granted to it by virtue of its status as an Intergovernmental Organization
+// or submit itself to any jurisdiction.
+
+/// \file JetTaggingUtilities.h
+/// \brief Jet tagging related utilities
+///
+/// \author Nima Zardoshti <[email protected]>
+
+#ifndef PWGJE_CORE_JETTAGGINGUTILITIES_H_
+#define PWGJE_CORE_JETTAGGINGUTILITIES_H_
+
+#include <cmath>
+#include <limits>
+#include <numeric>
+#include <tuple>
+#include <vector>
+
+#include "Framework/Logger.h"
+#include "Common/Core/RecoDecay.h"
+#include "PWGJE/Core/JetUtilities.h"
+
+enum JetTaggingSpecies {
+  none = 0,
+  charm = 1,
+  beauty = 2,
+  lightflavour = 3,
+  lightquark = 4,
+  gluon = 5
+};
+
+namespace JetTaggingUtilities
+{
+/**
+ * returns the globalIndex of the earliest mother of a particle in the shower. returns -1 if a suitable mother is not found
+ *
+ * @param particle MCParticle whose mother is to be found
+ */
+template <typename T>
+int getOriginalMotherIndex(const typename T::iterator& particle)
+{
+
+  // if (!particle) {
+  //   return -1.0;
+  // }
+  auto mother = particle;
+
+  while (mother.has_mothers()) {
+
+    mother = mother.template mothers_first_as<T>();
+
+    int motherStatusCode = std::abs(mother.getGenStatusCode());
+
+    if (motherStatusCode == 23 || motherStatusCode == 33 || motherStatusCode == 43 || motherStatusCode == 63) {
+      return mother.globalIndex();
+    }
+  }
+  return -1.0;
+}
+
+/**
+ * returns the globalIndex of the earliest HF mother of a particle in the shower. returns -1 if a suitable mother is not found. Should be used only on already identified HF particles
+ *
+ * @param hfparticle MCParticle whose mother is to be found
+ */
+
+template <typename T>
+int getOriginalHFMotherIndex(const typename T::iterator& hfparticle)
+{
+
+  // if (!hfparticle) {
+  //   return -1.0;
+  // }
+  auto mother = hfparticle;
+
+  while (mother.has_mothers()) {
+
+    mother = mother.template mothers_first_as<T>();
+
+    int motherStatusCode = std::abs(mother.getGenStatusCode());
+
+    if (motherStatusCode == 23 || motherStatusCode == 33 || motherStatusCode == 43 || motherStatusCode == 63 || (motherStatusCode == 51 && mother.template mothers_first_as<T>().pdgCode() == 21)) {
+      return mother.globalIndex();
+    }
+  }
+  return -1.0;
+}
+
+/**
+ * checks if atrack in a reco level jet originates from a HF shower. 0:no HF shower, 1:charm shower, 2:beauty shower. The first track originating from an HF shower can be extracted by reference
+ *
+ * @param jet
+ * @param particles table of generator level particles to be searched through
+ * @param hftrack track passed as reference which is then replaced by the first track that originated from an HF shower
+ */
+template <typename T, typename U, typename V>
+int jetTrackFromHFShower(T const& jet, U const& tracks, V const& particles, typename U::iterator& hftrack)
+{
+
+  bool hasMcParticle = false;
+  int origin = -1;
+  for (auto& track : jet.template tracks_as<U>()) {
+    if (!track.has_mcParticle()) {
+      continue;
+    }
+    hasMcParticle = true;
+    auto const& particle = track.template mcParticle_as<V>();
+    origin = RecoDecay::getCharmHadronOrigin(particles, particle, true);
+    if (origin == 1 || origin == 2) { // 1=charm , 2=beauty
+      hftrack = track;
+      if (origin == 1) {
+        return JetTaggingSpecies::charm;
+      }
+      if (origin == 2) {
+        return JetTaggingSpecies::beauty;
+      }
+    }
+  }
+  if (hasMcParticle) {
+    return JetTaggingSpecies::lightflavour;
+  } else {
+    return JetTaggingSpecies::none;
+  }
+}
+
+/**
+ * checks if a particle in a generator level jet originates from a HF shower. 0:no HF shower, 1:charm shower, 2:beauty shower. The first particle originating from an HF shower can be extracted by reference
+ *
+ * @param jet
+ * @param particles table of generator level particles to be searched through
+ * @param hfparticle particle passed as reference which is then replaced by the first track that originated from an HF shower
+ */
+template <typename T, typename U>
+int jetParticleFromHFShower(T const& jet, U const& particles, typename U::iterator& hfparticle)
+{
+
+  int origin = -1;
+  for (auto& particle : jet.template tracks_as<U>()) {
+    origin = RecoDecay::getCharmHadronOrigin(particles, particle, true);
+    if (origin == 1 || origin == 2) { // 1=charm , 2=beauty
+      hfparticle = particle;
+      if (origin == 1) {
+        return JetTaggingSpecies::charm;
+      }
+      if (origin == 2) {
+        return JetTaggingSpecies::beauty;
+      }
+    }
+  }
+  return JetTaggingSpecies::lightflavour;
+}
+
+/**
+ * returns if a reco level jet originates from a HF shower. 0:no HF shower, 1:charm shower, 2:beauty shower. The requirement is that the jet contains a particle from an HF shower and that the original HF quark is within dRMax of the jet axis in eta-phi
+ *
+ * @param jet
+ * @param particles table of generator level particles to be searched through
+ * @param dRMax maximum distance in eta-phi of initiating heavy-flavour quark from the jet axis
+ */
+
+template <typename T, typename U, typename V>
+int mcdJetFromHFShower(T const& jet, U const& tracks, V const& particles, float dRMax = 0.25)
+{
+
+  typename U::iterator hftrack;
+  int origin = jetTrackFromHFShower(jet, tracks, particles, hftrack);
+  if (origin == JetTaggingSpecies::charm || origin == JetTaggingSpecies::beauty) {
+    if (!hftrack.has_mcParticle()) {
+      return JetTaggingSpecies::none;
+    }
+    auto const& hfparticle = hftrack.template mcParticle_as<V>();
+
+    int originalHFMotherIndex = getOriginalHFMotherIndex<V>(hfparticle);
+    if (originalHFMotherIndex > -1.0) {
+
+      if (JetUtilities::deltaR(jet, particles.iteratorAt(originalHFMotherIndex)) < dRMax) {
+
+        return origin;
+
+      } else {
+        return JetTaggingSpecies::none;
+      }
+
+    } else {
+      return JetTaggingSpecies::none;
+    }
+
+  } else {
+
+    return JetTaggingSpecies::lightflavour;
+  }
+}
+
+/**
+ * checks if a generator level jet originates from a HF shower. 0:no HF shower, 1:charm shower, 2:beauty shower. The requirement is that the jet contains a particle from an HF shower and that the original HF quark is within dRMax of the jet axis in eta-phi
+ *
+ * @param jet
+ * @param particles table of generator level particles to be searched through
+ * @param dRMax maximum distance in eta-phi of initiating heavy-flavour quark from the jet axis
+ */
+
+template <typename T, typename U, typename V>
+int mcpJetFromHFShower(T const& jet, U const& particles, float dRMax = 0.25)
+{
+
+  typename U::iterator hfparticle;
+  int origin = jetParticleFromHFShower(jet, particles, hfparticle);
+  if (origin == JetTaggingSpecies::charm || origin == JetTaggingSpecies::beauty) {
+
+    int originalHFMotherIndex = getOriginalHFMotherIndex<U>(hfparticle);
+    if (originalHFMotherIndex > -1.0) {
+
+      if (JetUtilities::deltaR(jet, particles.iteratorAt(originalHFMotherIndex)) < dRMax) {
+
+        return origin;
+
+      } else {
+        return JetTaggingSpecies::none;
+      }
+
+    } else {
+      return JetTaggingSpecies::none;
+    }
+
+  } else {
+
+    return JetTaggingSpecies::lightflavour;
+  }
+}
+
+/**
+ * returns the pdg code of the original scattered parton closest to the jet axis, with the restriction that the parton and jet axis must be within dRMax in eta-phi
+ *
+ * @param jet
+ * @param particles table of generator level particles to be searched through
+ * @param dRMax maximum distance in eta-phi of initiating heavy-flavour quark from the jet axis
+ */
+
+template <typename T, typename U>
+int jetOrigin(T const& jet, U const& particles, float dRMax = 0.25)
+{
+  bool firstPartonFound = false;
+  typename U::iterator parton1;
+  typename U::iterator parton2;
+  for (auto& particle : particles) {
+    if (std::abs(particle.getGenStatusCode() == 23)) {
+      if (!firstPartonFound) {
+        parton1 = particle;
+        firstPartonFound = true;
+      } else {
+        parton2 = particle;
+      }
+    }
+  }
+
+  float dR1 = JetUtilities::deltaR(jet, parton1);
+  float dR2 = JetUtilities::deltaR(jet, parton2);
+
+  if (dR1 <= dR2 && dR1 < dRMax) {
+
+    return parton1.pdgCode();
+  }
+  if (dR2 <= dR1 && dR2 < dRMax) {
+
+    return parton2.pdgCode();
+  }
+
+  return 0;
+}
+
+}; // namespace JetTaggingUtilities
+
+#endif // PWGJE_CORE_JETTAGGINGUTILITIES_H_

PWGJE/Core/JetUtilities.h

@@ -13,6 +13,7 @@
 /// \brief Jet related utilities
 ///
 /// \author Raymond Ehlers <[email protected]>, ORNL
+/// \author Nima Zardoshti <[email protected]>
 
 #ifndef PWGJE_CORE_JETUTILITIES_H_
 #define PWGJE_CORE_JETUTILITIES_H_
@@ -26,6 +27,7 @@
 #include <TKDTree.h>
 
 #include "Framework/Logger.h"
+#include "Common/Core/RecoDecay.h"
 
 namespace JetUtilities
 {
@@ -384,6 +386,15 @@ std::tuple<std::vector<std::vector<int>>, std::vector<std::vector<int>>> MatchCl
   return std::make_tuple(matchIndexTrack, matchIndexCluster);
 }
 
+template <typename T, typename U>
+float deltaR(T const& A, U const& B)
+{
+  float dPhi = RecoDecay::constrainAngle(A.phi() - B.phi(), -M_PI);
+  float dEta = A.eta() - B.eta();
+
+  return TMath::Sqrt(dEta * dEta + dPhi * dPhi);
+}
+
 }; // namespace JetUtilities
 
 #endif // PWGJE_CORE_JETUTILITIES_H_

PWGJE/Core/PWGJECoreLinkDef.h

@@ -19,5 +19,6 @@
 #pragma link C++ class JetFinder + ;
 #pragma link C++ namespace JetUtilities + ;
 #pragma link C++ namespace FastJetUtilities + ;
+#pragma link C++ namespace JetTaggingUtilities + ;
 
 #endif // PWGJE_CORE_PWGJECORELINKDEF_H_

PWGJE/DataModel/JetTagging.h

@@ -0,0 +1,54 @@
+// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
+// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
+// All rights not expressly granted are reserved.
+//
+// This software is distributed under the terms of the GNU General Public
+// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
+//
+// In applying this license CERN does not waive the privileges and immunities
+// granted to it by virtue of its status as an Intergovernmental Organization
+// or submit itself to any jurisdiction.
+
+///
+/// \brief Table definitions for hf jet tagging
+///
+/// \author Nima Zardoshti
+
+#ifndef PWGJE_DATAMODEL_JETTAGGING_H_
+#define PWGJE_DATAMODEL_JETTAGGING_H_
+
+#include <cmath>
+#include "Framework/AnalysisDataModel.h"
+#include "PWGJE/DataModel/Jet.h"
+
+using namespace o2::analysis;
+
+namespace o2::aod
+{
+
+// Defines the jet substrcuture table definition
+#define JETTAGGING_TABLE_DEF(_jet_type_, _name_, _description_) \
+  namespace _name_##tagging                                     \
+  {                                                             \
+    DECLARE_SOA_COLUMN(Origin, origin, int);                    \
+    DECLARE_SOA_COLUMN(Algorithm1, algorithm1, int);            \
+    DECLARE_SOA_COLUMN(Algorithm2, algorithm2, int);            \
+    DECLARE_SOA_COLUMN(Algorithm3, algorithm3, int);            \
+  }                                                             \
+  DECLARE_SOA_TABLE(_jet_type_##Tags, "AOD", _description_ "Tags", _name_##tagging::Origin, _name_##tagging::Algorithm1, _name_##tagging::Algorithm2, _name_##tagging::Algorithm3);
+
+#define JETTAGGING_TABLES_DEF(_jet_type_, _description_)                                                    \
+  JETTAGGING_TABLE_DEF(_jet_type_##Jet, _jet_type_##jet, _description_)                                     \
+  JETTAGGING_TABLE_DEF(_jet_type_##MCDetectorLevelJet, _jet_type_##mcdetectorleveljet, _description_ "MCD") \
+  JETTAGGING_TABLE_DEF(_jet_type_##MCParticleLevelJet, _jet_type_##mcparticleleveljet, _description_ "MCP")
+
+JETTAGGING_TABLES_DEF(Charged, "C");
+JETTAGGING_TABLES_DEF(Full, "F");
+JETTAGGING_TABLES_DEF(Neutral, "N");
+JETTAGGING_TABLES_DEF(D0Charged, "D0");
+JETTAGGING_TABLES_DEF(LcCharged, "Lc");
+JETTAGGING_TABLES_DEF(BplusCharged, "BPL");
+
+} // namespace o2::aod
+
+#endif // PWGJE_DATAMODEL_JETTAGGING_H_

PWGJE/TableProducer/CMakeLists.txt

@@ -60,6 +60,11 @@ o2physics_add_dpl_workflow(jet-track-derived
                     SOURCES jettrackderived.cxx
                     PUBLIC_LINK_LIBRARIES O2::Framework O2Physics::PWGJECore O2Physics::AnalysisCore
                     COMPONENT_NAME Analysis)
+
+o2physics_add_dpl_workflow(jet-taggerhf
+                    SOURCES jettaggerhf.cxx
+                    PUBLIC_LINK_LIBRARIES O2::Framework O2Physics::PWGJECore O2Physics::AnalysisCore
+                    COMPONENT_NAME Analysis)
 endif()