FPGATrackSimGNNGraphConstructionTool.cxx

Go to the documentation of this file.

// Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
 
#include "FPGATrackSimGNNGraphConstructionTool.h"
 
#include <TFile.h>
#include <TTree.h>
#include "FourMomUtils/P4Helpers.h"
 
// AthAlgTool
 
FPGATrackSimGNNGraphConstructionTool::FPGATrackSimGNNGraphConstructionTool(const std::string& algname, const std::string &name, const IInterface *ifc) 
    : AthAlgTool(algname, name, ifc) {}
 
StatusCode FPGATrackSimGNNGraphConstructionTool::initialize()
{
    if(m_graphTool == "ModuleMap") {
        if(m_moduleMapPath == "") { // Require a path provided for the Module Map
                ATH_MSG_FATAL("ERROR! No Module Map provided. Please provide a valid path to a ROOT file."); 
                return StatusCode::FAILURE;
        }
        if(m_moduleMapType == "doublet") {
            loadDoubletModuleMap(); // Load the doublet module map and store entry branches in vectors
        }
    }
 
    return StatusCode::SUCCESS;
}
 
// Functions
 
StatusCode FPGATrackSimGNNGraphConstructionTool::getEdges(const std::vector<std::shared_ptr<FPGATrackSimGNNHit>> & hits, std::vector<std::shared_ptr<FPGATrackSimGNNEdge>> & edges)
{
    if(m_graphTool == "ModuleMap") {
        doModuleMap(hits, edges);
    }
 
    return StatusCode::SUCCESS;
}
 
void FPGATrackSimGNNGraphConstructionTool::loadDoubletModuleMap()
{
    std::unique_ptr<TFile> file(TFile::Open(m_moduleMapPath.value().c_str()));
    std::unique_ptr<TTree> tree(static_cast<TTree*>(file->Get("TreeModuleDoublet")));
 
    unsigned int mid1_value = 0;
    unsigned int mid2_value = 0;
    float z0min_12_value = 0.0;
    float dphimin_12_value = 0.0;
    float phislopemin_12_value = 0.0;
    float detamin_12_value = 0.0;
    float z0max_12_value = 0.0;
    float dphimax_12_value = 0.0;
    float phislopemax_12_value = 0.0;
    float detamax_12_value = 0.0;
 
    tree->SetBranchAddress("Module1", &mid1_value);
    tree->SetBranchAddress("Module2", &mid2_value);
    tree->SetBranchAddress("z0min_12", &z0min_12_value);
    tree->SetBranchAddress("dphimin_12", &dphimin_12_value);
    tree->SetBranchAddress("phiSlopemin_12", &phislopemin_12_value);
    tree->SetBranchAddress("detamin_12", &detamin_12_value);
    tree->SetBranchAddress("z0max_12", &z0max_12_value);
    tree->SetBranchAddress("dphimax_12", &dphimax_12_value);
    tree->SetBranchAddress("phiSlopemax_12", &phislopemax_12_value);
    tree->SetBranchAddress("detamax_12", &detamax_12_value);
 
    int64_t nEntries = tree->GetEntries();
    for (int64_t i = 0; i < nEntries; ++i) {
        tree->GetEntry(i);
        m_mid1.emplace_back(mid1_value);
        m_mid2.emplace_back(mid2_value);
        m_z0min_12.emplace_back(z0min_12_value);
        m_dphimin_12.emplace_back(dphimin_12_value);
        m_phislopemin_12.emplace_back(phislopemin_12_value);
        m_detamin_12.emplace_back(detamin_12_value);
        m_z0max_12.emplace_back(z0max_12_value);
        m_dphimax_12.emplace_back(dphimax_12_value);
        m_phislopemax_12.emplace_back(phislopemax_12_value);
        m_detamax_12.emplace_back(detamax_12_value);
    }
}
 
void FPGATrackSimGNNGraphConstructionTool::doModuleMap(const std::vector<std::shared_ptr<FPGATrackSimGNNHit>> & hits, std::vector<std::shared_ptr<FPGATrackSimGNNEdge>> & edges)
{
    // Use Module Map method for edge building
    // Two types of module maps: Doublet and Triplet
    // For each type of module map there is three functions: minmax, meanrms, and hybrid
    // Use the proper configuration set by the input script and passed as Gaudi::Property variables
    // Currently only Doublet Module Map with minmax cuts exist, but others can be implemented later on as desired
 
    if(m_moduleMapType == "doublet") {
        getDoubletEdges(hits, edges);
    }
}
 
void FPGATrackSimGNNGraphConstructionTool::getDoubletEdges(const std::vector<std::shared_ptr<FPGATrackSimGNNHit>> & hits, std::vector<std::shared_ptr<FPGATrackSimGNNEdge>> & edges)
{
    // Take the list of hits and use the doublet module map to generate all the edges between hits that pass the doublet cuts
 
    for (size_t i = 0; i < m_mid2.size(); i++) {
        std::vector<std::shared_ptr<FPGATrackSimGNNHit>> hit1_matches;
        std::vector<std::shared_ptr<FPGATrackSimGNNHit>> hit2_matches;
    
        std::vector<int> hit1_indices;
        std::vector<int> hit2_indices;
 
        for (size_t j = 0; j < hits.size(); j++) {
            if (hits[j]->getIdentifier() == m_mid1[i]) {
                hit1_matches.emplace_back(hits[j]);
                hit1_indices.emplace_back(j);
            }
            if (hits[j]->getIdentifier() == m_mid2[i]) {
                hit2_matches.emplace_back(hits[j]);
                hit2_indices.emplace_back(j);
            }
        }
 
        for (size_t h1 = 0; h1 < hit1_matches.size(); h1++) {
            for (size_t h2 = 0; h2 < hit2_matches.size(); h2++) {
                applyDoubletCuts(hit1_matches[h1], hit2_matches[h2], edges, hit1_indices[h1], hit2_indices[h2], i);
            }
        }
    }
}
 
void FPGATrackSimGNNGraphConstructionTool::applyDoubletCuts(const std::shared_ptr<FPGATrackSimGNNHit> & hit1, const std::shared_ptr<FPGATrackSimGNNHit> & hit2, std::vector<std::shared_ptr<FPGATrackSimGNNEdge>> & edges, int hit1_index, int hit2_index, unsigned int modulemap_id)
{
    // Four types of doublet cuts (dEta, z0, dPhi, phiSlope)
    // If an edge passes all four, then it is a valid edge and can be stored
 
    // delta_eta cuts
    float deta = hit1->getEta() - hit2->getEta();
    if(!doMask(deta, m_detamin_12[modulemap_id], m_detamax_12[modulemap_id])) return;
 
    // z0 cuts
    float dz = hit2->getZ() - hit1->getZ();
    float dr = hit2->getR() - hit1->getR();
    float z0 = dr==0. ? 0. : hit1->getZ() - (hit1->getR() * dz / dr);
    if(!doMask(z0, m_z0min_12[modulemap_id], m_z0max_12[modulemap_id])) return;
 
    // delta_phi cuts
    float dphi = P4Helpers::deltaPhi(hit2->getPhi(),hit1->getPhi());
    if(!doMask(dphi, m_dphimin_12[modulemap_id], m_dphimax_12[modulemap_id])) return;
 
    // phislope cuts
    float phislope = dr==0. ? 0. : dphi / dr;
    if(!doMask(phislope, m_phislopemin_12[modulemap_id], m_phislopemax_12[modulemap_id])) return;
 
    // if pass all doublet cuts, then record the edge information
    std::shared_ptr<FPGATrackSimGNNEdge> edge = std::make_shared<FPGATrackSimGNNEdge>();
    edge->setEdgeIndex1(hit1_index);
    edge->setEdgeIndex2(hit2_index);
    edge->setEdgeDR(dr);
    edge->setEdgeDPhi(dphi);
    edge->setEdgeDZ(dz);
    edge->setEdgeDEta(deta);
    edge->setEdgePhiSlope(phislope);
    edge->setEdgeRPhiSlope(0.5 * (hit2->getR() + hit1->getR()) * phislope);
    edges.emplace_back(edge);
}
 
bool FPGATrackSimGNNGraphConstructionTool::doMask(float val, float min, float max)
{
    bool mask = false;
    if(m_moduleMapFunc == "minmax") {
        mask = doMinMaxMask(val, min, max);
    }
 
    return mask;
}
 
bool FPGATrackSimGNNGraphConstructionTool::doMinMaxMask(float val, float min, float max)
{
    bool mask = false;
 
    if((val <= max * (1.0 + featureSign(max) * m_moduleMapTol)) &&
       (val >= min * (1.0 - featureSign(min) * m_moduleMapTol))) {
        mask = true;
    }
 
    return mask;
}
 
float FPGATrackSimGNNGraphConstructionTool::featureSign(float feature)
{
    if(feature < 0.0) { return -1.0; }
    else if(feature > 0.0) { return 1.0; }
    else { return 0.0; }
}