MlMsTrackSeeder.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2026 CERN for the benefit of the ATLAS collaboration
*/
#ifndef MUONTRACKFINDINGALGS_MLMSTRACKSEEDER_H
#define MUONTRACKFINDINGALGS_MLMSTRACKSEEDER_H
 
#include "AthenaBaseComps/AthAlgTool.h"
#include "CxxUtils/checker_macros.h"
#include "MuonRecToolInterfacesR4/ISegmentSelectionTool.h"
#include "MuonRecToolInterfacesR4/ITrackSeedingDiagnosticsTool.h"
#include "ActsEvent/ContextUtility.h"
#include "StoreGate/ReadDecorHandleKey.h"
#include "StoreGate/ReadHandleKey.h"
#include "xAODMuon/MuonSegmentContainer.h"
 
#include <array>
#include <atomic>
#include <chrono>
#include <cstdint>
#include <mutex>
#include <unordered_map>
#include <vector>
#include <span>
#include <string>
 
namespace MuonR4 {

class MlMsTrackSeeder final :
    public extends<AthAlgTool, ITrackSeedingTool> {
 public:
  using base_class::base_class;
 
  StatusCode initialize() override final;
  StatusCode finalize() override final;

  StatusCode findTrackSeeds(const EventContext& ctx,
                            std::vector<MsTrackSeed>& outSeeds) const override final;

  Acts::Result<Acts::BoundTrackParameters> estimateStartParameters(
      const EventContext& ctx,
      const MsTrackSeed& seed) const override final;
 
 private:
  ToolHandle<ITrackSeedingDiagnosticsTool> m_seedParameterEstimator{
      this, "BaselineSeeder", "",
      "Standard seeder: initial-parameter estimates for ML-built seeds and the projection "
      "onto its seeding cylinder used by PointingWindow"};
  ActsTrk::ContextUtility m_ctxProvider{this};
  SG::ReadHandleKey<xAOD::MuonSegmentContainer> m_segmentKey{
      this, "SegmentContainer", "MsSegments"};
  SG::ReadDecorHandleKey<xAOD::MuonSegmentContainer> m_candidateDecorKey{
      this, "CandidateDecoration", m_segmentKey, "trackCandidateIds",
      "Scheduler dependency on the ML candidate decoration"};
  SG::ReadDecorHandleKey<xAOD::MuonSegmentContainer> m_componentScoreKey{
      this, "ComponentScoreDecoration", m_segmentKey, "mlComponentScore",
      "ML scores per segment: [component's strongest cross-layer edge, segment's strongest cross-layer edge]"};
  Gaudi::Property<unsigned int> m_minSegmentsPerCandidate{
      this, "MinSegmentsPerCandidate", 2};
  Gaudi::Property<unsigned int> m_maxSegmentsPerCandidate{
      this, "MaxSegmentsPerCandidate", 0,
      "Discard ML components with more than this many segments; zero disables the cap"};
  Gaudi::Property<unsigned int> m_minLayers{
      this, "MinLayers", 2, "Minimum number of distinct muon layers a seed must span"};
  Gaudi::Property<double> m_minCosConsistency{
      this, "MinCosConsistency", 0.819,  // cos(35 deg), matching SegmentEdgeClassifierTool's default MaxDeltaThetaDeg
      "Minimum cos(angle) between a segment's direction and its seed anchor's "
      "direction for the segment to be included in the seed"};
  Gaudi::Property<double> m_pointingWindow{
      this, "PointingWindow", 250.,
      "Keep only seed segments whose position projected onto the standard seeder's cylinder "
      "lies within this distance [mm] of the window centre (the standard seeder uses "
      "SeedHalfLength = 250 mm); zero disables"};
  Gaudi::Property<bool> m_clusterWindowCentre{
      this, "ClusterWindowCentre", true,
      "PointingWindow: centre the window on the segment whose window covers the most distinct "
      "layers (ties: segment quality) instead of the anchor"};
  Gaudi::Property<unsigned int> m_extraClusterSeeds{
      this, "ExtraClusterSeeds", 2,
      "PointingWindow: after the first seed of an anchor, build up to this many further window seeds "
      "from the component segments not covered by an earlier window; zero disables"};
  Gaudi::Property<double> m_extraClusterMinScore{
      this, "ExtraClusterMinScore", 0.8,
      "ExtraClusterSeeds: an extra seed needs a segment whose own strongest selected cross-layer ML "
      "edge (ComponentScoreDecoration) reaches this value"};
  Gaudi::Property<double> m_innerPointingMaxDeg{
      this, "InnerPointingMaxDeg", 2.,
      "Remove from the seed pools the endcap inner segments whose |theta(direction) - theta(position)| "
      "exceeds this value [deg], i.e. that do not point back to the IP; zero disables"};
  Gaudi::Property<double> m_innerRescueWindow{
      this, "InnerRescueWindow", 0.,
      "After the MinLayers check and the 2-segment rule: fill a missing Inner layer with the closest "
      "IP-pointing (InnerRescueMaxDeg), direction-compatible inner segment of any component that "
      "projects within this distance [mm] of the window centre; zero disables"};
  Gaudi::Property<double> m_innerRescueMaxDeg{
      this, "InnerRescueMaxDeg", 1.,
      "InnerRescueWindow: maximum |theta(direction) - theta(position)| [deg] of the filled segment"};
  Gaudi::Property<bool> m_variantSeeds{
      this, "VariantSeeds", true,
      "PointingWindow: after a seed passing the Variant* gates, emit a second seed in which every "
      "layer takes its runner-up segment from the window (or, without any alternative, the weakest "
      "non-inner member is dropped), as a fallback if the fit of the first seed fails"};
  Gaudi::Property<unsigned int> m_variantMinLayers{
      this, "VariantMinLayers", 3,
      "VariantSeeds: minimum number of station layers of the seed"};
  Gaudi::Property<double> m_variantMaxPointingDeg{
      this, "VariantMaxPointingDeg", 1.,
      "VariantSeeds: endcap seeds need an inner segment with |theta(dir) - theta(pos)| below this [deg]"};
  Gaudi::Property<double> m_variantMinScore{
      this, "VariantMinScore", 0.8,
      "VariantSeeds: minimum strongest cross-layer ML edge of the seed's component (ComponentScoreDecoration); "
      "zero disables"};
  Gaudi::Property<double> m_twoSegMinScore{
      this, "TwoSegmentMinComponentScore", 0.3,
      "Keep a 2-segment seed only if its component's strongest selected cross-layer edge "
      "(ComponentScoreDecoration) reaches this value; zero disables"};
  ToolHandle<ISegmentSelectionTool> m_segSelector{
      this, "SegmentSelectionTool", "",
      "Seeding-quality check of the standard seeder, used by the truth diagnostics"};
  Gaudi::Property<bool> m_truthDiag{
      this, "EnableTruthDiagnostics", false,
      "MC only, and only at DEBUG level: truth-classify the seeds and time their fits"};

  struct SeedingCounters {
    std::atomic<std::size_t> events{0};
    std::atomic<std::size_t> components{0};
    std::atomic<std::size_t> rejectedTooSmall{0};
    std::atomic<std::size_t> rejectedTooLarge{0};
    std::atomic<std::size_t> rejectedNoAnchor{0};
    std::atomic<std::size_t> acceptedComponents{0};
    std::atomic<std::size_t> anchorsTooFewLayers{0};     
    std::atomic<std::size_t> windowCentreNotAnchor{0};   
    std::atomic<std::size_t> windowCentreOutOfBounds{0}; 
    std::atomic<std::size_t> windowDropped{0};           
    std::atomic<std::size_t> windowOutOfBounds{0};       
    std::atomic<std::size_t> twoSegLowScore{0};          
    std::atomic<std::size_t> twoSegNoScore{0};           
    std::atomic<std::size_t> extraSeeds{0};              
    std::atomic<std::size_t> extraNoStrong{0};           
    std::atomic<std::size_t> pointingRemoved{0};         
    std::atomic<std::size_t> pointingCentreSkipped{0};   
    std::atomic<std::size_t> innerRescued{0};            
    std::atomic<std::size_t> variantSeeds{0};            
    std::atomic<std::size_t> variantNone{0};             
    std::atomic<std::size_t> seeds{0};
    std::array<std::atomic<std::size_t>, 4> seedSegments{};  
    std::atomic<std::size_t> startParCalls{0};
    std::atomic<std::size_t> startParFailures{0};
    std::atomic<std::uint64_t> findSeedsNs{0};   
    std::atomic<std::uint64_t> startParNs{0};    
  };
  mutable SeedingCounters m_counters ATLAS_THREAD_SAFE;

  struct SeedAnatomy {
    static constexpr std::size_t kClasses = 4;
    std::array<std::atomic<std::uint64_t>, kClasses * 2> fitNs{};
    std::array<std::atomic<std::size_t>, kClasses * 2> fitCount{};
    std::atomic<std::size_t> untimedLastSeeds{0};   
    std::array<std::atomic<std::size_t>, kClasses * 3> nSegments{};     
    std::array<std::atomic<std::size_t>, kClasses * 4> thetaSpread{};   
    static constexpr std::size_t kPBins = 5;
    std::array<std::atomic<std::size_t>, kClasses * kPBins> seedsByP{};      
    std::array<std::atomic<std::uint64_t>, kClasses * kPBins> fitNsByP{};    
    std::array<std::atomic<std::size_t>, kClasses * kPBins> fitCountByP{};
    std::array<std::atomic<std::size_t>, 2 * 2> fitCountBySegs{};           
    std::array<std::atomic<std::uint64_t>, 2 * 2> fitNsBySegs{};
    std::array<std::atomic<std::size_t>, 2> pureNoStartParsQuality{};
    std::array<std::atomic<std::size_t>, 2> pureNoStartParsNsw{};
    std::atomic<std::size_t> pureNoStartPars{0};

    static constexpr std::size_t kGroups = 3;
    static constexpr std::size_t kDistBins = 4;   // projected distance <50 / 50-100 / 100-250 / >=250 or n.a. [mm]
    static constexpr std::size_t kQualBins = 3;   // 0 / 1 / 2 segments passing the seeding quality
    static constexpr std::size_t kScoreBins = 5;  // <0.05 / 0.05-0.3 / 0.3-0.8 / >=0.8 / n.a.
    static constexpr std::size_t kPairBins = 5;   // layer pair: I-M / M-O / I-O / with BEE-EE / other
    std::array<std::atomic<std::size_t>, kGroups * kDistBins> twoSegDistN{};
    std::array<std::atomic<std::uint64_t>, kGroups * kDistBins> twoSegDistNs{};
    std::array<std::atomic<std::size_t>, kGroups * kQualBins> twoSegQualN{};
    std::array<std::atomic<std::uint64_t>, kGroups * kQualBins> twoSegQualNs{};
    std::array<std::atomic<std::size_t>, kGroups * kScoreBins> twoSegCompScoreN{};   
    std::array<std::atomic<std::uint64_t>, kGroups * kScoreBins> twoSegCompScoreNs{};
    std::array<std::atomic<std::size_t>, kGroups * kScoreBins> twoSegPairScoreN{};   
    std::array<std::atomic<std::uint64_t>, kGroups * kScoreBins> twoSegPairScoreNs{};
    std::array<std::atomic<std::size_t>, kGroups * kPairBins> twoSegLayersN{};
    std::array<std::atomic<std::uint64_t>, kGroups * kPairBins> twoSegLayersNs{};
    std::array<std::atomic<std::size_t>, kGroups * 2> twoSegRegionN{};               
    std::array<std::atomic<std::uint64_t>, kGroups * 2> twoSegRegionNs{};
    std::array<std::atomic<std::size_t>, kGroups * kDistBins * kScoreBins> twoSegJointN{};
    std::array<std::atomic<std::uint64_t>, kGroups * kDistBins * kScoreBins> twoSegJointNs{};
    static constexpr std::size_t kMeasBins = 5;
    std::array<std::atomic<std::size_t>, kGroups * kMeasBins> measN{};
    std::array<std::atomic<std::uint64_t>, kGroups * kMeasBins> measNs{};
    std::array<std::atomic<std::size_t>, kGroups> twoSegDropped{};
    static constexpr std::size_t kPointBins = 6;
    static constexpr std::size_t kPointCats = 4;
    std::array<std::atomic<std::size_t>, kGroups * kPointCats * kPointBins> pointN{};
    std::array<std::atomic<std::uint64_t>, kGroups * kPointCats * kPointBins> pointNs{};
    std::array<std::atomic<std::size_t>, 3> extraSeedClass{};
    std::array<std::atomic<std::size_t>, 2> pureTwoSegOwnerHasLonger{};
  };
  mutable SeedAnatomy m_anatomy ATLAS_THREAD_SAFE;

  struct WindowCounters {
    static constexpr std::size_t kDistBins = 6;  // <100 / 100-250 / 250-500 / 500-1000 / >=1000 mm / out of bounds
    std::array<std::atomic<std::size_t>, 2 * kDistBins> distance{};
    std::array<std::atomic<std::size_t>, 2> dropped{};   
    std::array<std::atomic<std::size_t>, 3> seedsWithout{};  
    std::array<std::atomic<std::size_t>, 3> seedsLost{};     
    std::array<std::atomic<std::size_t>, 3> seedsPureAfter{}; 
    std::atomic<std::size_t> particlesSeeded{0};   
    std::atomic<std::size_t> particlesLost{0};     
    std::atomic<std::size_t> ownerSeeds{0};        
    std::atomic<std::size_t> centreOnOwner{0};     
    std::atomic<std::size_t> oracleKeeps{0};       
    std::atomic<std::size_t> actualKeeps{0};       
  };
  mutable WindowCounters m_window ATLAS_THREAD_SAFE;

  struct EtaFunnel {
    static constexpr std::size_t kRegions = 4;      // |eta| <1.05 / 1.05-1.3 / 1.3-2.0 / >=2.0
    static constexpr std::size_t kStages = 9;
    static constexpr std::size_t kMissReasons = 5;
    std::array<std::atomic<std::size_t>, kRegions * kStages> stage{};
    std::array<std::atomic<std::size_t>, kRegions> notSeedable{};  
    std::array<std::atomic<std::size_t>, kRegions * kMissReasons> missingLayers{};
    std::array<std::atomic<std::size_t>, kRegions> withMissing{};  
    std::array<std::atomic<std::size_t>, kRegions * 2> fakes{};     
    std::array<std::atomic<std::size_t>, kRegions> withFake{};      
  };
  mutable EtaFunnel m_etaFunnel ATLAS_THREAD_SAFE;

  struct RescueCounters {
    std::array<std::atomic<std::size_t>, 2> pointingRemoved{};  
    std::array<std::atomic<std::size_t>, 3> filled{};
  };
  mutable RescueCounters m_rescue ATLAS_THREAD_SAFE;

  struct VariantCounters {
    std::array<std::atomic<std::size_t>, 3> seedClass{};  
    std::atomic<std::size_t> pureMainNotPure{0};          
  };
  mutable VariantCounters m_variant ATLAS_THREAD_SAFE;

  struct SlotState {
    const MsTrackSeed* base{nullptr};
    std::vector<std::uint8_t> seedClass{};
    std::vector<std::uint8_t> startOk{};
    std::vector<std::uint8_t> pBin{};
    std::size_t lastNSeg{0};
    std::vector<std::array<std::uint8_t, 8>> desc{};
    long lastIndex{-1};
    std::chrono::steady_clock::time_point lastEnd{};
  };
  mutable std::mutex m_slotMutex;
  mutable std::unordered_map<std::size_t, SlotState> m_slots ATLAS_THREAD_SAFE;
};
 
}  // namespace MuonR4
 
#endif