diff --git a/Reconstruction/Jet/JetRecTools/JetRecTools/JetRecToolsDict.h b/Reconstruction/Jet/JetRecTools/JetRecTools/JetRecToolsDict.h
index 25c75195c70fb9157d393deb7fb75614f6be096b..201383c1b63db3e6755c49134f0a25bba915c121 100644
--- a/Reconstruction/Jet/JetRecTools/JetRecTools/JetRecToolsDict.h
+++ b/Reconstruction/Jet/JetRecTools/JetRecTools/JetRecToolsDict.h
@@ -25,5 +25,7 @@
#include "JetRecTools/CorrectPFOTool.h"
#include "JetRecTools/PuppiWeightTool.h"
#include "JetRecTools/ChargedHadronSubtractionTool.h"
+#include "JetRecTools/TARJetTool.h"
+#include "JetRecTools/SATScaleTool.h"
#endif
diff --git a/Reconstruction/Jet/JetRecTools/JetRecTools/SATScaleTool.h b/Reconstruction/Jet/JetRecTools/JetRecTools/SATScaleTool.h
new file mode 100644
index 0000000000000000000000000000000000000000..2253c4442a47cfc8e9cd14e622f2bc131c54363c
--- /dev/null
+++ b/Reconstruction/Jet/JetRecTools/JetRecTools/SATScaleTool.h
@@ -0,0 +1,69 @@
+/*
+ Copyright (C) 2002-2018 CERN for the benefit of the ATLAS collaboration
+*/
+
+// SATScaleTool.h
+
+#ifndef SATSCALETOOL_H
+#define SATSCALETOOL_H
+
+// Jason Veatch (jason.veatch@cern.ch)
+// November 2018
+//
+// Tool to produce rescaled tracks to be used as inputs to SAT jets
+// The output is a xAOD::TrackParticle collection
+
+#include "JetInterface/IJetExecuteTool.h"
+#include "AsgTools/AsgTool.h"
+
+#include "AthContainers/ConstDataVector.h"
+
+#include "JetRecTools/TrackAssistTool.h"
+
+class SATScaleTool
+: public IJetExecuteTool,
+ public asg::AsgTool,
+ public TrackAssistTool
+{
+ ASG_TOOL_CLASS(SATScaleTool, IJetExecuteTool)
+
+ public:
+
+ // Constructor
+ SATScaleTool(const std::string& myname = "SATScaleTool" );
+
+ // Initialize function
+ StatusCode initialize() override;
+
+ // Print all configurable parameters
+ void print() const override;
+
+ // Execute the tool
+ int execute() const override;
+
+ private:
+
+ // Produce collection of rescaled tracks
+ StatusCode makeSATTracks() const;
+
+ // Local method to retrieve the track vertex association
+ StatusCode getTrackVertexAssociation( const jet::TrackVertexAssociation *&tva ) const override;
+
+ // Local method to retrieve the primary vertex
+ StatusCode getPrimaryVertex( const xAOD::Vertex *&pvx ) const override;
+
+ // If jet merging is turned on, flag jets that meet the merging criteria
+ StatusCode flagJetsToMerge( const xAOD::JetContainer *jets ) const;
+
+ // Configurable parameters
+ std::string m_inTrackColl;
+ std::string m_inJetColl;
+ std::string m_outTrackColl;
+ float m_dRmatch;
+ bool m_doMerge;
+ float m_dRmerge;
+ float m_pTdRmerge;
+
+};
+
+#endif // SATSCALETOOL_H
diff --git a/Reconstruction/Jet/JetRecTools/JetRecTools/TARJetTool.h b/Reconstruction/Jet/JetRecTools/JetRecTools/TARJetTool.h
new file mode 100644
index 0000000000000000000000000000000000000000..d5681a3b09399b634def936505a5c47008a47cc6
--- /dev/null
+++ b/Reconstruction/Jet/JetRecTools/JetRecTools/TARJetTool.h
@@ -0,0 +1,60 @@
+/*
+ Copyright (C) 2002-2018 CERN for the benefit of the ATLAS collaboration
+*/
+
+// TARJetTool.h
+
+#ifndef TARJETTOOL_H
+#define TARJETTOOL_H
+
+// Jason Veatch (jason.veatch@cern.ch)
+// November 2018
+//
+// Tool to add rescaled tracks to reclustered jets in order to add TAR jet substructure moments
+// The output is a new collection of xAOD::TrackParticle and ElementLinks for each jet
+
+#include "JetInterface/IJetModifier.h"
+#include "AsgTools/AsgTool.h"
+
+#include "AthContainers/ConstDataVector.h"
+
+#include "JetRecTools/TrackAssistTool.h"
+
+class TARJetTool
+: public IJetModifier,
+ public asg::AsgTool,
+ public TrackAssistTool
+{
+ ASG_TOOL_CLASS(TARJetTool, IJetModifier)
+
+ public:
+
+ // Constructor
+ TARJetTool(const std::string& myname = "TARJetTool" );
+
+ // Initialize function
+ StatusCode initialize() override;
+
+ // Print all configurable parameters
+ void print() const override;
+
+ // Modify jet collection
+ int modify( xAOD::JetContainer& inJets ) const override;
+
+ private:
+
+ // Local method to retrieve the track vertex association
+ StatusCode getTrackVertexAssociation( const jet::TrackVertexAssociation *&tva ) const override;
+
+ // Local method to retrieve the primary vertex
+ StatusCode getPrimaryVertex( const xAOD::Vertex *&pvx ) const override;
+
+ // Configurable parameters
+ std::string m_inTrackColl;
+ std::string m_outTrackColl;
+ std::string m_assocTracksOutName;
+ float m_dRmatch;
+
+};
+
+#endif // TARJETTOOL_H
diff --git a/Reconstruction/Jet/JetRecTools/JetRecTools/TrackAssistTool.h b/Reconstruction/Jet/JetRecTools/JetRecTools/TrackAssistTool.h
new file mode 100644
index 0000000000000000000000000000000000000000..c3494727888327277018e83443a9f7616db9797c
--- /dev/null
+++ b/Reconstruction/Jet/JetRecTools/JetRecTools/TrackAssistTool.h
@@ -0,0 +1,84 @@
+/*
+ Copyright (C) 2002-2018 CERN for the benefit of the ATLAS collaboration
+*/
+
+// TrackAssist.h
+
+#ifndef TRACKASSISTTOOL_H
+#define TRACKASSISTTOOL_H
+
+// Jason Veatch (jason.veatch@cern.ch)
+// November 2018
+//
+// Base class for tools that implement track-assisted jet substructure algorithms
+
+#include "AsgTools/ToolHandle.h"
+#include "AsgTools/Check.h"
+#include "AsgTools/MessageCheck.h"
+
+#include "xAODCore/ShallowCopy.h"
+
+#include "xAODJet/Jet.h"
+#include "xAODJet/JetContainer.h"
+#include "xAODJet/JetAuxContainer.h"
+#include "xAODTracking/TrackParticle.h"
+#include "xAODTracking/TrackParticleContainer.h"
+#include "xAODTracking/TrackParticleAuxContainer.h"
+
+#include "xAODTracking/VertexContainer.h"
+#include "JetInterface/IJetTrackSelector.h"
+#include "JetEDM/TrackVertexAssociation.h"
+
+#include "xAODBase/IParticleHelpers.h"
+
+using namespace asg::msgUserCode;
+
+class TrackAssistTool
+{
+
+ public:
+
+ // Destructor.
+ virtual ~TrackAssistTool() { }
+
+ // Rescale tracks using jets with pre-determined scaling weights
+ StatusCode rescaleTracks( const xAOD::JetContainer *jets, xAOD::TrackParticleContainer *tracks) const;
+
+ protected:
+
+ // Configurable parameters
+ std::string m_assocTracksInName;
+ std::string m_vertexColl;
+ std::string m_trackVtxAssoc;
+ bool m_doTrackVtxAssoc;
+ ToolHandle<IJetTrackSelector> m_trackSelTool;
+
+ // Local method to check if track passes selection and vertex association criteria
+ bool isGoodTrack( const xAOD::TrackParticle &track, const xAOD::Vertex &pvx, const jet::TrackVertexAssociation &tva ) const;
+
+ // Constructor
+ TrackAssistTool();
+
+ // Initialize function
+ virtual StatusCode initialize();
+
+ // Print configured parameters
+ virtual void print() const;
+
+ // Local method to retrieve all vertex information
+ StatusCode getVertexInfo( const xAOD::Vertex *&pvx, const jet::TrackVertexAssociation *&tva ) const;
+
+ // Local method to retrieve the track vertex association
+ virtual StatusCode getTrackVertexAssociation( const jet::TrackVertexAssociation *&tva ) const = 0;
+
+ // Local method to retrieve the primary vertex
+ virtual StatusCode getPrimaryVertex( const xAOD::Vertex *&pvx ) const = 0;
+
+ private:
+
+ // Local method to calculate the factors to rescale each track
+ StatusCode getRescaleFactors( const xAOD::Jet &jet, xAOD::TrackParticleContainer &tracks) const;
+
+};
+
+#endif // TRACKASSISTTOOL_H
diff --git a/Reconstruction/Jet/JetRecTools/JetRecTools/selection.xml b/Reconstruction/Jet/JetRecTools/JetRecTools/selection.xml
index 624f405217aba358433a07d15e055304d8550fad..dbd65294113501eefc89b64a5ede3f0588891593 100644
--- a/Reconstruction/Jet/JetRecTools/JetRecTools/selection.xml
+++ b/Reconstruction/Jet/JetRecTools/JetRecTools/selection.xml
@@ -15,5 +15,7 @@
<class name="TrackPseudoJetGetter"/>
<class name="TrackVertexAssociationTool"/>
<class name="VoronoiWeightTool"/>
+<class name="TARJetTool"/>
+<class name="SATScaleTool"/>
</lcgdict>
diff --git a/Reconstruction/Jet/JetRecTools/Root/SATScaleTool.cxx b/Reconstruction/Jet/JetRecTools/Root/SATScaleTool.cxx
new file mode 100644
index 0000000000000000000000000000000000000000..f31c142bc7396716ec2969b1e045bf328dee4d54
--- /dev/null
+++ b/Reconstruction/Jet/JetRecTools/Root/SATScaleTool.cxx
@@ -0,0 +1,318 @@
+/*
+ Copyright (C) 2002-2018 CERN for the benefit of the ATLAS collaboration
+*/
+
+// SATScaleTool.cxx
+
+#include "JetRecTools/SATScaleTool.h"
+
+SATScaleTool::SATScaleTool(const std::string& myname)
+: AsgTool(myname),
+ TrackAssistTool(),
+ m_inTrackColl(""),
+ m_inJetColl(""),
+ m_outTrackColl(""),
+ m_dRmatch(0.2),
+ m_doMerge(true),
+ m_dRmerge(0.35),
+ m_pTdRmerge(0.3)
+{
+ declareProperty("InputTrackContainer", m_inTrackColl = "InDetTrackParticles");
+ declareProperty("InputJetContainer", m_inJetColl);
+ declareProperty("OutputTrackContainer", m_outTrackColl);
+ declareProperty("MatchDeltaR", m_dRmatch);
+ declareProperty("DoMergeJets", m_doMerge);
+ declareProperty("MergeDeltaR", m_dRmerge);
+ declareProperty("MergePtOverDeltaR", m_pTdRmerge);
+
+ declareProperty("InputAssociatedTracks", m_assocTracksInName = "GhostTrack");
+ declareProperty("VertexContainer", m_vertexColl = "PrimaryVertices");
+ declareProperty("TrackVertexAssociation", m_trackVtxAssoc);
+ declareProperty("DoTrackVertexAssociation", m_doTrackVtxAssoc = true);
+ declareProperty("TrackSelTool", m_trackSelTool);
+}
+
+StatusCode SATScaleTool::initialize()
+{
+
+ ATH_MSG_INFO("Initializing SATScaleTool " << name() << ".");
+ ATH_CHECK( TrackAssistTool::initialize() );
+ print();
+
+ return StatusCode::SUCCESS;
+}
+
+int SATScaleTool::execute() const {
+
+ if( makeSATTracks().isFailure() ) return 1;
+
+ return 0;
+}
+
+StatusCode SATScaleTool::makeSATTracks() const {
+
+ std::vector<int> matchedTrackIndices; // List of ghost matched track indices
+ std::vector<int> unmatchedTrackIndices; // List of non ghost matched track indices
+ std::vector<int> allGoodTrackIndices; // List of all good track indices
+
+ std::map<int, std::vector<int>> JetTrackMap; // Map of track indices to jet indices
+ std::map<int, std::vector<int> > JetMergeMap; // Map of jet indices for merging
+
+ // Get input jets
+ const xAOD::JetContainer *inJets = nullptr;
+ ATH_CHECK( evtStore()->retrieve(inJets,m_inJetColl) );
+
+ // Get input tracks
+ const xAOD::TrackParticleContainer *inTracks = nullptr;
+ ATH_CHECK( evtStore()->retrieve(inTracks,m_inTrackColl) );
+
+ // Get vertex objects
+ const xAOD::Vertex *pvx = nullptr;
+ const jet::TrackVertexAssociation *tva = nullptr;
+ ATH_CHECK( getVertexInfo(pvx,tva) );
+
+ // Make shallow copy of tracks
+ auto trackShallowCopy = xAOD::shallowCopyContainer(*inTracks);
+ std::unique_ptr<xAOD::TrackParticleContainer> outTracks (trackShallowCopy.first);
+ std::unique_ptr<xAOD::ShallowAuxContainer> outTracksAux (trackShallowCopy.second);
+
+ // Get list of all tracks
+ for( auto track : *inTracks ) {
+ // Make sure only tracks that pass selection are used
+ if( !isGoodTrack(*track,*pvx,*tva) ) continue;
+ allGoodTrackIndices.push_back( track->index() );
+ }
+
+ // Flag jets for merging
+ ATH_CHECK( flagJetsToMerge(inJets) );
+
+ // Loop over input jets
+ for( auto jet : *inJets ) {
+
+ // Get ghost-associated tracks
+ std::vector<const xAOD::TrackParticle*> myMatchedTracks = jet->getAssociatedObjects<xAOD::TrackParticle>(m_assocTracksInName);
+
+ // Loop over ghost-associated tracks
+ for( auto track : myMatchedTracks ) {
+
+ // Make sure only tracks that pass selection are used
+ if( !isGoodTrack(*track,*pvx,*tva) ) continue;
+
+ // Add to list of matched tracks
+ matchedTrackIndices.push_back(track->index());
+
+ // Record which jet the track belongs to
+ JetTrackMap[jet->index()].push_back(track->index());
+ }
+
+ // Fill map of jet indices to merge
+ int mergeIdx = jet->auxdata< int >("mergeIndex");
+ if( mergeIdx >= 0 ) {
+ int jetIdx = jet->index();
+ // Loop up the chain in case a nearest neighbor is also a jet to merge
+ while(inJets->at(mergeIdx)->auxdata< int >("mergeIndex") >= 0 ) {
+ mergeIdx = inJets->at(mergeIdx)->auxdata< int >("mergeIndex");
+ }
+ JetMergeMap[mergeIdx].push_back(jetIdx);
+ }
+ }
+
+ // Sort track indices so the difference can be taken
+ std::sort( allGoodTrackIndices.begin(), allGoodTrackIndices.end() ); // Sort matched track indices
+ std::sort( matchedTrackIndices.begin(), matchedTrackIndices.end() ); // Sort all track indices
+
+ // Get list of unmatched track indices
+ std::set_difference(allGoodTrackIndices.begin(), allGoodTrackIndices.end(), matchedTrackIndices.begin(), matchedTrackIndices.end(), std::inserter(unmatchedTrackIndices,unmatchedTrackIndices.begin()));
+
+ // dR-match all remaining tracks
+ for( auto trackIdx : unmatchedTrackIndices ) {
+
+ // Get input track
+ auto track = inTracks->at(trackIdx);
+
+ float dRmin = 9999999;
+ int jetIdx = -1;
+ for(auto jet : *inJets) {
+ // Find nearest jet to track
+ if(track->p4().DeltaR(jet->p4()) < dRmin) {
+ dRmin = track->p4().DeltaR(jet->p4());
+ jetIdx = jet->index();
+ }
+ }
+
+ // Check if dR matching criteria is met
+ if(dRmin < m_dRmatch) {
+ // Record which jet the track belongs to
+ JetTrackMap[jetIdx].push_back(track->index());
+ }
+ }
+
+ // Done matching tracks to small-R jets
+
+ // Container for combined jets (and jets that survive the combination)
+ auto combinedJets = std::make_unique<xAOD::JetContainer>();
+ auto combinedJetsAux = std::make_unique<xAOD::JetAuxContainer>();
+ combinedJets->setStore( combinedJetsAux.get() );
+
+ for(auto jet : *inJets) {
+
+ int mergeIndex = jet->auxdata< int >("mergeIndex");
+
+ // Make sure jet doesn't need to be merged
+ if( mergeIndex < 0 ) {
+
+ // Copy jet into combined jet container
+ xAOD::Jet* combinedJet = new xAOD::Jet();
+ combinedJets->push_back (combinedJet);
+ *combinedJet = *jet;
+
+ for( auto trackIdx : JetTrackMap[jet->index()] ) {
+ // Set jet associations for associated tracks
+ // Nominally, each track is associated to a single jet, but this construction
+ // allows each track to be associated to multiple jets with a weight for each
+ // association that is used in the rescaling
+ outTracks->at(trackIdx)->auxdata< std::vector< std::pair<int, float> > >("JetAssociations").emplace_back(combinedJet->index(),1.0);
+ }
+
+ // Check if there are any jets to merge
+ for( auto mergeIdx : JetMergeMap[jet->index()]) {
+
+ // Add jet 4-vectors together
+ combinedJet->setJetP4(combinedJet->jetP4() + inJets->at(mergeIdx)->jetP4());
+
+ for( auto trackIdx : JetTrackMap[mergeIdx] ) {
+ // Set merged jet track associations to combined jet
+ outTracks->at(trackIdx)->auxdata< std::vector< std::pair<int, float> > >("JetAssociations").emplace_back(combinedJet->index(),1.0);
+ }
+
+ }
+
+ }
+ }
+
+ // Rescale tracks
+ ATH_CHECK( rescaleTracks(inJets,outTracks.get()) );
+
+ // Create view container to store only selected tracks
+ std::unique_ptr< ConstDataVector<xAOD::TrackParticleContainer> > outSelTracks = std::make_unique< ConstDataVector< xAOD::TrackParticleContainer> >(SG::VIEW_ELEMENTS);
+
+ // Fill view container if track passes selection requirements
+ for(auto goodIdx : allGoodTrackIndices) {
+ outSelTracks->push_back(outTracks->at(goodIdx));
+ }
+
+ ATH_MSG_DEBUG("New SATTrack container size " << outSelTracks->size());
+
+ ATH_CHECK( evtStore()->record(outSelTracks.release(),"Sel"+m_outTrackColl) );
+
+ // Record output tracks
+ ATH_CHECK( evtStore()->record(outTracks.release(),m_outTrackColl) );
+ ATH_CHECK( evtStore()->record(outTracksAux.release(),m_outTrackColl+"Aux.") );
+
+ return StatusCode::SUCCESS;
+}
+
+StatusCode SATScaleTool::flagJetsToMerge( const xAOD::JetContainer *jets ) const {
+
+ ATH_MSG_DEBUG("In SATScaleTool::flagJetsToMerge");
+
+ for( auto *jet : *jets ) {
+
+ // Don't merge if flag is set to false
+ if(!m_doMerge) {
+ jet->auxdecor< int >("mergeIndex") = -1;
+ continue;
+ }
+
+ float dRmin = 999;
+ float ptRatioDR = -1;
+ int neighborIndex = -1;
+
+ for ( auto *jet2 : *jets ) {
+
+ // Don't use jet as its own nearest neighbor
+ if( jet == jet2 ) continue;
+
+ float dR = jet->p4().DeltaR(jet2->p4());
+
+ // Find nearest neighbor jet
+ if( dR < dRmin ) {
+ dRmin = dR;
+ // Calculate merging variables
+ ptRatioDR = (jet->pt() / jet2->pt()) / dR;
+ neighborIndex = jet2->index();
+ }
+
+ }
+
+ // Store index of jet to merge into if criteria are met
+ if( dRmin < m_dRmerge && ptRatioDR < m_pTdRmerge ) jet->auxdecor< int >("mergeIndex") = neighborIndex;
+ else jet->auxdecor< int >("mergeIndex") = -1;
+
+ }
+
+ return StatusCode::SUCCESS;
+}
+
+StatusCode SATScaleTool::getTrackVertexAssociation( const jet::TrackVertexAssociation *&tva ) const {
+
+ tva = nullptr;
+
+ // Get TrackVertexAssociation
+ ATH_CHECK( evtStore()->retrieve(tva,m_trackVtxAssoc) );
+
+ // Fail if none is found
+ if ( !tva ) {
+ ANA_MSG_ERROR("Could not retrieve the TrackVertexAssociation");
+ return StatusCode::FAILURE;
+ }
+
+ return StatusCode::SUCCESS;
+}
+
+StatusCode SATScaleTool::getPrimaryVertex( const xAOD::Vertex *&pvx ) const {
+
+ // Set pvx to nullptr to make sure we are actually retrieving the one we want
+ pvx = nullptr;
+
+ // Get vertex container
+ const xAOD::VertexContainer *vtxContainer = nullptr;
+ ATH_CHECK( evtStore()->retrieve(vtxContainer,m_vertexColl) );
+
+ // Fail if no vertex container or an empty vertex container is found
+ if ( vtxContainer == 0 || vtxContainer->size() == 0 ) {
+ ANA_MSG_ERROR("Failed to retrieve PrimaryVertices collection");
+ return StatusCode::FAILURE;
+ }
+
+ // Find the first primary vertex in the container
+ for ( const auto& vx : *vtxContainer ) {
+ if ( vx->vertexType() == xAOD::VxType::PriVtx ) {
+ pvx = vx;
+ break;
+ }
+ }
+
+ // Make sure the primary vertex is found
+ if( !pvx ) {
+ ANA_MSG_ERROR("No primary vertex assigned" );
+ return StatusCode::FAILURE;
+ }
+
+ return StatusCode::SUCCESS;
+}
+
+void SATScaleTool::print() const {
+ std::string smerge = m_doMerge ? "true" : "false";
+ ATH_MSG_INFO("Properties for SATScaleTool " << name());
+ ATH_MSG_INFO(" InputTrackContainer: " << m_inTrackColl);
+ ATH_MSG_INFO(" InputJetContainer: " << m_inJetColl);
+ ATH_MSG_INFO(" OutputTrackContainer: " << m_outTrackColl);
+ ATH_MSG_INFO(" MatchDeltaR: " << m_dRmatch);
+ ATH_MSG_INFO(" DoMergeJets: " << smerge);
+ ATH_MSG_INFO(" MergeDeltaR: " << m_dRmerge);
+ ATH_MSG_INFO(" MergePtOverDeltaR: " << m_pTdRmerge);
+ TrackAssistTool::print();
+ return;
+}
+
diff --git a/Reconstruction/Jet/JetRecTools/Root/TARJetTool.cxx b/Reconstruction/Jet/JetRecTools/Root/TARJetTool.cxx
new file mode 100644
index 0000000000000000000000000000000000000000..37c9a03b1bd257dd30a7a5a505f9a6e5bba81234
--- /dev/null
+++ b/Reconstruction/Jet/JetRecTools/Root/TARJetTool.cxx
@@ -0,0 +1,262 @@
+/*
+ Copyright (C) 2002-2018 CERN for the benefit of the ATLAS collaboration
+*/
+
+// TARJetTool.cxx
+
+#include "JetRecTools/TARJetTool.h"
+
+TARJetTool::TARJetTool(const std::string& myname)
+: AsgTool(myname),
+ TrackAssistTool(),
+ m_inTrackColl(""),
+ m_outTrackColl(""),
+ m_assocTracksOutName(""),
+ m_dRmatch(0.3)
+{
+ declareProperty("InputTrackContainer", m_inTrackColl = "InDetTrackParticles");
+ declareProperty("OutputTrackContainer", m_outTrackColl);
+ declareProperty("OutputAssociatedTracks", m_assocTracksOutName = "TARTracks");
+ declareProperty("MatchDeltaR", m_dRmatch);
+
+ declareProperty("InputAssociatedTracks", m_assocTracksInName = "GhostTrack");
+ declareProperty("VertexContainer", m_vertexColl = "PrimaryVertices");
+ declareProperty("TrackVertexAssociation", m_trackVtxAssoc);
+ declareProperty("DoTrackVertexAssociation", m_doTrackVtxAssoc = true);
+ declareProperty("TrackSelTool", m_trackSelTool);
+}
+
+StatusCode TARJetTool::initialize()
+{
+
+ ATH_MSG_INFO("Initializing TARJetTool " << name() << ".");
+ ATH_CHECK( TrackAssistTool::initialize() );
+ print();
+
+ return StatusCode::SUCCESS;
+}
+
+int TARJetTool::modify( xAOD::JetContainer& inJets ) const {
+
+ std::vector<int> matchedTrackIndices; // List of ghost matched track indices
+ std::vector<int> unmatchedTrackIndices; // List of non ghost matched track indices
+ std::vector<int> allGoodTrackIndices; // List of all good track indices
+
+ std::map<int, std::vector<int>> JetTrackMap; // Map of track indices to RC jet indices
+ std::map<int, std::vector<int>> ConstitTrackMap; // Map of track indices to constit jet indices
+ std::map<int, int> ConstitJetMap; // Map of all constituent jets to RC jets
+
+ // Get input tracks
+ const xAOD::TrackParticleContainer *inTracks = nullptr;
+ ATH_CHECK( evtStore()->retrieve(inTracks,m_inTrackColl) );
+
+ // Get vertex objects
+ const xAOD::Vertex *pvx = nullptr;
+ const jet::TrackVertexAssociation *tva = nullptr;
+ ATH_CHECK( getVertexInfo(pvx,tva) );
+
+ // Make container for constituent jets
+ auto constitJets = std::make_unique<xAOD::JetContainer>();
+ auto constitJetsAux = std::make_unique<xAOD::JetAuxContainer>();
+ constitJets->setStore( constitJetsAux.get() ); //< Connect the two
+
+ // Make shallow copy of tracks
+ auto trackShallowCopy = xAOD::shallowCopyContainer(*inTracks);
+ std::unique_ptr<xAOD::TrackParticleContainer> outTracks (trackShallowCopy.first);
+ std::unique_ptr<xAOD::ShallowAuxContainer> outTracksAux (trackShallowCopy.second);
+
+ // Get list of all tracks
+ for( auto track : *inTracks ) {
+ // Make sure only tracks that pass selection are used
+ if( !isGoodTrack(*track,*pvx,*tva) ) continue;
+ allGoodTrackIndices.push_back( track->index() );
+ }
+
+ // Loop over RC jets
+ for(auto jet : inJets) {
+
+ // Loop over constituent jets
+ for (unsigned int iConstit = 0; iConstit < jet->numConstituents(); iConstit++){
+
+ // Get constituent jet
+ const xAOD::Jet *constit = dynamic_cast< const xAOD::Jet* >( jet->rawConstituent(iConstit) );
+
+ // Copy constituent jet into temporary small-R jet collection
+ xAOD::Jet* constitJet = new xAOD::Jet();
+ constitJets->push_back (constitJet);
+ *constitJet = *constit;
+
+ // Record which RC jet constituent jet comes from
+ ConstitJetMap[constitJet->index()] = jet->index();
+
+ // Get ghost-associated tracks
+ std::vector<const xAOD::TrackParticle*> myMatchedTracks = constitJet->getAssociatedObjects<xAOD::TrackParticle>(m_assocTracksInName);
+
+ // Loop over ghost-associated tracks
+ for( auto track : myMatchedTracks ) {
+
+ // Make sure only tracks that pass selection are used
+ if( !isGoodTrack(*track,*pvx,*tva) ) continue;
+
+ // Add to list of matched tracks
+ matchedTrackIndices.push_back(track->index());
+
+ // Record which constituent jet and RC jet the track belongs to
+ ConstitTrackMap[constitJet->index()].push_back(track->index());
+ JetTrackMap[jet->index()].push_back(track->index());
+ }
+
+ }
+ }
+
+ // Sort track indices so the difference can be taken
+ std::sort( allGoodTrackIndices.begin(), allGoodTrackIndices.end() ); // Sort matched track indices
+ std::sort( matchedTrackIndices.begin(), matchedTrackIndices.end() ); // Sort all track indices
+
+ // Get list of unmatched track indices
+ std::set_difference(allGoodTrackIndices.begin(), allGoodTrackIndices.end(), matchedTrackIndices.begin(), matchedTrackIndices.end(), std::inserter(unmatchedTrackIndices,unmatchedTrackIndices.begin()));
+
+ // Loop over all remaining tracks
+ for( auto trackIdx : unmatchedTrackIndices ) {
+
+ // Get input track
+ auto track = inTracks->at(trackIdx);
+
+ float dRmin = 9999999;
+ int jetIdx = -1;
+
+ for(auto constitJet : *constitJets) {
+ // Find nearest small-R jet to track
+ if(track->p4().DeltaR(constitJet->p4()) < dRmin) {
+ dRmin = track->p4().DeltaR(constitJet->p4());
+ jetIdx = constitJet->index();
+ }
+ }
+
+ // Check if dR matching criteria is met
+ if(dRmin < m_dRmatch) {
+ // Record which constituent jet and RC jet track belongs to
+ ConstitTrackMap[jetIdx].push_back(track->index());
+ JetTrackMap[ConstitJetMap[jetIdx]].push_back(track->index());
+ }
+ }
+
+ // Loop over all constituent jets to assign jet associations to tracks
+ // NB: This is being done in a separate loop to make future modifications easier
+ for(auto constitJet : *constitJets) {
+
+ for( auto trackIdx : ConstitTrackMap[constitJet->index()] ) {
+ // Set jet associations for associated tracks
+ // Nominally, each track is associated to a single jet, but this construction
+ // allows each track to be associated to multiple jets with a weight for each
+ // association that is used in the rescaling
+ outTracks->at(trackIdx)->auxdata< std::vector< std::pair<int, float> > >("JetAssociations").emplace_back(constitJet->index(),1.0);
+ }
+
+ }
+
+ // Rescale tracks
+ ATH_CHECK( rescaleTracks(constitJets.get(),outTracks.get()) );
+
+ // Create view container to store only selected tracks
+ std::unique_ptr< ConstDataVector<xAOD::TrackParticleContainer> > outSelTracks = std::make_unique< ConstDataVector< xAOD::TrackParticleContainer> >(SG::VIEW_ELEMENTS);
+
+ // Fill view container if track passes selection requirements
+ for(auto goodIdx : allGoodTrackIndices) {
+ outSelTracks->push_back(outTracks->at(goodIdx));
+ }
+
+ ATH_MSG_DEBUG("New TARTrack container size " << outSelTracks->size());
+
+ ATH_CHECK( evtStore()->record(outSelTracks.release(),"Sel"+m_outTrackColl) );
+
+ // Get bare pointer for post-record operations
+ auto pOutTracks = outTracks.get();
+
+ // Record output tracks
+ ATH_CHECK( evtStore()->record(outTracks.release(), m_outTrackColl) );
+ ATH_CHECK( evtStore()->record(outTracksAux.release(), m_outTrackColl+"Aux.") );
+
+ // Loop over input jets to link tracks
+ for(auto jet : inJets) {
+
+ std::vector<const xAOD::TrackParticle*> TARTracks;
+
+ // Sum of TAR jets for mTAR
+ TLorentzVector TARJet;
+
+ // Collect all tracks associated to RC jet
+ for( auto trackIdx : JetTrackMap[jet->index()] ) {
+ TARTracks.push_back(pOutTracks->at(trackIdx));
+ TARJet += pOutTracks->at(trackIdx)->p4();
+ }
+
+ // Add associated TAR tracks
+ jet->setAssociatedObjects< xAOD::TrackParticle >(m_assocTracksOutName,TARTracks);
+
+ // Add mTAR decoration
+ jet->setAttribute("mTAR", TARJet.M());
+
+ }
+
+ return 0;
+}
+
+StatusCode TARJetTool::getTrackVertexAssociation( const jet::TrackVertexAssociation *&tva ) const {
+
+ tva = nullptr;
+
+ // Get TrackVertexAssociation
+ ATH_CHECK( evtStore()->retrieve(tva,m_trackVtxAssoc) );
+
+ // Fail if none is found
+ if ( !tva ) {
+ ANA_MSG_ERROR("Could not retrieve the TrackVertexAssociation");
+ return StatusCode::FAILURE;
+ }
+
+ return StatusCode::SUCCESS;
+}
+
+StatusCode TARJetTool::getPrimaryVertex( const xAOD::Vertex *&pvx ) const {
+
+ // Set pvx to nullptr to make sure we are actually retrieving the one we want
+ pvx = nullptr;
+
+ // Get vertex container
+ const xAOD::VertexContainer *vtxContainer = nullptr;
+ ATH_CHECK( evtStore()->retrieve(vtxContainer,m_vertexColl) );
+
+ // Fail if no vertex container or an empty vertex container is found
+ if ( vtxContainer == 0 || vtxContainer->size() == 0 ) {
+ ANA_MSG_ERROR("Failed to retrieve PrimaryVertices collection");
+ return StatusCode::FAILURE;
+ }
+
+ // Find the first primary vertex in the container
+ for ( const auto& vx : *vtxContainer ) {
+ if ( vx->vertexType() == xAOD::VxType::PriVtx ) {
+ pvx = vx;
+ break;
+ }
+ }
+
+ // Make sure the primary vertex is found
+ if( !pvx ) {
+ ANA_MSG_ERROR("No primary vertex assigned" );
+ return StatusCode::FAILURE;
+ }
+
+ return StatusCode::SUCCESS;
+}
+
+void TARJetTool::print() const {
+ ATH_MSG_INFO("Properties for TARJetTool " << name());
+ ATH_MSG_INFO(" InputTrackContainer: " << m_inTrackColl);
+ ATH_MSG_INFO(" OutputTrackContainer: " << m_outTrackColl);
+ ATH_MSG_INFO(" OutputAssociatedTracks: " << m_assocTracksOutName);
+ ATH_MSG_INFO(" MatchDeltaR: " << m_dRmatch);
+ TrackAssistTool::print();
+ return;
+}
+
diff --git a/Reconstruction/Jet/JetRecTools/Root/TrackAssistTool.cxx b/Reconstruction/Jet/JetRecTools/Root/TrackAssistTool.cxx
new file mode 100644
index 0000000000000000000000000000000000000000..03ac8ab1f7a61ff92fef2e8356e1e9df2c04e832
--- /dev/null
+++ b/Reconstruction/Jet/JetRecTools/Root/TrackAssistTool.cxx
@@ -0,0 +1,161 @@
+/*
+ Copyright (C) 2002-2018 CERN for the benefit of the ATLAS collaboration
+*/
+
+// TrackAssist.cxx
+
+#include "JetRecTools/TrackAssistTool.h"
+
+TrackAssistTool::TrackAssistTool()
+: m_assocTracksInName(""),
+ m_vertexColl(""),
+ m_trackVtxAssoc(""),
+ m_doTrackVtxAssoc(true),
+ m_trackSelTool("")
+{
+}
+
+StatusCode TrackAssistTool::initialize()
+{
+
+ // If no VertexContainer is given, give warning but proceed without it
+ if( m_vertexColl.empty() ) {
+ ANA_MSG_WARNING("No VertexContainer defined. The TrackVertexAssociation requirement will be turned off to avoid undefined behavior.");
+ m_doTrackVtxAssoc = false;
+ }
+
+ // If no TrackVertexAssociation is given, give warning but proceed without it
+ if( m_trackVtxAssoc.empty() ) {
+ ANA_MSG_WARNING("No TrackVertexAssociation defined. The TrackVertexAssociation requirement will be turned off to avoid undefined behavior.");
+ m_doTrackVtxAssoc = false;
+ }
+
+ // If TrackVertexAssociation is turned off, give warning but proceed
+ if( !m_doTrackVtxAssoc ) {
+ ANA_MSG_WARNING("TrackVertexAssociation retquirement will not be applied. This is not recommended for general use.");
+ }
+
+ // If no TrackSelector is given, give warning but proceed without it
+ if( m_trackSelTool.empty() ) {
+ ANA_MSG_WARNING("No TrackSelector defined. No track selection will be applied. This is not recommended for general use.");
+ }
+
+ return StatusCode::SUCCESS;
+}
+
+StatusCode TrackAssistTool::rescaleTracks( const xAOD::JetContainer *jets, xAOD::TrackParticleContainer *tracks) const {
+
+ for(auto jet : *jets) {
+ // Get track rescale weights
+ ANA_CHECK( getRescaleFactors(*jet, *tracks) );
+ }
+
+ for(auto track : *tracks) {
+
+ // Get the scales
+ float scale = track->auxdata< float >("TAScale");
+
+ // Include this to preserve the pT of any unmatched tracks
+ if(!scale) scale = 1.0;
+
+ // Rescale track pT
+ track->setDefiningParameters( track->d0(), track->z0(), track->phi0(), track->theta(), track->qOverP() / scale );
+
+ }
+
+ return StatusCode::SUCCESS;
+}
+
+StatusCode TrackAssistTool::getRescaleFactors( const xAOD::Jet &jet, xAOD::TrackParticleContainer &tracks) const {
+
+ float trackPt = 0;
+
+ // Get weighted sum of track pTs for jet
+ for(auto track : tracks) {
+
+ // Get associations and weights
+ std::vector< std::pair<int, float> > associations = track->auxdata< std::vector< std::pair<int, float> > >("JetAssociations");
+
+ // Loop over associations
+ for(auto association : associations) {
+
+ // Skip track if it isn't associated to the jet
+ if(unsigned(association.first) != jet.index()) continue;
+
+ // Add weighted pT to total
+ float weight = association.second;
+ trackPt += track->pt() * weight;
+
+ // Only add each track once
+ break;
+
+ }
+
+ }
+
+ // Save rescale factor for each track
+ for(auto track : tracks) {
+
+ float scale = 0.0;
+
+ // Get associations and weights
+ std::vector< std::pair<int, float> > associations = track->auxdata< std::vector< std::pair<int, float> > >("JetAssociations");
+
+ // Loop over associations
+ for(auto association : associations) {
+
+ // Skip track if it isn't associated to the jet
+ if(unsigned(association.first) != jet.index()) continue;
+
+ // Get factor to scale
+ float weight = association.second;
+ scale = jet.pt() * weight / trackPt;
+
+ // Break loop once association to jet is found
+ break;
+
+ }
+
+ // Save the sum of scales
+ track->auxdata< float >("TAScale") += scale;
+
+ }
+
+ return StatusCode::SUCCESS;
+}
+
+bool TrackAssistTool::isGoodTrack( const xAOD::TrackParticle &track, const xAOD::Vertex &pvx, const jet::TrackVertexAssociation &tva ) const {
+
+ // Do TrackVertexAssociation check if configured to do so
+ if ( m_doTrackVtxAssoc ) {
+ // Check if track is associated to primary vertex
+ if( &pvx != tva.associatedVertex(&track) ) return false;
+ }
+
+ // Do track selection check if configured to do so
+ if ( !m_trackSelTool.empty() ) {
+ // Check if track passes selection criteria
+ if( !m_trackSelTool->keep(track) ) return false;
+ }
+
+ return true;
+}
+
+StatusCode TrackAssistTool::getVertexInfo( const xAOD::Vertex *&pvx, const jet::TrackVertexAssociation *&tva ) const {
+
+ // Get primary vertex and vertex association if configured to do so
+ if( m_doTrackVtxAssoc ) {
+ ANA_CHECK( getPrimaryVertex(pvx) );
+ ANA_CHECK( getTrackVertexAssociation(tva) );
+ }
+ return StatusCode::SUCCESS;
+}
+
+void TrackAssistTool::print() const {
+ ANA_MSG_INFO("Properties inherited from TrackAssistTool");
+ ANA_MSG_INFO(" InputAssociatedTracks: " << m_assocTracksInName);
+ ANA_MSG_INFO(" VertexContainer: " << m_vertexColl);
+ ANA_MSG_INFO(" TrackVertexAssociation: " << m_trackVtxAssoc);
+ ANA_MSG_INFO(" TrackSelector: " << m_trackSelTool);
+ return;
+}
diff --git a/Reconstruction/Jet/JetRecTools/src/components/JetRecTools_entries.cxx b/Reconstruction/Jet/JetRecTools/src/components/JetRecTools_entries.cxx
index 217a76402e0429502aa7455a75dd9d23b1fee07a..04d1f92bc889f5510e80257f156cee7457054ec0 100755
--- a/Reconstruction/Jet/JetRecTools/src/components/JetRecTools_entries.cxx
+++ b/Reconstruction/Jet/JetRecTools/src/components/JetRecTools_entries.cxx
@@ -18,6 +18,8 @@
#include "JetRecTools/CorrectPFOTool.h"
#include "JetRecTools/ChargedHadronSubtractionTool.h"
#include "JetRecTools/PuppiWeightTool.h"
+#include "JetRecTools/TARJetTool.h"
+#include "JetRecTools/SATScaleTool.h"
DECLARE_COMPONENT( JetTrackSelectionTool )
DECLARE_COMPONENT( SimpleJetTrackSelectionTool )
@@ -36,4 +38,6 @@ DECLARE_COMPONENT( CorrectPFOTool )
DECLARE_COMPONENT( ChargedHadronSubtractionTool )
DECLARE_COMPONENT( PuppiWeightTool )
DECLARE_COMPONENT( ConstitTimeCutTool )
+DECLARE_COMPONENT( TARJetTool )
+DECLARE_COMPONENT( SATScaleTool )