diff --git a/Configuration/ProcessModifiers/python/ticlv5_TrackLinkingGNN_cff.py b/Configuration/ProcessModifiers/python/ticlv5_TrackLinkingGNN_cff.py new file mode 100644 index 0000000000000..d387226097369 --- /dev/null +++ b/Configuration/ProcessModifiers/python/ticlv5_TrackLinkingGNN_cff.py @@ -0,0 +1,7 @@ +import FWCore.ParameterSet.Config as cms + +from Configuration.ProcessModifiers.ticl_v5_cff import ticl_v5 + +# This modifier is for running TICL v5 with GNN track-trackster linking. +ticl_v5_TrackLinkingGNN = cms.Modifier() +ticlv5_TrackLinkingGNN = cms.ModifierChain(ticl_v5, ticl_v5_TrackLinkingGNN) diff --git a/Configuration/PyReleaseValidation/python/relval_Run4.py b/Configuration/PyReleaseValidation/python/relval_Run4.py index 58c8f7e0dac32..c4189376dbcd5 100644 --- a/Configuration/PyReleaseValidation/python/relval_Run4.py +++ b/Configuration/PyReleaseValidation/python/relval_Run4.py @@ -74,6 +74,7 @@ numWFIB.extend([prefixDet+34.751]) # HLTTiming75e33, alpaka numWFIB.extend([prefixDet+34.7511])# HLTTiming75e33, phase2CAExtension numWFIB.extend([prefixDet+34.752]) # HLTTiming75e33, ticl_v5 +numWFIB.extend([prefixDet+34.7521])# HLTTiming75e33, ticl_v5, ticlv5TrackLinkingGNN numWFIB.extend([prefixDet+34.753]) # HLTTiming75e33, alpaka,singleIterPatatrack numWFIB.extend([prefixDet+34.754]) # HLTTiming75e33, alpaka,singleIterPatatrack,trackingLST numWFIB.extend([prefixDet+34.755]) # HLTTiming75e33, alpaka,trackingLST diff --git a/Configuration/PyReleaseValidation/python/upgradeWorkflowComponents.py b/Configuration/PyReleaseValidation/python/upgradeWorkflowComponents.py index d44d5a71bfc12..c5e50acac1eff 100644 --- a/Configuration/PyReleaseValidation/python/upgradeWorkflowComponents.py +++ b/Configuration/PyReleaseValidation/python/upgradeWorkflowComponents.py @@ -1075,6 +1075,38 @@ def condition(self, fragment, stepList, key, hasHarvest): upgradeWFs['ticl_barrel_CPfromPU'].step3 = {'--procModifiers': 'ticl_barrel,enableCPfromPU'} upgradeWFs['ticl_barrel_CPfromPU'].step4 = {'--procModifiers': 'ticl_barrel,enableCPfromPU'} +class UpgradeWorkflow_ticlv5_TrackLinkingGNN(UpgradeWorkflow): + def setup_(self, step, stepName, stepDict, k, properties): + if ('Digi' in step and 'NoHLT' not in step) or ('HLTOnly' in step): + stepDict[stepName][k] = merge([self.step2, stepDict[step][k]]) + if 'RecoGlobal' in step: + stepDict[stepName][k] = merge([self.step3, stepDict[step][k]]) + if 'HARVESTGlobal' in step: + stepDict[stepName][k] = merge([self.step4, stepDict[step][k]]) + def condition(self, fragment, stepList, key, hasHarvest): + selected_fragments = ["TTbar_14TeV", "CloseByP", "Eta1p7_2p7", "ZEE_14"] + return any(sf in fragment for sf in selected_fragments) and 'Run4' in key + +upgradeWFs['ticlv5_TrackLinkingGNN'] = UpgradeWorkflow_ticlv5_TrackLinkingGNN( + steps = [ + 'HLTOnly', + 'DigiTrigger', + 'RecoGlobal', + 'HARVESTGlobal' + ], + PU = [ + 'HLTOnly', + 'DigiTrigger', + 'RecoGlobal', + 'HARVESTGlobal' + ], + suffix = '_ticlv5_TrackLinkGNN', + offset = 0.211, +) +upgradeWFs['ticlv5_TrackLinkingGNN'].step2 = {'--procModifiers': 'ticlv5_TrackLinkingGNN'} +upgradeWFs['ticlv5_TrackLinkingGNN'].step3 = {'--procModifiers': 'ticlv5_TrackLinkingGNN'} +upgradeWFs['ticlv5_TrackLinkingGNN'].step4 = {'--procModifiers': 'ticlv5_TrackLinkingGNN'} + # L3 Tracker Muon Outside-In reconstruction first class UpgradeWorkflow_phase2L3MuonsOIFirst(UpgradeWorkflow): def setup_(self, step, stepName, stepDict, k, properties): @@ -1956,6 +1988,20 @@ def condition(self, fragment, stepList, key, hasHarvest): '-s':'HARVESTING:@hltValidation' } +upgradeWFs['HLTTiming75e33TiclV5TrackLinkingGNN'] = deepcopy(upgradeWFs['HLTTiming75e33']) +upgradeWFs['HLTTiming75e33TiclV5TrackLinkingGNN'].suffix = '_HLT75e33TimingTiclV5TrackLinkGNN' +upgradeWFs['HLTTiming75e33TiclV5TrackLinkingGNN'].offset = 0.7521 +upgradeWFs['HLTTiming75e33TiclV5TrackLinkingGNN'].step2 = { + '-s':'DIGI:pdigi_valid,L1TrackTrigger,L1,L1P2GT,DIGI2RAW,HLT:75e33_timing,VALIDATION:@hltValidation', + '--procModifiers': 'ticlv5_TrackLinkingGNN', + '--datatier':'GEN-SIM-DIGI-RAW,DQMIO', + '--eventcontent':'FEVTDEBUGHLT,DQMIO' +} +upgradeWFs['HLTTiming75e33TiclV5TrackLinkingGNN'].step3 = { + '-s':'HARVESTING:@hltValidation' +} + + upgradeWFs['HLTTiming75e33AlpakaSingleIter'] = deepcopy(upgradeWFs['HLTTiming75e33']) upgradeWFs['HLTTiming75e33AlpakaSingleIter'].suffix = '_HLT75e33TimingAlpakaSingleIter' upgradeWFs['HLTTiming75e33AlpakaSingleIter'].offset = 0.753 diff --git a/Configuration/PyReleaseValidation/scripts/runTheMatrix.py b/Configuration/PyReleaseValidation/scripts/runTheMatrix.py index f5232e5bc2ef9..0ce4c8fa17635 100755 --- a/Configuration/PyReleaseValidation/scripts/runTheMatrix.py +++ b/Configuration/PyReleaseValidation/scripts/runTheMatrix.py @@ -163,6 +163,7 @@ def runSelected(opt): prefixDet+34.751, # HLT phase-2 timing menu Alpaka variant prefixDet+34.7511, # HLT phase-2 timing menu Phase2CAExtension variant prefixDet+34.752, # HLT phase-2 timing menu ticl_v5 variant + prefixDet+34.7521, # HLT phase-2 timing menu ticlv5TrackLinkGNN variant prefixDet+34.753, # HLT phase-2 timing menu Alpaka, single tracking iteration variant prefixDet+34.754, # HLT phase-2 timing menu Alpaka, single tracking iteration, LST building variant prefixDet+34.755, # HLT phase-2 timing menu Alpaka, LST building variant diff --git a/HLTrigger/Configuration/python/HLT_75e33/modules/hltTiclCandidate_cfi.py b/HLTrigger/Configuration/python/HLT_75e33/modules/hltTiclCandidate_cfi.py index ccb847a54447f..9a1318caf0363 100644 --- a/HLTrigger/Configuration/python/HLT_75e33/modules/hltTiclCandidate_cfi.py +++ b/HLTrigger/Configuration/python/HLT_75e33/modules/hltTiclCandidate_cfi.py @@ -47,3 +47,15 @@ useTimingAverage = cms.bool(False) ) +from Configuration.ProcessModifiers.ticlv5_TrackLinkingGNN_cff import ticl_v5_TrackLinkingGNN +ticl_v5_TrackLinkingGNN.toModify(hltTiclCandidate, + interpretationDescPSet = cms.PSet( + onnxTrkLinkingModelFirstDisk = cms.FileInPath('RecoHGCal/TICL/data/ticlv5/onnx_models/TrackLinking_GNN/FirstDiskPropGNN_v0.onnx'), + onnxTrkLinkingModelInterfaceDisk = cms.FileInPath('RecoHGCal/TICL/data/ticlv5/onnx_models/TrackLinking_GNN/InterfaceDiskPropGNN_v0.onnx'), + inputNames = cms.vstring('x', 'edge_index', 'edge_attr'), + output = cms.vstring('output'), + delta_tk_ts = cms.double(0.1), + thr_gnn = cms.double(0.5), + type = cms.string('GNNLink') + ) +) diff --git a/RecoHGCal/TICL/interface/TICLInterpretationAlgoBase.h b/RecoHGCal/TICL/interface/TICLInterpretationAlgoBase.h index 3c6a830561203..389b680af4f13 100644 --- a/RecoHGCal/TICL/interface/TICLInterpretationAlgoBase.h +++ b/RecoHGCal/TICL/interface/TICLInterpretationAlgoBase.h @@ -23,6 +23,7 @@ #include "FWCore/Framework/interface/ConsumesCollector.h" #include "PhysicsTools/TensorFlow/interface/TensorFlow.h" #include "DataFormats/Common/interface/MultiSpan.h" +#include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h" namespace edm { class Event; diff --git a/RecoHGCal/TICL/plugins/GNNInterpretationAlgo.cc b/RecoHGCal/TICL/plugins/GNNInterpretationAlgo.cc new file mode 100644 index 0000000000000..c59cfdce4324e --- /dev/null +++ b/RecoHGCal/TICL/plugins/GNNInterpretationAlgo.cc @@ -0,0 +1,572 @@ +// Author: Mohamed Darwish +#include "RecoHGCal/TICL/interface/TICLInterpretationAlgoBase.h" +#include "RecoHGCal/TICL/plugins/GNNInterpretationAlgo.h" + +#include "RecoParticleFlow/PFProducer/interface/PFMuonAlgo.h" +#include "TrackingTools/TrajectoryState/interface/TrajectoryStateTransform.h" + +using namespace ticl; +using namespace cms::Ort; + +using Vector = ticl::Trackster::Vector; + +GNNInterpretationAlgo::~GNNInterpretationAlgo() {} + +GNNInterpretationAlgo::GNNInterpretationAlgo(const edm::ParameterSet& conf, edm::ConsumesCollector cc) + : TICLInterpretationAlgoBase(conf, cc), + onnxLinkingRuntimeFirstDisk_(std::make_unique( + conf.getParameter("onnxTrkLinkingModelFirstDisk").fullPath().c_str())), + onnxLinkingRuntimeInterfaceDisk_(std::make_unique( + conf.getParameter("onnxTrkLinkingModelInterfaceDisk").fullPath().c_str())), + inputNames_(conf.getParameter>("inputNames")), + output_(conf.getParameter>("output")), + del_tk_ts_(conf.getParameter("delta_tk_ts")), + threshold_(conf.getParameter("thr_gnn")) { + onnxLinkingSessionFirstDisk_ = onnxLinkingRuntimeFirstDisk_.get(); + onnxLinkingSessionInterfaceDisk_ = onnxLinkingRuntimeInterfaceDisk_.get(); +} + +// Initialization +void GNNInterpretationAlgo::initialize(const HGCalDDDConstants* hgcons, + const hgcal::RecHitTools rhtools, + const edm::ESHandle bfieldH, + const edm::ESHandle propH) { + hgcons_ = hgcons; + rhtools_ = rhtools; + bfield_ = bfieldH; + propagator_ = propH; + + buildLayers(); +} +// Geometry construction +void GNNInterpretationAlgo::buildLayers() { + // Build propagation disks at: + // - HGCal front face + // - CE-E CE-H interface + + const float z_front = hgcons_->waferZ(1, true); + const auto r_front = hgcons_->rangeR(z_front, true); + + const float z_interface = rhtools_.getPositionLayer(rhtools_.lastLayerEE()).z(); + const auto r_interface = hgcons_->rangeR(z_interface, true); + + for (int side = 0; side < 2; ++side) { + const float sign = (side == 0 ? -1.f : 1.f); + + firstDisk_[side] = std::make_unique( + Disk::build(Disk::PositionType(0, 0, sign * z_front), + Disk::RotationType(), + SimpleDiskBounds(r_front.first, r_front.second, sign * z_front - 0.5f, sign * z_front + 0.5f)) + .get()); + + interfaceDisk_[side] = std::make_unique( + Disk::build(Disk::PositionType(0, 0, sign * z_interface), + Disk::RotationType(), + SimpleDiskBounds( + r_interface.first, r_interface.second, sign * z_interface - 0.5f, sign * z_interface + 0.5f)) + .get()); + } +} + +// Trackster propagation +Vector GNNInterpretationAlgo::propagateTrackster(const Trackster& t, + unsigned idx, + float zVal, + std::array& tracksterTiles) { + // needs only the positive Z co-ordinate of the surface to propagate to + // the correct sign is calculated inside according to the barycenter of trackster + + // Propagation direction + const Vector& barycenter = t.barycenter(); + + // NOTE: + // barycenter as direction for tracksters w/ poor PCA + // propagation still done to get the cartesian coords + // which are anyway converted to eta, phi in linking + // -> can be simplified later + + //FP: disable PCA propagation for the moment and fallback to barycenter position + // if (t.eigenvalues()[0] / t.eigenvalues()[1] < 20) + + Vector direction = barycenter.unit(); + + // Ensure correct Z-side propagation + zVal *= (barycenter.Z() > 0.f ? 1.f : -1.f); + + const float scale = (zVal - barycenter.Z()) / direction.Z(); + const Vector propPoint(scale * direction.X() + barycenter.X(), scale * direction.Y() + barycenter.Y(), zVal); + + // Fill spatial tiles for fast lookup + const bool isPositiveZ = (propPoint.Eta() > 0.f); + tracksterTiles[isPositiveZ].fill(propPoint.Eta(), propPoint.Phi(), idx); + + return propPoint; +} + +std::pair GNNInterpretationAlgo::calculateTrackstersError(const Trackster& trackster) { + const auto& barycenter = trackster.barycenter(); + const float x = barycenter.x(), y = barycenter.y(), z = barycenter.z(); + + const auto& s = trackster.sigmasPCA(); + const float s1 = s[0] * s[0], s2 = s[1] * s[1], s3 = s[2] * s[2]; + + const auto& v1 = trackster.eigenvectors()[0]; + const auto& v2 = trackster.eigenvectors()[1]; + const auto& v3 = trackster.eigenvectors()[2]; + + // Covariance in XY from 3D + const float cxx = s1 * v1.x() * v1.x() + s2 * v2.x() * v2.x() + s3 * v3.x() * v3.x(); + const float cxy = s1 * v1.x() * v1.y() + s2 * v2.x() * v2.y() + s3 * v3.x() * v3.y(); + const float cyy = s1 * v1.y() * v1.y() + s2 * v2.y() * v2.y() + s3 * v3.y() * v3.y(); + + // Geometry helpers + const float r2 = x * x + y * y; + const float denom_eta = r2 * (r2 + z * z); + const float sqrt_term = std::sqrt(r2 / (z * z) + 1); + + // Jacobian elements + const float J00 = -(x * z * z * sqrt_term) / denom_eta; + const float J01 = -(y * z * z * sqrt_term) / denom_eta; + const float J10 = -y / r2; + const float J11 = x / r2; + + // CovEtaPhi = J * CovXY * J^T + const float cee = J00 * (J00 * cxx + J01 * cxy) + J01 * (J00 * cxy + J01 * cyy); + const float cpp = J10 * (J10 * cxx + J11 * cxy) + J11 * (J10 * cxy + J11 * cyy); + + return {std::sqrt(std::abs(cee)), std::sqrt(std::abs(cpp))}; +} + +void GNNInterpretationAlgo::constructNodeFromWindow( + const edm::MultiSpan& tracksters, + const std::vector>& seeding, + const std::array& tracksterTiles, + const std::vector& tracksterPropPoints, + float delta2, + unsigned trackstersSize, + std::vector& graph) { + const float delta = 0.5f * delta2; + + for (const auto& [seedPos, seedIdx, _] : seeding) { + const float seedEta = seedPos.Eta(); + const float seedPhi = seedPos.Phi(); + const bool isPositiveZ = (seedEta > 0.0f); + + const TICLLayerTile& tile = tracksterTiles[isPositiveZ]; + + const float etaMin = std::max(std::abs(seedEta) - delta, static_cast(TileConstants::minEta)); + const float etaMax = std::min(std::abs(seedEta) + delta, static_cast(TileConstants::maxEta)); + + const auto searchBox = tile.searchBoxEtaPhi(etaMin, etaMax, seedPhi - delta, seedPhi + delta); + + ticl::Node node(seedIdx, false); + + for (int iEta = searchBox[0]; iEta <= searchBox[1]; ++iEta) { + for (int iPhi = searchBox[2]; iPhi <= searchBox[3]; ++iPhi) { + const int globalBin = tile.globalBin(iEta, iPhi % TileConstants::nPhiBins); + + const auto& candidates = tile[globalBin]; + + for (const unsigned tsIdx : candidates) { + if (tsIdx >= trackstersSize) + continue; + + const float sep2 = + reco::deltaR2(tracksterPropPoints[tsIdx].Eta(), tracksterPropPoints[tsIdx].Phi(), seedEta, seedPhi); + if (sep2 < delta2) { + node.addOuterNeighbour(tsIdx); + } + } + } + } + graph.emplace_back(std::move(node)); + } +} + +std::vector GNNInterpretationAlgo::padFeatures(const std::vector& core_feats, + size_t track_block_size, + size_t trackster_block_size, + bool isTrack) { + std::vector out; + out.reserve(track_block_size + trackster_block_size); + + if (isTrack) { + out.insert(out.end(), core_feats.begin(), core_feats.end()); + out.insert(out.end(), trackster_block_size, 0.f); + } else { + out.insert(out.end(), track_block_size, 0.f); + out.insert(out.end(), core_feats.begin(), core_feats.end()); + } + + return out; +} +void GNNInterpretationAlgo::buildGraphFromNodes(const std::tuple& TrackInfo, + const reco::Track& track, + const edm::MultiSpan& tracksters, + const std::vector& clusters, + const std::vector& nodeVec, + GraphData& outGraphData) { + outGraphData = {}; // clear previous data + + std::unordered_map> track_node_features; + std::unordered_map> trackster_node_features; + + const auto& [pos, localErrMatrix, track_idx] = TrackInfo; + + // Track coordinates and covariances + const float eta = pos.Eta(), phi = pos.Phi(); + const float x = pos.X(), y = pos.Y(), z = pos.Z(); + + AlgebraicMatrix22 covMatrixXY; + covMatrixXY(0, 0) = localErrMatrix(3, 3); + covMatrixXY(0, 1) = localErrMatrix(3, 4); + covMatrixXY(1, 0) = localErrMatrix(3, 4); + covMatrixXY(1, 1) = localErrMatrix(4, 4); + + const float sqrt_term = std::sqrt((x * x + y * y) / (z * z) + 1); + const float denom_eta = (x * x + y * y) * (x * x + y * y + z * z); + const float denom_phi = x * x + y * y; + + AlgebraicMatrix22 jacobian; + jacobian(0, 0) = -(x * z * z * sqrt_term) / denom_eta; + jacobian(0, 1) = -(y * z * z * sqrt_term) / denom_eta; + jacobian(1, 0) = -y / denom_phi; + jacobian(1, 1) = x / denom_phi; + + AlgebraicMatrix22 covMatrixEtaPhi = ROOT::Math::Transpose(jacobian) * covMatrixXY * jacobian; + const float track_etaErr = std::sqrt(covMatrixEtaPhi(0, 0)); + const float track_phiErr = std::sqrt(covMatrixEtaPhi(1, 1)); + + const float track_p = track.p(); + const float track_pt = track.pt(); + const float trackHits = track.recHitsSize(); + + std::vector trk_feats = { + std::abs(eta), phi, track_etaErr, track_phiErr, x, y, std::abs(z), track_p, track_pt, trackHits}; + trk_feats.resize(track_block_size); + + // Fill node features + for (const auto& node : nodeVec) { + if (!node.isTrackster() && static_cast(node.getId()) == track_idx) { + track_node_features[node.getId()] = padFeatures(trk_feats, track_block_size, trackster_block_size, true); + + for (unsigned ts_idx : node.getOuterNeighbours()) { + if (ts_idx >= tracksters.size()) + continue; + + if (!trackster_node_features.count(ts_idx)) { + const auto& ts = tracksters[ts_idx]; + auto [errEta, errPhi] = calculateTrackstersError(ts); + + std::vector ts_feats = {std::abs(ts.barycenter().eta()), + ts.barycenter().phi(), + errEta, + errPhi, + ts.barycenter().x(), + ts.barycenter().y(), + std::abs(ts.barycenter().z()), + ts.raw_energy(), + ts.time(), + ts.timeError(), + ts.raw_em_energy(), + ts.raw_em_pt(), + ts.raw_pt()}; + ts_feats.resize(trackster_block_size); + trackster_node_features[ts_idx] = padFeatures(ts_feats, track_block_size, trackster_block_size, false); + } + outGraphData.edge_index.emplace_back(node.getId(), ts_idx); + } + } + } + + // Insert nodes + size_t row_idx = 0; + for (const auto& [idx, feats] : track_node_features) { + outGraphData.nodeIndexToRow[{false, idx}] = row_idx++; + outGraphData.node_features.push_back(feats); + } + for (const auto& [idx, feats] : trackster_node_features) { + outGraphData.nodeIndexToRow[{true, idx}] = row_idx++; + outGraphData.node_features.push_back(feats); + } + outGraphData.num_nodes = outGraphData.node_features.size(); + + // Fill edge attributes + for (const auto& edge : outGraphData.edge_index) { + const NodeKey src_key{false, edge.first}; + const NodeKey dst_key{true, edge.second}; + + if (!outGraphData.nodeIndexToRow.count(src_key) || !outGraphData.nodeIndexToRow.count(dst_key)) + continue; + + const auto& src_feats = outGraphData.node_features[outGraphData.nodeIndexToRow[src_key]]; + const auto& dst_feats = outGraphData.node_features[outGraphData.nodeIndexToRow[dst_key]]; + + const int trkster_offset = track_block_size; + + const float trk_eta = src_feats[0], trk_phi = src_feats[1]; + const float ts_eta = dst_feats[trkster_offset], ts_phi = dst_feats[trkster_offset + 1]; + + const float delta_eta = trk_eta - ts_eta; + const float delta_phi = reco::deltaPhi(trk_phi, ts_phi); + const float deta_sig = delta_eta / std::sqrt(dst_feats[trkster_offset + 2] * dst_feats[trkster_offset + 2] + + src_feats[2] * src_feats[2] + 1e-8f); + const float dphi_sig = delta_phi / std::sqrt(dst_feats[trkster_offset + 3] * dst_feats[trkster_offset + 3] + + src_feats[3] * src_feats[3] + 1e-8f); + const float deltaR = std::sqrt(delta_eta * delta_eta + delta_phi * delta_phi); + + const float dx = dst_feats[trkster_offset + 4] - src_feats[4]; + const float dy = dst_feats[trkster_offset + 5] - src_feats[5]; + const float dz = dst_feats[trkster_offset + 6] - src_feats[6]; + const float dist3D = std::sqrt(dx * dx + dy * dy + dz * dz); + const float distXY = std::sqrt(dx * dx + dy * dy); + + const float dE = dst_feats[trkster_offset + 7] - src_feats[7]; + const float E_ratio = dst_feats[trkster_offset + 7] / (src_feats[7] + 1e-8f); + + const auto& ts = tracksters[edge.second]; + const auto& vertices = ts.vertices(); + float min_dist = std::numeric_limits::max(); + float max_dist = 0.f; + + for (const auto& vtx : vertices) { + const auto& cl = clusters[vtx]; + const float dist = std::sqrt(std::pow(cl.x() - src_feats[4], 2) + std::pow(cl.y() - src_feats[5], 2) + + std::pow(std::abs(cl.z()) - src_feats[6], 2)); + min_dist = std::min(min_dist, dist); + max_dist = std::max(max_dist, dist); + } + + outGraphData.edge_attr.push_back( + {delta_eta, delta_phi, deta_sig, dphi_sig, deltaR, dist3D, distXY, dE, E_ratio, min_dist, max_dist}); + } +} + +void GNNInterpretationAlgo::makeCandidates(const Inputs& input, + edm::Handle inputTiming_h, + std::vector& resultTracksters, + std::vector& resultCandidate) { + const auto& tracks = *input.tracksHandle; + const auto& maskTracks = input.maskedTracks; + const auto& tracksters = input.tracksters; + const auto& clusters = input.layerClusters; + + const auto bFieldProd = bfield_.product(); + const Propagator& prop = (*propagator_); + + // propagated point collections + // elements in the propagated points collecions are used + // to look for potential linkages in the appropriate tiles + // Track propagation + using TrackPropInfo = std::tuple; + + std::vector tkPropFront; // propagated to first disk + std::vector tkPropInt; // propagated to interface disk + tkPropFront.reserve(tracks.size()); + tkPropInt.reserve(tracks.size()); + + std::vector candidateTrackIds; + candidateTrackIds.reserve(tracks.size()); + + for (unsigned i = 0; i < tracks.size(); ++i) { + if (maskTracks[i]) + candidateTrackIds.push_back(i); + } + std::sort(candidateTrackIds.begin(), candidateTrackIds.end(), [&](unsigned i, unsigned j) { + return tracks[i].p() > tracks[j].p(); + }); + + for (unsigned trkId : candidateTrackIds) { + const auto& tk = tracks[trkId]; + const int side = (tk.eta() > 0); + + const auto& fts = trajectoryStateTransform::outerFreeState(tk, bFieldProd); + // to the HGCal front + const auto tsosFront = prop.propagate(fts, firstDisk_[side]->surface()); + if (tsosFront.isValid()) { + tkPropFront.emplace_back( + Vector(tsosFront.globalPosition().x(), tsosFront.globalPosition().y(), tsosFront.globalPosition().z()), + trkId, + tsosFront.localError().matrix()); + } + + // Interface disk + const auto tsosInt = prop.propagate(fts, interfaceDisk_[side]->surface()); + if (tsosInt.isValid()) { + tkPropInt.emplace_back( + Vector(tsosInt.globalPosition().x(), tsosInt.globalPosition().y(), tsosInt.globalPosition().z()), + trkId, + tsosInt.localError().matrix()); + } + } // Tracks + tkPropFront.shrink_to_fit(); + tkPropInt.shrink_to_fit(); + candidateTrackIds.shrink_to_fit(); + // Propagate tracksters + // Record postions of all tracksters propagated to layer 1 and lastLayerEE, + // to be used later for distance calculation in the link finding stage + // indexed by trackster index in event collection + std::array tsTilesFront = {}; + std::array tsTilesInt = {}; + + std::vector tsPropFront, tsPropInt; + tsPropFront.reserve(tracksters.size()); + tsPropInt.reserve(tracksters.size()); + + for (unsigned i = 0; i < tracksters.size(); ++i) { + const auto& ts = tracksters[i]; + + float zFront = hgcons_->waferZ(1, true); + tsPropFront.emplace_back(propagateTrackster(ts, i, zFront, tsTilesFront)); + + float zInt = rhtools_.getPositionLayer(rhtools_.lastLayerEE()).z(); + tsPropInt.emplace_back(propagateTrackster(ts, i, zInt, tsTilesInt)); + } + + // Step 1: Construct nodes from tracksters and tracks + std::vector nodesFront, nodesInt; + + constructNodeFromWindow( + tracksters, tkPropFront, tsTilesFront, tsPropFront, del_tk_ts_, tracksters.size(), nodesFront); + + constructNodeFromWindow(tracksters, tkPropInt, tsTilesInt, tsPropInt, del_tk_ts_, tracksters.size(), nodesInt); + + std::vector> trackToTracksters(tracks.size()); + std::vector>> trackToScores(tracks.size()); + std::vector tracksterAvailable(tracksters.size(), true); + + auto runInferenceForTrack = [&](unsigned trkId, + const std::vector& tkProps, + std::vector& nodes, + bool useInterfaceModel) { + auto it = + std::find_if(tkProps.begin(), tkProps.end(), [&](const auto& info) { return std::get<1>(info) == trkId; }); + + if (it == tkProps.end()) + return; + + GraphData graphData; + buildGraphFromNodes(std::make_tuple(std::get<0>(*it), std::get<2>(*it), trkId), + tracks[trkId], + tracksters, + clusters, + nodes, + graphData); + + // Prepare ONNX input + std::vector> inputData(3); + std::vector> inputShapes; + + const int64_t nNodes = graphData.node_features.size(); + const int64_t nNodeFeat = nNodes ? graphData.node_features.front().size() : 0; + + for (const auto& feat : graphData.node_features) + inputData[0].insert(inputData[0].end(), feat.begin(), feat.end()); + + inputShapes.push_back({nNodes, nNodeFeat}); + + std::vector src_nodes, dst_nodes; + for (const auto& edge : graphData.edge_index) { + NodeKey src_key = {false, edge.first}; + NodeKey dst_key = {true, edge.second}; + src_nodes.push_back(graphData.nodeIndexToRow.at(src_key)); + dst_nodes.push_back(graphData.nodeIndexToRow.at(dst_key)); + } + inputData[1].insert(inputData[1].end(), src_nodes.begin(), src_nodes.end()); + inputData[1].insert(inputData[1].end(), dst_nodes.begin(), dst_nodes.end()); + inputShapes.push_back({2, static_cast(graphData.edge_index.size())}); + + const int64_t nEdges = graphData.edge_attr.size(); + const int64_t nEdgeFeat = nEdges ? graphData.edge_attr.front().size() : 0; + + for (const auto& attr : graphData.edge_attr) + inputData[2].insert(inputData[2].end(), attr.begin(), attr.end()); + + inputShapes.push_back({nEdges, nEdgeFeat}); + + if (inputData[1].empty()) + return; + const auto& outputs = useInterfaceModel + ? onnxLinkingSessionInterfaceDisk_->run(inputNames_, inputData, inputShapes, output_) + : onnxLinkingSessionFirstDisk_->run(inputNames_, inputData, inputShapes, output_); + + const auto& scores = outputs[0]; + + for (size_t i = 0; i < graphData.edge_index.size(); ++i) { + if (scores[i] <= threshold_) + continue; + const auto& edge = graphData.edge_index[i]; + const float deltaR = graphData.edge_attr[i][4]; + const float score = std::log(tracks[trkId].pt() / (deltaR + 1e-5f)); + + trackToScores[trkId].emplace_back(edge.second, score); + } + }; + + for (unsigned trkId : candidateTrackIds) { + runInferenceForTrack(trkId, tkPropFront, nodesFront, false); //First disk + runInferenceForTrack(trkId, tkPropInt, nodesInt, true); //Interface disk + } + // Resolve global associations + std::vector> allLinks; + + for (unsigned trkId = 0; trkId < trackToScores.size(); ++trkId) { + for (const auto& [tsId, score] : trackToScores[trkId]) + allLinks.emplace_back(trkId, tsId, score); + } + + std::sort( + allLinks.begin(), allLinks.end(), [](const auto& a, const auto& b) { return std::get<2>(a) > std::get<2>(b); }); + for (const auto& [trkId, tsId, score] : allLinks) { + if (tracksterAvailable[tsId]) { + trackToTracksters[trkId].push_back(tsId); + tracksterAvailable[tsId] = false; + } + } + // Build output tracksters + + for (unsigned trkId = 0; trkId < trackToTracksters.size(); ++trkId) { + if (trackToTracksters[trkId].empty()) + continue; + + resultCandidate[trkId] = resultTracksters.size(); + + if (trackToTracksters[trkId].size() == 1) { + resultTracksters.push_back(tracksters[trackToTracksters[trkId][0]]); + } else { + Trackster merged; + merged.mergeTracksters(tracksters, trackToTracksters[trkId]); + + bool hasHadron = false; + for (auto tsId : trackToTracksters[trkId]) + hasHadron |= tracksters[tsId].isHadronic(); + merged.setIdProbability(hasHadron ? Trackster::ParticleType::charged_hadron : Trackster::ParticleType::electron, + 1.f); + + resultTracksters.push_back(std::move(merged)); + } + } + + // Add unlinked tracksters + for (unsigned i = 0; i < tracksters.size(); ++i) { + if (tracksterAvailable[i]) + resultTracksters.push_back(tracksters[i]); + } +} + +void GNNInterpretationAlgo::fillPSetDescription(edm::ParameterSetDescription& desc) { + desc.add( + "onnxTrkLinkingModelFirstDisk", + edm::FileInPath("RecoHGCal/TICL/data/ticlv5/onnx_models/TrackLinking_GNN/FirstDiskPropGNN_v0.onnx")) + ->setComment("Path to ONNX tracks tracksters linking model at first disk "); + desc.add( + "onnxTrkLinkingModelInterfaceDisk", + edm::FileInPath("RecoHGCal/TICL/data/ticlv5/onnx_models/TrackLinking_GNN/InterfaceDiskPropGNN_v0.onnx")) + ->setComment("Path to ONNX tracks tracksters linking model at interface disk "); + desc.add>("inputNames", {"x", "edge_index", "edge_attr"}); + desc.add>("output", {"output"}); + desc.add("delta_tk_ts", 0.1); + desc.add("thr_gnn", 0.5); + + TICLInterpretationAlgoBase::fillPSetDescription(desc); +} diff --git a/RecoHGCal/TICL/plugins/GNNInterpretationAlgo.h b/RecoHGCal/TICL/plugins/GNNInterpretationAlgo.h new file mode 100644 index 0000000000000..56f7e947040bb --- /dev/null +++ b/RecoHGCal/TICL/plugins/GNNInterpretationAlgo.h @@ -0,0 +1,108 @@ +#ifndef RecoHGCal_TICL_GNNInterpretationAlgo_H_ +#define RecoHGCal_TICL_GNNInterpretationAlgo_H_ + +#include "FWCore/Framework/interface/Frameworkfwd.h" +#include "FWCore/Framework/interface/ESHandle.h" +#include "FWCore/Framework/interface/MakerMacros.h" +#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h" +#include "RecoHGCal/TICL/interface/TICLInterpretationAlgoBase.h" +#include "DataFormats/TrackReco/interface/Track.h" +#include "DataFormats/GeometrySurface/interface/BoundDisk.h" +#include "DataFormats/Math/interface/deltaR.h" +#include "RecoHGCal/TICL/plugins/TICLGraph.h" +#include "TMatrixDSym.h" +#include "TMatrixD.h" + +namespace ticl { + + using NodeKey = std::pair; + + struct pair_hash { + template + std::size_t operator()(const std::pair &p) const { + auto h1 = std::hash{}(p.first); + auto h2 = std::hash{}(p.second); + return h1 ^ (h2 << 1); + } + }; + + struct GraphData { + std::vector> node_features; + std::vector> edge_index; + std::vector> edge_attr; + std::unordered_map nodeIndexToRow; + int num_nodes; + }; + + class GNNInterpretationAlgo : public TICLInterpretationAlgoBase { + public: + GNNInterpretationAlgo(const edm::ParameterSet &conf, edm::ConsumesCollector iC); + + ~GNNInterpretationAlgo() override; + + void makeCandidates(const Inputs &input, + edm::Handle inputTiming_h, + std::vector &resultTracksters, + std::vector &resultCandidate) override; + + void initialize(const HGCalDDDConstants *hgcons, + const hgcal::RecHitTools rhtools, + const edm::ESHandle bfieldH, + const edm::ESHandle propH) override; + + static void fillPSetDescription(edm::ParameterSetDescription &iDesc); + + private: + void buildLayers(); + const std::unique_ptr onnxLinkingRuntimeFirstDisk_; + const cms::Ort::ONNXRuntime *onnxLinkingSessionFirstDisk_; + const std::unique_ptr onnxLinkingRuntimeInterfaceDisk_; + const cms::Ort::ONNXRuntime *onnxLinkingSessionInterfaceDisk_; + const std::vector inputNames_; + const std::vector output_; + + Vector propagateTrackster(const Trackster &t, + const unsigned idx, + float zVal, + std::array &tracksterTiles); + + std::pair calculateTrackstersError(const Trackster &trackster); + std::vector padFeatures(const std::vector &core_feats, + size_t track_block_size, + size_t trackster_block_size, + bool isTrack); + void constructNodeFromWindow(const edm::MultiSpan &tracksters, + const std::vector> &seeding, + const std::array &tracksterTiles, + const std::vector &tracksterPropPoints, + float delta, + unsigned trackstersSize, + std::vector &graph); + void printGraphSummary(const GraphData &graphData); + void buildGraphFromNodes(const std::tuple &TrackInfo, + const reco::Track &track, + const edm::MultiSpan &tracksters, + const std::vector &clusters, + const std::vector &nodeVec, + GraphData &outGraphData); + + const float del_tk_ts_; + const float threshold_; + + const size_t track_block_size = 10; // number of track features + const size_t trackster_block_size = 13; // number of trackster features + + const HGCalDDDConstants *hgcons_; + + std::unique_ptr firstDisk_[2]; + std::unique_ptr interfaceDisk_[2]; + + hgcal::RecHitTools rhtools_; + + edm::ESHandle bfield_; + edm::ESHandle propagator_; + }; + +} // namespace ticl + +#endif diff --git a/RecoHGCal/TICL/plugins/TICLCandidateProducer.cc b/RecoHGCal/TICL/plugins/TICLCandidateProducer.cc index 6231c6ad64759..dc88b0526f0f1 100644 --- a/RecoHGCal/TICL/plugins/TICLCandidateProducer.cc +++ b/RecoHGCal/TICL/plugins/TICLCandidateProducer.cc @@ -31,6 +31,7 @@ #include "RecoHGCal/TICL/interface/TICLInterpretationAlgoBase.h" #include "RecoHGCal/TICL/plugins/TICLInterpretationPluginFactory.h" #include "RecoHGCal/TICL/plugins/GeneralInterpretationAlgo.h" +#include "RecoHGCal/TICL/plugins/GNNInterpretationAlgo.h" #include "RecoParticleFlow/PFProducer/interface/PFMuonAlgo.h" diff --git a/RecoHGCal/TICL/plugins/TICLGraph.h b/RecoHGCal/TICL/plugins/TICLGraph.h index c78123bdf6577..6f2e282a12bf6 100644 --- a/RecoHGCal/TICL/plugins/TICLGraph.h +++ b/RecoHGCal/TICL/plugins/TICLGraph.h @@ -25,6 +25,7 @@ namespace ticl { return findInner != innerNeighboursId_.end(); } inline bool alreadyVisited() const { return alreadyVisited_; } + inline bool isTrackster() const { return isTrackster_; } ~Node() = default; diff --git a/RecoHGCal/TICL/plugins/TICLInterpretationPluginFactory.cc b/RecoHGCal/TICL/plugins/TICLInterpretationPluginFactory.cc index 00681086d3a58..20a3aa62510c3 100644 --- a/RecoHGCal/TICL/plugins/TICLInterpretationPluginFactory.cc +++ b/RecoHGCal/TICL/plugins/TICLInterpretationPluginFactory.cc @@ -4,8 +4,10 @@ #include "FWCore/ParameterSet/interface/ValidatedPluginMacros.h" #include "RecoHGCal/TICL/plugins/TICLInterpretationPluginFactory.h" #include "GeneralInterpretationAlgo.h" +#include "GNNInterpretationAlgo.h" EDM_REGISTER_VALIDATED_PLUGINFACTORY(TICLGeneralInterpretationPluginFactory, "TICLGeneralInterpretationPluginFactory"); EDM_REGISTER_VALIDATED_PLUGINFACTORY(TICLEGammaInterpretationPluginFactory, "TICLEGammaInterpretationPluginFactory"); DEFINE_EDM_VALIDATED_PLUGIN(TICLGeneralInterpretationPluginFactory, ticl::GeneralInterpretationAlgo, "General"); +DEFINE_EDM_VALIDATED_PLUGIN(TICLGeneralInterpretationPluginFactory, ticl::GNNInterpretationAlgo, "GNNLink"); // DEFINE_EDM_VALIDATED_PLUGIN(TICLEGammaInterpretationPluginFactory, ticl::EGammaInterpretation, "EGamma"); diff --git a/RecoHGCal/TICL/python/iterativeTICL_cff.py b/RecoHGCal/TICL/python/iterativeTICL_cff.py index 27b1ee2df6c3d..1dc53b7123876 100644 --- a/RecoHGCal/TICL/python/iterativeTICL_cff.py +++ b/RecoHGCal/TICL/python/iterativeTICL_cff.py @@ -23,6 +23,8 @@ from RecoHGCal.TICL.mtdSoAProducer_cfi import mtdSoAProducer as _mtdSoAProducer from Configuration.ProcessModifiers.ticl_v5_cff import ticl_v5 +from Configuration.ProcessModifiers.ticlv5_TrackLinkingGNN_cff import ticl_v5_TrackLinkingGNN + from Configuration.ProcessModifiers.ticl_superclustering_dnn_cff import ticl_superclustering_dnn from Configuration.ProcessModifiers.ticl_superclustering_mustache_pf_cff import ticl_superclustering_mustache_pf from Configuration.ProcessModifiers.ticl_superclustering_mustache_ticl_cff import ticl_superclustering_mustache_ticl @@ -144,6 +146,17 @@ pfTICL = _pfTICLProducer.clone() ticl_v5.toModify(pfTICL, ticlCandidateSrc = cms.InputTag('ticlCandidate'), isTICLv5 = cms.bool(True), useTimingAverage=True) +ticl_v5_TrackLinkingGNN.toModify(ticlCandidate, + interpretationDescPSet = cms.PSet( + onnxTrkLinkingModelFirstDisk = cms.FileInPath('RecoHGCal/TICL/data/ticlv5/onnx_models/TrackLinking_GNN/FirstDiskPropGNN_v0.onnx'), + onnxTrkLinkingModelInterfaceDisk = cms.FileInPath('RecoHGCal/TICL/data/ticlv5/onnx_models/TrackLinking_GNN/InterfaceDiskPropGNN_v0.onnx'), + inputNames = cms.vstring('x', 'edge_index', 'edge_attr'), + output = cms.vstring('output'), + delta_tk_ts = cms.double(0.1), + thr_gnn = cms.double(0.5), + type = cms.string('GNNLink') + ) + ) ticlPFTask = cms.Task(pfTICL)