Merge branch 'MuonMemMan' into 'master'

Better memory managment in combined muon algorithms

See merge request atlas/athena!16596

Former-commit-id: 3dd57363

MuonSpectrometer/MuonReconstruction/MuonRecTools/MuonRecToolInterfaces/MuonRecToolInterfaces/IMuonRecoValidationTool.h

@@ -60,7 +60,7 @@ namespace Muon {
 
     /** add new MuonCandidate */
     virtual bool addMuonCandidate( const xAOD::TrackParticle& indetTrackParticle, const MuonCandidate* candidate, 
-                                   const Trk::Track* combinedTrack, int ntimes, float beta, float chi2ndof, int stage ) = 0;
+                                   Trk::Track* combinedTrack, int ntimes, float beta, float chi2ndof, int stage ) = 0;
 
     /** add StauHits to ntuple */
     virtual bool addTimeMeasurements( const xAOD::TrackParticle& indetTrackParticle, const MuGirlNS::StauHits& stauHits ) = 0;

MuonSpectrometer/MuonTruthAlgs/src/MuonDetailedTrackTruthMaker.cxx

@@ -7,6 +7,7 @@
 
 #include "MuonTruthAlgs/MuonDetailedTrackTruthMaker.h"
 #include <iterator>
+#include <map>
 
 //================================================================
 MuonDetailedTrackTruthMaker::MuonDetailedTrackTruthMaker(const std::string &name, ISvcLocator *pSvcLocator) :
@@ -33,10 +34,17 @@ StatusCode MuonDetailedTrackTruthMaker::initialize()
 
   if(m_useNSW) m_PRD_TruthNames={"sTGC_TruthMap","MM_TruthMap","RPC_TruthMap","TGC_TruthMap","MDT_TruthMap"};
 
-  m_detailedTrackTruthNames.reserve ( m_trackCollectionNames.size());
-  for(unsigned int i=0;i<m_trackCollectionNames.size();i++){
-    m_detailedTrackTruthNames.emplace_back(m_trackCollectionNames.at(i).key()+"DetailedTruth");
-    ATH_MSG_INFO("process "<<m_trackCollectionNames.at(i).key()<<" for detailed truth collection "<<m_detailedTrackTruthNames.at(i).key());
+  if(m_detailedTrackTruthNames.empty()){
+    m_detailedTrackTruthNames.reserve ( m_trackCollectionNames.size());
+    for(unsigned int i=0;i<m_trackCollectionNames.size();i++){
+      m_detailedTrackTruthNames.emplace_back(m_trackCollectionNames.at(i).key()+"DetailedTruth");
+      ATH_MSG_INFO("process "<<m_trackCollectionNames.at(i).key()<<" for detailed truth collection "<<m_detailedTrackTruthNames.at(i).key());
+    }
+  }
+  else{
+    for(unsigned int i=0;i<m_detailedTrackTruthNames.size();i++){
+      m_detailedTrackTruthNames.at(i)=m_detailedTrackTruthNames.at(i).key()+"DetailedTruth";
+    }
   }
 
   ATH_CHECK(m_trackCollectionNames.initialize());
@@ -72,6 +80,21 @@ StatusCode MuonDetailedTrackTruthMaker::execute() {
 
   //----------------------------------------------------------------
   // Retrieve track collections
+
+  std::map<std::string,DetailedTrackTruthCollection*> dttcMap;
+  for(SG::WriteHandle<DetailedTrackTruthCollection>& h_dttc : m_detailedTrackTruthNames.makeHandles()){
+    ATH_CHECK(h_dttc.record(std::make_unique<DetailedTrackTruthCollection>()));
+    if(h_dttc.key().find("ExtrapolatedMuonTracks")!=std::string::npos) dttcMap.insert(std::pair<std::string,DetailedTrackTruthCollection*>("METracks",h_dttc.ptr()));
+    else if(h_dttc.key().find("CombinedMuonTracks")!=std::string::npos) dttcMap.insert(std::pair<std::string,DetailedTrackTruthCollection*>("CombinedTracks",h_dttc.ptr()));
+    else if(h_dttc.key().find("MSOnlyExtrapolated")!=std::string::npos) dttcMap.insert(std::pair<std::string,DetailedTrackTruthCollection*>("MSOnlyExtrapolated",h_dttc.ptr()));
+    else{
+      std::string cname=h_dttc.key();
+      int pos=cname.find("DetailedTruth");
+      cname.erase(pos,cname.length()-pos);
+      dttcMap.insert(std::pair<std::string,DetailedTrackTruthCollection*>(cname,h_dttc.ptr()));
+    }
+  }
+
   int i=0;
   for(SG::ReadHandle<TrackCollection>& tcol : m_trackCollectionNames.makeHandles()){
     if(!tcol.isValid()){
@@ -83,10 +106,19 @@ StatusCode MuonDetailedTrackTruthMaker::execute() {
     //----------------------------------------------------------------
     // Produce and store the output.
 
-    SG::WriteHandle<DetailedTrackTruthCollection> dttc(m_detailedTrackTruthNames.at(i));
-    ATH_CHECK(dttc.record(std::make_unique<DetailedTrackTruthCollection>()));
-    dttc->setTrackCollection(tcol.cptr());
-    m_truthTool->buildDetailedTrackTruth(dttc.ptr(), *(tcol.cptr()), prdCollectionVector);
+    DetailedTrackTruthCollection* dttc=0;
+    for(auto entry : dttcMap){
+      if(tcol.key().find(entry.first)!=std::string::npos){
+	dttc=entry.second;
+	break;
+      }
+    }
+    if(!dttc){
+      ATH_MSG_WARNING("no detailed track collection found!");
+      continue;
+    }
+    if(!dttc->trackCollectionLink().isValid()) dttc->setTrackCollection(tcol.cptr());
+    m_truthTool->buildDetailedTrackTruth(dttc, *(tcol.cptr()), prdCollectionVector);
     i++;
   }
   return StatusCode::SUCCESS;

MuonSpectrometer/MuonValidation/MuonRecValidation/MuonTrackPerformance/src/MuonRecoValidationTool.cxx

@@ -198,7 +198,7 @@ namespace Muon {
 
     auto pos = std::find(m_trackParticles.begin(),m_trackParticles.end(),&indetTrackParticle);
     if( pos == m_trackParticles.end() ) {
-      ATH_MSG_WARNING("addMuonCandidate: indetTrackParticle not found ");
+      ATH_MSG_WARNING("addTimeMeasurement: indetTrackParticle not found ");
       return false;
     }
     int index = std::distance(m_trackParticles.begin(),pos);
@@ -502,7 +502,7 @@ namespace Muon {
   }
 
   bool MuonRecoValidationTool::addMuonCandidate( const xAOD::TrackParticle& indetTrackParticle, const MuonCandidate* candidate, 
-                                                 const Trk::Track* combinedTrack, int ntimes, float beta, float chi2ndof, int stage ) {
+                                                 Trk::Track* combinedTrack, int ntimes, float beta, float chi2ndof, int stage ) {
 
     auto pos = std::find(m_trackParticles.begin(),m_trackParticles.end(),&indetTrackParticle);
     if( pos == m_trackParticles.end() ) {
@@ -515,8 +515,9 @@ namespace Muon {
     int ntrigPhi = 0;
     int ntrigEta = 0;
     int nseg = candidate ? candidate->layerIntersections.size() : 0;
+    const Trk::Track* track=combinedTrack;
     if( combinedTrack ){
-      IMuonHitSummaryTool::CompactSummary summary = m_hitSummaryTool->summary(*combinedTrack);
+      IMuonHitSummaryTool::CompactSummary summary = m_hitSummaryTool->summary(*track);
       nprec = summary.nprecisionLayers;
       ntrigPhi = summary.nphiLayers;  
       ntrigEta = summary.ntrigEtaLayers;  

MuonSpectrometer/MuonValidation/MuonRecValidation/MuonTrackPerformance/src/MuonRecoValidationTool.h

@@ -58,30 +58,30 @@ namespace Muon {
     StatusCode finalize();
     
     /** add a new TrackParticle with it's muon system extension */
-    bool addTrackParticle(  const xAOD::TrackParticle& indetTrackParticle, const MuonSystemExtension& muonSystemExtention );
+    bool addTrackParticle(  const xAOD::TrackParticle& indetTrackParticle, const MuonSystemExtension& muonSystemExtention ) override;
 
     /** add a new segment */
-    bool add( const MuonSystemExtension::Intersection& intersection, const MuonSegment& segment, int stage );
+    bool add( const MuonSystemExtension::Intersection& intersection, const MuonSegment& segment, int stage ) override;
 
     /** add a new hough maximum */
-    bool add( const MuonSystemExtension::Intersection& intersection, const MuonHough::MuonLayerHough::Maximum& maximum );
+    bool add( const MuonSystemExtension::Intersection& intersection, const MuonHough::MuonLayerHough::Maximum& maximum ) override;
 
     /** add a new prd */
-    bool add( const MuonSystemExtension::Intersection& intersection, const Trk::PrepRawData& prd, float expos, float expos_err );
+    bool add( const MuonSystemExtension::Intersection& intersection, const Trk::PrepRawData& prd, float expos, float expos_err ) override;
 
     /** add a new time measurement */
-    bool addTimeMeasurement( const MuonSystemExtension::Intersection& intersection, const Trk::MeasurementBase& meas );
+    bool addTimeMeasurement( const MuonSystemExtension::Intersection& intersection, const Trk::MeasurementBase& meas ) override;
 
     /** add a new time measurement */
     bool addTimeMeasurement( const MuonSystemExtension::Intersection& intersection, const Identifier& id,
-                             const Amg::Vector3D& gpos, float time, float errorTime );
+                             const Amg::Vector3D& gpos, float time, float errorTime ) override;
 
     /** add StauHits to ntuple */
-    bool addTimeMeasurements( const xAOD::TrackParticle& indetTrackParticle, const MuGirlNS::StauHits& stauHits );
+    bool addTimeMeasurements( const xAOD::TrackParticle& indetTrackParticle, const MuGirlNS::StauHits& stauHits ) override;
 
     /** add a new muon candidate */
     bool addMuonCandidate( const xAOD::TrackParticle& indetTrackParticle, const MuonCandidate* candidate, 
-                           const Trk::Track* combinedTrack, int ntimes, float beta, float chi2ndof, int stage );
+                           Trk::Track* combinedTrack, int ntimes, float beta, float chi2ndof, int stage ) override;
 
     /**  incident service handle for EndEvent */
     void handle(const Incident& inc);

Reconstruction/MuonIdentification/MuGirl/src/MuGirlRecoTool.cxx

@@ -474,7 +474,7 @@ bool MuGirlRecoTool::RunFromID(MuGirlNS::CandidateSummaryList& summaryList) {
 
             std::vector<const Muon::MuonSegment*> stauSegmentList = *m_pStauTool->newMdtSegments();
             m_pSegmentManager->setInStore(stauSegmentList, true);
-
+	    /*
             std::unique_ptr<MuonCombined::MuGirlLowBetaTag> lowbetatag{new MuonCombined::MuGirlLowBetaTag(stauRefittedTrack,
                                                                                             stauSegmentList)};
 
@@ -483,7 +483,7 @@ bool MuGirlRecoTool::RunFromID(MuGirlNS::CandidateSummaryList& summaryList) {
             lowbetatag->setMuBeta(Beta);
             lowbetatag->setStauSummary(stauSummary.release());
             lowbetatag->setStauExtras(stauExtras.release());
-
+	    */
             //pParticle->pInDetCandidate->addTag(*lowbetatag);
           } else {
             ATH_MSG_DEBUG("stauRefittedTrack is NULL");
@@ -509,10 +509,11 @@ bool MuGirlRecoTool::RunFromID(MuGirlNS::CandidateSummaryList& summaryList) {
           ATH_MSG_DEBUG("RunFromID: doRH == true");
           std::unique_ptr<MuGirlNS::RHExtras> rhExtras(new MuGirlNS::RHExtras);
           if (m_pStauTool->fillRHExtras(rhExtras.get()).isFailure()) ATH_MSG_DEBUG("RunFromID: Cannot fill rhExtras");
-          if (rhExtras->numCaloCells > 0) {
+          if (rhExtras->numCaloCells > 0) {/*
             std::unique_ptr<MuonCombined::MuGirlLowBetaTag> lowbetatag
               {new MuonCombined::MuGirlLowBetaTag(pTrackParticle->track(), std::vector<const Muon::MuonSegment*>())};
             lowbetatag->setRHExtras(rhExtras.release());
+					   */
             //pParticle->pInDetCandidate->addTag(*lowbetatag);
           }
           ATH_MSG_DEBUG("RunFromID: doRH == true - done");
@@ -527,7 +528,7 @@ bool MuGirlRecoTool::RunFromID(MuGirlNS::CandidateSummaryList& summaryList) {
       //Generate the MuGirlTag according to the new running schema: MuGirlTag will delete the segments
       //Otherwise put the list of MuGirl Segments into Storegate: Storegate will delete the segments
 
-
+      /*
       ATH_MSG_DEBUG("RunFromID: delivering the muon pTrkRefitted track to the MuGirl tag");
       std::vector<const Muon::MuonSegment*> muonSegmentList = summary->muonSegmentList;
       std::unique_ptr<MuonCombined::MuGirlTag> tag{};
@@ -536,9 +537,10 @@ bool MuGirlRecoTool::RunFromID(MuGirlNS::CandidateSummaryList& summaryList) {
         tag.reset(new MuonCombined::MuGirlTag(muonSegmentList));
       } else tag.reset(new MuonCombined::MuGirlTag(summary->pTrkRefitted, muonSegmentList));
       tag->setUpdatedExtrapolatedTrack(std::move(summary->pTrkMSRefitted));
+      */
       //pParticle->pInDetCandidate->addTag(*tag);
       //set the segment into SegmentManager undeletable
-      m_pSegmentManager->setInStore(muonSegmentList, true);
+      //m_pSegmentManager->setInStore(muonSegmentList, true);
       summary->saveStau = 0;
       //m_pCandidate->releaseMuonSegments(false);   //the segments released from Candidate ownership
       //summaryList.push_back(summary.release());
@@ -547,14 +549,15 @@ bool MuGirlRecoTool::RunFromID(MuGirlNS::CandidateSummaryList& summaryList) {
 
       //  m_pCandidate->releaseMuonSegments(true);
       //Generate the MuGirlTag according to the new running schema: MuGirlTag will delete the segments
-
+      /*
       ATH_MSG_DEBUG("RunFromID: delivering the pTrkLowBeta to the MuGirl tag");
       std::vector<const Muon::MuonSegment*> muonSegmentList = summary->muonSegmentList;
       MuonCombined::MuGirlTag* tag = new MuonCombined::MuGirlTag(summary->pTrkLowBeta, muonSegmentList);
       tag->setUpdatedExtrapolatedTrack(nullptr);
+      */
       //pParticle->pInDetCandidate->addTag(*tag);
       //set the segment into SegmentManager undeletable
-      m_pSegmentManager->setInStore(muonSegmentList, true);
+      //m_pSegmentManager->setInStore(muonSegmentList, true);
     }
 
     if (m_doNTuple) {

Reconstruction/MuonIdentification/MuidInterfaces/MuidInterfaces/IMuidMuonRecovery.h

@@ -45,8 +45,8 @@ public:
 
     /**IMuidMuonRecovery interface:
        algorithmic code for recovering muon spectrometer using the inner detector track */    
-    virtual const Trk::Track*           recoverableMatch (const Trk::Track& indetTrack, 
-							  const Trk::Track& spectrometerTrack) const = 0;
+    virtual Trk::Track*           recoverableMatch (const Trk::Track& indetTrack, 
+						    const Trk::Track& spectrometerTrack) const = 0;
 };
  
 }	// end of namespace

Reconstruction/MuonIdentification/MuidTrackBuilder/MuidTrackBuilder/MuidMuonRecovery.h

@@ -54,8 +54,8 @@ public:
     
     /**IMuidMuonRecovery interface:
        algorithmic code for recovering muon spectrometer using the inner detector track */    
-    const Trk::Track*	recoverableMatch (const Trk::Track&	indetTrack, 
-					  const Trk::Track&	spectrometerTrack) const;
+    Trk::Track*	recoverableMatch (const Trk::Track&	indetTrack, 
+				  const Trk::Track&	spectrometerTrack) const;
 
  private:
     // helpers, managers, tools

Reconstruction/MuonIdentification/MuidTrackBuilder/src/MuidMuonRecovery.cxx

@@ -36,6 +36,7 @@
 #include "TrkToolInterfaces/IResidualPullCalculator.h"
 #include "TrkTrack/Track.h"
 #include "TrkTrackSummary/TrackSummary.h"
+#include "TrkSegment/SegmentCollection.h"
 
 namespace Rec
 {
@@ -147,7 +148,7 @@ MuidMuonRecovery::finalize()
 }
 
 
-const Trk::Track*
+Trk::Track*
 MuidMuonRecovery::recoverableMatch (const Trk::Track& indetTrack,
 				    const Trk::Track& spectrometerTrack) const
 {
@@ -390,7 +391,7 @@ if (msgLvl(MSG::DEBUG))
     }
 
     // fit the combined track
-    const Trk::Track* combinedTrack = 0;
+    Trk::Track* combinedTrack = 0;
     if (! m_trackBuilder.empty())
     {
 	combinedTrack	= m_trackBuilder->indetExtension(indetTrack,

Reconstruction/MuonIdentification/MuonCombinedAlgs/src/MuonCombinedAlg.cxx

@@ -16,6 +16,8 @@ StatusCode MuonCombinedAlg::initialize()
   ATH_CHECK(m_indetCandidateCollectionName.initialize());
   ATH_CHECK(m_muonCandidateCollectionName.initialize());
   ATH_CHECK(m_combTagMaps.initialize());
+  ATH_CHECK(m_muidCombinedTracks.initialize());
+  ATH_CHECK(m_muidMETracks.initialize());
 
   return StatusCode::SUCCESS; 
 }
@@ -50,9 +52,16 @@ StatusCode MuonCombinedAlg::execute()
     maps.push_back(h.ptr());
   }
 
+  SG::WriteHandle<TrackCollection> muidCombTracks(m_muidCombinedTracks);
+  ATH_CHECK(muidCombTracks.record(std::make_unique<TrackCollection>()));
+
+  SG::WriteHandle<TrackCollection> muidMETracks(m_muidMETracks);
+  ATH_CHECK(muidMETracks.record(std::make_unique<TrackCollection>()));
+
   if(inDetCandidateCollection->empty() || muonCandidateCollection->empty()) return StatusCode::SUCCESS;
 
-  m_muonCombinedTool->combine(*muonCandidateCollection,*inDetCandidateCollection,maps);
+  //note that STACO does not create new Trk::Tracks so it doesn't need collections here
+  m_muonCombinedTool->combine(*muonCandidateCollection,*inDetCandidateCollection,maps,muidCombTracks.ptr(),muidMETracks.ptr());
 
   return StatusCode::SUCCESS;
 }

Reconstruction/MuonIdentification/MuonCombinedAlgs/src/MuonCombinedAlg.h

@@ -14,6 +14,7 @@
 #include "MuonCombinedEvent/MuonCandidate.h"
 #include "MuonCombinedEvent/InDetCandidateToTagMap.h"
 #include "MuonCombinedToolInterfaces/IMuonCombinedTool.h"
+#include "TrkTrack/TrackCollection.h"
 #include <string>
 #include <vector>
 
@@ -39,6 +40,8 @@ class MuonCombinedAlg : public AthAlgorithm
   SG::ReadHandleKey<InDetCandidateCollection> m_indetCandidateCollectionName{this,"InDetCandidateLocation","InDetCandidates","name of ID candidate collection"};
   SG::ReadHandleKey<MuonCandidateCollection> m_muonCandidateCollectionName{this,"MuonCandidateLocation","MuonCandidates","name of muon candidate collection"};
   SG::WriteHandleKeyArray<MuonCombined::InDetCandidateToTagMap> m_combTagMaps{this,"CombinedTagMaps",{"muidcoTagMap","stacoTagMap"},"combined muon tag maps"};
+  SG::WriteHandleKey<TrackCollection> m_muidCombinedTracks{this,"MuidCombinedTracksLocation","MuidCombinedTracks","Muidco combined Tracks"};
+  SG::WriteHandleKey<TrackCollection> m_muidMETracks{this,"MuidMETracksLocation","MuidMETracks","Muidco ME Tracks"};
 
 };
 

Reconstruction/MuonIdentification/MuonCombinedAlgs/src/MuonCombinedInDetExtensionAlg.cxx

@@ -28,6 +28,9 @@ StatusCode MuonCombinedInDetExtensionAlg::initialize()
   ATH_CHECK(m_sTGC_ContainerName.initialize(m_usePRDs && m_useNSW));
   ATH_CHECK(m_MM_ContainerName.initialize(m_usePRDs && m_useNSW));
   ATH_CHECK(m_tagMap.initialize());
+  ATH_CHECK(m_combTracks.initialize(m_combTracks.key()!=""));
+  ATH_CHECK(m_METracks.initialize(m_METracks.key()!=""));
+  ATH_CHECK(m_segments.initialize(m_segments.key()!=""));
 
   return StatusCode::SUCCESS; 
 }
@@ -44,23 +47,50 @@ StatusCode MuonCombinedInDetExtensionAlg::execute()
   SG::WriteHandle<MuonCombined::InDetCandidateToTagMap> tagMap(m_tagMap);
   ATH_CHECK( tagMap.record (std::make_unique<MuonCombined::InDetCandidateToTagMap>()) );
 
+  TrackCollection* combTracks=nullptr;
+  TrackCollection* meTracks=nullptr;
+  Trk::SegmentCollection* segments=nullptr;
+
+  if(m_combTracks.key()!=""){
+    SG::WriteHandle<TrackCollection> wh_combTracks(m_combTracks);
+    ATH_CHECK(wh_combTracks.record(std::make_unique<TrackCollection>()));
+    combTracks=wh_combTracks.ptr();
+  }
+  if(m_METracks.key()!=""){
+    SG::WriteHandle<TrackCollection> wh_meTracks(m_METracks);
+    ATH_CHECK(wh_meTracks.record(std::make_unique<TrackCollection>()));
+    meTracks=wh_meTracks.ptr();
+  }
+
+  if(m_segments.key()!=""){
+    SG::WriteHandle<Trk::SegmentCollection> wh_segs(m_segments);
+    ATH_CHECK(wh_segs.record(std::make_unique<Trk::SegmentCollection>()));
+    segments=wh_segs.ptr();
+  }
+
   if(m_usePRDs){
+    MuonCombined::IMuonCombinedInDetExtensionTool::MuonPrdData prdData;
     SG::ReadHandle<Muon::MdtPrepDataContainer> mdtPRDContainer(m_MDT_ContainerName);
-    SG::ReadHandle<Muon::CscPrepDataContainer> cscPRDContainer(m_CSC_ContainerName);
-    SG::ReadHandle<Muon::RpcPrepDataContainer> rpcPRDContainer(m_RPC_ContainerName);
-    SG::ReadHandle<Muon::TgcPrepDataContainer> tgcPRDContainer(m_TGC_ContainerName);
-    if (m_useNSW) {
+    prdData.mdtPrds=mdtPRDContainer.cptr();
+    if(!m_useNSW){
+      SG::ReadHandle<Muon::CscPrepDataContainer> cscPRDContainer(m_CSC_ContainerName);
+      prdData.cscPrds=cscPRDContainer.cptr();
+    }
+    else{
       SG::ReadHandle<Muon::sTgcPrepDataContainer> stgcPRDContainer(m_sTGC_ContainerName);
       SG::ReadHandle<Muon::MMPrepDataContainer> mmPRDContainer(m_MM_ContainerName);
-      for(auto& tool : m_muonCombinedInDetExtensionTools)
-        tool->extendWithPRDs(*indetCandidateCollection,tagMap.ptr(),mdtPRDContainer.cptr(),cscPRDContainer.cptr(),rpcPRDContainer.cptr(),tgcPRDContainer.cptr(),stgcPRDContainer.cptr(),mmPRDContainer.cptr());
-     } else {
-       for(auto& tool : m_muonCombinedInDetExtensionTools)
-         tool->extendWithPRDs(*indetCandidateCollection,tagMap.ptr(),mdtPRDContainer.cptr(),cscPRDContainer.cptr(),rpcPRDContainer.cptr(),tgcPRDContainer.cptr(),0,0);
-     }
+      prdData.stgcPrds=stgcPRDContainer.cptr();
+      prdData.mmPrds=mmPRDContainer.cptr();
+    }
+    SG::ReadHandle<Muon::RpcPrepDataContainer> rpcPRDContainer(m_RPC_ContainerName);
+    prdData.rpcPrds=rpcPRDContainer.cptr();
+    SG::ReadHandle<Muon::TgcPrepDataContainer> tgcPRDContainer(m_TGC_ContainerName);
+    prdData.tgcPrds=tgcPRDContainer.cptr();
+    for(auto& tool : m_muonCombinedInDetExtensionTools)
+      tool->extendWithPRDs(*indetCandidateCollection,tagMap.ptr(),prdData,combTracks,meTracks,segments); 
   } else{
     for(auto& tool : m_muonCombinedInDetExtensionTools)
-      tool->extend(*indetCandidateCollection,tagMap.ptr());
+      tool->extend(*indetCandidateCollection,tagMap.ptr(),combTracks,meTracks,segments);
   }
   
   return StatusCode::SUCCESS;

Reconstruction/MuonIdentification/MuonCombinedAlgs/src/MuonCombinedInDetExtensionAlg.h

@@ -13,6 +13,8 @@
 #include "MuonPrepRawData/MuonPrepDataContainer.h"
 #include "MuonPrepRawData/MMPrepDataContainer.h"
 #include "MuonPrepRawData/sTgcPrepDataContainer.h"
+#include "TrkTrack/TrackCollection.h"
+#include "TrkSegment/SegmentCollection.h"
 #include <string>
 
 namespace MuonCombined {
@@ -40,6 +42,9 @@ class MuonCombinedInDetExtensionAlg : public AthAlgorithm
   SG::ReadHandleKey<Muon::sTgcPrepDataContainer> m_sTGC_ContainerName{this,"sTGCPrepDataLocation","sTGC_Clusters","sTGC prep data"};
   SG::ReadHandleKey<Muon::MMPrepDataContainer> m_MM_ContainerName{this,"MMPrepDataLocation","MM_Clusters","MM prep data"};
   SG::WriteHandleKey<MuonCombined::InDetCandidateToTagMap> m_tagMap{this,"TagMap","muGirlTagMap","tag map"};
+  SG::WriteHandleKey<TrackCollection> m_combTracks{this,"CombinedTrackCollection","MuGirlCombinedTracks","combined track collection"};
+  SG::WriteHandleKey<TrackCollection> m_METracks{this,"METrackCollection","MuGirlMETracks","ME track collection"};
+  SG::WriteHandleKey<Trk::SegmentCollection> m_segments{this,"SegmentCollection","MuGirlSegments","Segment collection"};
   bool m_usePRDs;
   bool m_useNSW;
 };

Reconstruction/MuonIdentification/MuonCombinedAlgs/src/MuonCombinedMuonCandidateAlg.cxx

@@ -16,6 +16,7 @@ StatusCode MuonCombinedMuonCandidateAlg::initialize()
   ATH_CHECK(m_muonCandidateTool.retrieve());
   ATH_CHECK(m_muonTrackParticleLocation.initialize());
   ATH_CHECK(m_candidateCollectionName.initialize());
+  ATH_CHECK(m_msOnlyTracks.initialize());
   return StatusCode::SUCCESS; 
 }
 
@@ -35,7 +36,9 @@ StatusCode MuonCombinedMuonCandidateAlg::execute()
 
   SG::WriteHandle<MuonCandidateCollection> muonCandidates(m_candidateCollectionName);
   ATH_CHECK(muonCandidates.record(std::make_unique<MuonCandidateCollection>()));
-  m_muonCandidateTool->create(*muonTrackParticles,*muonCandidates);
+  SG::WriteHandle<TrackCollection> msOnlyTracks(m_msOnlyTracks);
+  ATH_CHECK(msOnlyTracks.record(std::make_unique<TrackCollection>()));
+  m_muonCandidateTool->create(*muonTrackParticles,*muonCandidates,*msOnlyTracks);
 
   return StatusCode::SUCCESS;
 }

Reconstruction/MuonIdentification/MuonCombinedAlgs/src/MuonCombinedMuonCandidateAlg.h

@@ -12,6 +12,7 @@
 #include "xAODTracking/TrackParticleContainer.h"
 #include "MuonCombinedEvent/MuonCandidateCollection.h"
 #include "MuonCombinedToolInterfaces/IMuonCandidateTool.h"
+#include "TrkTrack/TrackCollection.h"
 #include <string>
 
 namespace MuonCombined {
@@ -33,7 +34,7 @@ class MuonCombinedMuonCandidateAlg : public AthAlgorithm
   ToolHandle<MuonCombined::IMuonCandidateTool> m_muonCandidateTool{this,"MuonCandidateTool","MuonCombined::MuonCandidateTool/MuonCandidateTool","Muon Candidate Tool"};
   SG::ReadHandleKey<xAOD::TrackParticleContainer> m_muonTrackParticleLocation{this,"MuonSpectrometerTrackParticleLocation","MuonSpectrometerTrackParticles","MS Track Particle collection"};
   SG::WriteHandleKey<MuonCandidateCollection> m_candidateCollectionName{this,"MuonCandidateLocation","MuonCandidates","Muon candidate collection"};
-
+  SG::WriteHandleKey<TrackCollection> m_msOnlyTracks{this,"MSOnlyExtrapolatedTrackLocation","MSOnlyExtrapolatedMuonTracks","MS extrapolated muon tracks"};
 };
 
 

Reconstruction/MuonIdentification/MuonCombinedAlgs/src/MuonCreatorAlg.cxx

@@ -40,22 +40,13 @@ StatusCode MuonCreatorAlg::initialize()
   ATH_CHECK(m_muonCandidateCollectionName.initialize(!m_buildSlowMuon));
   //Can't use a flag in intialize for an array of keys
   if(!m_doSA) ATH_CHECK(m_tagMaps.initialize());
-  m_segTrkContainerName = "Trk"+m_segContainerName.key();
-  m_segContainerName = "xaod"+m_segContainerName.key();
   ATH_CHECK(m_segContainerName.initialize());
-  ATH_CHECK(m_segTrkContainerName.initialize());
-  m_combinedTrkCollectionName = m_combinedCollectionName.key()+"Tracks";
   m_combinedCollectionName = m_combinedCollectionName.key()+"TrackParticles";
   ATH_CHECK(m_combinedCollectionName.initialize());
-  ATH_CHECK(m_combinedTrkCollectionName.initialize());
-  m_extrapolatedTrkCollectionName = m_extrapolatedCollectionName.key()+"Tracks";
   m_extrapolatedCollectionName = m_extrapolatedCollectionName.key()+"TrackParticles";
   ATH_CHECK(m_extrapolatedCollectionName.initialize());
-  ATH_CHECK(m_extrapolatedTrkCollectionName.initialize());
-  m_msOnlyExtrapolatedTrkCollectionName = m_msOnlyExtrapolatedCollectionName.key()+"Tracks";
   m_msOnlyExtrapolatedCollectionName = m_msOnlyExtrapolatedCollectionName.key()+"TrackParticles";
   ATH_CHECK(m_msOnlyExtrapolatedCollectionName.initialize());
-  ATH_CHECK(m_msOnlyExtrapolatedTrkCollectionName.initialize());
   ATH_CHECK(m_clusterContainerName.initialize(m_makeClusters));
   m_clusterContainerLinkName = m_clusterContainerName.key()+"_links";
   ATH_CHECK(m_clusterContainerLinkName.initialize(m_makeClusters));
@@ -88,43 +79,27 @@ StatusCode MuonCreatorAlg::execute()
   
   MuonCombined::IMuonCreatorTool::OutputData output(*(wh_muons.ptr()));
 
-
-
   // Create and record track particles:
   //combined tracks
   SG::WriteHandle<xAOD::TrackParticleContainer> wh_combtp(m_combinedCollectionName);
-  SG::WriteHandle<TrackCollection> wh_combtrk(m_combinedTrkCollectionName);
   ATH_CHECK(wh_combtp.record(std::make_unique<xAOD::TrackParticleContainer>(), std::make_unique<xAOD::TrackParticleAuxContainer>()));
-  ATH_CHECK(wh_combtrk.record(std::make_unique<TrackCollection>()));
   output.combinedTrackParticleContainer = wh_combtp.ptr();
-  output.combinedTrackCollection = wh_combtrk.ptr();
 
   //extrapolated tracks
   SG::WriteHandle<xAOD::TrackParticleContainer> wh_extrtp(m_extrapolatedCollectionName);
-  SG::WriteHandle<TrackCollection> wh_extrtrk(m_extrapolatedTrkCollectionName);
   ATH_CHECK(wh_extrtp.record(std::make_unique<xAOD::TrackParticleContainer>(), std::make_unique<xAOD::TrackParticleAuxContainer>()));
-  ATH_CHECK(wh_extrtrk.record(std::make_unique<TrackCollection>()));
   output.extrapolatedTrackParticleContainer = wh_extrtp.ptr();
-  output.extrapolatedTrackCollection = wh_extrtrk.ptr();
 
   //msonly tracks
   SG::WriteHandle<xAOD::TrackParticleContainer> wh_msextrtp(m_msOnlyExtrapolatedCollectionName);
-  SG::WriteHandle<TrackCollection> wh_msextrtrk(m_msOnlyExtrapolatedTrkCollectionName);
   ATH_CHECK(wh_msextrtp.record(std::make_unique<xAOD::TrackParticleContainer>(), std::make_unique<xAOD::TrackParticleAuxContainer>()));
-  ATH_CHECK(wh_msextrtrk.record(std::make_unique<TrackCollection>()));
   output.msOnlyExtrapolatedTrackParticleContainer = wh_msextrtp.ptr();
-  output.msOnlyExtrapolatedTrackCollection = wh_msextrtrk.ptr();
-
 
   //segments
   SG::WriteHandle<xAOD::MuonSegmentContainer> wh_segment(m_segContainerName);
   ATH_CHECK(wh_segment.record(std::make_unique<xAOD::MuonSegmentContainer>(),std::make_unique<xAOD::MuonSegmentAuxContainer>()));
   output.xaodSegmentContainer=wh_segment.ptr();
 
-  SG::WriteHandle<Trk::SegmentCollection> wh_segmentTrk(m_segTrkContainerName);
-  ATH_CHECK(wh_segmentTrk.record(std::make_unique<Trk::SegmentCollection>()));
-  output.muonSegmentCollection=wh_segmentTrk.ptr();
-
   // calo clusters
   SG::WriteHandle<xAOD::CaloClusterContainer> wh_clusters;
   SG::WriteHandle<CaloClusterCellLinkContainer> wh_clusterslink;

Reconstruction/MuonIdentification/MuonCombinedAlgs/src/MuonCreatorAlg.h

@@ -42,16 +42,12 @@ class MuonCreatorAlg : public AthAlgorithm
   SG::WriteHandleKey<xAOD::MuonContainer> m_muonCollectionName{this,"MuonContainerLocation", "Muons", "Muon Container"};
   SG::WriteHandleKey<xAOD::SlowMuonContainer> m_slowMuonCollectionName{this, "SlowMuonContainerLocation", "SlowMuons", "Slow Muon Container"};
   SG::WriteHandleKey<xAOD::TrackParticleContainer> m_combinedCollectionName{this, "CombinedLocation", "CombinedMuon", "Combined muons"};
-  SG::WriteHandleKey<TrackCollection> m_combinedTrkCollectionName{this,"CombinedTrackLocation","CombinedMuon","Combined muon tracks"};
   SG::WriteHandleKey<xAOD::TrackParticleContainer> m_extrapolatedCollectionName{this, "ExtrapolatedLocation", "ExtrapolatedMuon", "Extrapolated muons"};
-  SG::WriteHandleKey<TrackCollection> m_extrapolatedTrkCollectionName{this,"ExtrapolatedTrackLocation","ExtrapolatedMuon","Extrapolated muon tracks"};
   SG::WriteHandleKey<xAOD::TrackParticleContainer> m_msOnlyExtrapolatedCollectionName{this, "MSOnlyExtrapolatedLocation", "MSOnlyExtrapolatedMuon", "MS Extrapolated muons"};
-  SG::WriteHandleKey<TrackCollection> m_msOnlyExtrapolatedTrkCollectionName{this,"MSOnlyExtrapolatedTrackLocation","MSOnlyExtrapolatedMuon","MS extrapolated muon tracks"};
   SG::ReadHandleKey<InDetCandidateCollection> m_indetCandidateCollectionName{this,"InDetCandidateLocation","InDetCandidates","ID candidates"};
   SG::ReadHandleKey<MuonCandidateCollection> m_muonCandidateCollectionName{this,"MuonCandidateLocation","MuonCandidates","Muon candidates"};
   SG::ReadHandleKeyArray<MuonCombined::InDetCandidateToTagMap> m_tagMaps{this,"TagMaps",{"muidcoTagMap","stacoTagMap","muGirlTagMap","caloTagMap","segmentTagMap"},"ID candidate to tag maps"};
-  SG::WriteHandleKey<xAOD::MuonSegmentContainer> m_segContainerName{this, "SegmentContainerName", "MuonSegments", "Segments"};
-  SG::WriteHandleKey<Trk::SegmentCollection> m_segTrkContainerName{this,"TrackSegmentContainerName","MuonSegments","Track segments"};
+  SG::WriteHandleKey<xAOD::MuonSegmentContainer> m_segContainerName{this, "SegmentContainerName", "MuGirlMuonSegments", "New segments"};
   SG::WriteHandleKey<xAOD::CaloClusterContainer> m_clusterContainerName{this, "ClusterContainerName", "MuonClusterCollection", "Clusters"};
   SG::WriteHandleKey<CaloClusterCellLinkContainer> m_clusterContainerLinkName{this,"CaloClusterCellLinkName","MuonClusterCollection","Cluster links"};
   bool m_buildSlowMuon;

Reconstruction/MuonIdentification/MuonCombinedBaseTools/src/MuGirlTagTool.cxx

@@ -50,18 +50,18 @@ namespace MuonCombined {
     return StatusCode::SUCCESS;
   }
 
-  void MuGirlTagTool::extend( const InDetCandidateCollection& inDetCandidates, InDetCandidateToTagMap* tagMap ) {
-    if(!tagMap) ATH_MSG_DEBUG("We don't run this tool any more anyway");
+  void MuGirlTagTool::extend( const InDetCandidateCollection& inDetCandidates, InDetCandidateToTagMap* tagMap, TrackCollection* combTracks, TrackCollection* meTracks, 
+			      Trk::SegmentCollection* segments ) {
+    if(!tagMap || !combTracks || !meTracks || !segments) ATH_MSG_DEBUG("We don't run this tool any more anyway");
     if (m_pMuGirlReconstruction->MuGirlReco(inDetCandidates).isFailure()) {
       ATH_MSG_ERROR("MuGirl reconstruction failed!");
     }
   }
 
-  void MuGirlTagTool::extendWithPRDs( const InDetCandidateCollection& inDetCandidates, InDetCandidateToTagMap* tagMap, 
-				      const Muon::MdtPrepDataContainer* mdtPRDs, const Muon::CscPrepDataContainer* cscPRDs, const Muon::RpcPrepDataContainer* rpcPRDs, 
-				      const Muon::TgcPrepDataContainer *tgcPRDs, const Muon::sTgcPrepDataContainer* stgcPRDs, const Muon::MMPrepDataContainer* mmPRDs ) {
+  void MuGirlTagTool::extendWithPRDs( const InDetCandidateCollection<