MlMsTrackSeeder.cxx

Go to the documentation of this file.

/*
  Copyright (C) 2002-2026 CERN for the benefit of the ATLAS collaboration
*/
#include "MlMsTrackSeeder.h"
 
#include "AthenaBaseComps/AthCheckMacros.h"
#include "MuonTrackEvent/MsTrackSeed.h"
#include "MuonTrackEvent/TrackingHelpers.h"
#include "MuonStationIndex/MuonStationIndex.h"
#include "StoreGate/ReadDecorHandle.h"
#include "StoreGate/ReadHandle.h"
 
#include <map>
#include <utility>
#include <vector>
 
namespace MuonR4 {
 
namespace {
 
struct MlComponent {
  std::vector<const xAOD::MuonSegment*> segments;
  std::vector<const xAOD::MuonSegment*> anchors;
};
 
}  // namespace
 
StatusCode MlMsTrackSeeder::initialize() {
  ATH_CHECK(m_seedParameterEstimator.retrieve());
  ATH_CHECK(m_segmentKey.initialize());
  ATH_CHECK(m_candidateDecorKey.initialize());
  return StatusCode::SUCCESS;
}
 
StatusCode MlMsTrackSeeder::findTrackSeeds(
    const EventContext& ctx,
    std::vector<MsTrackSeed>& outSeeds) const {
  const xAOD::MuonSegmentContainer* segments{nullptr};
  ATH_CHECK(SG::get(segments, m_segmentKey, ctx));
 
  // The decoration payload is [componentId, isSeedAnchor].
  SG::ReadDecorHandle<xAOD::MuonSegmentContainer, std::vector<unsigned int>>
      componentAcc{m_candidateDecorKey, ctx};
  if (!componentAcc.isAvailable()) {
    ATH_MSG_DEBUG("No ML candidate decoration found");
    return StatusCode::SUCCESS;
  }
 
  std::map<unsigned int, MlComponent> components;
 
  for (const xAOD::MuonSegment* segment : *segments) {
    if (!segment) continue;
    const std::vector<unsigned int>& payload = componentAcc(*segment);
    if (payload.size() != 2 || payload[0] == 0) continue;
    MlComponent& component = components[payload[0]];
    component.segments.push_back(segment);
    if (payload[1] != 0) component.anchors.push_back(segment);
  }
 
  std::size_t acceptedComponents = 0;
  std::size_t constructedSeeds = 0;
  for (const auto& [componentId, component] : components) {
    if (component.segments.size() < m_minSegmentsPerCandidate.value() ||
        (m_maxSegmentsPerCandidate.value() != 0 &&
         component.segments.size() > m_maxSegmentsPerCandidate.value()) ||
        component.anchors.empty()) {
      continue;
    }
    ++acceptedComponents;
 
    for (const xAOD::MuonSegment* anchor : component.anchors) {
      const MsTrackSeed::Location location =
          Muon::MuonStationIndex::isBarrel(anchor->chamberIndex())
              ? MsTrackSeed::Location::Barrel
              : MsTrackSeed::Location::Endcap;
      MsTrackSeed seed{location, ExpandedSector{anchor->position().phi()}};
 
      // One segment per layer, best aligned with the anchor: the baseline seeder
      // only builds seeds that span several layers, and the fit rejects single-layer tracks.
      const Muon::MuonStationIndex::LayerIndex anchorLayer =
          Muon::MuonStationIndex::toLayerIndex(anchor->chamberIndex());
      std::map<Muon::MuonStationIndex::LayerIndex, std::pair<double, const xAOD::MuonSegment*>> bestPerLayer;
      for (const xAOD::MuonSegment* segment : component.segments) {
        if (segment == anchor) continue;
        const Muon::MuonStationIndex::LayerIndex layer =
            Muon::MuonStationIndex::toLayerIndex(segment->chamberIndex());
        if (layer == anchorLayer) continue;
        const double cosAngle = anchor->direction().dot(segment->direction());
        if (cosAngle < m_minCosConsistency.value()) continue;
        auto& best = bestPerLayer[layer];
        if (!best.second || cosAngle > best.first) best = {cosAngle, segment};
      }
      if (bestPerLayer.size() + 1 < m_minLayers.value()) continue;
 
      seed.addSegment(anchor);
      for (const auto& [layer, choice] : bestPerLayer) seed.addSegment(choice.second);
      seed.setPosition(anchor->position());
      outSeeds.push_back(std::move(seed));
      ++constructedSeeds;
    }
    ATH_MSG_VERBOSE("Constructed " << component.anchors.size()
                    << " seed(s) from ML component " << componentId
                    << " with " << component.segments.size()
                    << " segments");
  }
 
  ATH_MSG_DEBUG("Constructed " << constructedSeeds << " seed(s) from "
                << acceptedComponents << "/" << components.size()
                << " decorated ML component(s)");
  return StatusCode::SUCCESS;
}
 
Acts::Result<Acts::BoundTrackParameters>
MlMsTrackSeeder::estimateStartParameters(
    const EventContext& ctx,
    const MsTrackSeed& seed) const {
  return m_seedParameterEstimator->estimateStartParameters(ctx, seed);
}
 
}  // namespace MuonR4