MsTrackTester.cxx

Go to the documentation of this file.

/*
  Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
*/
#include "MsTrackTester.h"
 
#include "StoreGate/ReadHandle.h"
#include "MuonTrackEvent/TrackingHelpers.h"
#include "MuonTrackEvent/HitSummary.h"
#include "MuonTruthHelpers/MuonSimHitHelpers.h"
#include "MuonDetDescrUtils/MuonSectorMapping.h"
#include "MuonPRDTestR4/TrackContainerModule.h"
#include "MuonTesterTree/MuonTesterTreeDict.h"
#include "xAODTruth/xAODTruthHelpers.h"
 
#include "Acts/Definitions/Units.hpp"
 
#include <format>
 
using namespace MuonVal;
using namespace MuonPRDTest;
using namespace MuonR4;
using namespace Acts::UnitLiterals;
 
namespace {
    using SegLink_t = ElementLink<xAOD::MuonSegmentContainer>;
    using SegLinkVec_t = std::vector<SegLink_t>;
    static const SG::ConstAccessor<SegLinkVec_t> acc_truthSegLinks{"truthSegmentLinks"};
    constexpr double MeVtoGeV = 1.e-3;
    
    using Location = MsTrackSeeder::Location;
 
 
    constexpr int overlapSector(const int sec1 , const int sec2) {
        constexpr int nSec = Muon::MuonStationIndex::numberOfSectors();
        if (sec2 > sec1) return overlapSector(sec2, sec1);
        if (sec1 == 1 && sec2 == nSec) return 0;
        return sec1 + sec2;
    }   
 
}
 
namespace MuonValR4 {
    std::optional<MsTrackSeed> MsTrackTester::makeSeedFromTruth(const ActsTrk::GeometryContext& gctx,
                                                                const xAOD::TruthParticle& truthMuon) const {
        int secMax{-1}, secMin{100};
        std::vector<const xAOD::MuonSegment*> matchedSegs = MuonR4::getTruthSegments(truthMuon);
        if (matchedSegs.empty()) {
            return std::nullopt;
        }
        for (const xAOD::MuonSegment* seg : matchedSegs) {
            secMax = std::max(secMax, seg->sector());
            secMin = std::min(secMin, seg->sector());
        }
        const int orSec = overlapSector(secMax, secMin);
        MsTrackSeed barrelSeed{Location::Barrel, orSec};
        MsTrackSeed endcapSeed{Location::Endcap, orSec};
        for (const xAOD::MuonSegment* seg : matchedSegs) {
            barrelSeed.addSegment(seg);
            endcapSeed.addSegment(seg);
        }
        barrelSeed.setPosition(matchedSegs[0]->position());
        endcapSeed.setPosition(matchedSegs[0]->position());
        const auto [barrelLength, barrelTheta] = calcSeedLength(gctx, barrelSeed);
        const auto [endcapLength, endcapTheta] = calcSeedLength(gctx, endcapSeed);
        ATH_MSG_VERBOSE(__func__<<"() "<<__LINE__<<" - Constructed new seed from truth muon wih pT:"
                <<(truthMuon.pt()/ Gaudi::Units::GeV)<<" [GeV], eta: "<<truthMuon.eta()
                <<", phi: "<<(truthMuon.phi() / 1._degree)<<", q: "<<truthMuon.charge()
                <<", matchedSeg: "<<matchedSegs.size()<< " barrel (L/theta): "<<barrelLength
                <<"/"<<(barrelTheta / 1._degree)<<" - endcap (L/theta): "
                <<endcapLength<<"/"<<(endcapTheta / 1._degree)<<"\n"<<barrelSeed);
        if (barrelLength < 0 && endcapLength < 0) {
            ATH_MSG_WARNING(__func__<<"() "<<__LINE__<<" - Invalid seed");
            return std::nullopt;
        }
        return barrelLength < 0 || std::abs(endcapLength) < barrelLength 
              ? endcapSeed : barrelSeed;
    }
 
    std::pair<double, double> MsTrackTester::calcSeedLength(const ActsTrk::GeometryContext& gctx, 
                                                            const MuonR4::MsTrackSeed& seed) const {
        double maxL{-1.*Gaudi::Units::km}, minL{1.*Gaudi::Units::km},
               maxTheta{-M_PI}, minTheta{M_PI};
        for (const xAOD::MuonSegment* seg : seed.segments()) {
            const auto secProj = m_seeder->projectorFromSeed(*seg, seed);
            const Amg::Vector2D projPos{m_seeder->expressOnCylinder(gctx, *seg, seed.location(), secProj)};
            if (!m_seeder->withinBounds(projPos, seed.location())) {
                continue;
            }
            const double projected = projPos[seed.location()==Location::Barrel];
            const double theta = seg->direction().theta();
            minL = std::min(minL, projected);
            maxL = std::max(maxL, projected);
            minTheta = std::min(minTheta, theta);
            maxTheta = std::max(maxTheta, theta);
        }
        return std::make_pair(maxL - minL, maxTheta - minTheta);
    }
    StatusCode MsTrackTester::initialize() {
        ATH_CHECK(m_truthSegmentKey.initialize(m_isMC));
        ATH_CHECK(m_truthKey.initialize(m_isMC));
        ATH_CHECK(m_msTrkSeedKey.initialize());
        ATH_CHECK(m_recoSegmentKey.initialize());
        ATH_CHECK(m_segSelector.retrieve());
        ATH_CHECK(m_geoCtxKey.initialize());
        ATH_CHECK(m_fieldCacheKey.initialize());
        ATH_CHECK(m_trackKey.initialize());
        ATH_CHECK(m_summaryTool.retrieve());
 
        ATH_CHECK(m_legacyTrackKey.initialize(!m_legacyTrackKey.empty()));
        ATH_CHECK(m_legacySegmentKey.initialize(!m_legacySegmentKey.empty()));
        ATH_CHECK(m_legacyMuonKey.initialize(!m_legacyMuonKey.empty()));
 
        ATH_CHECK(detStore()->retrieve(m_detMgr));
 
        MsTrackSeeder::Config seederCfg{};
        seederCfg.detMgr = m_detMgr;
        seederCfg.seedHalfLength = 2.*Gaudi::Units::m;
        seederCfg.endcapDiscRadius = 40.*Gaudi::Units::m;
        m_seeder =  std::make_unique<MuonR4::MsTrackSeeder>(name(), std::move(seederCfg));
 
        int evOpts{0};
 
        m_recoSegs = std::make_unique<SegmentVariables>(m_tree, m_recoSegmentKey.key(), "Segments", msgLevel());
        m_recoSegs->addVariable(std::make_unique<MuonVal::GenericAuxDecorationBranch<unsigned short>>(m_tree, 
            "Segments_passSeedQual",[this](const SG::AuxElement* aux){
            const auto* seg = static_cast<const xAOD::MuonSegment*>(aux);
            return m_segSelector->passSeedingQuality(Gaudi::Hive::currentContext(),
                                                     *MuonR4::detailedSegment(*seg)); }));
        m_recoSegs->addVariable(std::make_unique<MuonVal::GenericAuxDecorationBranch<unsigned short>>(m_tree, 
                                "Segments_passTrackQual",[this](const SG::AuxElement* aux){
            const auto* seg = static_cast<const xAOD::MuonSegment*>(aux);
            return m_segSelector->passTrackQuality(Gaudi::Hive::currentContext(),
                                                   *MuonR4::detailedSegment(*seg)); 
        }));
 
 
        if (m_isMC) {
            evOpts |= EventInfoBranch::isMC;
            m_recoSegs->addVariable(std::make_unique<MuonVal::GenericAuxDecorationBranch<unsigned short>>(m_tree, 
                "Segments_truthSegLink",[this](const SG::AuxElement* aux){
                const auto* seg = static_cast<const xAOD::MuonSegment*>(aux);
                const xAOD::MuonSegment* truthS = MuonR4::getMatchedTruthSegment(*seg);
                const unsigned linkIdx = truthS ? m_truthSegs->push_back(*truthS) : -1;
                if (truthS) {
                    m_truthSegToRecoLink.push_back(linkIdx, m_recoSegs->push_back(*seg));
                    const xAOD::TruthParticle* truthMuon = MuonR4::getTruthMatchedParticle(*truthS);
                    if (truthMuon){
                        m_truthTrks->push_back(truthMuon);
                        m_truthMuRecoSegLinks[m_truthTrks->find(truthMuon)].push_back(m_recoSegs->push_back(*seg));
                    }
                }
                return linkIdx;
            }));
 
            m_truthSegs = std::make_unique<SegmentVariables>(m_tree, m_truthSegmentKey.key(), "TruthSegments", msgLevel());
            
            m_truthSegs->addVariable(std::make_unique<MuonVal::GenericAuxDecorationBranch<unsigned short>>(m_tree, 
                    "TruthSegments_truthLink",[this](const SG::AuxElement* aux){
                    const auto* seg = static_cast<const xAOD::MuonSegment*>(aux);
                    const xAOD::TruthParticle* truthP = MuonR4::getTruthMatchedParticle(*seg);
                    m_truthTrks->push_back(truthP);
                    unsigned short linkIdx = m_truthTrks->find(truthP);                    
                    return linkIdx;
                }));
            m_truthSegs->addVariable(std::make_unique<MuonVal::GenericAuxDecorationBranch<unsigned short>>(m_tree,
                "TruthSegments_hasBarrelProj", [this](const SG::AuxElement* aux){
                    const auto* seg = static_cast<const xAOD::MuonSegment*>(aux);
                    const ActsTrk::GeometryContext* gctx{nullptr};
                    SG::get(gctx, m_geoCtxKey, Gaudi::Hive::currentContext()).ignore();
                    using enum MuonR4::MsTrackSeeder::SectorProjector;
                    for (const auto proj : {leftOverlap, center, rightOverlap}) {
                        const Amg::Vector2D projPos{m_seeder->expressOnCylinder(*gctx, *seg, Location::Barrel, proj)};
                        if (m_seeder->withinBounds(projPos, Location::Barrel)) {
                            return 1;
                        }
                    }
                    return 0;
                }));
 
            m_truthSegs->addVariable(std::make_unique<MuonVal::GenericAuxDecorationBranch<unsigned short>>(m_tree,
                "TruthSegments_hasEndcapProj", [this](const SG::AuxElement* aux){
                    const auto* seg = static_cast<const xAOD::MuonSegment*>(aux);
                    const ActsTrk::GeometryContext* gctx{nullptr};
                    SG::get(gctx, m_geoCtxKey, Gaudi::Hive::currentContext()).ignore();
                    using enum MuonR4::MsTrackSeeder::SectorProjector;
                    for (const auto proj : {leftOverlap, center, rightOverlap}) {
                        const Amg::Vector2D projPos{m_seeder->expressOnCylinder(*gctx, *seg, Location::Endcap, proj)};
                        if (m_seeder->withinBounds(projPos, Location::Endcap)) {
                            return 1;
                        }
                    }
                    return 0;
                }));
            m_tree.addBranch(m_truthSegs);
 
            m_truthTrks = std::make_unique<IParticleFourMomBranch>(m_tree, "TruthMuons");
            m_truthTrks->addVariable<int>(-1, "truthOrigin");
            m_truthTrks->addVariable<int>(-1, "truthType");
            m_tree.addBranch(m_truthTrks);
            m_trkTruthLinks.emplace_back(m_truthSegmentKey, "truthParticleLink");
            m_trkTruthLinks.emplace_back(m_truthKey, "truthSegmentLinks");
            m_trkTruthLinks.emplace_back(m_recoSegmentKey, "truthSegmentLink");
        }
 
        m_tree.addBranch(m_recoSegs);
        m_tree.addBranch(std::make_unique<EventInfoBranch>(m_tree, evOpts));
        m_tree.addBranch(std::make_unique<TrackContainerModule>(m_tree, m_trackKey.key(), msgLevel()));
 
        if(!m_legacyTrackKey.empty()) {
            m_legacyTrks = std::make_unique<IParticleFourMomBranch>(m_tree, "LegacyMSTrks");
            m_legacyTrks->addVariable(std::make_unique<TrackChi2Branch>(*m_legacyTrks));
            if (m_isMC) {
                BilateralLinkerBranch::connectCollections(m_legacyTrks, m_truthTrks, [](const xAOD::IParticle* trk){ 
                                                      return xAOD::TruthHelpers::getTruthParticle(*trk); }, "truth", "LegacyMS");
            }
            for (const auto& summary :{
                 // Inner
                "innerSmallHits", "innerLargeHits", "innerSmallHoles", "innerLargeHoles",
                // Middle
                "middleSmallHits", "middleLargeHits", "middleSmallHoles",
                "middleLargeHoles",
                // Outer
                "outerSmallHits", "outerLargeHits", "outerSmallHoles", "outerLargeHoles",
                // Extended
                "extendedSmallHits", "extendedLargeHits", "extendedSmallHoles",
                "extendedLargeHoles",
                "etaLayer1Hits", "phiLayer1Hits", "etaLayer2Hits", "phiLayer2Hits",
                "etaLayer3Hits", "phiLayer3Hits", "etaLayer4Hits", "phiLayer4Hits",
                "etaLayer1Holes", "phiLayer1Holes", "etaLayer2Holes", "phiLayer2Holes",
                "etaLayer3Holes", "phiLayer3Holes", "etaLayer4Holes", "phiLayer4Holes" }) {
                m_legacyTrks->addVariable<uint8_t>(-1, summary); 
            }
 
            m_tree.addBranch(m_legacyTrks);
 
 
            m_legacyRecoSegs = std::make_unique<SegmentVariables>(m_tree, m_legacySegmentKey.key(), "LegacyRecoSegments", msgLevel());
            m_tree.addBranch(m_legacyRecoSegs);
        } 
 
 
        m_seedSummary = std::make_shared<TrackSummaryModule>(m_tree, "MsTrkSeed", m_summaryTool.get());
        m_trackSummary = std::make_shared<TrackSummaryModule>(m_tree, "ActsMsTracks", m_summaryTool.get());
        m_tree.addBranch(m_seedSummary);
        m_tree.addBranch(m_trackSummary);
        
        ATH_CHECK(m_trkTruthLinks.initialize());
        ATH_CHECK(m_tree.init(this));
        return StatusCode::SUCCESS;
    }
    StatusCode MsTrackTester::execute() {
        const EventContext& ctx{Gaudi::Hive::currentContext()};
 
        const xAOD::TrackParticleContainer* legacyTrks{nullptr};
        ATH_CHECK(SG::get(legacyTrks, m_legacyTrackKey, ctx));
 
        //This for now is to be able to retrieve the matching between the legacy segments and tracks ...
        const xAOD::MuonContainer* legacyMuons{nullptr};
        ATH_CHECK(SG::get(legacyMuons, m_legacyMuonKey, ctx));
 
        //Dump also legacy segments
        const xAOD::MuonSegmentContainer* legacyRecoSegs{nullptr};
        ATH_CHECK(SG::get(legacyRecoSegs, m_legacySegmentKey, ctx));
 
 
        std::unordered_map<const xAOD::MuonSegment*, unsigned short> segIndex;
        
        if (legacyTrks){
            unsigned short iTrk = 0;
            for (const xAOD::TrackParticle* track : *legacyTrks) {
                ATH_MSG_VERBOSE("Legacy track "<< iTrk << ": pT: "<<(track->pt() *MeVtoGeV)<<" [GeV], eta: "<<track->eta()
                                <<", phi: "<<(track->phi() / 1._degree)<<", q: "<<track->charge());
                m_summaryTool->copySummary(m_summaryTool->makeSummary(ctx, *track->track()),
                                           *track);
                m_legacyTrks->push_back(track);
 
                //Change navigation from down to particle track? Initially only tracks were stored/looked at, whereas the following code is to navigate to associated muon segments ... in principle can change the order and thus avoid the matching part and some extra looping
                unsigned iMuon = 0;
                for (const xAOD::Muon* muon : *legacyMuons){
                    ATH_MSG_VERBOSE("iMuon " << iMuon << " pT: "<<(muon->pt() *MeVtoGeV)<<" [GeV], eta: "<<muon->eta()
                                <<", phi: "<<(muon->phi() / 1._degree)<<", q: "<<muon->charge()  <<", nSegments: "<<muon->nMuonSegments());
 
                   ++iMuon;
                  
                   //MS or ME? MS as we probably don't have yet material and extrapolated to IP 
                   const xAOD::TrackParticle* msTrack = muon->trackParticle( xAOD::Muon::MuonSpectrometerTrackParticle);
                   if ( !msTrack || (msTrack != track) ) { continue; }
 
                   // retrieve associated segments for this track
                   for (size_t s = 0; s < muon->nMuonSegments(); ++s) {
                        const xAOD::MuonSegment* segment = muon->muonSegment(s);
 
                        ATH_MSG_VERBOSE(std::format( "Legacy muon-segment link: segment {:}  @{:}, eta: {:.2f}, phi {:.2f}", 
                                        printID(*segment), Amg::toString(segment->position()),
                                        segment->direction().eta(), segment->direction().phi() / 1._degree));
 
                         segIndex[segment] = iTrk;
                   }
 
                }
                ++iTrk;
            }
        }
 
        if (legacyRecoSegs) {
            unsigned short iSeg = 0;
            m_legacySegToTrkLinks[legacyRecoSegs->size()-1];
            for (const xAOD::MuonSegment* seg : *legacyRecoSegs) {
                //Store all segments
                m_legacyRecoSegs->push_back(*seg);
 
                //Get the ordering in the legacy segment container to be able to link to the legacy tracks later on
                m_legacySegToTrkLinks[iSeg] = segIndex.find(seg) != segIndex.end() ? segIndex[seg] : -1;
 
                ATH_MSG_VERBOSE(std::format( "Dump legacy reco segment index {:} {:}  @{:}, eta: {:.2f}, phi {:.2f} associated muon: {:}", 
                                iSeg, printID(*seg), Amg::toString(seg->position()),
                                seg->direction().eta(), seg->direction().phi() / 1._degree, m_legacySegToTrkLinks[iSeg]));
 
                ++iSeg;
            }
        }
        
        const xAOD::MuonSegmentContainer* recoSegments{nullptr};
        ATH_CHECK(SG::get(recoSegments, m_recoSegmentKey, ctx));
        const MuonR4::MsTrackSeedContainer* trkSeeds{nullptr};
        ATH_CHECK(SG::get(trkSeeds, m_msTrkSeedKey, ctx));
        const ActsTrk::GeometryContext* gctx{nullptr};
        ATH_CHECK(SG::get(gctx, m_geoCtxKey, ctx));
        const AtlasFieldCacheCondObj* magCache{nullptr};
        ATH_CHECK(SG::get(magCache, m_fieldCacheKey, ctx));
        
        std::map<const xAOD::TruthParticle*, std::vector<unsigned>> truthToSeedMatchCounter{};
        for (const MuonR4::MsTrackSeed& seed : *trkSeeds) {
            unsigned int seedIdx = m_seedPos.size();
            m_seedPos += seed.position();
            m_seedType+= Acts::toUnderlying(seed.location());
            m_seedSector += seed.sector();
            m_seedSummary->push_back(ctx, seed);
            ATH_MSG_VERBOSE(" Dump new seed: "<<seed);
            for (const xAOD::MuonSegment* seg : seed.segments()){
                m_seedRecoSegMatch[seedIdx].push_back(m_recoSegs->push_back(*seg));
                if (const xAOD::MuonSegment* truthSeg =  MuonR4::getMatchedTruthSegment(*seg);
                    truthSeg != nullptr) {
                    std::vector<unsigned>& matchCounter = truthToSeedMatchCounter[MuonR4::getTruthMatchedParticle(*truthSeg)];
                    if (seedIdx >= matchCounter.size()) {
                        matchCounter.resize(seedIdx +1);
                    }
                    ++matchCounter[seedIdx];
                }
            }
            const auto[seedLength, theta] = calcSeedLength(*gctx, seed);
            m_seedLength+= seedLength;
            m_seedThetaCone+=theta;
            m_seedQP += m_seeder->estimateQtimesP(*gctx, *magCache, seed) / Gaudi::Units::GeV; 
        }
        for (const xAOD::MuonSegment* seg : *recoSegments) {
            m_recoSegs->push_back(*seg);
        }
        const xAOD::MuonSegmentContainer* truthSegs{nullptr};
        ATH_CHECK(SG::get(truthSegs, m_truthSegmentKey, ctx));
        if (truthSegs) {
            for (const xAOD::MuonSegment* seg : *truthSegs) {
                ATH_MSG_VERBOSE(std::format( "Dump truth segment {:}  @{:}, eta: {:.2f}, phi {:.2f}", 
                                printID(*seg), Amg::toString(seg->position()),
                                seg->direction().eta(), seg->direction().phi() / 1._degree));
                m_truthSegs->push_back(*seg);
            }
            if (truthSegs->size()) {
                m_truthSegToRecoLink[truthSegs->size()-1];
            }
        }
        const xAOD::TruthParticleContainer* truthMuons{nullptr};
        ATH_CHECK(SG::get(truthMuons, m_truthKey, ctx));
        if (truthMuons && truthMuons->size()) {
            m_truthMuToSeedIdx[truthMuons->size() -1];
            m_truthMuToSeedCounter[truthMuons->size() -1];
            m_truthMuTruthSegLinks[truthMuons->size() -1];
            m_truthMuRecoSegLinks[truthMuons->size() -1];
 
            for (const xAOD::TruthParticle* truth : *truthMuons) {
                ATH_MSG_DEBUG("Truth muon: pT: "<<(truth->pt() *MeVtoGeV)<<", eta: "<<truth->eta()
                             <<", phi: "<<(truth->phi() / 1._degree)<<", q: "<<truth->charge());
 
                m_truthTrks->push_back(*truth);
                unsigned truthIdx = m_truthTrks->find(truth);
                std::vector<unsigned>& matchCounter = truthToSeedMatchCounter[truth];                
                std::vector<unsigned>::iterator maxSeed = matchCounter.begin();
                while ( (maxSeed = std::ranges::max_element(matchCounter))!=matchCounter.end() && (*maxSeed) > 0) {
                    m_truthMuToSeedIdx[truthIdx].push_back(std::distance(matchCounter.begin(), maxSeed));
                    m_truthMuToSeedCounter[truthIdx].push_back(*maxSeed);
                    (*maxSeed) = 0;
                }
                
                std::vector<unsigned short>& truthSegLinks = m_truthMuTruthSegLinks[truthIdx];
                const std::vector<const xAOD::MuonSegment*> truthSegs = getTruthSegments(*truth);
                m_truthMuTruthNSegs += truthSegs.size();
                for (const xAOD::MuonSegment* truthSeg: truthSegs) {
                    truthSegLinks.push_back(m_truthSegs->push_back(*truthSeg));
                }
 
                auto truthSeed = makeSeedFromTruth(*gctx, *truth);
                if (!truthSeed) {
                    m_truthMuonsSeedLength[truthIdx] = -1;
                    m_truthMuonsSeedCone[truthIdx] = -1;
                    m_truthMuonQP[truthIdx] =0;
                } else {
                    const auto [length, cone] = calcSeedLength(*gctx, *truthSeed);
                    m_truthMuonsSeedLength[truthIdx] = length;
                    m_truthMuonsSeedCone [truthIdx]  = cone;
                    m_truthMuonQP[truthIdx]  =  m_seeder->estimateQtimesP(*gctx, *magCache, *truthSeed) / Gaudi::Units::GeV;
                }
                 
            }
        }
        const ActsTrk::TrackContainer* msTracks{nullptr};
        ATH_CHECK(SG::get(msTracks, m_trackKey, ctx));
        for (const auto trk : *msTracks) {
            m_trackSummary->push_back(ctx, trk);
        }
 
        ATH_CHECK(m_tree.fill(ctx));
        return StatusCode::SUCCESS;
    }
    StatusCode MsTrackTester::finalize() {
        ATH_CHECK(m_tree.write());
        return StatusCode::SUCCESS;
    }
}