Update truth decorations

BTagTrainingPreprocessing/CMakeLists.txt

@@ -81,6 +81,7 @@ atlas_add_library(dataset-dumper
   src/SingleBTagTools.cxx
   src/JetLeptonDecayLabelDecorator.cxx
   src/LeptonTruthDecorator.cxx
+  src/JetTruthSummaryDecorator.cxx
   PUBLIC_HEADERS src
   LINK_LIBRARIES ${LINK_LIBRARIES}
   )

BTagTrainingPreprocessing/src/JetTruthAssociator.cxx

@@ -39,7 +39,7 @@ namespace {
         if (x.status() != 1) return false;
         return truth::getParentHadron(&x);
       };
-    case p::leptonFromWZ:
+    case p::overlapLepton:
       return [](const TP& x) {
         if (!x.isElectron() && !x.isMuon()) return false;
         if (x.status() != 1) return false;
@@ -102,7 +102,6 @@ void JetTruthAssociator::decorate(const xAOD::Jet& jet,
 bool JetTruthAssociator::passed_cuts(const xAOD::TruthParticle& part,
                                     const xAOD::Jet& jet) const
 {
-
   if ( part.pt() < m_kinematics.pt_minimum ) return false;
   if ( std::abs(part.eta()) > m_kinematics.abs_eta_maximum ) return false;
   if (jet.p4().DeltaR(part.p4()) > m_kinematics.dr_maximum ) return false;

BTagTrainingPreprocessing/src/JetTruthAssociator.hh

-#ifndef JETTRUTHDECORATOR_HH
-#define JETTRUTHDECORATOR_HH
+#ifndef JET_TRUTH_ASSOCIATOR_HH
+#define JET_TRUTH_ASSOCIATOR_HH
 
 #include "TruthSelectorConfig.hh"
 

BTagTrainingPreprocessing/src/JetTruthSummaryDecorator.cxx

+#include "JetTruthSummaryDecorator.hh"
+#include "xAODJet/Jet.h"
+#include "xAODTruth/TruthParticle.h"
+
+
+// the constructor just builds the decorator
+JetTruthSummaryDecorator::JetTruthSummaryDecorator(const std::string& link_name, const std::string& out_name)
+    : m_truth_particle_links(link_name),
+      m_deco_n_truth("n_truth_" + out_name),
+      m_deco_min_dr("min_dr_truth_" + out_name) {}
+
+// return min deltaR(jet,particles), or infinity if no particles exist
+float JetTruthSummaryDecorator::getMinimumDeltaR(
+    const xAOD::Jet& jet,
+    const PartLinks& part_links) const {
+  
+  std::vector<float> delta_r = {INFINITY};
+  for (const PartLink& tl : part_links) {
+
+    // cast to truth particle
+    const auto* truth_particle = dynamic_cast<const xAOD::TruthParticle*>(*tl);
+    if (!truth_particle) {
+      throw std::runtime_error("Can't cast xAOD::IParticle to xAOD::TruthParticle");
+    }
+    
+    // get minimum dR
+    float dr = jet.p4().DeltaR(truth_particle->p4());
+    delta_r.push_back(dr);
+  }
+
+  return *std::min_element(delta_r.begin(), delta_r.end());
+}
+
+// this call actually does the work on the jet
+void JetTruthSummaryDecorator::decorate(
+    const xAOD::Jet& jet) const {
+
+  // get linked truth particles
+  auto part_links = m_truth_particle_links(jet);
+
+  // decorate
+  m_deco_n_truth(jet) = part_links.size();
+  m_deco_min_dr(jet) = getMinimumDeltaR(jet, part_links);
+}

BTagTrainingPreprocessing/src/JetTruthSummaryDecorator.hh

+#ifndef TRUTH_SUMMARY_DECORATOR_HH
+#define TRUTH_SUMMARY_DECORATOR_HH
+
+#include "xAODBTagging/BTaggingContainerFwd.h"
+#include "xAODJet/JetContainerFwd.h"
+#include "xAODTruth/TruthParticleFwd.h"
+
+#include "AthLinks/ElementLink.h"
+
+
+class JetTruthSummaryDecorator {
+ public:
+  using AE = SG::AuxElement;
+  typedef ElementLink<xAOD::IParticleContainer> PartLink;
+  typedef std::vector<PartLink> PartLinks;
+
+  JetTruthSummaryDecorator(const std::string& link_name, const std::string& out_name);
+
+  float getMinimumDeltaR(const xAOD::Jet& jet, const PartLinks& part_links) const;
+  void decorate(const xAOD::Jet& jet) const;
+
+ private:
+  AE::ConstAccessor<PartLinks> m_truth_particle_links;
+  AE::Decorator<int> m_deco_n_truth;
+  AE::Decorator<float> m_deco_min_dr;
+};
+#endif

BTagTrainingPreprocessing/src/LeptonTruthDecorator.cxx

@@ -4,94 +4,31 @@
 #include "xAODTruth/TruthParticleContainer.h"
 
 // the constructor just builds the decorator
-LeptonTruthDecorator::LeptonTruthDecorator(const std::string& link_name)
+LeptonTruthDecorator::LeptonTruthDecorator()
     : m_acc_softMuon_link("softMuon_link"),
-      m_acc_softMuon_trackLeptonID("leptonID"),
-      m_acc_classifierParticleType("classifierParticleType"),
       m_acc_softMuon_truthOrigin("truthOrigin"),
       m_acc_softMuon_truthType("truthType"),
       m_deco_softMuon_truthOrigin("softMuon_truthOrigin"),
-      m_deco_softMuon_truthType("softMuon_truthType"),
-      m_truth_particle_links(link_name),
-      m_deco_dr_closest_prompt_truth_lepton("dRClosestPromptTruthLepton"),
-      m_deco_dr_closest_non_prompt_truth_lepton("dRClosestNonPromptTruthLepton"),
-      m_deco_n_truth_muons("nTruthMuons"),
-      m_deco_n_prompt_leptons("nPromptLeptons"),
-      m_deco_n_non_prompt_leptons("nNonPromptLeptons") {}
-
-// return smallest element in vector, default to -1 if empty
-float defaultOrMinValue(const std::vector<float>& vector) {
-  if (vector.size() > 0) {
-    return *std::min_element(vector.begin(), vector.end());
-  } else {
-    return -1;
-  }
-}
+      m_deco_softMuon_truthType("softMuon_truthType") {}
 
 // this call actually does the work on the jet
 void LeptonTruthDecorator::decorate(
     const xAOD::BTagging& btag,
     const xAOD::Jet& jet) const {
 
-  // check for an associated soft muon
   int softMuon_truthOrigin = -1;
   int softMuon_truthType = -1;
+
+  // check for an associated soft muon and get info
   auto softMuon_link = m_acc_softMuon_link(btag);
   if (softMuon_link.isValid()) {
-    // get jet, muon and muon truths
     const xAOD::Muon& muon = **softMuon_link;
     softMuon_truthOrigin = m_acc_softMuon_truthOrigin(muon);
     softMuon_truthType = m_acc_softMuon_truthType(muon);
   }
+
+  // decorate
   m_deco_softMuon_truthOrigin(jet) = softMuon_truthOrigin;
   m_deco_softMuon_truthType(jet) = softMuon_truthType;
 
-  // store info on nearby truth leptons
-  std::vector<float> dr_prompt_truth_leptons;
-  std::vector<float> dr_non_prompt_truth_leptons;
-  int n_truth_muons = 0;
-  int n_prompt_leptons = 0;
-  int n_non_prompt_leptons = 0;
-  for (const PartLink& tl : m_truth_particle_links(jet)) {
-
-    const auto* truth_particle = dynamic_cast<const xAOD::TruthParticle*>(*tl);
-    if (!truth_particle) {
-      throw std::runtime_error("Can't cast to a truth particle");
-    }
-
-    float dr = jet.p4().DeltaR(truth_particle->p4());
-    unsigned int truth_particle_type = m_acc_classifierParticleType(
-      *truth_particle);
-
-    // prompt (iso) leptons
-    if (truth_particle_type == MCTruthPartClassifier::IsoElectron or
-        truth_particle_type == MCTruthPartClassifier::IsoMuon) {
-      dr_prompt_truth_leptons.push_back(dr);
-      n_prompt_leptons++;
-    }
-
-    //  non-prompt (non-iso) leptons
-    if (truth_particle_type == MCTruthPartClassifier::NonIsoElectron or
-        truth_particle_type == MCTruthPartClassifier::NonIsoMuon) {
-      dr_non_prompt_truth_leptons.push_back(dr);
-      n_non_prompt_leptons++;
-    }
-
-    // prompt and non-prompt muons
-    if (truth_particle_type == MCTruthPartClassifier::IsoMuon or
-        truth_particle_type == MCTruthPartClassifier::NonIsoMuon) {
-      n_truth_muons++;
-    }
-  }
-
-  // Default or closest dr
-  float dr_closest_prompt_truth_lepton = defaultOrMinValue(dr_prompt_truth_leptons);
-  float dr_closest_non_prompt_truth_lepton = defaultOrMinValue(dr_non_prompt_truth_leptons);
-
-  // decorate
-  m_deco_dr_closest_prompt_truth_lepton(jet) = dr_closest_prompt_truth_lepton;
-  m_deco_dr_closest_non_prompt_truth_lepton(jet) = dr_closest_non_prompt_truth_lepton;
-  m_deco_n_truth_muons(jet) = n_truth_muons;
-  m_deco_n_prompt_leptons(jet) = n_prompt_leptons;
-  m_deco_n_non_prompt_leptons(jet) = n_non_prompt_leptons;
 }

BTagTrainingPreprocessing/src/LeptonTruthDecorator.hh

-#ifndef MUNO_INFO_DECORATOR_HH
-#define MUNO_INFO_DECORATOR_HH
+#ifndef LEPTON_INFO_DECORATOR_HH
+#define LEPTON_INFO_DECORATOR_HH
 
 #include "xAODMuon/MuonContainer.h"
 
@@ -13,26 +13,19 @@ class LeptonTruthDecorator {
   using PartLink = ElementLink<xAOD::IParticleContainer>;
   using PartLinks = std::vector<PartLink>;
 
-  LeptonTruthDecorator(const std::string& link_name);
+  LeptonTruthDecorator();
 
   void decorate(const xAOD::BTagging& btag,
                 const xAOD::Jet& jet) const;
 
  private:
   AE::ConstAccessor<ElementLink<xAOD::MuonContainer>> m_acc_softMuon_link;
-  AE::ConstAccessor<int> m_acc_softMuon_trackLeptonID;
-  AE::ConstAccessor<unsigned int> m_acc_classifierParticleType;
+
   AE::ConstAccessor<int> m_acc_softMuon_truthOrigin;
   AE::ConstAccessor<int> m_acc_softMuon_truthType;
 
   AE::Decorator<int> m_deco_softMuon_truthOrigin;
   AE::Decorator<int> m_deco_softMuon_truthType;
 
-  AE::ConstAccessor<PartLinks> m_truth_particle_links;
-  AE::Decorator<float> m_deco_dr_closest_prompt_truth_lepton;
-  AE::Decorator<float> m_deco_dr_closest_non_prompt_truth_lepton;
-  AE::Decorator<int> m_deco_n_truth_muons;
-  AE::Decorator<int> m_deco_n_prompt_leptons;
-  AE::Decorator<int> m_deco_n_non_prompt_leptons;
 };
 #endif

BTagTrainingPreprocessing/src/SingleBTagConfig.cxx

@@ -7,6 +7,7 @@
 #include <filesystem>
 #include <set>
 
+
 namespace {
 
   TrackSortOrder get_track_sort_order(const boost::property_tree::ptree& pt,
@@ -94,7 +95,7 @@ namespace {
     if (name == "truth_leptons") return p::lepton;
     if (name == "truth_fromBC") return p::fromBC;
 #define TRY(KEY) if (name == #KEY) return p::KEY
-    TRY(leptonFromWZ);
+    TRY(overlapLepton);
     TRY(promptLepton);
     TRY(nonPromptLepton);
     TRY(muon);
@@ -111,31 +112,45 @@ namespace {
 
     // the configuration can have either a selection object or a merge
     // list, not both
-    Opts selection(json.get("association"));
+    Opts association(json.get("association"));
     const nlohmann::ordered_json& merge = json.get("merge");
-    if (!selection.empty() && !merge.empty()) {
+    if (!association.empty() && !merge.empty()) {
       throw std::runtime_error(
-        "Truth config can't contain both selection and merging");
+        "Truth config can't contain both association and merging");
     }
 
-    if (!selection.empty()) {
+    std::string particles = "";
+    if (!association.empty()) {
       TruthSelectorConfig sel;
-      sel.particle = get_truth_particle(
-        selection.require<std::string>("particles"));
+      particles = association.require<std::string>("particles");
+      sel.particle = get_truth_particle(particles);
       auto& kin = sel.kinematics;
-      kin.pt_minimum = selection.get("pt_minimum", 0.0f);
-      kin.abs_eta_maximum = selection.get("abs_eta_maximum", INFINITY);
-      kin.dr_maximum = selection.require<float>("dr_maximum");
+      kin.pt_minimum = association.get("pt_minimum", 0.0f);
+      kin.abs_eta_maximum = association.get("abs_eta_maximum", INFINITY);
+      kin.dr_maximum = association.require<float>("dr_maximum");
       cfg.selection = sel;
     }
-    selection.throw_if_unused("selection");
+    association.throw_if_unused("selection");
 
     if (!merge.empty()) {
       cfg.merge = merge.get<std::vector<std::string>>();
     }
 
     cfg.overlap_dr = json.get("overlap_dr", 0.0);
-    cfg.association_name = json.require<std::string>("association_name");
+
+    // use association name if specified, otherwise reuse particles name
+    std::string association_name = json.get("association_name", "");
+    if (not association_name.empty()) {
+      cfg.association_name = association_name;
+    }
+    else if (not particles.empty()) {
+      cfg.association_name = particles;
+    }
+    else {
+      throw std::runtime_error(
+        "If not specifying \"association_name\" you need to specify some"
+        " \"particles\" to associate");
+    }
 
     Opts output(json.get("output"));
     if (!output.empty()) {
@@ -148,13 +163,7 @@ namespace {
     }
     output.throw_if_unused("output");
 
-    Opts decorate_json(json.get("decorate"));
-    if (!decorate_json.empty()) {
-      TruthConfig::Decorate decorate;
-      decorate.lepton = decorate_json.get("lepton", false);
-      cfg.decorate = decorate;
-    }
-    decorate_json.throw_if_unused("decorate option");
+    cfg.decorate = json.get("decorate_summary", false);
 
     json.throw_if_unused("truth config");
 

BTagTrainingPreprocessing/src/SingleBTagConfig.hh

@@ -40,10 +40,7 @@ struct TruthConfig {
   float overlap_dr;
   std::string association_name;
   std::optional<TruthOutputConfig> output;
-  struct Decorate {
-    bool lepton;
-  };
-  std::optional<Decorate> decorate;
+  std::optional<bool> decorate;
 };
 
 struct DecorateConfig {

BTagTrainingPreprocessing/src/SingleBTagTools.cxx

@@ -126,9 +126,10 @@ SingleBTagTools::SingleBTagTools(const SingleBTagConfig& jobcfg):
         });
     }
     if (cfg.decorate) {
-      if (cfg.decorate->lepton) {
-        lepTruthDecorators.emplace_back(cfg.association_name);
-      }
+      std::string out_name;
+      if (cfg.output) { out_name = cfg.output->name; }
+      else { out_name = cfg.association_name; }
+      jetTruthSummaryDecorators.emplace_back(cfg.association_name, out_name);
     }
   }
   if (jobcfg.hits) {

BTagTrainingPreprocessing/src/SingleBTagTools.hh

@@ -23,6 +23,7 @@
 #include "TruthSelectorConfig.hh"
 #include "JetLeptonDecayLabelDecorator.hh"
 #include "LeptonTruthDecorator.hh"
+#include "JetTruthSummaryDecorator.hh"
 
 #include "FlavorTagDiscriminants/BTagJetAugmenter.h"
 #include "FlavorTagDiscriminants/BTagTrackIpAccessor.h"
@@ -99,20 +100,24 @@ struct SingleBTagTools {
   Accessors acc;
   Decorators dec;
 
+  // truth tools
   std::vector<JetTruthAssociator> truth_associators;
   std::vector<JetTruthMerger> jet_truth_mergers;
+  std::vector<JetTruthSummaryDecorator> jetTruthSummaryDecorators;
+  std::vector<std::function<bool(const xAOD::Jet&)>> overlap_checks;
 
+  // jet decorators
+  DecoratorExample example_decorator;  
+  JetLeptonDecayLabelDecorator jet_lepton_decay_label_decorator;
+  LeptonTruthDecorator lepTruthDecorator;
+  
+  // track decorators
   TrackTruthDecorator trkTruthDecorator;
   TrackVertexDecorator trkVertexDecorator;
   TrackLeptonDecorator trkLeptonDecorator;
-  DecoratorExample example_decorator;
-  JetLeptonDecayLabelDecorator jet_lepton_decay_label_decorator;
-  std::vector<LeptonTruthDecorator> lepTruthDecorators;
 
+  // hits decorators
   std::unique_ptr<HitDecorator> hit_decorator;
-
-  std::vector<std::function<bool(const xAOD::Jet&)>> overlap_checks;
-
 };
 
 struct SingleBTagOutputs {
@@ -124,7 +129,6 @@ struct SingleBTagOutputs {
   std::unique_ptr<HitWriter> hits;
 
   TruthCorruptionCounter truth_counts;
-
 };
 
 #endif

BTagTrainingPreprocessing/src/TruthSelectorConfig.hh

@@ -10,7 +10,7 @@ struct TruthKinematicConfig
 
 struct TruthSelectorConfig
 {
-  enum class Particle { hadron, lepton, fromBC, leptonFromWZ,
+  enum class Particle { hadron, lepton, fromBC, overlapLepton,
     promptLepton, nonPromptLepton, muon, stableNonGeant };
   Particle particle;
   TruthKinematicConfig kinematics;

BTagTrainingPreprocessing/src/processSingleBTagEvent.cxx

@@ -218,9 +218,14 @@ void processSingleBTagEventImpl(Event& event,
     // get the b-tagging object
     const xAOD::BTagging* btag = tools.acc.btaggingLink(*calib_jet);
 
-    // decorate jet with lepton truth info
-    for (const auto& dec: tools.lepTruthDecorators) {
-      dec.decorate(*btag, *calib_jet);
+    // decorate jet with lepton MCTC truth info
+    if (jobcfg.decorate.soft_muon) {
+      tools.lepTruthDecorator.decorate(*btag, *calib_jet);
+    }
+
+    // decorate jet with summary info about associated truth collections
+    for (const auto& dec: tools.jetTruthSummaryDecorators) {
+      dec.decorate(*calib_jet);
     }
 
     tools.dec.jet_rank(*calib_jet) = rank++;

configs/single-b-tag/EMPFlow.json

@@ -49,12 +49,9 @@
         {
             "association": {
                 "file": "fragments/baseline-truth-kinematics.json",
-                "particles": "stableNonGeant"
+                "particles": "promptLepton"
             },
-            "association_name": "stable_nongeant",
-            "decorate": {
-                "lepton": true
-            }
+            "decorate_summary": true
         }
     ],
     "decorate": {

configs/single-b-tag/EMPFlowTruth.json

@@ -6,7 +6,6 @@
                 "file": "fragments/baseline-truth-kinematics.json",
                 "particles": "truth_hadrons"
             },
-            "association_name": "truth_hadrons",
             "output": {
                 "n_to_save": 5,
                 "sort_order": "pt"
@@ -17,7 +16,6 @@
                 "file": "fragments/baseline-truth-kinematics.json",
                 "particles": "truth_leptons"
             },
-            "association_name": "truth_leptons",
             "output": {
                 "n_to_save": 5,
                 "sort_order": "pt"
@@ -28,7 +26,6 @@
                 "file": "fragments/baseline-truth-kinematics.json",
                 "particles": "truth_fromBC"
             },
-            "association_name": "truth_fromBC",
             "output": {
                 "n_to_save": 5,
                 "sort_order": "pt"

configs/single-b-tag/fragments/overlap-leptons.json

 {
     "association": {
         "file": "baseline-truth-kinematics.json",
-        "particles": "leptonFromWZ"
+        "particles": "overlapLepton"
     },
-    "association_name": "overlapLeptons",
     "overlap_dr": 0.3
 }

configs/single-b-tag/fragments/single-btag-variables-all.json

@@ -6,7 +6,7 @@
     ],
     "jet_int_variables": [
         "LeptonDecayLabel",
-        "nPromptLeptons"
+        "n_truth_promptLepton"
     ],
     "jet_floats": [
         "HadronConeExclTruthLabelPt"

docs/configuration.md

@@ -59,7 +59,7 @@ sel & merge  --> assoc(association_name)
 subgraph outputs [" "]
 or(overlap_dr)
 writer(output)
-dec(decorate)
+dec(decorate_summary)
 end
 assoc --> or & dec & writer
 ```
@@ -69,15 +69,15 @@ or none of the first two options can be used:
 
 - `association`: Runs $\Delta R$ particle association if specified. It
   can be further configured:
-    - `particles`: The type of particles to select.
-    - `pt_minimum`, `abs_eta_maximum`: should be self-explanatory
-    - `dr_maximum`: maximum $\Delta R$ between jet and an associated
+    - `particles`: The type of particles to select, e.g. `fromBC` or `promptLepton`.
+    - `pt_minimum`, `abs_eta_maximum`: should be self-explanatory.
+    - `dr_maximum`: maximum $\Delta R$ between jet and an associated.
       particle
 
     Example:
     ```
     "association": {
-      "particles": "leptonsFromWZ",
+      "particles": "overlapLepton",
       "pt_minimum": 2e3,
       "abs_eta_maximum": 2.5,
       "dr_maximum": 0.4
@@ -99,9 +99,10 @@ If neither of the above options is given, we assume the particles are
 already associated in your input xAOD. What happens with them can be
 further configured:
 
-- `association_name`: The only required field. It specifies the name
-  of the `std::vector<ElementLink<xAOD::IParticleContainer>>` auxiliary
-  data that is saved on the jet. Example:
+- `association_name`: It specifies the name of the 
+  `std::vector<ElementLink<xAOD::IParticleContainer>>` auxiliary
+  data that is saved on the jet. If not specified we'll use `particles`
+  from the `association` block. Example:
   `"association_name": "ConeExclHadrons"`.
 - `overlap_dr`: The minimum $\Delta R(\text{jet},x)$, where $x$ is any
   particle in this collection. Jets are vetoed if closer particles are
@@ -122,9 +123,9 @@ further configured:
     }
 
     ```
-- `decorate`: Enable other decorations on the jet.
-    - `lepton`: decorate with lepton counts and $\min \Delta
-      R(\text{jet}, \text{lepton})$.
+- `decorate_summary`: Enable summary decorations on the jet.
+   This will add `n_truth_<name>` and `min_dr_truth_<name>`
+   decorations to the jet, which can be written out.
 
 ### DL2 Config