Adding a cut on IP for TreeCreator and SV reconstruction

Author: hahassan7

PWGJE/Core/JetTaggingUtilities.h

@@ -211,12 +211,12 @@ int jetTrackFromHFShower(T const& jet, U const& /*tracks*/, V const& particles,
     hasMcParticle = true;
     auto const& particle = track.template mcParticle_as<V>();
     origin = RecoDecay::getParticleOrigin(particles, particle, searchUpToQuark);
-    if (origin == 1 || origin == 2) { // 1=charm , 2=beauty
+    if (origin == RecoDecay::OriginType::Prompt || origin == RecoDecay::OriginType::NonPrompt) { // 1=charm , 2=beauty
       hftrack = track;
-      if (origin == 1) {
+      if (origin == RecoDecay::OriginType::Prompt) {
         return JetTaggingSpecies::charm;
       }
-      if (origin == 2) {
+      if (origin == RecoDecay::OriginType::NonPrompt) {
         return JetTaggingSpecies::beauty;
       }
     }
@@ -244,12 +244,12 @@ int jetParticleFromHFShower(T const& jet, U const& particles, typename U::iterat
   for (const auto& particle : jet.template tracks_as<U>()) {
     hfparticle = particle; // for init if origin is 1 or 2, the particle is not hfparticle
     origin = RecoDecay::getParticleOrigin(particles, particle, searchUpToQuark);
-    if (origin == 1 || origin == 2) { // 1=charm , 2=beauty
+    if (origin == RecoDecay::OriginType::Prompt || origin == RecoDecay::OriginType::NonPrompt) { // 1=charm , 2=beauty
       hfparticle = particle;
-      if (origin == 1) {
+      if (origin == RecoDecay::OriginType::Prompt) {
         return JetTaggingSpecies::charm;
       }
-      if (origin == 2) {
+      if (origin == RecoDecay::OriginType::NonPrompt) {
         return JetTaggingSpecies::beauty;
       }
     }
@@ -993,11 +993,11 @@ void analyzeJetSVInfo4ML(AnalysisJet const& myJet, AnyTracks const& /*allTracks*
 
 // Looping over the track info and putting them in the input vector
 template <typename AnalysisJet, typename AnyTracks, typename SecondaryVertices>
-void analyzeJetTrackInfo4ML(AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, SecondaryVertices const& /*allSVs*/, std::vector<BJetTrackParams>& tracksParams, float trackPtMin = 0.5)
+void analyzeJetTrackInfo4ML(AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, SecondaryVertices const& /*allSVs*/, std::vector<BJetTrackParams>& tracksParams, float trackPtMin = 0.5, float trackDcaXYMax = 10.0, float trackDcaZMax = 10.0)
 {
   for (const auto& constituent : analysisJet.template tracks_as<AnyTracks>()) {
 
-    if (constituent.pt() < trackPtMin) {
+    if (constituent.pt() < trackPtMin || !trackAcceptanceWithDca(constituent, trackDcaXYMax, trackDcaZMax)) {
       continue;
     }
 
@@ -1026,11 +1026,11 @@ void analyzeJetTrackInfo4ML(AnalysisJet const& analysisJet, AnyTracks const& /*a
 
 // Looping over the track info and putting them in the input vector without using any SV info
 template <typename AnalysisJet, typename AnyTracks>
-void analyzeJetTrackInfo4MLnoSV(AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, std::vector<BJetTrackParams>& tracksParams, float trackPtMin = 0.5)
+void analyzeJetTrackInfo4MLnoSV(AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, std::vector<BJetTrackParams>& tracksParams, float trackPtMin = 0.5, float trackDcaXYMax = 10.0, float trackDcaZMax = 10.0)
 {
   for (const auto& constituent : analysisJet.template tracks_as<AnyTracks>()) {
 
-    if (constituent.pt() < trackPtMin) {
+    if (constituent.pt() < trackPtMin || !trackAcceptanceWithDca(constituent, trackDcaXYMax, trackDcaZMax)) {
       continue;
     }
 

PWGJE/Core/MlResponseHfTagging.h

@@ -113,6 +113,11 @@ class MlResponseHfTagging : public MlResponse<TypeOutputScore>
   /// @return A vector of input shapes
   std::vector<std::vector<int64_t>> getInputShape() const { return this->mModels[0].getInputShapes(); }
 
+  /// @brief Method to get the output shape of a model
+  /// \param imod is the index of the model
+  /// @return number of output nodes
+  int getOutputNodes(int imod = 0) const { return this->mModels[imod].getNumOutputNodes(); }
+
   /// Method to fill the inputs of jet, tracks and secondary vertices
   /// \param jet is the b-jet candidate
   /// \param tracks is the vector of tracks associated to the jet

PWGJE/TableProducer/jetTaggerHF.cxx

@@ -94,6 +94,7 @@ struct JetTaggerHFTask {
   Configurable<LabeledArray<double>> cutsMl{"cutsMl", {DefaultCutsMl[0], 1, 2, {"pT bin 0"}, {"score for default b-jet tagging", "uncer 1"}}, "ML selections per pT bin"};
   Configurable<int> nClassesMl{"nClassesMl", 2, "Number of classes in ML model"};
   Configurable<std::vector<std::string>> namesInputFeatures{"namesInputFeatures", std::vector<std::string>{"feature1", "feature2"}, "Names of ML model input features"};
+  Configurable<bool> useDb{"useDb", false, "Flag to use DB for ML model instead of the score"};
 
   Configurable<std::string> ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"};
   Configurable<std::vector<std::string>> modelPathsCCDB{"modelPathsCCDB", std::vector<std::string>{"Users/h/hahassan"}, "Paths of models on CCDB"};
@@ -507,7 +508,7 @@ struct JetTaggerHFTask {
       std::vector<jettaggingutilities::BJetSVParams> svsParams;
 
       jettaggingutilities::analyzeJetSVInfo4ML(analysisJet, allTracks, allSVs, svsParams, svPtMin, svReductionFactor);
-      jettaggingutilities::analyzeJetTrackInfo4ML(analysisJet, allTracks, allSVs, tracksParams, trackPtMin);
+      jettaggingutilities::analyzeJetTrackInfo4ML(analysisJet, allTracks, allSVs, tracksParams, trackPtMin, trackDcaXYMax, trackDcaZMax);
 
       int nSVs = analysisJet.template secondaryVertices_as<SecondaryVertices>().size();
 
@@ -525,7 +526,15 @@ struct JetTaggerHFTask {
         bMlResponse.isSelectedMl(inputML, analysisJet.pt(), output);
       }
 
-      scoreML[analysisJet.globalIndex()] = output[0];
+      if (bMlResponse.getOutputNodes() > 1) {
+        auto Db = [](std::vector<float> scores, float fC) {
+          return std::log(scores[2] / (fC * scores[1] + (1 - fC) * scores[0]));
+        };
+
+        scoreML[analysisJet.globalIndex()] = useDb ? Db(output, fC) : output[2]; // 2 is the b-jet index
+      } else {
+        scoreML[analysisJet.globalIndex()] = output[0];
+      }
     }
   }
 
@@ -537,7 +546,7 @@ struct JetTaggerHFTask {
       std::vector<jettaggingutilities::BJetTrackParams> tracksParams;
       std::vector<jettaggingutilities::BJetSVParams> svsParams;
 
-      jettaggingutilities::analyzeJetTrackInfo4MLnoSV(analysisJet, allTracks, tracksParams, trackPtMin);
+      jettaggingutilities::analyzeJetTrackInfo4MLnoSV(analysisJet, allTracks, tracksParams, trackPtMin, trackDcaXYMax, trackDcaZMax);
 
       jettaggingutilities::BJetParams jetparam = {analysisJet.pt(), analysisJet.eta(), analysisJet.phi(), static_cast<int>(tracksParams.size()), 0, analysisJet.mass()};
       tracksParams.resize(nJetConst); // resize to the number of inputs of the ML

PWGJE/TableProducer/secondaryVertexReconstruction.cxx

@@ -72,6 +72,8 @@ struct SecondaryVertexReconstruction {
   Configurable<float> ptMinTrack{"ptMinTrack", -1., "min. track pT"};
   Configurable<float> etaMinTrack{"etaMinTrack", -99999., "min. pseudorapidity"};
   Configurable<float> etaMaxTrack{"etaMaxTrack", 4., "max. pseudorapidity"};
+  Configurable<float> maxIPxy{"maxIPxy", 10, "maximum track DCA in xy plane"};
+  Configurable<float> maxIPz{"maxIPz", 10, "maximum track DCA in z direction"};
   Configurable<bool> fillHistograms{"fillHistograms", true, "do validation plots"};
 
   Configurable<std::string> ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"};
@@ -288,7 +290,7 @@ struct SecondaryVertexReconstruction {
     for (size_t iprong = prongIndex; iprong < particles.size(); ++iprong) {
 
       const auto& testTrack = particles[iprong].template track_as<OriginalTracks>();
-      if (testTrack.pt() < ptMinTrack || testTrack.eta() < etaMinTrack || testTrack.eta() > etaMaxTrack) {
+      if (testTrack.pt() < ptMinTrack || testTrack.eta() < etaMinTrack || testTrack.eta() > etaMaxTrack || std::abs(testTrack.dcaXY()) > maxIPxy || std::abs(testTrack.dcaZ()) > maxIPz) {
         continue;
       }