BucketDumperAlg.h

Go to the documentation of this file.

/*
   Copyright (C) 2002-2024 CERN for the benefit of the ATLAS collaboration
*/
#ifndef MUONCSVDUMP_BucketDumperAlg_H
#define MUONCSVDUMP_BucketDumperAlg_H
 
 
#include "AthenaBaseComps/AthHistogramAlgorithm.h"
 
#include "StoreGate/ReadHandleKeyArray.h"
#include "StoreGate/ReadDecorHandleKeyArray.h"
 
#include <MuonPatternEvent/MuonPatternContainer.h>
#include <MuonSpacePoint/SpacePointContainer.h>
#include <ActsGeometryInterfaces/GeometryContext.h>
#include <MuonReadoutGeometryR4/MuonDetectorManager.h>
#include "xAODMuon/MuonSegmentContainer.h"
#include "xAODTruth/TruthParticle.h"
 
#include "MuonTesterTree/MuonTesterTree.h"
#include "MuonTesterTree/ThreeVectorBranch.h"
#include "MuonTesterTree/IdentifierBranch.h"
 
#include "MuonRecToolInterfacesR4/IPatternVisualizationTool.h"
 
#include "AthenaKernel/IAthRNGSvc.h"
#include "CLHEP/Random/RandomEngine.h"
 
#include "MuonInferenceInterfaces/IGraphInferenceTool.h"
 
#include <unordered_map>
 
 
namespace MuonR4{
class BucketDumperAlg: public AthHistogramAlgorithm {
 
   public:
    using AthHistogramAlgorithm::AthHistogramAlgorithm;
    ~BucketDumperAlg() = default;
 
    virtual StatusCode initialize() override final;
    virtual StatusCode finalize() override final;
    virtual StatusCode execute() override final;
 
   private:
      StatusCode dumpContainer(const EventContext& ctx,
                               const SG::ReadHandleKey<SpacePointContainer>& spacePointKey,
                               const SG::ReadHandleKey<xAOD::MuonSegmentContainer>& segmentKey);

      StatusCode computeAllBucketScores(const EventContext& ctx,
                                        const SpacePointContainer* spContainer,
                                        std::unordered_map<const SpacePointBucket*, std::vector<float>>& bucketScoreMap) const;
 
    SG::ReadHandleKeyArray<SpacePointContainer> m_spacePointKeys{this, "SpacePointKeys", {"MuonSpacePoints"}, 
                                                     "Key to the space point container"};
 
    SG::ReadHandleKeyArray<xAOD::MuonSegmentContainer> m_inSegmentKeys{this, "SegmentKey", {"MuonSegmentsFromR4"}};
 
    SG::ReadDecorHandleKeyArray<xAOD::MuonSegmentContainer> m_truthDecorKeys{this, "TruthDecorLinks", {}};
    SG::ReadHandleKey<ActsTrk::GeometryContext> m_geoCtxKey{this, "AlignmentKey", "ActsAlignment", "cond handle key"};
    
    Gaudi::Property<bool> m_isMC{this, "isMC", true};
    Gaudi::Property<double> m_fracToKeep{this,"dataFracToKeep", 1.}; // 0.055 to balanced dataset without MC
    Gaudi::Property<std::string> m_streamName{this, "StreamName", ""};
    Gaudi::Property<bool> m_doMLBucketScore{this, "DoMLBucketScore", false, 
        "Enable ML bucket score computation and dumping (requires InferenceTool to be configured)"};
    Gaudi::Property<bool> m_doMLBucketFilter{this, "DoMLBucketFilter", false,
        "Skip dumping buckets classified as class 0 (reject) by ML model. Requires DoMLBucketScore=true."};
    ServiceHandle<IAthRNGSvc> m_rndmSvc{this, "RndmSvc", "AthRNGSvc", ""};

    ToolHandle<MuonML::IGraphInferenceTool> m_inferenceTool{this, "InferenceTool", "",
        "Optional ML inference tool for computing bucket scores. If not provided, scores will be empty."};

    ToolHandle<MuonValR4::IPatternVisualizationTool> m_visionTool{this, "VisualizationTool", ""};
    CLHEP::HepRandomEngine* getRandomEngine(const EventContext&ctx) const;
    
    MuonVal::MuonTesterTree m_tree{"MuonBucketDump","MuonBucketDump"};
 
    MuonVal::ScalarBranch<float>&           m_bucket_min{m_tree.newScalar<float>("bucket_min", -1)};
    MuonVal::ScalarBranch<float>&           m_bucket_max{m_tree.newScalar<float>("bucket_max", -1)};
    MuonVal::ScalarBranch<uint16_t>&        m_bucket_spacePoints{m_tree.newScalar<uint16_t>("bucket_spacePoints", 0)};
    MuonVal::ScalarBranch<uint16_t>&        m_bucket_segments{m_tree.newScalar<uint16_t>("bucket_segments", 0)};
    MuonVal::ScalarBranch<uint16_t>&        m_bucket_layers{m_tree.newScalar<uint16_t>("bucket_layers", 0)};
    MuonVal::ScalarBranch<uint8_t>&         m_bucket_truthHit{m_tree.newScalar<uint8_t>("bucket_hasTruth", 0)};
    MuonVal::ScalarBranch<float>&           m_bucket_posX{m_tree.newScalar<float>("bucket_positionX")};
    MuonVal::ScalarBranch<float>&           m_bucket_posY{m_tree.newScalar<float>("bucket_positionY")};
    MuonVal::ScalarBranch<float>&           m_bucket_posZ{m_tree.newScalar<float>("bucket_positionZ")};
 
    MuonVal::ScalarBranch<uint8_t>&         m_bucket_sector{m_tree.newScalar<uint8_t>("bucket_sector")};
    MuonVal::ScalarBranch<uint8_t>&         m_bucket_chamberIdx{m_tree.newScalar<uint8_t>("bucket_chamberIndex")};
    MuonVal::ScalarBranch<Char_t>&          m_bucket_side{m_tree.newScalar<Char_t>("bucket_side")};

    MuonVal::VectorBranch<float>&           m_bucket_ml_score_class0{m_tree.newVector<float>("bucket_ml_score_class0", false)};
    MuonVal::VectorBranch<float>&           m_bucket_ml_score_class1{m_tree.newVector<float>("bucket_ml_score_class1", false)};
    MuonVal::VectorBranch<float>&           m_bucket_ml_score_class2{m_tree.newVector<float>("bucket_ml_score_class2", false)};
 
    MuonVal::ThreeVectorBranch              m_spoint_localPosition{m_tree, "localPosition"}; 
    MuonVal::ThreeVectorBranch              m_spoint_globalPosition{m_tree, "globalPosition"}; 
 
    MuonVal::MuonIdentifierBranch           m_spoint_id{m_tree, "id"};
    MuonVal::VectorBranch<uint16_t>&        m_spoint_layer{m_tree.newVector<uint16_t>("Layer")};
    MuonVal::VectorBranch<unsigned short>&            m_spoint_isStrip{m_tree.newVector<unsigned short>("isStrip", false)};
    MuonVal::VectorBranch<unsigned short>&            m_spoint_isMdt{m_tree.newVector<unsigned short>("isMdt", false)};
 
    MuonVal::VectorBranch<uint16_t>&        m_spoint_adc{m_tree.newVector<uint16_t>("adc")};
    MuonVal::VectorBranch<uint16_t>&        m_spoint_tdc{m_tree.newVector<uint16_t>("tdc")};
 
    MuonVal::VectorBranch<float>&           m_spoint_covX{m_tree.newVector<float>("covX")};
    MuonVal::VectorBranch<float>&           m_spoint_covY{m_tree.newVector<float>("covY")};
    MuonVal::VectorBranch<float>&           m_spoint_driftR{m_tree.newVector<float>("driftR")};
 
    MuonVal::VectorBranch<unsigned short>&  m_spoint_measuresEta{m_tree.newVector<unsigned short>("measuresEta")};
    MuonVal::VectorBranch<unsigned short>&  m_spoint_measuresPhi{m_tree.newVector<unsigned short>("measuresPhi")};
    MuonVal::VectorBranch<unsigned short>&  m_spoint_trueLabel{m_tree.newVector<unsigned short>("trueLabel")};
    MuonVal::VectorBranch<unsigned int>&    m_spoint_nEtaInstances{m_tree.newVector<unsigned int>("nEtaInUse")};
    MuonVal::VectorBranch<unsigned int>&    m_spoint_nPhiInstances{m_tree.newVector<unsigned int>("nPhiInUse")};
    MuonVal::VectorBranch<unsigned int>&    m_spoint_dimension{m_tree.newVector<unsigned int>("dimension")};
 
    MuonVal::VectorBranch<uint16_t>&        m_spoint_nSegments{m_tree.newVector<uint16_t>("nSegments")};
    MuonVal::MatrixBranch<int16_t>&         m_spoint_mat{m_tree.newMatrix<int16_t>("sp_seg_matching",-1)};
    MuonVal::MatrixBranch<int16_t>&         m_spoint_trueSeg{m_tree.newMatrix<int16_t>("spacePoint_truthSegLink",-1)};
    
    
    
    MuonVal::VectorBranch<float>&           m_segmentLocX{m_tree.newVector<float>("segmentLocalX")};
    MuonVal::VectorBranch<float>&           m_segmentLocY{m_tree.newVector<float>("segmentLocalY")};
    MuonVal::VectorBranch<float>&           m_segmentLocTheta{m_tree.newVector<float>("segmentLocalTheta")};
    MuonVal::VectorBranch<float>&           m_segmentLocPhi{m_tree.newVector<float>("segmentLocalPhi")};
    
    MuonVal::VectorBranch<uint16_t>&        m_segmentTruthIdx{m_tree.newVector<uint16_t>("segmentTruthPart")};
    
    MuonVal::ThreeVectorBranch              m_segmentPos{m_tree, "segmentPosition"}; 
    MuonVal::ThreeVectorBranch              m_segmentDir{m_tree, "segmentDirection"};
    MuonVal::VectorBranch<float>&           m_segment_chiSquared{m_tree.newVector<float>("segment_chiSquared")};
    MuonVal::VectorBranch<float>&           m_segment_numberDoF{m_tree.newVector<float>("segment_numberDoF")};
 
   MuonVal::VectorBranch<float>&            m_truthSegLocX{m_tree.newVector<float>("truthSegmentLocalX")};
   MuonVal::VectorBranch<float>&            m_truthSegLocY{m_tree.newVector<float>("truthSegmentLocalY")};
   MuonVal::VectorBranch<float>&            m_truthSegLocTheta{m_tree.newVector<float>("truthSegmentLocalTheta")};
   MuonVal::VectorBranch<float>&            m_truthSegLocPhi{m_tree.newVector<float>("truthSegmentLocalPhi")};
   
   
};
}
#endif