Merge branch 'TrkTruthMatch' into 'main'

MuonR4 - Add algorithm to perform track matching to truth particle

See merge request !77876

MuonSpectrometer/MuonPhaseII/MuonPatternRecognition/MuonTruthAlgsR4/CMakeLists.txt

@@ -10,7 +10,7 @@ atlas_add_component( MuonTruthAlgsR4
                      src/components/*.cxx src/*.cxx
                      LINK_LIBRARIES  xAODMuonSimHit AthenaBaseComps GaudiKernel xAODMuon MuonIdHelpersLib
                                      MuonReadoutGeometryR4 xAODMuonPrepData MuonPatternEvent xAODMuonViews
-                                     TrkTruthData MuonTruthHelpers)
+                                     TrkTruthData MuonTruthHelpers MuonRecHelperToolsLib)
 
 # Install files from the package:
 atlas_install_python_modules( python/*.py POST_BUILD_CMD ${ATLAS_FLAKE8} )

MuonSpectrometer/MuonPhaseII/MuonPatternRecognition/MuonTruthAlgsR4/python/MuonTruthAlgsConfig.py

@@ -87,24 +87,33 @@ def PrdMultiTruthMakerCfg(flags):
         result.addEventAlgo(the_alg) 
     return result
 
-def TruthSegmentToTruthPartCfg(flags, name="MuonTruthSegmentToTruthAssocAlg", **kwargs):
+def TruthSegmentToTruthPartAssocCfg(flags, name="MuonTruthSegmentToTruthAssocAlg", **kwargs):
     result = ComponentAccumulator()
     hitDecors = []
-    if flags.Detector.GeometryMDT:
-        hitDecors+=["truthMdtHits"]
-    if flags.Detector.GeometryRPC:
-        hitDecors+=["truthRpcHits"]
-    if flags.Detector.GeometryTGC:
-        hitDecors+=["truthTgcHits"]
-    if flags.Detector.GeometryMM:
-        hitDecors+=["truthMMHits"]
-    if flags.Detector.GeometrysTGC:
-        hitDecors+=["truthStgcHits"]
+    if flags.Detector.GeometryMDT: hitDecors+=["truthMdtHits"]
+    if flags.Detector.GeometryRPC: hitDecors+=["truthRpcHits"]
+    if flags.Detector.GeometryTGC: hitDecors+=["truthTgcHits"]
+    if flags.Detector.GeometryMM: hitDecors+=["truthMMHits"]
+    if flags.Detector.GeometrysTGC: hitDecors+=["truthStgcHits"]
     kwargs.setdefault("SimHitIds", hitDecors)
     the_alg = CompFactory.MuonR4.TruthSegToTruthPartAssocAlg(name, **kwargs)
     result.addEventAlgo(the_alg, primary = True)
     return result
 
+def TrackToTruthPartAssocCfg(flags, **kwargs):
+    result = ComponentAccumulator()
+    hitDecors = []
+    if flags.Detector.GeometryMDT: hitDecors+=["truthMdtHits"]
+    if flags.Detector.GeometryRPC: hitDecors+=["truthRpcHits"]
+    if flags.Detector.GeometryTGC: hitDecors+=["truthTgcHits"]
+    if flags.Detector.GeometryMM: hitDecors+=["truthMMHits"]
+    if flags.Detector.GeometrysTGC: hitDecors+=["truthStgcHits"]
+    kwargs.setdefault("SimHitIds", hitDecors)
+    the_alg = CompFactory.MuonR4.TrackToTruthPartAssocAlg(**kwargs)
+    result.addEventAlgo(the_alg, primary = True)
+    return result
+
+
 @AccumulatorCache
 def MuonTruthAlgsCfg(flags):
     result = ComponentAccumulator()
@@ -123,5 +132,5 @@ def MuonTruthAlgsCfg(flags):
     #### Disable for the moment because tracking geometry explodes for R4
     ### from MuonConfig.MuonTruthAlgsConfig import MuonTruthAddTrackRecordsAlgCfg
     ### result.merge(MuonTruthAddTrackRecordsAlgCfg(flags))
-    result.merge(TruthSegmentToTruthPartCfg(flags))
+    result.merge(TruthSegmentToTruthPartAssocCfg(flags))
     return result

MuonSpectrometer/MuonPhaseII/MuonPatternRecognition/MuonTruthAlgsR4/src/TrackToTruthPartAssocAlg.cxx

+/*
+  Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
+*/
+
+#include "TrackToTruthPartAssocAlg.h"
+
+#include "StoreGate/ReadDecorHandleKeyArray.h"
+#include "StoreGate/WriteDecorHandle.h"
+#include "xAODTruth/xAODTruthHelpers.h"
+
+#include <unordered_set>
+namespace {
+    using IdSet_t = std::unordered_set<Identifier>;
+    using TruthLink_t = ElementLink<xAOD::TruthParticleContainer>;
+    using TruthPartWithIds_t = std::tuple<const xAOD::TruthParticle*, IdSet_t>;
+    using IdDecorHandle_t = SG::ReadDecorHandle<xAOD::TruthParticleContainer, std::vector<unsigned long long>>;
+}
+
+namespace MuonR4{
+    StatusCode TrackToTruthPartAssocAlg::initialize() {
+        
+        
+        ATH_CHECK(m_trkKey.initialize());
+        ATH_CHECK(m_originWriteKey.initialize());
+        ATH_CHECK(m_typeWriteKey.initialize());
+        ATH_CHECK(m_linkWriteKey.initialize());
+        ATH_CHECK(m_truthMuonKey.initialize());
+
+        ATH_CHECK(m_idHelperSvc.retrieve());
+        ATH_CHECK(m_edmHelperSvc.retrieve());
+
+        for (const std::string& hitIds  : m_simHitIds) {
+            m_simHitKeys.emplace_back(m_truthMuonKey, hitIds);
+        }
+        ATH_CHECK(m_simHitKeys.initialize());
+        ATH_CHECK(m_truMuOriginKey.initialize());
+        ATH_CHECK(m_truMuTypeKey.initialize());
+        return StatusCode::SUCCESS;
+    }
+    StatusCode TrackToTruthPartAssocAlg::execute(const EventContext& ctx) const {
+        SG::ReadHandle tracks{m_trkKey, ctx};
+        ATH_CHECK(tracks.isPresent());
+
+
+        SG::WriteDecorHandle<xAOD::TrackParticleContainer, int> acc_truthOrigin{m_originWriteKey, ctx};
+        SG::WriteDecorHandle<xAOD::TrackParticleContainer, int> acc_truthType{m_typeWriteKey, ctx};
+        SG::WriteDecorHandle<xAOD::TrackParticleContainer, TruthLink_t> acc_truthLink{m_linkWriteKey, ctx};
+        ///
+        /// Initialize the Identifier decorators
+        std::vector<IdDecorHandle_t> idDecorHandles{};
+        for (const SG::ReadDecorHandleKey<xAOD::TruthParticleContainer>& hitKey : m_simHitKeys) {
+            idDecorHandles.emplace_back(hitKey, ctx);
+        }
+        
+        std::vector<TruthPartWithIds_t> truthWithIds{};
+        ///
+        SG::ReadHandle truthMuonCont{m_truthMuonKey, ctx};
+        ATH_CHECK(truthMuonCont.isPresent());
+        for (const xAOD::TruthParticle* truthMuon : *truthMuonCont) {
+            IdSet_t assocIds{};
+            ATH_MSG_DEBUG("Truth muon: pT:"<<truthMuon->pt()<<" [GeV], eta: "<<truthMuon->eta()<<", phi: "
+                        <<truthMuon->phi()<<", q: "<<truthMuon->charge()<<", truthType: "<<xAOD::TruthHelpers::getParticleTruthType(*truthMuon)
+                        <<", origin: "<<xAOD::TruthHelpers::getParticleTruthOrigin(*truthMuon));
+            for (const IdDecorHandle_t& hitDecor : idDecorHandles) {
+                std::ranges::transform(hitDecor(*truthMuon), std::inserter(assocIds, assocIds.begin()),
+                                [this](unsigned long long rawId) {
+                                    const Identifier id{rawId};
+                                    ATH_MSG_VERBOSE(" --- associated hit id: "<<m_idHelperSvc->toString(id));
+                                    return id; 
+                                });
+            }
+            truthWithIds.emplace_back(std::make_tuple(truthMuon, std::move(assocIds)));
+        }
+
+        for (const xAOD::TrackParticle* trackPart : *tracks){
+            const Trk::Track* track = trackPart->track();
+            if (!track) {
+                ATH_MSG_ERROR("Associated reconstructed track is not available for "<<m_trkKey);
+                return StatusCode::FAILURE;
+            }
+            /// Assemble the track identifiers
+            IdSet_t trackIds{};
+            for (const Trk::TrackStateOnSurface* tsos : *track->trackStateOnSurfaces()) {
+                const Trk::MeasurementBase* meas = tsos->measurementOnTrack();
+                /// Probably scatterer
+                if (!meas) {
+                    continue;
+                }
+                const Identifier measId = m_edmHelperSvc->getIdentifier(*meas);
+                if (!measId.is_valid() || !m_idHelperSvc->isMuon(measId)){
+                    ATH_MSG_VERBOSE("Measurement is not a muon one");
+                    continue;
+                }
+                trackIds.insert(measId);
+            }
+            /// Start matching by finding the particle with the largest fraction
+            int bestMatchFrac{-1};
+            const xAOD::TruthParticle* bestMatch{nullptr};
+            for (const auto& [truthPart, truthIds]: truthWithIds) {
+                const int matchedReco = std::ranges::count_if(trackIds,[&truthIds](const Identifier& recoId){
+                                                                return truthIds.count(recoId);
+                                                            });
+                ///  If the particle has no matched hits or it's worse than the best matched
+                if (!matchedReco || matchedReco < bestMatchFrac){
+                    continue;
+                }
+                bestMatchFrac = matchedReco;
+                bestMatch = truthPart;
+            }
+            /// Track could not be matched
+            if (!bestMatch) {
+                continue;
+            }
+            acc_truthOrigin(*trackPart) = xAOD::TruthHelpers::getParticleTruthOrigin(*bestMatch);
+            acc_truthType(*trackPart) = xAOD::TruthHelpers::getParticleTruthType(*bestMatch);
+            const auto bestMatchCont = static_cast<const xAOD::TruthParticleContainer*>(bestMatch->container());
+            acc_truthLink(*trackPart) = TruthLink_t{bestMatchCont, bestMatch->index()};
+        }
+        return StatusCode::SUCCESS;
+    }
+}
\ No newline at end of file

MuonSpectrometer/MuonPhaseII/MuonPatternRecognition/MuonTruthAlgsR4/src/TrackToTruthPartAssocAlg.h

+/*
+  Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
+*/
+#ifndef MUONTRUTHSEGMENTMAKER_TrackToTruthPartAssocAlg_H
+#define MUONTRUTHSEGMENTMAKER_TrackToTruthPartAssocAlg_H
+
+#include <AthenaBaseComps/AthReentrantAlgorithm.h>
+
+#include <xAODTruth/TruthParticleContainer.h>
+#include <xAODTracking/TrackParticleContainer.h>
+
+#include <StoreGate/ReadHandleKey.h>
+#include <StoreGate/ReadDecorHandleKeyArray.h>
+#include <StoreGate/WriteDecorHandleKey.h>
+
+#include <MuonIdHelpers/IMuonIdHelperSvc.h>
+#include <MuonRecHelperTools/IMuonEDMHelperSvc.h>
+
+namespace MuonR4{
+    /** @brief The TrackToTruthPartAssocAlg matches the reconstructed tracks to truth muons.
+     *         As baseline, the 
+     */
+    class TrackToTruthPartAssocAlg : public AthReentrantAlgorithm {
+        public:
+            using AthReentrantAlgorithm::AthReentrantAlgorithm;
+
+            StatusCode initialize() override final;
+            StatusCode execute(const EventContext& ctx) const override final;
+        private:
+            
+        
+            /** @brief IdHelperSvc to decode the Identifiers */
+            ServiceHandle<Muon::IMuonIdHelperSvc> m_idHelperSvc{this, "IdHelperSvc",  "Muon::MuonIdHelperSvc/MuonIdHelperSvc"};
+            /** @brief Helper service to handle the Identifiers of measurements */
+            ServiceHandle<Muon::IMuonEDMHelperSvc> m_edmHelperSvc{this, "EdmHelperSvc", "Muon::MuonEDMHelperSvc/MuonEDMHelperSvc"};
+        
+
+            using TrkReadKey_t = SG::ReadHandleKey<xAOD::TrackParticleContainer>;
+            using TrkWriteDecorKey_t = SG::WriteDecorHandleKey<xAOD::TrackParticleContainer>;
+
+            TrkReadKey_t m_trkKey{this, "TrackCollection", "MSTrks"};
+            /** @brief Decorations to be written to the TrackParticle truthOrigin/truthType/truthParticleLink */
+            TrkWriteDecorKey_t m_originWriteKey{this, "TruthOriginWriteKey", m_trkKey, "truthOrigin"};
+            TrkWriteDecorKey_t m_typeWriteKey{this, "TruthTypeWriteKey", m_trkKey, "truthType"};
+            TrkWriteDecorKey_t m_linkWriteKey{this, "TruthLinkWriteKey", m_trkKey, "truthParticleLink"};
+            /** @brief Input truth particle keys */
+            using TruthReadKey_t = SG::ReadHandleKey<xAOD::TruthParticleContainer>;
+            using TruthDecor_t = SG::ReadDecorHandleKey<xAOD::TruthParticleContainer>;
+            using TruthDecorArr_t =  SG::ReadDecorHandleKeyArray<xAOD::TruthParticleContainer>;
+
+            TruthReadKey_t m_truthMuonKey{this, "TruthMuonKey", "MuonTruthParticles"};
+
+            /** @brief List of simHit id decorations to read from the truth particle */
+            Gaudi::Property<std::vector<std::string>> m_simHitIds{this, "SimHitIds", {}};
+            /** @brief Declaration of the dependency on the simHit decorations */
+            TruthDecorArr_t m_simHitKeys{this, "TruthSimHitIdKeys", {}};
+
+            TruthDecor_t m_truMuOriginKey{this, "TruthMuonOriginKey", m_truthMuonKey, "truthOrigin"};
+            TruthDecor_t m_truMuTypeKey{this, "TruthMuonTypeKey", m_truthMuonKey, "truthType"};
+
+    };
+}
+
+#endif
\ No newline at end of file

MuonSpectrometer/MuonPhaseII/MuonPatternRecognition/MuonTruthAlgsR4/src/components/MuonTruthAlgsR4_entries.cxx

@@ -7,7 +7,9 @@
 #include "../PrepDataToSimHitAssocAlg.h"
 #include "../PrdMultiTruthMaker.h"
 #include "../TruthSegToTruthPartAssocAlg.h"
+#include "../TrackToTruthPartAssocAlg.h"
 DECLARE_COMPONENT(MuonR4::TruthSegmentMaker)
 DECLARE_COMPONENT(MuonR4::PrepDataToSimHitAssocAlg)
 DECLARE_COMPONENT(MuonR4::PrdMultiTruthMaker)
-DECLARE_COMPONENT(MuonR4::TruthSegToTruthPartAssocAlg)
\ No newline at end of file
+DECLARE_COMPONENT(MuonR4::TruthSegToTruthPartAssocAlg)
+DECLARE_COMPONENT(MuonR4::TrackToTruthPartAssocAlg)
\ No newline at end of file

MuonSpectrometer/MuonPhaseII/MuonValidation/MuonPatternRecognitionTest/python/PatternTestConfig.py

@@ -161,22 +161,11 @@ def TrackTruthMatchCfg(flags):
     result = ComponentAccumulator()
     if not flags.Input.isMC:
         return result
-    from MuonConfig.MuonTruthAlgsConfig import MuonDetailedTrackTruthMakerCfg
+    from MuonTruthAlgsR4.MuonTruthAlgsConfig import TrackToTruthPartAssocCfg
     
-    track_cols = ["MuonTracksR4", "MuonTracksFromHoughR4", "MuonSpectrometerTracks"]
     track_colstp = ["MuonSpectrometerTrackParticlesR4", "MuonSpectrometerTrackParticlesFromHoughR4", "MuonSpectrometerTrackParticles"]
 
-    for trk in track_cols:
-        result.merge(MuonDetailedTrackTruthMakerCfg(flags, name=f"MuonDetailedTruthTrkMaker{trk}",
-                                                        TrackCollectionNames=[trk]))
+    for trk in track_colstp:
+        result.merge(TrackToTruthPartAssocCfg(flags, name=f"TrackToTruth{trk}",  TrackCollection=trk))
 
-        from MuonConfig.MuonTruthAlgsConfig import MuonTruthAssociationAlgCfg
-        result.merge(MuonTruthAssociationAlgCfg(flags, TrackContainers=[]))
-        for i in range(len(track_cols)):
-            from TrkConfig.TrkTruthAlgsConfig import TrackTruthSelectorCfg, TrackParticleTruthAlgCfg
-            result.merge(TrackTruthSelectorCfg(flags, tracks=track_cols[i]))
-
-            result.merge(TrackParticleTruthAlgCfg(flags, tracks=track_cols[i],
-                                                  TrackParticleName=track_colstp[i]))
     return result
-