PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/CMakeLists.txt

@@ -27,6 +27,7 @@ atlas_add_library( FlavorTagDiscriminants
   Root/DL2Tool.cxx
   Root/customGetter.cxx
   Root/FlipTagEnums.cxx
+  Root/AssociationEnums.cxx
   Root/VRJetOverlapDecorator.cxx
   Root/VRJetOverlapDecoratorTool.cxx
   Root/HbbTag.cxx
@@ -44,6 +45,7 @@ atlas_add_library( FlavorTagDiscriminants
 if (NOT XAOD_STANDALONE)
   atlas_add_component( FlavorTagDiscriminantsLib
     src/BTagDecoratorAlg.cxx
+    src/JetTagDecoratorAlg.cxx
     src/BTagToJetLinkerAlg.cxx
     src/JetToBTagLinkerAlg.cxx
     src/BTagTrackLinkCopyAlg.cxx

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/FlavorTagDiscriminants/AssociationEnums.h

+/*
+  Copyright (C) 2002-2021 CERN for the benefit of the ATLAS collaboration
+*/
+
+
+#ifndef ASSOCIATION_ENUMS_HH
+#define ASSOCIATION_ENUMS_HH
+
+#include <string>
+
+namespace FlavorTagDiscriminants {
+  enum class TrackLinkType {
+    TRACK_PARTICLE,
+    IPARTICLE
+  };
+  TrackLinkType trackLinkTypeFromString(const std::string&);
+}
+
+#endif

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/FlavorTagDiscriminants/BTagDecoratorAlg.h

@@ -5,41 +5,26 @@
 #ifndef BTAG_DECORATOR_ALG_H
 #define BTAG_DECORATOR_ALG_H
 
-#include "AthenaBaseComps/AthReentrantAlgorithm.h"
+
+#include "FlavorTagDiscriminants/DecoratorAlg.h"
 #include "FlavorTagDiscriminants/IBTagDecorator.h"
 
 #include "xAODBTagging/BTaggingContainer.h"
 #include "xAODTracking/TrackParticleContainer.h"
-#include "StoreGate/WriteDecorHandleKey.h"
-#include "StoreGate/ReadDecorHandleKey.h"
+
+namespace detail {
+  using BTag_t = FlavorTagDiscriminants::DecoratorAlg<
+    xAOD::BTaggingContainer,
+    IBTagDecorator,
+    xAOD::TrackParticleContainer
+    >;
+}
 
 namespace FlavorTagDiscriminants {
-  class BTagDecoratorAlg : public AthReentrantAlgorithm
+  class BTagDecoratorAlg : public detail::BTag_t
   {
   public:
     BTagDecoratorAlg(const std::string& name, ISvcLocator* svcloc);
-    virtual StatusCode initialize() override;
-    virtual StatusCode execute(const EventContext& cxt) const override;
-    virtual StatusCode finalize() override;
-  private:
-    SG::ReadHandleKey<xAOD::BTaggingContainer> m_btagContainerKey {
-      this, "btagContainer", "", "Key for the input btag collection"};
-    SG::ReadHandleKey<xAOD::TrackParticleContainer> m_trackContainerKey {
-      this, "trackContainer", "InDetTrackParticles",
-        "Key for track particle container"};
-    Gaudi::Property<std::vector<std::string>> m_undeclaredReadDecorKeys {
-      this, "undeclaredReadDecorKeys", {},
-      "List of read handles that we don't read, e.g. static variables" };
-
-    ToolHandle<IBTagDecorator> m_decorator{
-      this, "decorator", "", "Decorator tool"};
-
-    // Keys to keep track of the inputs / outputs
-    std::vector<SG::ReadDecorHandleKey<xAOD::BTaggingContainer>> m_btagAux;
-    std::vector<SG::ReadDecorHandleKey<xAOD::TrackParticleContainer>> m_trkAux;
-    std::vector<SG::WriteDecorHandleKey<xAOD::BTaggingContainer>> m_btagDecor;
-
-
   };
 }
 

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/FlavorTagDiscriminants/DL2.h

@@ -8,6 +8,7 @@
 // local includes
 #include "FlavorTagDiscriminants/customGetter.h"
 #include "FlavorTagDiscriminants/FlipTagEnums.h"
+#include "FlavorTagDiscriminants/AssociationEnums.h"
 #include "FlavorTagDiscriminants/DL2DataDependencyNames.h"
 #include "xAODBTagging/ftagfloat_t.h"
 
@@ -67,6 +68,7 @@ namespace FlavorTagDiscriminants {
     FlipTagConfig flip;
     std::string track_link_name;
     std::map<std::string,std::string> remap_scalar;
+    TrackLinkType track_link_type;
   };
 
 
@@ -87,7 +89,7 @@ namespace FlavorTagDiscriminants {
                                  const xAOD::Jet&)> TrackSequenceFilter;
 
     // getter functions
-    typedef std::function<NamedVar(const BTagging&)> VarFromBTag;
+    typedef std::function<NamedVar(const SG::AuxElement&)> VarFromBTag;
     typedef std::function<NamedSeq(const Jet&, const Tracks&)> SeqFromTracks;
 
     // ___________________________________________________________________
@@ -112,7 +114,7 @@ namespace FlavorTagDiscriminants {
         m_name(name)
         {
         }
-      NamedVar operator()(const xAOD::BTagging& btag) const {
+      NamedVar operator()(const SG::AuxElement& btag) const {
         T ret_value = m_getter(btag);
         bool is_default = m_default_flag(btag);
         if constexpr (std::is_floating_point<T>::value) {
@@ -140,7 +142,7 @@ namespace FlavorTagDiscriminants {
         m_name(name)
         {
         }
-      NamedVar operator()(const xAOD::BTagging& btag) const {
+      NamedVar operator()(const SG::AuxElement& btag) const {
         T ret_value = m_getter(btag);
         if constexpr (std::is_floating_point<T>::value) {
           if (std::isnan(ret_value)) {
@@ -159,11 +161,15 @@ namespace FlavorTagDiscriminants {
     public:
       TracksFromJet(SortOrder, TrackSelection, const DL2Options&);
       Tracks operator()(const xAOD::Jet& jet,
-                        const xAOD::BTagging& btag) const;
+                        const SG::AuxElement& btag) const;
     private:
-      typedef SG::AuxElement AE;
-      typedef std::vector<ElementLink<xAOD::TrackParticleContainer>> TrackLinks;
-      AE::ConstAccessor<TrackLinks> m_trackAssociator;
+      using AE = SG::AuxElement;
+      using IPC = xAOD::IParticleContainer;
+      using TPC = xAOD::TrackParticleContainer;
+      using TrackLinks = std::vector<ElementLink<TPC>>;
+      using PartLinks = std::vector<ElementLink<IPC>>;
+      using TPV = std::vector<const xAOD::TrackParticle*>;
+      std::function<TPV(const SG::AuxElement&)> m_associator;
       TrackSortVar m_trackSortVar;
       TrackFilter m_trackFilter;
     };
@@ -199,6 +205,8 @@ namespace FlavorTagDiscriminants {
         const std::vector<DL2TrackSequenceConfig>& = {},
         const DL2Options& = DL2Options());
     void decorate(const xAOD::BTagging& btag) const;
+    void decorate(const xAOD::Jet& jet) const;
+    void decorate(const xAOD::Jet& jet, const SG::AuxElement& decorated) const;
 
     // functions to report data depdedencies
     DL2DataDependencyNames getDataDependencyNames() const;

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/FlavorTagDiscriminants/DL2HighLevel.h

@@ -6,10 +6,12 @@
 #define DL2_HIGH_LEVEL_HH
 
 #include "FlavorTagDiscriminants/FlipTagEnums.h"
+#include "FlavorTagDiscriminants/AssociationEnums.h"
 #include "FlavorTagDiscriminants/DL2DataDependencyNames.h"
 
 // EDM includes
 #include "xAODBTagging/BTaggingFwd.h"
+#include "xAODJet/JetFwd.h"
 
 #include <memory>
 #include <string>
@@ -24,10 +26,12 @@ namespace FlavorTagDiscriminants {
   public:
     DL2HighLevel(const std::string& nn_file_name,
                  FlipTagConfig = FlipTagConfig::STANDARD,
-                 std::map<std::string, std::string> remap_scalar = {});
+                 std::map<std::string, std::string> remap_scalar = {},
+                 TrackLinkType = TrackLinkType::TRACK_PARTICLE);
     DL2HighLevel(DL2HighLevel&&);
     ~DL2HighLevel();
     void decorate(const xAOD::BTagging& btag) const;
+    void decorate(const xAOD::Jet& jet) const;
     DL2DataDependencyNames getDataDependencyNames() const;
   private:
     std::unique_ptr<DL2> m_dl2;

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/FlavorTagDiscriminants/DL2Tool.h

 // for text editors: this file is -*- C++ -*-
 /*
-  Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration
+  Copyright (C) 2002-2021 CERN for the benefit of the ATLAS collaboration
 */
 
 #ifndef DL2_TOOL_H
@@ -8,6 +8,7 @@
 
 #include "AsgTools/AsgTool.h"
 #include "FlavorTagDiscriminants/IBTagDecorator.h"
+#include "FlavorTagDiscriminants/IJetTagDecorator.h"
 
 namespace FlavorTagDiscriminants {
 
@@ -17,11 +18,14 @@ namespace FlavorTagDiscriminants {
     std::string nnFile;
     std::string flipTagConfig;
     std::map<std::string,std::string> variableRemapping;
+    std::string trackLinkType;
   };
 
-  class DL2Tool : public asg::AsgTool, virtual public IBTagDecorator
+  class DL2Tool : public asg::AsgTool,
+                  virtual public IBTagDecorator,
+                  virtual public IJetTagDecorator
   {
-    ASG_TOOL_CLASS(DL2Tool, IBTagDecorator )
+    ASG_TOOL_CLASS2(DL2Tool, IBTagDecorator, IJetTagDecorator )
   public:
     DL2Tool(const std::string& name);
     ~DL2Tool();
@@ -29,7 +33,8 @@ namespace FlavorTagDiscriminants {
     StatusCode initialize() override;
 
     // returns 0 for success
-    virtual void decorate(const xAOD::BTagging& jet) const override;
+    virtual void decorate(const xAOD::BTagging& btag) const override;
+    virtual void decorate(const xAOD::Jet& jet) const override;
 
     virtual std::set<std::string> getDecoratorKeys() const override;
     virtual std::set<std::string> getAuxInputKeys() const override;

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/FlavorTagDiscriminants/DecoratorAlg.h

+/*
+  Copyright (C) 2002-2021 CERN for the benefit of the ATLAS collaboration
+*/
+
+#ifndef DECORATOR_ALG_H
+#define DECORATOR_ALG_H
+
+#include "AthenaBaseComps/AthReentrantAlgorithm.h"
+
+#include "StoreGate/WriteDecorHandleKey.h"
+#include "StoreGate/ReadDecorHandleKey.h"
+
+namespace FlavorTagDiscriminants {
+  template <typename CONTAINER, typename DECORATOR, typename CONSTITUENTS>
+  class DecoratorAlg : public AthReentrantAlgorithm
+  {
+  public:
+    DecoratorAlg(const std::string& name, ISvcLocator* svcloc);
+    virtual StatusCode initialize() override;
+    virtual StatusCode execute(const EventContext& cxt) const override;
+    virtual StatusCode finalize() override;
+  private:
+    SG::ReadHandleKey<CONTAINER> m_containerKey {
+      this, "container", "", "Key for the input collection"};
+    SG::ReadHandleKey<CONSTITUENTS> m_constituentKey {
+      this, "constituentContainer", "",
+        "Key for track inputs container"};
+    Gaudi::Property<std::vector<std::string>> m_undeclaredReadDecorKeys {
+      this, "undeclaredReadDecorKeys", {},
+      "List of read handles that we don't read, e.g. static variables" };
+
+    ToolHandle<DECORATOR> m_decorator{
+      this, "decorator", "", "Decorator tool"};
+
+    // Keys to keep track of the inputs / outputs
+    std::vector<SG::ReadDecorHandleKey<CONTAINER>> m_aux;
+    std::vector<SG::ReadDecorHandleKey<CONSTITUENTS>> m_constituentAux;
+    std::vector<SG::WriteDecorHandleKey<CONTAINER>> m_decor;
+  };
+}
+
+#include "DecoratorAlg.icc"
+
+#endif

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/FlavorTagDiscriminants/DecoratorAlg.icc

+/*
+  Copyright (C) 2002-2021 CERN for the benefit of the ATLAS collaboration
+*/
+
+#include "StoreGate/ReadDecorHandle.h"
+
+#include <exception>
+
+namespace FlavorTagDiscriminants {
+  template <typename C, typename D, typename N>
+  DecoratorAlg<C,D,N>::DecoratorAlg(
+    const std::string& name, ISvcLocator* svcloc):
+    AthReentrantAlgorithm(name, svcloc)
+  {
+  }
+
+
+  template <typename C, typename D, typename N>
+  StatusCode DecoratorAlg<C,D,N>::initialize() {
+    ATH_CHECK(m_containerKey.initialize());
+    ATH_CHECK(m_constituentKey.initialize());
+
+    ATH_CHECK(m_decorator.retrieve());
+
+    std::set<std::string> veto(
+      m_undeclaredReadDecorKeys.begin(),
+      m_undeclaredReadDecorKeys.end());
+
+    // now we build data dependencies from the internal tools. We have
+    // to reserve the vectors here to prevent a segfault since read /
+    // write handles aren't movable once declaired as a property.
+    m_aux.reserve(m_decorator->getAuxInputKeys().size());
+    for (const std::string& key: m_decorator->getAuxInputKeys()) {
+      const std::string full = m_containerKey.key() + "." + key;
+      if (veto.count(full)) {
+        ATH_MSG_DEBUG("Not declaring accessor: " + full);
+        continue;
+      }
+      ATH_MSG_DEBUG("Adding accessor: " + full);
+      m_aux.emplace_back(this, key, full, "");
+      ATH_CHECK(m_aux.back().initialize());
+    }
+    m_constituentAux.reserve(m_decorator->getConstituentAuxInputKeys().size());
+    for (const std::string& key: m_decorator->getConstituentAuxInputKeys()) {
+      const std::string full = m_constituentKey.key() + "." + key;
+      if (veto.count(full)) {
+        ATH_MSG_DEBUG("Not declaring accessor: " + full);
+        continue;
+      }
+      ATH_MSG_DEBUG("Adding constituent accessor: " + full);
+      m_constituentAux.emplace_back(this, key, full, "");
+      ATH_CHECK(m_constituentAux.back().initialize());
+    }
+    m_decor.reserve(m_decorator->getDecoratorKeys().size());
+    for (const std::string& key: m_decorator->getDecoratorKeys()) {
+      const std::string full = m_containerKey.key() + "." + key;
+      ATH_MSG_DEBUG("Adding decorator: " + full);
+      m_decor.emplace_back(this, key, full, "");
+      ATH_CHECK(m_decor.back().initialize());
+    }
+    ATH_MSG_DEBUG("Finished setting up");
+    return StatusCode::SUCCESS;
+  }
+
+
+  template <typename C, typename D, typename N>
+  StatusCode DecoratorAlg<C,D,N>::execute(const EventContext& cxt ) const {
+    SG::ReadHandle<C> container(
+      m_containerKey, cxt);
+    if (!container.isValid()) {
+      ATH_MSG_ERROR("no container " << container.key());
+      return StatusCode::FAILURE;
+    }
+    ATH_MSG_DEBUG(
+      "Decorating " + std::to_string(container->size()) + " elements");
+    for (const auto* element: *container) {
+      m_decorator->decorate(*element);
+    }
+    return StatusCode::SUCCESS;
+  }
+
+  template <typename C, typename D, typename N>
+  StatusCode DecoratorAlg<C,D,N>::finalize() {
+    return StatusCode::SUCCESS;
+  }
+
+}

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/FlavorTagDiscriminants/IBTagDecorator.h

@@ -6,10 +6,14 @@
 #ifndef I_BTAG_DECORATOR_H
 #define I_BTAG_DECORATOR_H
 
+#include "IDependencyReporter.h"
+
 #include "AsgTools/IAsgTool.h"
 #include "xAODBTagging/BTaggingFwd.h"
 
-class IBTagDecorator : virtual public asg::IAsgTool {
+class IBTagDecorator : virtual public asg::IAsgTool,
+                       virtual public IDependencyReporter
+{
 ASG_TOOL_INTERFACE(IBTagDecorator)
 
 public:
@@ -19,11 +23,6 @@ public:
 
   /// Method to decorate a jet.
   virtual void decorate(const xAOD::BTagging& btag) const = 0;
-
-  // Names of the decorations being added
-  virtual std::set<std::string> getDecoratorKeys() const = 0;
-  virtual std::set<std::string> getAuxInputKeys() const = 0;
-  virtual std::set<std::string> getConstituentAuxInputKeys() const = 0;
 };
 
 

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/FlavorTagDiscriminants/IDependencyReporter.h

+// for text editors: this file is -*- C++ -*-
+/*
+  Copyright (C) 2002-2021 CERN for the benefit of the ATLAS collaboration
+*/
+
+#ifndef I_DEPENDENCY_REPORTER
+#define I_DEPENDENCY_REPORTER
+
+#include <set>
+#include <string>
+
+class IDependencyReporter {
+public:
+  /// Destructor.
+  virtual ~IDependencyReporter() { };
+
+  // Names of the decorations being added
+  virtual std::set<std::string> getDecoratorKeys() const = 0;
+  virtual std::set<std::string> getAuxInputKeys() const = 0;
+  virtual std::set<std::string> getConstituentAuxInputKeys() const = 0;
+};
+
+#endif

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/FlavorTagDiscriminants/IJetTagDecorator.h

+// for text editors: this file is -*- C++ -*-
+/*
+  Copyright (C) 2002-2021 CERN for the benefit of the ATLAS collaboration
+*/
+
+#ifndef I_JETTAG_DECORATOR_H
+#define I_JETTAG_DECORATOR_H
+
+#include "IDependencyReporter.h"
+
+#include "AsgTools/IAsgTool.h"
+#include "xAODJet/JetFwd.h"
+
+class IJetTagDecorator : virtual public asg::IAsgTool,
+                         virtual public IDependencyReporter
+{
+ASG_TOOL_INTERFACE(IJetTagDecorator)
+
+public:
+
+  /// Destructor.
+  virtual ~IJetTagDecorator() { };
+
+  /// Method to decorate a jet.
+  virtual void decorate(const xAOD::Jet& jet) const = 0;
+
+};
+
+
+#endif

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/FlavorTagDiscriminants/JetTagDecoratorAlg.h

+/*
+  Copyright (C) 2002-2021 CERN for the benefit of the ATLAS collaboration
+*/
+
+#ifndef JET_TAG_DECORATOR_ALG_H
+#define JET_TAG_DECORATOR_ALG_H
+
+#include "FlavorTagDiscriminants/DecoratorAlg.h"
+#include "FlavorTagDiscriminants/IJetTagDecorator.h"
+
+#include "xAODJet/JetContainer.h"
+#include "xAODTracking/TrackParticleContainer.h"
+
+namespace detail {
+  using JetTag_t = FlavorTagDiscriminants::DecoratorAlg<
+    xAOD::JetContainer,
+    IJetTagDecorator,
+    xAOD::TrackParticleContainer
+    >;
+}
+
+namespace FlavorTagDiscriminants {
+  class JetTagDecoratorAlg : public detail::JetTag_t
+  {
+  public:
+    JetTagDecoratorAlg(const std::string& name, ISvcLocator* svcloc);
+  };
+}
+
+#endif

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/FlavorTagDiscriminants/customGetter.h

@@ -11,7 +11,7 @@
 // EDM includes
 #include "xAODJet/JetFwd.h"
 #include "xAODTracking/TrackParticleFwd.h"
-#include "xAODBTagging/BTaggingFwd.h"
+#include "AthContainers/AuxElement.h"
 
 #include <functional>
 #include <string>
@@ -54,7 +54,7 @@ namespace FlavorTagDiscriminants {
 
   // internal functions
   namespace internal {
-    std::function<std::pair<std::string, double>(const xAOD::BTagging&)>
+    std::function<std::pair<std::string, double>(const SG::AuxElement&)>
     customGetterAndName(const std::string&);
 
     std::pair<std::function<std::pair<std::string, std::vector<double>>(

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/Root/AssociationEnums.cxx

+/*
+  Copyright (C) 2002-2021 CERN for the benefit of the ATLAS collaboration
+*/
+
+
+#include "FlavorTagDiscriminants/AssociationEnums.h"
+
+#include <stdexcept>
+
+namespace FlavorTagDiscriminants{
+
+#define RETURN_CONFIG(cfg) \
+  if (name == std::string(#cfg)) return TrackLinkType::cfg
+
+  TrackLinkType trackLinkTypeFromString(const std::string& name) {
+    RETURN_CONFIG(TRACK_PARTICLE);
+    RETURN_CONFIG(IPARTICLE);
+    throw std::logic_error("DL2 association scheme '" + name + "' unknown");
+  }
+
+#undef RETURN_CONFIG
+
+}

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/Root/DL2.cxx

@@ -20,6 +20,7 @@ namespace FlavorTagDiscriminants {
     track_prefix = "btagIp_";
     flip = FlipTagConfig::STANDARD;
     track_link_name = "BTagTrackToJetAssociator";
+    track_link_type = TrackLinkType::TRACK_PARTICLE;
   }
 
   // DL2
@@ -109,6 +110,18 @@ namespace FlavorTagDiscriminants {
   }
 
   void DL2::decorate(const xAOD::BTagging& btag) const {
+    auto jetLink = m_jetLink(btag);
+    if (!jetLink.isValid()) {
+      throw std::runtime_error("invalid jetLink");
+    }
+    const xAOD::Jet& jet = **jetLink;
+    decorate(jet, btag);
+  }
+  void DL2::decorate(const xAOD::Jet& jet) const {
+    decorate(jet, jet);
+  }
+
+  void DL2::decorate(const xAOD::Jet& jet, const SG::AuxElement& btag) const {
     using namespace internal;
     std::vector<NamedVar> vvec;
     for (const auto& getter: m_varsFromBTag) {
@@ -126,11 +139,6 @@ namespace FlavorTagDiscriminants {
     }
 
     // add track sequences
-    auto jetLink = m_jetLink(btag);
-    if (!jetLink.isValid()) {
-      throw std::runtime_error("invalid jetLink");
-    }
-    const xAOD::Jet& jet = **jetLink;
     std::map<std::string,std::map<std::string, std::vector<double>>> seqs;
     for (const auto& builder: m_trackSequenceBuilders) {
       Tracks sorted_tracks = builder.tracksFromJet(jet, btag);
@@ -174,19 +182,51 @@ namespace FlavorTagDiscriminants {
     TracksFromJet::TracksFromJet(SortOrder order,
                                  TrackSelection selection,
                                  const DL2Options& options):
-      m_trackAssociator(options.track_link_name),
       m_trackSortVar(get::trackSortVar(order, options)),
       m_trackFilter(get::trackFilter(selection, options).first)
     {
+      // We have several ways to get tracks: either we retrieve an
+      // IParticleContainer and cast the pointers to TrackParticle, or
+      // we retrieve a TrackParticleContainer directly. Unfortunately
+      // the way tracks are stored isn't consistent across the EDM, so
+      // we allow configuration for both setups.
+      //
+      if (options.track_link_type == TrackLinkType::IPARTICLE) {
+        SG::AuxElement::ConstAccessor<PartLinks> acc(options.track_link_name);
+        m_associator = [acc](const SG::AuxElement& btag) -> TPV {
+          TPV tracks;
+          for (const ElementLink<IPC>& link: acc(btag)) {
+            if (!link.isValid()) {
+              throw std::logic_error("invalid particle link");
+            }
+            const auto* trk = dynamic_cast<const xAOD::TrackParticle*>(*link);
+            if (!trk) {
+              throw std::logic_error("iparticle does not cast to Track");
+            }
+            tracks.push_back(trk);
+          }
+          return tracks;
+        };
+      } else if (options.track_link_type == TrackLinkType::TRACK_PARTICLE){
+        SG::AuxElement::ConstAccessor<TrackLinks> acc(options.track_link_name);
+        m_associator = [acc](const SG::AuxElement& btag) -> TPV {
+          TPV tracks;
+          for (const ElementLink<TPC>& link: acc(btag)) {
+            if (!link.isValid()) {
+              throw std::logic_error("invalid track link");
+            }
+            tracks.push_back(*link);
+          }
+          return tracks;
+        };
+      } else {
+        throw std::logic_error("Unknown TrackLinkType");
+      }
     }
     Tracks TracksFromJet::operator()(const xAOD::Jet& jet,
-                                     const xAOD::BTagging& btagging) const {
+                                     const SG::AuxElement& btagging) const {
       std::vector<std::pair<double, const xAOD::TrackParticle*>> tracks;
-      for (const auto &link : m_trackAssociator(btagging)) {
-        if(!link.isValid()) {
-          throw std::logic_error("invalid track link");
-        }
-        const xAOD::TrackParticle *tp = *link;
+      for (const xAOD::TrackParticle *tp : m_associator(btagging)) {
         if (m_trackFilter(tp)) {
           tracks.push_back({m_trackSortVar(tp, jet), tp});
         };

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/Root/DL2HighLevel.cxx

@@ -31,7 +31,8 @@ namespace FlavorTagDiscriminants {
 
   DL2HighLevel::DL2HighLevel(const std::string& nn_file_name,
                              FlipTagConfig flip_config,
-                             std::map<std::string,std::string> remap_scalar):
+                             std::map<std::string,std::string> remap_scalar,
+                             TrackLinkType track_link_type):
     m_dl2(nullptr)
   {
     // get the graph
@@ -168,6 +169,7 @@ namespace FlavorTagDiscriminants {
     }
     options.flip = flip_config;
     options.remap_scalar = remap_scalar;
+    options.track_link_type = track_link_type;
 
     m_dl2.reset(
       new DL2(
@@ -184,6 +186,9 @@ namespace FlavorTagDiscriminants {
   void DL2HighLevel::decorate(const xAOD::BTagging& btag) const {
     m_dl2->decorate(btag);
   }
+  void DL2HighLevel::decorate(const xAOD::Jet& jet) const {
+    m_dl2->decorate(jet);
+  }
 
   DL2DataDependencyNames DL2HighLevel::getDataDependencyNames() const
   {

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/Root/DL2Tool.cxx

@@ -15,6 +15,7 @@ namespace FlavorTagDiscriminants {
     declareProperty("nnFile", m_props.nnFile);
     declareProperty("flipTagConfig", m_props.flipTagConfig);
     declareProperty("variableRemapping", m_props.variableRemapping);
+    declareProperty("trackLinkType", m_props.trackLinkType);
   }
   DL2Tool::~DL2Tool() {}
 
@@ -24,16 +25,26 @@ namespace FlavorTagDiscriminants {
     if (m_props.flipTagConfig.size() > 0) {
       flipConfig = flipTagConfigFromString(m_props.flipTagConfig);
     }
+    TrackLinkType trackLinkType = TrackLinkType::TRACK_PARTICLE;
+    if (m_props.trackLinkType.size() > 0) {
+      trackLinkType = trackLinkTypeFromString(m_props.trackLinkType);
+    }
     m_dl2.reset(
       new DL2HighLevel(
         m_props.nnFile,
         flipConfig,
-        m_props.variableRemapping)
+        m_props.variableRemapping,
+        trackLinkType
+        )
       );
     return StatusCode::SUCCESS;
   }
 
-  void DL2Tool::decorate(const xAOD::BTagging& jet) const {
+  void DL2Tool::decorate(const xAOD::BTagging& btag) const {
+    ATH_MSG_DEBUG("Decoration from: " + m_props.nnFile);
+    m_dl2->decorate(btag);
+  }
+  void DL2Tool::decorate(const xAOD::Jet& jet) const {
     ATH_MSG_DEBUG("Decoration from: " + m_props.nnFile);
     m_dl2->decorate(jet);
   }

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/Root/customGetter.cxx

@@ -4,9 +4,8 @@
 #include "FlavorTagDiscriminants/customGetter.h"
 #include "FlavorTagDiscriminants/BTagTrackIpAccessor.h"
 
-#include "xAODJet/Jet.h"
+#include "xAODJet/JetContainer.h"
 #include "xAODTracking/TrackParticle.h"
-#include "xAODBTagging/BTagging.h"
 
 #include <optional>
 
@@ -17,12 +16,12 @@ namespace {
   // this function is not at all optimized, but then it doesn't have
   // to be since it should only be called in the initialization stage.
   //
-  std::function<double(const xAOD::BTagging&)> customGetter(
+  std::function<double(const SG::AuxElement&)> customGetter(
     const std::string& name)
   {
     using JL = ElementLink<xAOD::JetContainer>;
     SG::AuxElement::ConstAccessor<JL> jl("jetLink");
-    auto jg = [jl](const xAOD::BTagging& btag) -> const xAOD::Jet& {
+    auto jg = [jl](const SG::AuxElement& btag) -> const xAOD::Jet& {
       auto link = jl(btag);
       if (!link.isValid()) {
         throw std::runtime_error("invalid jetLink");
@@ -30,10 +29,10 @@ namespace {
       return **link;
     };
     if (name == "pt") {
-      return [jg](const xAOD::BTagging& b) {return jg(b).pt();};
+      return [jg](const SG::AuxElement& b) {return jg(b).pt();};
     }
     if (name == "abs_eta") {
-      return [jg](const xAOD::BTagging& b) {return std::abs(jg(b).eta());};
+      return [jg](const SG::AuxElement& b) {return std::abs(jg(b).eta());};
     }
     throw std::logic_error("no match for custom getter " + name);
   }
@@ -185,10 +184,10 @@ namespace FlavorTagDiscriminants {
     // which returns the pair we wanted.
     //
     // Case for jet variables
-    std::function<std::pair<std::string, double>(const xAOD::BTagging&)>
+    std::function<std::pair<std::string, double>(const SG::AuxElement&)>
     customGetterAndName(const std::string& name) {
       auto getter = customGetter(name);
-      return [name, getter](const xAOD::BTagging& j) {
+      return [name, getter](const SG::AuxElement& j) {
                return std::make_pair(name, getter(j));
              };
     }

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/src/BTagDecoratorAlg.cxx

@@ -3,85 +3,11 @@
 */
 
 #include "FlavorTagDiscriminants/BTagDecoratorAlg.h"
-#include "xAODBTagging/BTaggingContainer.h"
-#include "xAODTracking/TrackParticleContainer.h"
-#include "StoreGate/ReadDecorHandle.h"
-
-#include <exception>
 
 namespace FlavorTagDiscriminants {
   BTagDecoratorAlg::BTagDecoratorAlg(
     const std::string& name, ISvcLocator* svcloc):
-    AthReentrantAlgorithm(name, svcloc)
+    detail::BTag_t(name, svcloc)
   {
   }
-
-
-  StatusCode BTagDecoratorAlg::initialize() {
-    ATH_CHECK(m_btagContainerKey.initialize());
-    ATH_CHECK(m_trackContainerKey.initialize());
-
-    ATH_CHECK(m_decorator.retrieve());
-
-    std::set<std::string> veto(
-      m_undeclaredReadDecorKeys.begin(),
-      m_undeclaredReadDecorKeys.end());
-
-    // now we build data dependencies from the internal tools. We have
-    // to reserve the vectors here to prevent a segfault since read /
-    // write handles aren't movable once declaired as a property.
-    m_btagAux.reserve(m_decorator->getAuxInputKeys().size());
-    for (const std::string& key: m_decorator->getAuxInputKeys()) {
-      const std::string full = m_btagContainerKey.key() + "." + key;
-      if (veto.count(full)) {
-        ATH_MSG_DEBUG("Not declaring accessor: " + full);
-        continue;
-      }
-      ATH_MSG_DEBUG("Adding btag accessor: " + full);
-      m_btagAux.emplace_back(this, key, full, "");
-      ATH_CHECK(m_btagAux.back().initialize());
-    }
-    m_trkAux.reserve(m_decorator->getConstituentAuxInputKeys().size());
-    for (const std::string& key: m_decorator->getConstituentAuxInputKeys()) {
-      const std::string full = m_trackContainerKey.key() + "." + key;
-      if (veto.count(full)) {
-        ATH_MSG_DEBUG("Not declaring accessor: " + full);
-        continue;
-      }
-      ATH_MSG_DEBUG("Adding track accessor: " + full);
-      m_trkAux.emplace_back(this, key, full, "");
-      ATH_CHECK(m_trkAux.back().initialize());
-    }
-    m_btagDecor.reserve(m_decorator->getDecoratorKeys().size());
-    for (const std::string& key: m_decorator->getDecoratorKeys()) {
-      const std::string full = m_btagContainerKey.key() + "." + key;
-      ATH_MSG_DEBUG("Adding btag decorator: " + full);
-      m_btagDecor.emplace_back(this, key, full, "");
-      ATH_CHECK(m_btagDecor.back().initialize());
-    }
-    ATH_MSG_DEBUG("Finished setting up");
-    return StatusCode::SUCCESS;
-  }
-
-
-  StatusCode BTagDecoratorAlg::execute(const EventContext& cxt ) const {
-    SG::ReadHandle<xAOD::BTaggingContainer> btagContainer(
-      m_btagContainerKey, cxt);
-    if (!btagContainer.isValid()) {
-      ATH_MSG_ERROR("no b-tagging container " << btagContainer.key());
-      return StatusCode::FAILURE;
-    }
-    ATH_MSG_DEBUG(
-      "Decorating " + std::to_string(btagContainer->size()) + " btags");
-    for (const auto* btag: *btagContainer) {
-      m_decorator->decorate(*btag);
-    }
-    return StatusCode::SUCCESS;
-  }
-
-
-  StatusCode BTagDecoratorAlg::finalize() {
-    return StatusCode::SUCCESS;
-  }
-
 }

PhysicsAnalysis/JetTagging/FlavorTagDiscriminants/src/JetTagDecoratorAlg.cxx

+/*
+  Copyright (C) 2002-2021 CERN for the benefit of the ATLAS collaboration
+*/
+
+#include "FlavorTagDiscriminants/JetTagDecoratorAlg.h"
+
+namespace FlavorTagDiscriminants {
+  JetTagDecoratorAlg::JetTagDecoratorAlg(
+    const std::string& name, ISvcLocator* svcloc):
+    detail::JetTag_t(name, svcloc)
+  {
+  }
+}